odbc_api/buffers/columnar.rs
1use std::{
2 collections::HashSet,
3 num::NonZeroUsize,
4 str::{Utf8Error, from_utf8},
5};
6
7use crate::{
8 Error, ResultSetMetadata, RowSetBuffer,
9 columnar_bulk_inserter::BoundInputSlice,
10 cursor::TruncationInfo,
11 fixed_sized::Pod,
12 handles::{CDataMut, Statement, StatementRef},
13 parameter::WithDataType,
14 result_set_metadata::utf8_display_sizes,
15};
16
17use super::{Indicator, TextColumn};
18
19impl<C: ColumnBuffer> ColumnarBuffer<C> {
20 /// Create a new instance from columns with unique indicies. Capacity of the buffer will be the
21 /// minimum capacity of the columns. The constructed buffer is always empty (i.e. the number of
22 /// valid rows is considered to be zero).
23 ///
24 /// You do not want to call this constructor directly unless you want to provide your own buffer
25 /// implentation. Most users of this crate may want to use the constructors like
26 /// [`crate::buffers::ColumnarAnyBuffer::from_descs`] or
27 /// [`crate::buffers::TextRowSet::from_max_str_lens`] instead.
28 pub fn new(columns: Vec<(u16, C)>) -> Self {
29 // Assert capacity
30 let capacity = columns
31 .iter()
32 .map(|(_, col)| col.capacity())
33 .min()
34 .unwrap_or(0);
35
36 // Assert uniqueness of indices
37 let mut indices = HashSet::new();
38 if columns
39 .iter()
40 .any(move |&(col_index, _)| !indices.insert(col_index))
41 {
42 panic!("Column indices must be unique.")
43 }
44
45 unsafe { Self::new_unchecked(capacity, columns) }
46 }
47
48 /// # Safety
49 ///
50 /// * Indices must be unique
51 /// * Columns all must have enough `capacity`.
52 pub unsafe fn new_unchecked(capacity: usize, columns: Vec<(u16, C)>) -> Self {
53 ColumnarBuffer {
54 num_rows: Box::new(0),
55 row_capacity: capacity,
56 columns,
57 }
58 }
59
60 /// Number of valid rows in the buffer.
61 pub fn num_rows(&self) -> usize {
62 *self.num_rows
63 }
64
65 /// Return the number of columns in the row set.
66 pub fn num_cols(&self) -> usize {
67 self.columns.len()
68 }
69}
70
71impl<C: Slice> ColumnarBuffer<C> {
72 /// Use this method to gain read access to the actual column data.
73 ///
74 /// # Parameters
75 ///
76 /// * `buffer_index`: Please note that the buffer index is not identical to the ODBC column
77 /// index. For one it is zero based. It also indexes the buffer bound, and not the columns of
78 /// the output result set. This is important, because not every column needs to be bound. Some
79 /// columns may simply be ignored. That being said, if every column of the output is bound in
80 /// the buffer, in the same order in which they are enumerated in the result set, the
81 /// relationship between column index and buffer index is `buffer_index = column_index - 1`.
82 pub fn column(&self, buffer_index: usize) -> C::Slice<'_> {
83 self.columns[buffer_index].1.slice(*self.num_rows)
84 }
85}
86
87unsafe impl<C> RowSetBuffer for ColumnarBuffer<C>
88where
89 C: ColumnBuffer,
90{
91 fn bind_type(&self) -> usize {
92 0 // Specify columnar binding
93 }
94
95 fn row_array_size(&self) -> usize {
96 self.row_capacity
97 }
98
99 fn mut_num_fetch_rows(&mut self) -> &mut usize {
100 self.num_rows.as_mut()
101 }
102
103 unsafe fn bind_colmuns_to_cursor(&mut self, mut cursor: StatementRef<'_>) -> Result<(), Error> {
104 unsafe {
105 for (col_number, column) in &mut self.columns {
106 cursor.bind_col(*col_number, column).into_result(&cursor)?;
107 }
108 }
109 Ok(())
110 }
111
112 fn find_truncation(&self) -> Option<TruncationInfo> {
113 self.columns
114 .iter()
115 .enumerate()
116 .find_map(|(buffer_index, (_col_index, col_buffer))| {
117 col_buffer
118 .has_truncated_values(*self.num_rows)
119 .map(|indicator| TruncationInfo {
120 indicator: indicator.length(),
121 buffer_index,
122 })
123 })
124 }
125}
126
127/// A columnar buffer intended to be bound with [crate::Cursor::bind_buffer] in order to obtain
128/// results from a cursor.
129///
130/// Binds to the result set column-wise. This is usually helpful in data engineering or data science
131/// tasks. This buffer type can be used in situations where the schema of the queried data is known
132/// at compile time, as well as for generic applications which work with a wide range of different
133/// data.
134///
135/// # Example: Fetching results column wise with `ColumnarDynBuffer`.
136///
137/// Consider querying a table with two columns `year` and `name`.
138///
139/// ```no_run
140/// use odbc_api::{
141/// Environment, Cursor, ConnectionOptions,
142/// buffers::{BufferDesc, Item, ColumnarDynBuffer},
143/// };
144///
145/// let env = Environment::new()?;
146///
147/// let batch_size = 1000; // Maximum number of rows in each row set
148/// let buffer_description = [
149/// // We know year to be a Nullable SMALLINT
150/// BufferDesc::I16 { nullable: true },
151/// // and name to be a required VARCHAR
152/// BufferDesc::Text { max_str_len: 255 },
153/// ];
154///
155/// /// Creates a columnar buffer fitting the buffer description with the capacity of `batch_size`.
156/// let mut buffer = ColumnarDynBuffer::from_descs(batch_size, buffer_description);
157///
158/// let mut conn = env.connect(
159/// "YourDatabase", "SA", "My@Test@Password1",
160/// ConnectionOptions::default(),
161/// )?;
162/// let query = "SELECT year, name FROM Birthdays;";
163/// let params = ();
164/// let timeout_sec = None;
165/// if let Some(cursor) = conn.execute(query, params, timeout_sec)? {
166/// // Bind buffer to cursor. We bind the buffer as a mutable reference here, which makes it
167/// // easier to reuse for other queries, but we could have taken ownership.
168/// let mut row_set_cursor = cursor.bind_buffer(&mut buffer)?;
169/// // Loop over row sets
170/// while let Some(row_set) = row_set_cursor.fetch()? {
171/// // Process years in row set
172/// let year_col = row_set.column(0);
173/// for year in year_col
174/// .as_nullable_slice::<i16>()
175/// .expect("Year column buffer expected to be nullable Int")
176/// {
177/// // Iterate over `Option<i16>` with it ..
178/// }
179/// // Process names in row set
180/// let name_col = row_set.column(1);
181/// for name in name_col
182/// .as_text()
183/// .expect("Name column buffer expected to be text")
184/// .iter()
185/// {
186/// // Iterate over `Option<&CStr> ..
187/// }
188/// }
189/// }
190/// # Ok::<(), odbc_api::Error>(())
191/// ```
192///
193/// This second example changes two things: we do not know the schema in advance and use the
194/// SQL DataType to determine the best fit for the buffers. Also, we want to do everything in a
195/// function and return a `Cursor` with an already bound buffer. This approach is best if you have
196/// few and very long queries, so the overhead of allocating buffers is negligible and you want to
197/// have an easier time with the borrow checker.
198///
199/// ```no_run
200/// use odbc_api::{
201/// Connection, BlockCursor, Error, Cursor, Nullability, ResultSetMetadata,
202/// buffers::{ BufferDesc, ColumnarDynBuffer }
203/// };
204///
205/// fn get_birthdays<'a>(conn: &'a mut Connection)
206/// -> Result<BlockCursor<impl Cursor + 'a, ColumnarDynBuffer>, Error>
207/// {
208/// let query = "SELECT year, name FROM Birthdays;";
209/// let params = ();
210/// let timeout_sec = None;
211/// let mut cursor = conn.execute(query, params, timeout_sec)?.unwrap();
212/// let mut column_description = Default::default();
213/// let buffer_description : Vec<_> = (0..cursor.num_result_cols()?).map(|index| {
214/// cursor.describe_col(index as u16 + 1, &mut column_description)?;
215/// let nullable = matches!(
216/// column_description.nullability,
217/// Nullability::Unknown | Nullability::Nullable
218/// );
219/// let desc = BufferDesc::from_data_type(
220/// column_description.data_type,
221/// nullable
222/// ).unwrap_or(BufferDesc::Text{ max_str_len: 255 });
223/// Ok(desc)
224/// }).collect::<Result<_, Error>>()?;
225///
226/// // Row set size of 5000 rows.
227/// let buffer = ColumnarDynBuffer::from_descs(5000, buffer_description);
228/// // Bind buffer and take ownership over it.
229/// cursor.bind_buffer(buffer)
230/// }
231/// ```
232pub struct ColumnarBuffer<C> {
233 /// A mutable pointer to num_rows_fetched is passed to the C-API. It is used to write back the
234 /// number of fetched rows. `num_rows` is heap allocated, so the pointer is not invalidated,
235 /// even if the `ColumnarBuffer` instance is moved in memory.
236 num_rows: Box<usize>,
237 /// aka: batch size, row array size
238 row_capacity: usize,
239 /// Column index and bound buffer
240 columns: Vec<(u16, C)>,
241}
242
243/// Access a safe view of the column buffer.
244///
245/// After a fetch operation buffers may only partially be filled with data, the rest of the buffer
246/// may contain uninitialized values. Also we must not permit any operation which would invalidate
247/// the addresses of the buffer. To make reading buffer contents after a fetch safe,
248/// [`ColumnBuffer`]s implement this trait to offer safe views.
249///
250/// # Safety
251///
252/// Views must not allow access to uninitialized / invalid rows.
253pub unsafe trait Slice {
254 /// Immutable view on the column data. Used in safe abstractions. User must not be able to
255 /// access uninitialized or invalid memory of the buffer through this interface.
256 type Slice<'a>
257 where
258 Self: 'a;
259
260 /// Num rows may not exceed the actual amount of valid num_rows filled by the ODBC API. The
261 /// column buffer does not know how many elements were in the last row group, and therefore can
262 /// not guarantee the accessed element to be valid and in a defined state. It also can not panic
263 /// on accessing an undefined element.
264 fn slice(&self, valid_rows: usize) -> Self::Slice<'_>;
265}
266
267/// A buffer for a single column intended to be used together with [`ColumnarBuffer`].
268///
269/// # Safety
270///
271/// Implementations must ensure that:
272///
273/// * Capacity must be correctly reported otherwise data may be written outside its bounds.
274/// * truncation must be correctly reported. Code which reuses the same column buffer for reading
275/// and inserting may rely on this in order to avoid passing values indicators of truncated values
276/// in bulk insertions. This could lead to out of bounds memory access.
277pub unsafe trait ColumnBuffer: CDataMut {
278 /// Current capacity of the column
279 fn capacity(&self) -> usize;
280
281 /// `Some` if any value is truncated in the range [0, num_rows).
282 ///
283 /// After fetching data we may want to know if any value has been truncated due to the buffer
284 /// not being able to hold elements of that size. This method checks the indicator buffer
285 /// element wise.
286 fn has_truncated_values(&self, num_rows: usize) -> Option<Indicator>;
287}
288
289unsafe impl<T> Slice for WithDataType<T>
290where
291 T: Slice,
292{
293 type Slice<'a>
294 = T::Slice<'a>
295 where
296 T: 'a;
297
298 fn slice(&self, valid_rows: usize) -> T::Slice<'_> {
299 self.value.slice(valid_rows)
300 }
301}
302
303unsafe impl<T> ColumnBuffer for WithDataType<T>
304where
305 T: ColumnBuffer,
306{
307 fn capacity(&self) -> usize {
308 self.value.capacity()
309 }
310
311 fn has_truncated_values(&self, num_rows: usize) -> Option<Indicator> {
312 self.value.has_truncated_values(num_rows)
313 }
314}
315
316unsafe impl<'a, T> BoundInputSlice<'a> for WithDataType<T>
317where
318 T: BoundInputSlice<'a>,
319{
320 type SliceMut = T::SliceMut;
321
322 unsafe fn as_view_mut(
323 &'a mut self,
324 parameter_index: u16,
325 stmt: StatementRef<'a>,
326 ) -> Self::SliceMut {
327 unsafe { self.value.as_view_mut(parameter_index, stmt) }
328 }
329}
330
331/// This row set binds a string buffer to each column, which is large enough to hold the maximum
332/// length string representation for each element in the row set at once.
333///
334/// # Example
335///
336/// ```no_run
337/// //! A program executing a query and printing the result as CSV to standard output. Requires
338/// //! `anyhow` and `csv` crate.
339///
340/// use anyhow::Error;
341/// use odbc_api::{buffers::TextRowSet, Cursor, Environment, ConnectionOptions, ResultSetMetadata};
342/// use std::{
343/// ffi::CStr,
344/// io::{stdout, Write},
345/// path::PathBuf,
346/// };
347///
348/// /// Maximum number of rows fetched with one row set. Fetching batches of rows is usually much
349/// /// faster than fetching individual rows.
350/// const BATCH_SIZE: usize = 5000;
351///
352/// fn main() -> Result<(), Error> {
353/// // Write csv to standard out
354/// let out = stdout();
355/// let mut writer = csv::Writer::from_writer(out);
356///
357/// // We know this is going to be the only ODBC environment in the entire process, so this is
358/// // safe.
359/// let environment = unsafe { Environment::new() }?;
360///
361/// // Connect using a DSN. Alternatively we could have used a connection string
362/// let mut connection = environment.connect(
363/// "DataSourceName",
364/// "Username",
365/// "Password",
366/// ConnectionOptions::default(),
367/// )?;
368///
369/// // Execute a one-off query without any parameters.
370/// let query = "SELECT * FROM TableName";
371/// let params = ();
372/// let timeout_sec = None;
373/// match connection.execute(query, params, timeout_sec)? {
374/// Some(mut cursor) => {
375/// // Write the column names to stdout
376/// let mut headline: Vec<String> = cursor.column_names()?.collect::<Result<_,_>>()?;
377/// writer.write_record(headline)?;
378///
379/// // Use schema in cursor to initialize a text buffer large enough to hold the largest
380/// // possible strings for each column up to an upper limit of 4KiB
381/// let mut buffers = TextRowSet::for_cursor(BATCH_SIZE, &mut cursor, Some(4096))?;
382/// // Bind the buffer to the cursor. It is now being filled with every call to fetch.
383/// let mut row_set_cursor = cursor.bind_buffer(&mut buffers)?;
384///
385/// // Iterate over batches
386/// while let Some(batch) = row_set_cursor.fetch()? {
387/// // Within a batch, iterate over every row
388/// for row_index in 0..batch.num_rows() {
389/// // Within a row iterate over every column
390/// let record = (0..batch.num_cols()).map(|col_index| {
391/// batch
392/// .at(col_index, row_index)
393/// .unwrap_or(&[])
394/// });
395/// // Writes the row as CSV
396/// writer.write_record(record)?;
397/// }
398/// }
399/// }
400/// None => {
401/// eprintln!(
402/// "Query came back empty. No output has been created."
403/// );
404/// }
405/// }
406///
407/// Ok(())
408/// }
409/// ```
410pub type TextRowSet = ColumnarBuffer<TextColumn<u8>>;
411
412impl TextRowSet {
413 /// The resulting text buffer is not in any way tied to the cursor, other than that its buffer
414 /// sizes a tailor fitted to result set the cursor is iterating over.
415 ///
416 /// This method performs fallible buffer allocations, if no upper bound is set, so you may see
417 /// a speedup, by setting an upper bound using `max_str_limit`.
418 ///
419 ///
420 /// # Parameters
421 ///
422 /// * `batch_size`: The maximum number of rows the buffer is able to hold.
423 /// * `cursor`: Used to query the display size for each column of the row set. For character
424 /// data the length in characters is multiplied by 4 in order to have enough space for 4 byte
425 /// utf-8 characters. This is a pessimization for some data sources (e.g. SQLite 3) which do
426 /// interpret the size of a `VARCHAR(5)` column as 5 bytes rather than 5 characters.
427 /// * `max_str_limit`: Some queries make it hard to estimate a sensible upper bound and
428 /// sometimes drivers are just not that good at it. This argument allows you to specify an
429 /// upper bound for the length of character data. Any size reported by the driver is capped to
430 /// this value. In case the upper bound can not inferred by the metadata reported by the
431 /// driver the element size is set to this upper bound, too.
432 pub fn for_cursor(
433 batch_size: usize,
434 cursor: &mut impl ResultSetMetadata,
435 max_str_limit: Option<usize>,
436 ) -> Result<TextRowSet, Error> {
437 let buffers = utf8_display_sizes(cursor)?
438 .enumerate()
439 .map(|(buffer_index, reported_len)| {
440 let buffer_index = buffer_index as u16;
441 let col_index = buffer_index + 1;
442 let max_str_len = reported_len?;
443 let buffer = if let Some(upper_bound) = max_str_limit {
444 let max_str_len = max_str_len
445 .map(NonZeroUsize::get)
446 .unwrap_or(upper_bound)
447 .min(upper_bound);
448 TextColumn::new(batch_size, max_str_len)
449 } else {
450 let max_str_len = max_str_len.map(NonZeroUsize::get).ok_or(
451 Error::TooLargeColumnBufferSize {
452 buffer_index,
453 num_elements: batch_size,
454 element_size: usize::MAX,
455 },
456 )?;
457 TextColumn::try_new(batch_size, max_str_len).map_err(|source| {
458 Error::TooLargeColumnBufferSize {
459 buffer_index,
460 num_elements: source.num_elements,
461 element_size: source.element_size,
462 }
463 })?
464 };
465
466 Ok::<_, Error>((col_index, buffer))
467 })
468 .collect::<Result<_, _>>()?;
469 Ok(TextRowSet {
470 row_capacity: batch_size,
471 num_rows: Box::new(0),
472 columns: buffers,
473 })
474 }
475
476 /// Creates a text buffer large enough to hold `batch_size` rows with one column for each item
477 /// `max_str_lengths` of respective size.
478 pub fn from_max_str_lens(
479 row_capacity: usize,
480 max_str_lengths: impl IntoIterator<Item = usize>,
481 ) -> Result<Self, Error> {
482 let buffers = max_str_lengths
483 .into_iter()
484 .enumerate()
485 .map(|(index, max_str_len)| {
486 Ok::<_, Error>((
487 (index + 1).try_into().unwrap(),
488 TextColumn::try_new(row_capacity, max_str_len)
489 .map_err(|source| source.add_context(index.try_into().unwrap()))?,
490 ))
491 })
492 .collect::<Result<_, _>>()?;
493 Ok(TextRowSet {
494 row_capacity,
495 num_rows: Box::new(0),
496 columns: buffers,
497 })
498 }
499
500 /// Access the element at the specified position in the row set.
501 pub fn at(&self, buffer_index: usize, row_index: usize) -> Option<&[u8]> {
502 assert!(row_index < *self.num_rows);
503 self.columns[buffer_index].1.value_at(row_index)
504 }
505
506 /// Access the element at the specified position in the row set.
507 pub fn at_as_str(&self, col_index: usize, row_index: usize) -> Result<Option<&str>, Utf8Error> {
508 self.at(col_index, row_index).map(from_utf8).transpose()
509 }
510
511 /// Indicator value at the specified position. Useful to detect truncation of data.
512 ///
513 /// # Example
514 ///
515 /// ```
516 /// use odbc_api::buffers::{Indicator, TextRowSet};
517 ///
518 /// fn is_truncated(buffer: &TextRowSet, col_index: usize, row_index: usize) -> bool {
519 /// match buffer.indicator_at(col_index, row_index) {
520 /// // There is no value, therefore there is no value not fitting in the column buffer.
521 /// Indicator::Null => false,
522 /// // The value did not fit into the column buffer, we do not even know, by how much.
523 /// Indicator::NoTotal => true,
524 /// Indicator::Length(total_length) => {
525 /// // If the maximum string length is shorter than the values total length, the
526 /// // has been truncated to fit into the buffer.
527 /// buffer.max_len(col_index) < total_length
528 /// }
529 /// }
530 /// }
531 /// ```
532 pub fn indicator_at(&self, buf_index: usize, row_index: usize) -> Indicator {
533 assert!(row_index < *self.num_rows);
534 self.columns[buf_index].1.indicator_at(row_index)
535 }
536
537 /// Maximum length in bytes of elements in a column.
538 pub fn max_len(&self, buf_index: usize) -> usize {
539 self.columns[buf_index].1.max_len()
540 }
541}
542
543unsafe impl<T> ColumnBuffer for Vec<T>
544where
545 T: Pod,
546{
547 fn capacity(&self) -> usize {
548 self.len()
549 }
550
551 fn has_truncated_values(&self, _num_rows: usize) -> Option<Indicator> {
552 None
553 }
554}
555
556unsafe impl<T> Slice for Vec<T>
557where
558 T: Pod,
559{
560 type Slice<'a> = &'a [T];
561
562 fn slice(&self, valid_rows: usize) -> &[T] {
563 &self[..valid_rows]
564 }
565}
566
567/// A column buffer which can be resized.
568///
569/// Resizing is useful if a column buffer is used for inserting parameters, rather than fetching.
570/// Imagine an application which inserts data from a stream with row groups of varying size. If it
571/// encounters a row group with a new maximum size, it may want to resize the parameter buffers to
572/// send the entire row group in one go.
573pub trait Resize {
574 /// Resize the buffer to the given capacity.
575 ///
576 /// # Parameters
577 ///
578 /// * `new_capacity`: The new capacity of the buffer.
579 fn resize(&mut self, new_capacity: usize);
580}
581
582impl<T> Resize for Vec<T>
583where
584 T: Default + Clone,
585{
586 fn resize(&mut self, new_capacity: usize) {
587 Vec::resize(self, new_capacity, T::default());
588 }
589}
590
591impl<T> Resize for WithDataType<T>
592where
593 T: Resize,
594{
595 fn resize(&mut self, new_capacity: usize) {
596 self.value.resize(new_capacity);
597 }
598}
599
600#[cfg(test)]
601mod tests {
602
603 use super::Resize;
604 use crate::buffers::{BufferDesc, ColumnarDynBuffer};
605
606 #[test]
607 #[should_panic(expected = "Column indices must be unique.")]
608 fn assert_unique_column_indices() {
609 let bd = BufferDesc::I32 { nullable: false };
610 ColumnarDynBuffer::from_descs_and_indices(1, [(1, bd), (2, bd), (1, bd)].iter().cloned());
611 }
612
613 /// Vec's can resize just fine without this library, yet it is important that they implement the
614 /// `Resize` trait, so that other generic types know about it.
615 #[test]
616 fn vec_is_resize() {
617 let mut my_int_column_buffer = vec![1, 2];
618
619 Resize::resize(&mut my_int_column_buffer, 4);
620
621 assert_eq!(my_int_column_buffer[0], 1);
622 assert_eq!(my_int_column_buffer[1], 2);
623 assert_eq!(my_int_column_buffer[2], 0);
624 assert_eq!(my_int_column_buffer[3], 0);
625 }
626}