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

proof_of_sql/base/database/
owned_column.rs
