proof_of_sql/base/database/
owned_column.rs

1/// A column of data, with type included. This is simply a wrapper around `Vec<T>` for enumerated `T`.
2/// This is primarily used as an internal result that is used before
3/// converting to the final result in either Arrow format or JSON.
4/// This is the analog of an arrow Array.
5use super::{Column, ColumnCoercionError, ColumnType, OwnedColumnError, OwnedColumnResult};
6use crate::base::{
7    math::{
8        decimal::Precision,
9        permutation::{Permutation, PermutationError},
10    },
11    posql_time::{PoSQLTimeUnit, PoSQLTimeZone},
12    scalar::Scalar,
13    slice_ops::{inner_product_ref_cast, inner_product_with_bytes},
14};
15use alloc::{
16    string::{String, ToString},
17    vec::Vec,
18};
19use itertools::Itertools;
20use serde::{Deserialize, Serialize};
21
22#[derive(Debug, PartialEq, Clone, Eq, Serialize, Deserialize)]
23#[non_exhaustive]
24#[cfg_attr(test, derive(proptest_derive::Arbitrary))]
25/// Supported types for [`OwnedColumn`]
26pub enum OwnedColumn<S: Scalar> {
27    /// Boolean columns
28    Boolean(Vec<bool>),
29    /// u8 columns
30    Uint8(Vec<u8>),
31    /// i8 columns
32    TinyInt(Vec<i8>),
33    /// i16 columns
34    SmallInt(Vec<i16>),
35    /// i32 columns
36    Int(Vec<i32>),
37    /// i64 columns
38    BigInt(Vec<i64>),
39    /// i128 columns
40    Int128(Vec<i128>),
41    /// String columns
42    VarChar(Vec<String>),
43    /// Decimal columns
44    #[cfg_attr(test, proptest(skip))]
45    Decimal75(Precision, i8, Vec<S>),
46    /// Timestamp columns
47    #[cfg_attr(test, proptest(skip))]
48    TimestampTZ(PoSQLTimeUnit, PoSQLTimeZone, Vec<i64>),
49    /// Scalar columns
50    #[cfg_attr(test, proptest(skip))]
51    Scalar(Vec<S>),
52    /// Variable length binary columns
53    VarBinary(Vec<Vec<u8>>),
54}
55
56impl<S: Scalar> OwnedColumn<S> {
57    /// Compute the inner product of the column with a vector of scalars.
58    pub(crate) fn inner_product(&self, vec: &[S]) -> S {
59        match self {
60            OwnedColumn::Boolean(col) => inner_product_ref_cast(col, vec),
61            OwnedColumn::Uint8(col) => inner_product_ref_cast(col, vec),
62            OwnedColumn::TinyInt(col) => inner_product_ref_cast(col, vec),
63            OwnedColumn::SmallInt(col) => inner_product_ref_cast(col, vec),
64            OwnedColumn::Int(col) => inner_product_ref_cast(col, vec),
65            OwnedColumn::BigInt(col) | OwnedColumn::TimestampTZ(_, _, col) => {
66                inner_product_ref_cast(col, vec)
67            }
68            OwnedColumn::VarChar(col) => inner_product_ref_cast(col, vec),
69            OwnedColumn::VarBinary(col) => inner_product_with_bytes(col, vec),
70            OwnedColumn::Int128(col) => inner_product_ref_cast(col, vec),
71            OwnedColumn::Decimal75(_, _, col) | OwnedColumn::Scalar(col) => {
72                inner_product_ref_cast(col, vec)
73            }
74        }
75    }
76
77    /// Returns the length of the column.
78    #[must_use]
79    pub fn len(&self) -> usize {
80        match self {
81            OwnedColumn::Boolean(col) => col.len(),
82            OwnedColumn::TinyInt(col) => col.len(),
83            OwnedColumn::Uint8(col) => col.len(),
84            OwnedColumn::SmallInt(col) => col.len(),
85            OwnedColumn::Int(col) => col.len(),
86            OwnedColumn::BigInt(col) | OwnedColumn::TimestampTZ(_, _, col) => col.len(),
87            OwnedColumn::VarChar(col) => col.len(),
88            OwnedColumn::VarBinary(col) => col.len(),
89            OwnedColumn::Int128(col) => col.len(),
90            OwnedColumn::Decimal75(_, _, col) | OwnedColumn::Scalar(col) => col.len(),
91        }
92    }
93
94    /// Returns the column with its entries permutated
95    pub fn try_permute(&self, permutation: &Permutation) -> Result<Self, PermutationError> {
96        Ok(match self {
97            OwnedColumn::Boolean(col) => OwnedColumn::Boolean(permutation.try_apply(col)?),
98            OwnedColumn::TinyInt(col) => OwnedColumn::TinyInt(permutation.try_apply(col)?),
99            OwnedColumn::Uint8(col) => OwnedColumn::Uint8(permutation.try_apply(col)?),
100            OwnedColumn::SmallInt(col) => OwnedColumn::SmallInt(permutation.try_apply(col)?),
101            OwnedColumn::Int(col) => OwnedColumn::Int(permutation.try_apply(col)?),
102            OwnedColumn::BigInt(col) => OwnedColumn::BigInt(permutation.try_apply(col)?),
103            OwnedColumn::VarChar(col) => OwnedColumn::VarChar(permutation.try_apply(col)?),
104            OwnedColumn::VarBinary(col) => OwnedColumn::VarBinary(permutation.try_apply(col)?),
105            OwnedColumn::Int128(col) => OwnedColumn::Int128(permutation.try_apply(col)?),
106            OwnedColumn::Decimal75(precision, scale, col) => {
107                OwnedColumn::Decimal75(*precision, *scale, permutation.try_apply(col)?)
108            }
109            OwnedColumn::Scalar(col) => OwnedColumn::Scalar(permutation.try_apply(col)?),
110            OwnedColumn::TimestampTZ(tu, tz, col) => {
111                OwnedColumn::TimestampTZ(*tu, *tz, permutation.try_apply(col)?)
112            }
113        })
114    }
115
116    /// Returns the sliced column.
117    #[must_use]
118    pub fn slice(&self, start: usize, end: usize) -> Self {
119        match self {
120            OwnedColumn::Boolean(col) => OwnedColumn::Boolean(col[start..end].to_vec()),
121            OwnedColumn::TinyInt(col) => OwnedColumn::TinyInt(col[start..end].to_vec()),
122            OwnedColumn::Uint8(col) => OwnedColumn::Uint8(col[start..end].to_vec()),
123            OwnedColumn::SmallInt(col) => OwnedColumn::SmallInt(col[start..end].to_vec()),
124            OwnedColumn::Int(col) => OwnedColumn::Int(col[start..end].to_vec()),
125            OwnedColumn::BigInt(col) => OwnedColumn::BigInt(col[start..end].to_vec()),
126            OwnedColumn::VarChar(col) => OwnedColumn::VarChar(col[start..end].to_vec()),
127            OwnedColumn::VarBinary(col) => OwnedColumn::VarBinary(col[start..end].to_vec()),
128            OwnedColumn::Int128(col) => OwnedColumn::Int128(col[start..end].to_vec()),
129            OwnedColumn::Decimal75(precision, scale, col) => {
130                OwnedColumn::Decimal75(*precision, *scale, col[start..end].to_vec())
131            }
132            OwnedColumn::Scalar(col) => OwnedColumn::Scalar(col[start..end].to_vec()),
133            OwnedColumn::TimestampTZ(tu, tz, col) => {
134                OwnedColumn::TimestampTZ(*tu, *tz, col[start..end].to_vec())
135            }
136        }
137    }
138
139    /// Returns true if the column is empty.
140    #[must_use]
141    pub fn is_empty(&self) -> bool {
142        match self {
143            OwnedColumn::Boolean(col) => col.is_empty(),
144            OwnedColumn::TinyInt(col) => col.is_empty(),
145            OwnedColumn::Uint8(col) => col.is_empty(),
146            OwnedColumn::SmallInt(col) => col.is_empty(),
147            OwnedColumn::Int(col) => col.is_empty(),
148            OwnedColumn::BigInt(col) | OwnedColumn::TimestampTZ(_, _, col) => col.is_empty(),
149            OwnedColumn::VarChar(col) => col.is_empty(),
150            OwnedColumn::VarBinary(col) => col.is_empty(),
151            OwnedColumn::Int128(col) => col.is_empty(),
152            OwnedColumn::Scalar(col) | OwnedColumn::Decimal75(_, _, col) => col.is_empty(),
153        }
154    }
155    /// Returns the type of the column.
156    #[must_use]
157    pub fn column_type(&self) -> ColumnType {
158        match self {
159            OwnedColumn::Boolean(_) => ColumnType::Boolean,
160            OwnedColumn::TinyInt(_) => ColumnType::TinyInt,
161            OwnedColumn::Uint8(_) => ColumnType::Uint8,
162            OwnedColumn::SmallInt(_) => ColumnType::SmallInt,
163            OwnedColumn::Int(_) => ColumnType::Int,
164            OwnedColumn::BigInt(_) => ColumnType::BigInt,
165            OwnedColumn::VarChar(_) => ColumnType::VarChar,
166            OwnedColumn::VarBinary(_) => ColumnType::VarBinary,
167            OwnedColumn::Int128(_) => ColumnType::Int128,
168            OwnedColumn::Scalar(_) => ColumnType::Scalar,
169            OwnedColumn::Decimal75(precision, scale, _) => {
170                ColumnType::Decimal75(*precision, *scale)
171            }
172            OwnedColumn::TimestampTZ(tu, tz, _) => ColumnType::TimestampTZ(*tu, *tz),
173        }
174    }
175
176    /// Convert a slice of scalars to a vec of owned columns
177    pub fn try_from_scalars(scalars: &[S], column_type: ColumnType) -> OwnedColumnResult<Self> {
178        match column_type {
179            ColumnType::Boolean => Ok(OwnedColumn::Boolean(
180                scalars
181                    .iter()
182                    .map(|s| -> Result<bool, _> { TryInto::<bool>::try_into(*s) })
183                    .collect::<Result<Vec<_>, _>>()
184                    .map_err(|_| OwnedColumnError::ScalarConversionError {
185                        error: "Overflow in scalar conversions".to_string(),
186                    })?,
187            )),
188            ColumnType::Uint8 => Ok(OwnedColumn::Uint8(
189                scalars
190                    .iter()
191                    .map(|s| -> Result<u8, _> { TryInto::<u8>::try_into(*s) })
192                    .collect::<Result<Vec<_>, _>>()
193                    .map_err(|_| OwnedColumnError::ScalarConversionError {
194                        error: "Overflow in scalar conversions".to_string(),
195                    })?,
196            )),
197            ColumnType::TinyInt => Ok(OwnedColumn::TinyInt(
198                scalars
199                    .iter()
200                    .map(|s| -> Result<i8, _> { TryInto::<i8>::try_into(*s) })
201                    .collect::<Result<Vec<_>, _>>()
202                    .map_err(|_| OwnedColumnError::ScalarConversionError {
203                        error: "Overflow in scalar conversions".to_string(),
204                    })?,
205            )),
206            ColumnType::SmallInt => Ok(OwnedColumn::SmallInt(
207                scalars
208                    .iter()
209                    .map(|s| -> Result<i16, _> { TryInto::<i16>::try_into(*s) })
210                    .collect::<Result<Vec<_>, _>>()
211                    .map_err(|_| OwnedColumnError::ScalarConversionError {
212                        error: "Overflow in scalar conversions".to_string(),
213                    })?,
214            )),
215            ColumnType::Int => Ok(OwnedColumn::Int(
216                scalars
217                    .iter()
218                    .map(|s| -> Result<i32, _> { TryInto::<i32>::try_into(*s) })
219                    .collect::<Result<Vec<_>, _>>()
220                    .map_err(|_| OwnedColumnError::ScalarConversionError {
221                        error: "Overflow in scalar conversions".to_string(),
222                    })?,
223            )),
224            ColumnType::BigInt => Ok(OwnedColumn::BigInt(
225                scalars
226                    .iter()
227                    .map(|s| -> Result<i64, _> { TryInto::<i64>::try_into(*s) })
228                    .collect::<Result<Vec<_>, _>>()
229                    .map_err(|_| OwnedColumnError::ScalarConversionError {
230                        error: "Overflow in scalar conversions".to_string(),
231                    })?,
232            )),
233            ColumnType::Int128 => Ok(OwnedColumn::Int128(
234                scalars
235                    .iter()
236                    .map(|s| -> Result<i128, _> { TryInto::<i128>::try_into(*s) })
237                    .collect::<Result<Vec<_>, _>>()
238                    .map_err(|_| OwnedColumnError::ScalarConversionError {
239                        error: "Overflow in scalar conversions".to_string(),
240                    })?,
241            )),
242            ColumnType::Scalar => Ok(OwnedColumn::Scalar(scalars.to_vec())),
243            ColumnType::Decimal75(precision, scale) => {
244                Ok(OwnedColumn::Decimal75(precision, scale, scalars.to_vec()))
245            }
246            ColumnType::TimestampTZ(tu, tz) => {
247                let raw_values: Vec<i64> = scalars
248                    .iter()
249                    .map(|s| -> Result<i64, _> { TryInto::<i64>::try_into(*s) })
250                    .collect::<Result<Vec<_>, _>>()
251                    .map_err(|_| OwnedColumnError::ScalarConversionError {
252                        error: "Overflow in scalar conversions".to_string(),
253                    })?;
254                Ok(OwnedColumn::TimestampTZ(tu, tz, raw_values))
255            }
256            // Can not convert scalars to VarChar
257            ColumnType::VarChar | ColumnType::VarBinary => Err(OwnedColumnError::TypeCastError {
258                from_type: ColumnType::Scalar,
259                to_type: ColumnType::VarChar,
260            }),
261        }
262    }
263
264    /// Convert a slice of option scalars to a vec of owned columns
265    pub fn try_from_option_scalars(
266        option_scalars: &[Option<S>],
267        column_type: ColumnType,
268    ) -> OwnedColumnResult<Self> {
269        let scalars = option_scalars
270            .iter()
271            .copied()
272            .collect::<Option<Vec<_>>>()
273            .ok_or(OwnedColumnError::Unsupported {
274                error: "NULL is not supported yet".to_string(),
275            })?;
276        Self::try_from_scalars(&scalars, column_type)
277    }
278    #[cfg(test)]
279    /// Returns an iterator over the raw data of the column
280    /// assuming the underlying type is [u8], panicking if it is not.
281    pub fn u8_iter(&self) -> impl Iterator<Item = &u8> {
282        match self {
283            OwnedColumn::Uint8(col) => col.iter(),
284            _ => panic!("Expected Uint8 column"),
285        }
286    }
287    #[cfg(test)]
288    /// Returns an iterator over the raw data of the column
289    /// assuming the underlying type is [i8], panicking if it is not.
290    pub fn i8_iter(&self) -> impl Iterator<Item = &i8> {
291        match self {
292            OwnedColumn::TinyInt(col) => col.iter(),
293            _ => panic!("Expected TinyInt column"),
294        }
295    }
296    #[cfg(test)]
297    /// Returns an iterator over the raw data of the column
298    /// assuming the underlying type is [i16], panicking if it is not.
299    pub fn i16_iter(&self) -> impl Iterator<Item = &i16> {
300        match self {
301            OwnedColumn::SmallInt(col) => col.iter(),
302            _ => panic!("Expected SmallInt column"),
303        }
304    }
305    #[cfg(test)]
306    /// Returns an iterator over the raw data of the column
307    /// assuming the underlying type is [i32], panicking if it is not.
308    pub fn i32_iter(&self) -> impl Iterator<Item = &i32> {
309        match self {
310            OwnedColumn::Int(col) => col.iter(),
311            _ => panic!("Expected Int column"),
312        }
313    }
314    #[cfg(test)]
315    /// Returns an iterator over the raw data of the column
316    /// assuming the underlying type is [i64], panicking if it is not.
317    pub fn i64_iter(&self) -> impl Iterator<Item = &i64> {
318        match self {
319            OwnedColumn::TimestampTZ(_, _, col) | OwnedColumn::BigInt(col) => col.iter(),
320            _ => panic!("Expected TimestampTZ or BigInt column"),
321        }
322    }
323    #[cfg(test)]
324    /// Returns an iterator over the raw data of the column
325    /// assuming the underlying type is [i128], panicking if it is not.
326    pub fn i128_iter(&self) -> impl Iterator<Item = &i128> {
327        match self {
328            OwnedColumn::Int128(col) => col.iter(),
329            _ => panic!("Expected Int128 column"),
330        }
331    }
332    #[cfg(test)]
333    /// Returns an iterator over the raw data of the column
334    /// assuming the underlying type is [bool], panicking if it is not.
335    pub fn bool_iter(&self) -> impl Iterator<Item = &bool> {
336        match self {
337            OwnedColumn::Boolean(col) => col.iter(),
338            _ => panic!("Expected Boolean column"),
339        }
340    }
341    #[cfg(test)]
342    /// Returns an iterator over the raw data of the column
343    /// assuming the underlying type is a [Scalar], panicking if it is not.
344    pub fn scalar_iter(&self) -> impl Iterator<Item = &S> {
345        match self {
346            OwnedColumn::Decimal75(_, _, col) | OwnedColumn::Scalar(col) => col.iter(),
347            _ => panic!("Expected Scalar or Decimal75 column"),
348        }
349    }
350    #[cfg(test)]
351    /// Returns an iterator over the raw data of the column
352    /// assuming the underlying type is [String], panicking if it is not.
353    pub fn string_iter(&self) -> impl Iterator<Item = &String> {
354        match self {
355            OwnedColumn::VarChar(col) => col.iter(),
356            _ => panic!("Expected VarChar column"),
357        }
358    }
359}
360
361impl<'a, S: Scalar> From<&Column<'a, S>> for OwnedColumn<S> {
362    fn from(col: &Column<'a, S>) -> Self {
363        match col {
364            Column::Boolean(col) => OwnedColumn::Boolean(col.to_vec()),
365            Column::TinyInt(col) => OwnedColumn::TinyInt(col.to_vec()),
366            Column::Uint8(col) => OwnedColumn::Uint8(col.to_vec()),
367            Column::SmallInt(col) => OwnedColumn::SmallInt(col.to_vec()),
368            Column::Int(col) => OwnedColumn::Int(col.to_vec()),
369            Column::BigInt(col) => OwnedColumn::BigInt(col.to_vec()),
370            Column::VarChar((col, _)) => {
371                OwnedColumn::VarChar(col.iter().map(ToString::to_string).collect())
372            }
373            Column::VarBinary((col, _)) => {
374                OwnedColumn::VarBinary(col.iter().map(|slice| slice.to_vec()).collect())
375            }
376            Column::Int128(col) => OwnedColumn::Int128(col.to_vec()),
377            Column::Decimal75(precision, scale, col) => {
378                OwnedColumn::Decimal75(*precision, *scale, col.to_vec())
379            }
380            Column::Scalar(col) => OwnedColumn::Scalar(col.to_vec()),
381            Column::TimestampTZ(tu, tz, col) => OwnedColumn::TimestampTZ(*tu, *tz, col.to_vec()),
382        }
383    }
384}
385
386impl<S: Scalar> OwnedColumn<S> {
387    /// Attempts to coerce a column of scalars to a numeric column of the specified type.
388    /// If the specified type is the same as the current column type, the function will return the column as is.
389    ///
390    /// # Arguments
391    ///
392    /// * `to_type` - The target numeric column type to coerce to.
393    ///
394    /// # Returns
395    ///
396    /// * `Ok(Self)` - If the coercion is successful.
397    /// * `Err(ColumnCoercionError)` - If the coercion fails due to type mismatch or overflow.
398    ///
399    /// # Errors
400    ///
401    /// If the specified type is the same as the current column type, the function will not error.
402    ///
403    /// Otherwise, this function will return an error if:
404    /// * The column type is not `Scalar`.
405    /// * The target type is not a numeric type.
406    /// * There is an overflow during the coercion.
407    pub(crate) fn try_coerce_scalar_to_numeric(
408        self,
409        to_type: ColumnType,
410    ) -> Result<Self, ColumnCoercionError> {
411        if self.column_type() == to_type {
412            Ok(self)
413        } else if let OwnedColumn::Scalar(vec) = self {
414            match to_type {
415                ColumnType::Uint8 => vec
416                    .into_iter()
417                    .map(TryInto::try_into)
418                    .try_collect()
419                    .map_err(|_| ColumnCoercionError::Overflow)
420                    .map(OwnedColumn::Uint8),
421                ColumnType::TinyInt => vec
422                    .into_iter()
423                    .map(TryInto::try_into)
424                    .try_collect()
425                    .map_err(|_| ColumnCoercionError::Overflow)
426                    .map(OwnedColumn::TinyInt),
427                ColumnType::SmallInt => vec
428                    .into_iter()
429                    .map(TryInto::try_into)
430                    .try_collect()
431                    .map_err(|_| ColumnCoercionError::Overflow)
432                    .map(OwnedColumn::SmallInt),
433                ColumnType::Int => vec
434                    .into_iter()
435                    .map(TryInto::try_into)
436                    .try_collect()
437                    .map_err(|_| ColumnCoercionError::Overflow)
438                    .map(OwnedColumn::Int),
439                ColumnType::BigInt => vec
440                    .into_iter()
441                    .map(TryInto::try_into)
442                    .try_collect()
443                    .map_err(|_| ColumnCoercionError::Overflow)
444                    .map(OwnedColumn::BigInt),
445                ColumnType::Int128 => vec
446                    .into_iter()
447                    .map(TryInto::try_into)
448                    .try_collect()
449                    .map_err(|_| ColumnCoercionError::Overflow)
450                    .map(OwnedColumn::Int128),
451                ColumnType::Decimal75(precision, scale) => {
452                    Ok(OwnedColumn::Decimal75(precision, scale, vec))
453                }
454                _ => Err(ColumnCoercionError::InvalidTypeCoercion),
455            }
456        } else {
457            Err(ColumnCoercionError::InvalidTypeCoercion)
458        }
459    }
460}
461
462#[cfg(test)]
463mod test {
464    use super::*;
465    use crate::base::{
466        math::decimal::Precision,
467        scalar::{test_scalar::TestScalar, ScalarExt},
468    };
469    use alloc::vec;
470    use bumpalo::Bump;
471
472    #[test]
473    fn we_can_slice_a_column() {
474        let col: OwnedColumn<TestScalar> = OwnedColumn::Int128(vec![1, 2, 3, 4, 5]);
475        assert_eq!(col.slice(1, 4), OwnedColumn::Int128(vec![2, 3, 4]));
476    }
477
478    #[test]
479    fn we_can_permute_a_column() {
480        let col: OwnedColumn<TestScalar> = OwnedColumn::Int128(vec![1, 2, 3, 4, 5]);
481        let permutation = Permutation::try_new(vec![1, 3, 4, 0, 2]).unwrap();
482        assert_eq!(
483            col.try_permute(&permutation).unwrap(),
484            OwnedColumn::Int128(vec![2, 4, 5, 1, 3])
485        );
486    }
487
488    #[test]
489    fn we_can_convert_columns_to_owned_columns_round_trip() {
490        let alloc = Bump::new();
491        // Integers
492        let col: Column<'_, TestScalar> = Column::Int128(&[1, 2, 3, 4, 5]);
493        let owned_col: OwnedColumn<TestScalar> = (&col).into();
494        assert_eq!(owned_col, OwnedColumn::Int128(vec![1, 2, 3, 4, 5]));
495        let new_col = Column::<TestScalar>::from_owned_column(&owned_col, &alloc);
496        assert_eq!(col, new_col);
497
498        // Booleans
499        let col: Column<'_, TestScalar> = Column::Boolean(&[true, false, true, false, true]);
500        let owned_col: OwnedColumn<TestScalar> = (&col).into();
501        assert_eq!(
502            owned_col,
503            OwnedColumn::Boolean(vec![true, false, true, false, true])
504        );
505        let new_col = Column::<TestScalar>::from_owned_column(&owned_col, &alloc);
506        assert_eq!(col, new_col);
507
508        // Strings
509        let strs = [
510            "Space and Time",
511            "מרחב וזמן",
512            "Χώρος και Χρόνος",
513            "Տարածություն և ժամանակ",
514            "ቦታ እና ጊዜ",
515            "სივრცე და დრო",
516        ];
517        let scalars = strs.iter().map(TestScalar::from).collect::<Vec<_>>();
518        let col: Column<'_, TestScalar> = Column::VarChar((&strs, &scalars));
519        let owned_col: OwnedColumn<TestScalar> = (&col).into();
520        assert_eq!(
521            owned_col,
522            OwnedColumn::VarChar(
523                strs.iter()
524                    .map(ToString::to_string)
525                    .collect::<Vec<String>>()
526            )
527        );
528        let new_col = Column::<TestScalar>::from_owned_column(&owned_col, &alloc);
529        assert_eq!(col, new_col);
530
531        // Decimals
532        let scalars: Vec<TestScalar> = [1, 2, 3, 4, 5].iter().map(TestScalar::from).collect();
533        let col: Column<'_, TestScalar> =
534            Column::Decimal75(Precision::new(75).unwrap(), -128, &scalars);
535        let owned_col: OwnedColumn<TestScalar> = (&col).into();
536        assert_eq!(
537            owned_col,
538            OwnedColumn::Decimal75(Precision::new(75).unwrap(), -128, scalars.clone())
539        );
540        let new_col = Column::<TestScalar>::from_owned_column(&owned_col, &alloc);
541        assert_eq!(col, new_col);
542    }
543
544    #[test]
545    fn we_can_convert_scalars_to_owned_columns() {
546        // Int
547        let scalars = [1, 2, 3, 4, 5]
548            .iter()
549            .map(TestScalar::from)
550            .collect::<Vec<_>>();
551        let column_type = ColumnType::Int128;
552        let owned_col = OwnedColumn::try_from_scalars(&scalars, column_type).unwrap();
553        assert_eq!(owned_col, OwnedColumn::Int128(vec![1, 2, 3, 4, 5]));
554
555        // Boolean
556        let scalars = [true, false, true, false, true]
557            .iter()
558            .map(TestScalar::from)
559            .collect::<Vec<_>>();
560        let column_type = ColumnType::Boolean;
561        let owned_col = OwnedColumn::try_from_scalars(&scalars, column_type).unwrap();
562        assert_eq!(
563            owned_col,
564            OwnedColumn::Boolean(vec![true, false, true, false, true])
565        );
566
567        // Decimal
568        let scalars = [1, 2, 3, 4, 5]
569            .iter()
570            .map(TestScalar::from)
571            .collect::<Vec<_>>();
572        let column_type = ColumnType::Decimal75(Precision::new(75).unwrap(), -128);
573        let owned_col = OwnedColumn::try_from_scalars(&scalars, column_type).unwrap();
574        assert_eq!(
575            owned_col,
576            OwnedColumn::Decimal75(Precision::new(75).unwrap(), -128, scalars)
577        );
578    }
579
580    #[test]
581    fn we_cannot_convert_scalars_to_owned_columns_if_varchar() {
582        let scalars = ["a", "b", "c", "d", "e"]
583            .iter()
584            .map(TestScalar::from)
585            .collect::<Vec<_>>();
586        let column_type = ColumnType::VarChar;
587        let res = OwnedColumn::try_from_scalars(&scalars, column_type);
588        assert!(matches!(res, Err(OwnedColumnError::TypeCastError { .. })));
589    }
590
591    #[test]
592    fn we_cannot_convert_scalars_to_owned_columns_if_overflow() {
593        // Int
594        let scalars = [i128::MAX, i128::MAX, i128::MAX, i128::MAX, i128::MAX]
595            .iter()
596            .map(TestScalar::from)
597            .collect::<Vec<_>>();
598        let column_type = ColumnType::BigInt;
599        let res = OwnedColumn::try_from_scalars(&scalars, column_type);
600        assert!(matches!(
601            res,
602            Err(OwnedColumnError::ScalarConversionError { .. })
603        ));
604
605        // Boolean
606        let scalars = [i128::MAX, i128::MAX, i128::MAX, i128::MAX, i128::MAX]
607            .iter()
608            .map(TestScalar::from)
609            .collect::<Vec<_>>();
610        let column_type = ColumnType::Boolean;
611        let res = OwnedColumn::try_from_scalars(&scalars, column_type);
612        assert!(matches!(
613            res,
614            Err(OwnedColumnError::ScalarConversionError { .. })
615        ));
616    }
617
618    #[test]
619    fn we_can_convert_option_scalars_to_owned_columns() {
620        // Int
621        let option_scalars = [Some(1), Some(2), Some(3), Some(4), Some(5)]
622            .iter()
623            .map(|s| s.map(TestScalar::from))
624            .collect::<Vec<_>>();
625        let column_type = ColumnType::Int128;
626        let owned_col = OwnedColumn::try_from_option_scalars(&option_scalars, column_type).unwrap();
627        assert_eq!(owned_col, OwnedColumn::Int128(vec![1, 2, 3, 4, 5]));
628
629        // Boolean
630        let option_scalars = [Some(true), Some(false), Some(true), Some(false), Some(true)]
631            .iter()
632            .map(|s| s.map(TestScalar::from))
633            .collect::<Vec<_>>();
634        let column_type = ColumnType::Boolean;
635        let owned_col = OwnedColumn::try_from_option_scalars(&option_scalars, column_type).unwrap();
636        assert_eq!(
637            owned_col,
638            OwnedColumn::Boolean(vec![true, false, true, false, true])
639        );
640
641        // Decimal
642        let option_scalars = [Some(1), Some(2), Some(3), Some(4), Some(5)]
643            .iter()
644            .map(|s| s.map(TestScalar::from))
645            .collect::<Vec<_>>();
646        let scalars = [1, 2, 3, 4, 5]
647            .iter()
648            .map(|&i| TestScalar::from(i))
649            .collect::<Vec<_>>();
650        let column_type = ColumnType::Decimal75(Precision::new(75).unwrap(), 127);
651        let owned_col = OwnedColumn::try_from_option_scalars(&option_scalars, column_type).unwrap();
652        assert_eq!(
653            owned_col,
654            OwnedColumn::Decimal75(Precision::new(75).unwrap(), 127, scalars)
655        );
656    }
657
658    #[test]
659    fn we_cannot_convert_option_scalars_to_owned_columns_if_varchar() {
660        let option_scalars = ["a", "b", "c", "d", "e"]
661            .iter()
662            .map(|s| Some(TestScalar::from(*s)))
663            .collect::<Vec<_>>();
664        let column_type = ColumnType::VarChar;
665        let res = OwnedColumn::try_from_option_scalars(&option_scalars, column_type);
666        assert!(matches!(res, Err(OwnedColumnError::TypeCastError { .. })));
667    }
668
669    #[test]
670    fn we_cannot_convert_option_scalars_to_owned_columns_if_overflow() {
671        // Int
672        let option_scalars = [
673            Some(i128::MAX),
674            Some(i128::MAX),
675            Some(i128::MAX),
676            Some(i128::MAX),
677            Some(i128::MAX),
678        ]
679        .iter()
680        .map(|s| s.map(TestScalar::from))
681        .collect::<Vec<_>>();
682        let column_type = ColumnType::BigInt;
683        let res = OwnedColumn::try_from_option_scalars(&option_scalars, column_type);
684        assert!(matches!(
685            res,
686            Err(OwnedColumnError::ScalarConversionError { .. })
687        ));
688
689        // Boolean
690        let option_scalars = [
691            Some(i128::MAX),
692            Some(i128::MAX),
693            Some(i128::MAX),
694            Some(i128::MAX),
695            Some(i128::MAX),
696        ]
697        .iter()
698        .map(|s| s.map(TestScalar::from))
699        .collect::<Vec<_>>();
700        let column_type = ColumnType::Boolean;
701        let res = OwnedColumn::try_from_option_scalars(&option_scalars, column_type);
702        assert!(matches!(
703            res,
704            Err(OwnedColumnError::ScalarConversionError { .. })
705        ));
706    }
707
708    #[test]
709    fn we_cannot_convert_option_scalars_to_owned_columns_if_none() {
710        // Int
711        let option_scalars = [Some(1), Some(2), None, Some(4), Some(5)]
712            .iter()
713            .map(|s| s.map(TestScalar::from))
714            .collect::<Vec<_>>();
715        let column_type = ColumnType::Int128;
716        let res = OwnedColumn::try_from_option_scalars(&option_scalars, column_type);
717        assert!(matches!(res, Err(OwnedColumnError::Unsupported { .. })));
718
719        // Boolean
720        let option_scalars = [Some(true), Some(false), None, Some(false), Some(true)]
721            .iter()
722            .map(|s| s.map(TestScalar::from))
723            .collect::<Vec<_>>();
724        let column_type = ColumnType::Boolean;
725        let res = OwnedColumn::try_from_option_scalars(&option_scalars, column_type);
726        assert!(matches!(res, Err(OwnedColumnError::Unsupported { .. })));
727    }
728
729    #[test]
730    fn we_can_coerce_scalar_to_numeric() {
731        let scalars = vec![
732            TestScalar::from(1),
733            TestScalar::from(2),
734            TestScalar::from(3),
735        ];
736        let col = OwnedColumn::Scalar(scalars.clone());
737
738        // Coerce to TinyInt
739        let coerced_col = col
740            .clone()
741            .try_coerce_scalar_to_numeric(ColumnType::TinyInt)
742            .unwrap();
743        assert_eq!(coerced_col, OwnedColumn::TinyInt(vec![1, 2, 3]));
744
745        // Coerce to SmallInt
746        let coerced_col = col
747            .clone()
748            .try_coerce_scalar_to_numeric(ColumnType::SmallInt)
749            .unwrap();
750        assert_eq!(coerced_col, OwnedColumn::SmallInt(vec![1, 2, 3]));
751
752        // Coerce to Int
753        let coerced_col = col
754            .clone()
755            .try_coerce_scalar_to_numeric(ColumnType::Int)
756            .unwrap();
757        assert_eq!(coerced_col, OwnedColumn::Int(vec![1, 2, 3]));
758
759        // Coerce to BigInt
760        let coerced_col = col
761            .clone()
762            .try_coerce_scalar_to_numeric(ColumnType::BigInt)
763            .unwrap();
764        assert_eq!(coerced_col, OwnedColumn::BigInt(vec![1, 2, 3]));
765
766        // Coerce to Int128
767        let coerced_col = col
768            .clone()
769            .try_coerce_scalar_to_numeric(ColumnType::Int128)
770            .unwrap();
771        assert_eq!(coerced_col, OwnedColumn::Int128(vec![1, 2, 3]));
772
773        // Coerce to Decimal75
774        let coerced_col = col
775            .clone()
776            .try_coerce_scalar_to_numeric(ColumnType::Decimal75(Precision::new(75).unwrap(), 0))
777            .unwrap();
778        assert_eq!(
779            coerced_col,
780            OwnedColumn::Decimal75(Precision::new(75).unwrap(), 0, scalars)
781        );
782    }
783
784    #[test]
785    fn we_cannot_coerce_scalar_to_invalid_type() {
786        let scalars = vec![
787            TestScalar::from(1),
788            TestScalar::from(2),
789            TestScalar::from(3),
790        ];
791        let col = OwnedColumn::Scalar(scalars);
792
793        // Attempt to coerce to VarChar
794        let res = col
795            .clone()
796            .try_coerce_scalar_to_numeric(ColumnType::VarChar);
797        assert!(matches!(res, Err(ColumnCoercionError::InvalidTypeCoercion)));
798
799        // Attempt to coerce non-scalar column
800        let col = OwnedColumn::<TestScalar>::Int(vec![1, 2, 3]);
801        let res = col.try_coerce_scalar_to_numeric(ColumnType::BigInt);
802        assert!(matches!(res, Err(ColumnCoercionError::InvalidTypeCoercion)));
803    }
804
805    #[test]
806    fn we_cannot_coerce_scalar_to_numeric_if_overflow() {
807        let scalars = vec![TestScalar::from(i128::MAX), -TestScalar::from(i128::MIN)];
808        let col = OwnedColumn::Scalar(scalars);
809
810        // Attempt to coerce to TinyInt
811        let res = col
812            .clone()
813            .try_coerce_scalar_to_numeric(ColumnType::TinyInt);
814        assert!(matches!(res, Err(ColumnCoercionError::Overflow)));
815
816        // Attempt to coerce to SmallInt
817        let res = col
818            .clone()
819            .try_coerce_scalar_to_numeric(ColumnType::SmallInt);
820        assert!(matches!(res, Err(ColumnCoercionError::Overflow)));
821
822        // Attempt to coerce to Int
823        let res = col.clone().try_coerce_scalar_to_numeric(ColumnType::Int);
824        assert!(matches!(res, Err(ColumnCoercionError::Overflow)));
825
826        // Attempt to coerce to BigInt
827        let res = col.clone().try_coerce_scalar_to_numeric(ColumnType::BigInt);
828        assert!(matches!(res, Err(ColumnCoercionError::Overflow)));
829
830        // Attempt to coerce to Int128
831        let res = col.try_coerce_scalar_to_numeric(ColumnType::Int128);
832        assert!(matches!(res, Err(ColumnCoercionError::Overflow)));
833    }
834
835    #[test]
836    fn we_can_slice_and_permute_varbinary_columns() {
837        let col = OwnedColumn::<TestScalar>::VarBinary(vec![
838            b"foo".to_vec(),
839            b"bar".to_vec(),
840            b"baz".to_vec(),
841            b"qux".to_vec(),
842        ]);
843        assert_eq!(
844            col.slice(1, 3),
845            OwnedColumn::VarBinary(vec![b"bar".to_vec(), b"baz".to_vec()])
846        );
847        let permutation = Permutation::try_new(vec![2, 0, 3, 1]).unwrap();
848        assert_eq!(
849            col.try_permute(&permutation).unwrap(),
850            OwnedColumn::VarBinary(vec![
851                b"baz".to_vec(),
852                b"foo".to_vec(),
853                b"qux".to_vec(),
854                b"bar".to_vec()
855            ])
856        );
857    }
858
859    #[test]
860    fn we_can_convert_varbinary_column_round_trip_using_hash() {
861        let raw_bytes = [b"abc".as_ref(), b"xyz".as_ref()];
862
863        let scalars: Vec<TestScalar> = raw_bytes
864            .iter()
865            .map(|&b| TestScalar::from_byte_slice_via_hash(b))
866            .collect();
867
868        let col: Column<'_, TestScalar> =
869            Column::VarBinary((raw_bytes.as_slice(), scalars.as_slice()));
870
871        let owned_col: OwnedColumn<TestScalar> = (&col).into();
872
873        assert_eq!(
874            owned_col,
875            OwnedColumn::VarBinary(vec![b"abc".to_vec(), b"xyz".to_vec()])
876        );
877
878        let bump = bumpalo::Bump::new();
879        let new_col = Column::<TestScalar>::from_owned_column(&owned_col, &bump);
880
881        assert_eq!(col, new_col);
882    }
883
884    #[test]
885    fn we_can_compute_inner_product_with_varbinary_columns_using_hash() {
886        let lhs = OwnedColumn::<TestScalar>::VarBinary(vec![
887            b"foo".to_vec(),
888            b"bar".to_vec(),
889            b"baz".to_vec(),
890        ]);
891
892        let scalars = vec![
893            TestScalar::from(10),
894            TestScalar::from(20),
895            TestScalar::from(30),
896        ];
897
898        let product = lhs.inner_product(&scalars);
899
900        let lhs_hashes: Vec<TestScalar> = [b"foo".as_ref(), b"bar".as_ref(), b"baz".as_ref()]
901            .iter()
902            .map(|&bytes| TestScalar::from_byte_slice_via_hash(bytes))
903            .collect();
904
905        let expected =
906            lhs_hashes[0] * scalars[0] + lhs_hashes[1] * scalars[1] + lhs_hashes[2] * scalars[2];
907
908        assert_eq!(product, expected);
909    }
910}
proof_of_sql/base/database/owned_column.rs

proof_of_sql/base/database/
owned_column.rs
