feat: add column_index_operation

crates/proof-of-sql/src/base/database/column_index_operation.rs

@@ -0,0 +1,150 @@
+use super::{slice_operation::apply_slice_to_indexes, Column, ColumnOperationResult, ColumnType};
+use crate::base::scalar::Scalar;
+use bumpalo::Bump;
+
+/// Apply a `Column` to a vector of indexes, returning a new `Column` with the
+/// values at the given indexes. Repetitions are allowed.
+///
+/// # Panics
+/// Panics if any of the indexes are out of bounds.
+#[allow(dead_code)]
+pub(crate) fn apply_column_to_indexes<'a, S>(
+    column: &Column<'a, S>,
+    alloc: &'a Bump,
+    indexes: &[usize],
+) -> ColumnOperationResult<Column<'a, S>>
+where
+    S: Scalar,
+{
+    match column.column_type() {
+        ColumnType::Boolean => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_boolean().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::Boolean(alloc.alloc_slice_copy(&raw_values) as &[_]))
+        }
+        ColumnType::TinyInt => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_tinyint().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::TinyInt(alloc.alloc_slice_copy(&raw_values) as &[_]))
+        }
+        ColumnType::SmallInt => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_smallint().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::SmallInt(alloc.alloc_slice_copy(&raw_values) as &[_]))
+        }
+        ColumnType::Int => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_int().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::Int(alloc.alloc_slice_copy(&raw_values) as &[_]))
+        }
+        ColumnType::BigInt => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_bigint().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::BigInt(alloc.alloc_slice_copy(&raw_values) as &[_]))
+        }
+        ColumnType::Int128 => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_int128().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::Int128(alloc.alloc_slice_copy(&raw_values) as &[_]))
+        }
+        ColumnType::Scalar => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_scalar().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::Scalar(alloc.alloc_slice_copy(&raw_values) as &[_]))
+        }
+        ColumnType::Decimal75(precision, scale) => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_decimal75().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::Decimal75(
+                precision,
+                scale,
+                alloc.alloc_slice_copy(&raw_values) as &[_],
+            ))
+        }
+        ColumnType::VarChar => {
+            let (raw_values, raw_scalars) = column.as_varchar().expect("Column types should match");
+            let raw_values = apply_slice_to_indexes(raw_values, indexes)?;
+            let scalars = apply_slice_to_indexes(raw_scalars, indexes)?;
+            Ok(Column::VarChar((
+                alloc.alloc_slice_clone(&raw_values) as &[_],
+                alloc.alloc_slice_copy(&scalars) as &[_],
+            )))
+        }
+        ColumnType::TimestampTZ(tu, tz) => {
+            let raw_values = apply_slice_to_indexes(
+                column.as_timestamptz().expect("Column types should match"),
+                indexes,
+            )?;
+            Ok(Column::TimestampTZ(
+                tu,
+                tz,
+                alloc.alloc_slice_copy(&raw_values) as &[_],
+            ))
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::base::{database::ColumnOperationError, scalar::test_scalar::TestScalar};
+
+    #[test]
+    fn test_apply_index_op() {
+        let bump = Bump::new();
+        let column: Column<TestScalar> = Column::Int(&[1, 2, 3, 4, 5]);
+        let indexes = [1, 3, 1, 2];
+        let result = apply_column_to_indexes(&column, &bump, &indexes).unwrap();
+        assert_eq!(result, Column::Int(&[2, 4, 2, 3]));
+
+        let scalars = [1, 2, 3].iter().map(TestScalar::from).collect::<Vec<_>>();
+        let column: Column<TestScalar> = Column::Scalar(&scalars);
+        let indexes = [1, 1, 1];
+        let result = apply_column_to_indexes(&column, &bump, &indexes).unwrap();
+        let expected_scalars = [2, 2, 2].iter().map(TestScalar::from).collect::<Vec<_>>();
+        assert_eq!(result, Column::Scalar(&expected_scalars));
+
+        let strings = vec!["a", "b", "c"];
+        let scalars = strings.iter().map(TestScalar::from).collect::<Vec<_>>();
+        let column: Column<TestScalar> = Column::VarChar((&strings, &scalars));
+        let indexes = [2, 1, 1];
+        let result = apply_column_to_indexes(&column, &bump, &indexes).unwrap();
+        let expected_strings = vec!["c", "b", "b"];
+        let expected_scalars = expected_strings
+            .iter()
+            .map(TestScalar::from)
+            .collect::<Vec<_>>();
+        assert_eq!(
+            result,
+            Column::VarChar((&expected_strings, &expected_scalars))
+        );
+    }
+
+    #[test]
+    fn test_apply_index_op_out_of_bound() {
+        let bump = Bump::new();
+        let column: Column<TestScalar> = Column::Int(&[1, 2, 3, 4, 5]);
+        let indexes = [1, 3, 1, 2, 5];
+        let result = apply_column_to_indexes(&column, &bump, &indexes);
+        assert!(matches!(
+            result,
+            Err(ColumnOperationError::IndexOutOfBounds { .. })
+        ));
+    }
+}

crates/proof-of-sql/src/base/database/column_operation_error.rs

@@ -63,6 +63,15 @@ pub enum ColumnOperationError {
         /// The type of the column that caused the error
         actual_type: ColumnType,
     },
+
+    /// Errors related to index out of bounds
+    #[snafu(display("Index out of bounds: {index} >= {len}"))]
+    IndexOutOfBounds {
+        /// The index that caused the error
+        index: usize,
+        /// The length of the column
+        len: usize,
+    },
 }
 
 /// Result type for column operations

crates/proof-of-sql/src/base/database/mod.rs

@@ -25,6 +25,8 @@ pub(super) use column_comparison_operation::{
     ComparisonOp, EqualOp, GreaterThanOrEqualOp, LessThanOrEqualOp,
 };
 
+mod column_index_operation;
+
 mod column_repetition_operation;
 pub(super) use column_repetition_operation::{ColumnRepeatOp, ElementwiseRepeatOp, RepetitionOp};
 

crates/proof-of-sql/src/base/database/slice_operation.rs

@@ -242,6 +242,28 @@ pub(super) fn repeat_elementwise<S: Clone>(slice: &[S], n: usize) -> impl Iterat
         .flat_map(move |s| core::iter::repeat(s).take(n).cloned())
 }
 
+/// Apply a slice to a slice of indexes.
+///
+/// e.g. `apply_slice_to_indexes(&[1, 2, 3], &[0, 0, 1, 0]).unwrap()` -> `vec![1, 1, 2, 1]`
+/// Note that the function will return an error if any index is out of bounds.
+pub(super) fn apply_slice_to_indexes<S: Clone>(
+    slice: &[S],
+    indexes: &[usize],
+) -> ColumnOperationResult<Vec<S>> {
+    let max_index = slice.len();
+    indexes
+        .iter()
+        .map(|&i| {
+            (i < max_index).then(|| slice[i].clone()).ok_or(
+                ColumnOperationError::IndexOutOfBounds {
+                    index: i,
+                    len: max_index,
+                },
+            )
+        })
+        .collect()
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -662,4 +684,54 @@ mod test {
         let expected: Vec<i32> = vec![];
         assert_eq!(expected, actual);
     }
+
+    #[test]
+    fn we_can_apply_slices_to_indexes() {
+        // We can apply a slice to indexes
+        let slice = [1_i16, 2, 3];
+        let indexes = [0, 2];
+        let actual = apply_slice_to_indexes(&slice, &indexes).unwrap();
+        let expected = vec![1_i16, 3];
+        assert_eq!(expected, actual);
+
+        // We can apply an empty slice to no indexes
+        let slice = [0; 0];
+        let indexes = [];
+        let actual = apply_slice_to_indexes(&slice, &indexes).unwrap();
+        let expected: Vec<i32> = vec![];
+        assert_eq!(expected, actual);
+
+        // We can apply a slice to no indexes
+        let slice = [1_i16, 2, 3];
+        let indexes = [];
+        let actual = apply_slice_to_indexes(&slice, &indexes).unwrap();
+        let expected: Vec<i16> = vec![];
+        assert_eq!(expected, actual);
+
+        // Repetition in indexes is allowed
+        let slice = [1_i16, 2, 3];
+        let indexes = [0, 0, 2, 2];
+        let actual = apply_slice_to_indexes(&slice, &indexes).unwrap();
+        let expected = vec![1_i16, 1, 3, 3];
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn we_cannot_apply_slices_to_indexes_if_out_of_bounds() {
+        // We cannot apply a slice to out-of-bounds indexes
+        let slice = [1_i16, 2, 3];
+        let indexes = [0, 3];
+        assert!(matches!(
+            apply_slice_to_indexes(&slice, &indexes),
+            Err(ColumnOperationError::IndexOutOfBounds { .. })
+        ));
+
+        // We can not apply an empty slice to any non-empty indexes
+        let slice = [0; 0];
+        let indexes = [0];
+        assert!(matches!(
+            apply_slice_to_indexes(&slice, &indexes),
+            Err(ColumnOperationError::IndexOutOfBounds { .. })
+        ));
+    }
 }