First pass on IN/Not In

Cargo.toml

@@ -24,6 +24,7 @@ version = "0.2.0"
 arrow = { version = "^52.0" }
 arrow-arith = { version = "^52.0" }
 arrow-array = { version = "^52.0" }
+arrow-buffer = { version = "^52.0" }
 arrow-cast = { version = "^52.0" }
 arrow-data = { version = "^52.0" }
 arrow-ord = { version = "^52.0" }

kernel/Cargo.toml

@@ -37,6 +37,8 @@ delta_kernel_derive = { path = "../derive-macros", version = "0.2.0" }
 visibility = "0.1.0"
 
 # Used in default engine
+arrow-array = { workspace = true, optional = true }
+arrow-buffer = { workspace = true, optional = true }
 arrow-array = { workspace = true, optional = true, features = ["chrono-tz"] }
 arrow-select = { workspace = true, optional = true }
 arrow-arith = { workspace = true, optional = true }
@@ -74,6 +76,7 @@ default-engine = [
   "arrow-conversion",
   "arrow-expression",
   "arrow-array",
+  "arrow-buffer",
   "arrow-cast",
   "arrow-json",
   "arrow-schema",

kernel/src/engine/arrow_expression.rs

@@ -4,25 +4,32 @@ use std::sync::Arc;
 use arrow_arith::boolean::{and_kleene, is_null, not, or_kleene};
 use arrow_arith::numeric::{add, div, mul, sub};
 use arrow_array::cast::AsArray;
+use arrow_array::types::*;
 use arrow_array::{
     Array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Datum, Decimal128Array, Float32Array,
     Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, ListArray, RecordBatch,
     StringArray, StructArray, TimestampMicrosecondArray,
 };
+use arrow_buffer::OffsetBuffer;
 use arrow_ord::cmp::{distinct, eq, gt, gt_eq, lt, lt_eq, neq};
+use arrow_ord::comparison::{in_list, in_list_utf8};
 use arrow_schema::{
-    ArrowError, DataType as ArrowDataType, Field as ArrowField, Fields, Schema as ArrowSchema,
+    ArrowError, DataType as ArrowDataType, Field as ArrowField, Fields, IntervalUnit,
+    Schema as ArrowSchema, TimeUnit,
 };
+use arrow_select::concat::concat;
 use itertools::Itertools;
 
-use super::arrow_conversion::LIST_ARRAY_ROOT;
 use crate::engine::arrow_data::ArrowEngineData;
+use crate::engine::arrow_utils::prim_array_cmp;
 use crate::engine::arrow_utils::ensure_data_types;
 use crate::error::{DeltaResult, Error};
 use crate::expressions::{BinaryOperator, Expression, Scalar, UnaryOperator, VariadicOperator};
 use crate::schema::{DataType, PrimitiveType, SchemaRef};
 use crate::{EngineData, ExpressionEvaluator, ExpressionHandler};
 
+use super::arrow_conversion::LIST_ARRAY_ROOT;
+
 // TODO leverage scalars / Datum
 
 fn downcast_to_bool(arr: &dyn Array) -> DeltaResult<&BooleanArray> {
@@ -67,6 +74,20 @@ impl Scalar {
                     .try_collect()?;
                 Arc::new(StructArray::try_new(fields, arrays, None)?)
             }
+            Array(data) => {
+                let values = data.array_elements();
+                let vecs: Vec<_> = values.iter().map(|v| v.to_array(num_rows)).try_collect()?;
+                let values: Vec<_> = vecs.iter().map(|x| x.as_ref()).collect();
+                let offsets: Vec<_> = vecs.iter().map(|v| v.len()).collect();
+                let offset_buffer = OffsetBuffer::from_lengths(offsets);
+                let field = ArrowField::try_from(data.array_type())?;
+                Arc::new(ListArray::new(
+                    Arc::new(field),
+                    offset_buffer,
+                    concat(values.as_slice())?,
+                    None,
+                ))
+            }
             Null(data_type) => match data_type {
                 DataType::Primitive(primitive) => match primitive {
                     PrimitiveType::Byte => Arc::new(Int8Array::new_null(num_rows)),
@@ -168,7 +189,6 @@ fn evaluate_expression(
 ) -> DeltaResult<ArrayRef> {
     use BinaryOperator::*;
     use Expression::*;
-
     match (expression, result_type) {
         (Literal(scalar), _) => Ok(scalar.to_array(batch.num_rows())?),
         (Column(name), _) => {
@@ -216,6 +236,69 @@ fn evaluate_expression(
                 UnaryOperator::IsNull => Arc::new(is_null(&arr)?),
             })
         }
+        (
+            BinaryOperation {
+                op: In,
+                left,
+                right,
+            },
+            _,
+        ) => {
+            let left_arr = evaluate_expression(left.as_ref(), batch, None)?;
+            let right_arr = evaluate_expression(right.as_ref(), batch, None)?;
+            if let Some(string_arr) = left_arr.as_string_opt::<i32>() {
+                return in_list_utf8(string_arr, right_arr.as_list::<i32>())
+                    .map(wrap_comparison_result)
+                    .map_err(Error::generic_err);
+            }
+            prim_array_cmp! {
+                left_arr, right_arr,
+                (ArrowDataType::Int8, Int8Type),
+                (ArrowDataType::Int16, Int16Type),
+                (ArrowDataType::Int32, Int32Type),
+                (ArrowDataType::Int64, Int64Type),
+                (ArrowDataType::UInt8, UInt8Type),
+                (ArrowDataType::UInt16, UInt16Type),
+                (ArrowDataType::UInt32, UInt32Type),
+                (ArrowDataType::UInt64, UInt64Type),
+                (ArrowDataType::Float16, Float16Type),
+                (ArrowDataType::Float32, Float32Type),
+                (ArrowDataType::Float64, Float64Type),
+                (ArrowDataType::Timestamp(TimeUnit::Second, _), TimestampSecondType),
+                (ArrowDataType::Timestamp(TimeUnit::Millisecond, _), TimestampMillisecondType),
+                (ArrowDataType::Timestamp(TimeUnit::Microsecond, _), TimestampMicrosecondType),
+                (ArrowDataType::Timestamp(TimeUnit::Nanosecond, _), TimestampNanosecondType),
+                (ArrowDataType::Date32, Date32Type),
+                (ArrowDataType::Date64, Date64Type),
+                (ArrowDataType::Time32(TimeUnit::Second), Time32SecondType),
+                (ArrowDataType::Time32(TimeUnit::Millisecond), Time32MillisecondType),
+                (ArrowDataType::Time64(TimeUnit::Microsecond), Time64MicrosecondType),
+                (ArrowDataType::Time64(TimeUnit::Nanosecond), Time64NanosecondType),
+                (ArrowDataType::Duration(TimeUnit::Second), DurationSecondType),
+                (ArrowDataType::Duration(TimeUnit::Millisecond), DurationMillisecondType),
+                (ArrowDataType::Duration(TimeUnit::Microsecond), DurationMicrosecondType),
+                (ArrowDataType::Duration(TimeUnit::Nanosecond), DurationNanosecondType),
+                (ArrowDataType::Interval(IntervalUnit::DayTime), IntervalDayTimeType),
+                (ArrowDataType::Interval(IntervalUnit::YearMonth), IntervalYearMonthType),
+                (ArrowDataType::Interval(IntervalUnit::MonthDayNano), IntervalMonthDayNanoType),
+                (ArrowDataType::Decimal128(_, _), Decimal128Type),
+                (ArrowDataType::Decimal256(_, _), Decimal256Type)
+            }
+        }
+        (
+            BinaryOperation {
+                op: NotIn,
+                left,
+                right,
+            },
+            _,
+        ) => {
+            let reverse_op = Expression::binary(In, *left.clone(), *right.clone());
+            let reverse_expr = evaluate_expression(&reverse_op, batch, None)?;
+            not(reverse_expr.as_boolean())
+                .map(wrap_comparison_result)
+                .map_err(Error::generic_err)
+        }
         (BinaryOperation { op, left, right }, _) => {
             let left_arr = evaluate_expression(left.as_ref(), batch, None)?;
             let right_arr = evaluate_expression(right.as_ref(), batch, None)?;
@@ -233,6 +316,7 @@ fn evaluate_expression(
                 Equal => |l, r| eq(l, r).map(wrap_comparison_result),
                 NotEqual => |l, r| neq(l, r).map(wrap_comparison_result),
                 Distinct => |l, r| distinct(l, r).map(wrap_comparison_result),
+                _ => return Err(Error::generic("Invalid expression given")),
             };
 
             eval(&left_arr, &right_arr).map_err(Error::generic_err)
@@ -321,11 +405,81 @@ impl ExpressionEvaluator for DefaultExpressionEvaluator {
 
 #[cfg(test)]
 mod tests {
+    use std::ops::{Add, Div, Mul, Sub};
 
-    use super::*;
-    use arrow_array::Int32Array;
+    use arrow_array::{GenericStringArray, Int32Array};
+    use arrow_buffer::ScalarBuffer;
     use arrow_schema::{DataType, Field, Fields, Schema};
-    use std::ops::{Add, Div, Mul, Sub};
+
+    use crate::expressions::*;
+
+    use super::*;
+
+    #[test]
+    fn test_array_column() {
+        let values = Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8]);
+        let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 3, 6, 9]));
+        let field = Arc::new(Field::new("item", DataType::Int32, true));
+        let arr_field = Arc::new(Field::new("item", DataType::List(field.clone()), true));
+
+        let schema = Schema::new(vec![arr_field.clone()]);
+
+        let array = ListArray::new(field.clone(), offsets, Arc::new(values), None);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array.clone())]).unwrap();
+
+        let not_op = Expression::binary(
+            BinaryOperator::NotIn,
+            Expression::literal(5),
+            Expression::column("item"),
+        );
+
+        let in_op = Expression::binary(
+            BinaryOperator::NotIn,
+            Expression::literal(5),
+            Expression::column("item"),
+        );
+
+        let result = evaluate_expression(&not_op, &batch, None).unwrap();
+        let expected = BooleanArray::from(vec![true, false, true]);
+        assert_eq!(result.as_ref(), &expected);
+
+        let in_result = evaluate_expression(&in_op, &batch, None).unwrap();
+        let in_expected = BooleanArray::from(vec![true, false, true]);
+        assert_eq!(in_result.as_ref(), &in_expected);
+    }
+
+    #[test]
+    fn test_str_arrays() {
+        let values = GenericStringArray::<i32>::from(vec![
+            "hi", "bye", "hi", "hi", "bye", "bye", "hi", "bye", "hi",
+        ]);
+        let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 3, 6, 9]));
+        let field = Arc::new(Field::new("item", DataType::Utf8, true));
+        let arr_field = Arc::new(Field::new("item", DataType::List(field.clone()), true));
+        let schema = Schema::new(vec![arr_field.clone()]);
+        let array = ListArray::new(field.clone(), offsets, Arc::new(values), None);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array.clone())]).unwrap();
+
+        let str_not_op = Expression::binary(
+            BinaryOperator::NotIn,
+            Expression::literal("bye"),
+            Expression::column("item"),
+        );
+
+        let str_in_op = Expression::binary(
+            BinaryOperator::In,
+            Expression::literal("hi"),
+            Expression::column("item"),
+        );
+
+        let result = evaluate_expression(&str_in_op, &batch, None).unwrap();
+        let expected = BooleanArray::from(vec![true, true, true]);
+        assert_eq!(result.as_ref(), &expected);
+
+        let in_result = evaluate_expression(&str_not_op, &batch, None).unwrap();
+        let in_expected = BooleanArray::from(vec![false, false, false]);
+        assert_eq!(in_result.as_ref(), &in_expected);
+    }
 
     #[test]
     fn test_extract_column() {

kernel/src/engine/arrow_utils.rs

@@ -20,6 +20,39 @@ use itertools::Itertools;
 use parquet::{arrow::ProjectionMask, schema::types::SchemaDescriptor};
 use tracing::debug;
 
+macro_rules! prim_array_cmp {
+    ( $left_arr: ident, $right_arr: ident, $(($data_ty: pat, $prim_ty: ty)),+ ) => {
+
+        return match $left_arr.data_type() {
+        $(
+            $data_ty => {
+                let prim_array = $left_arr.as_primitive_opt::<$prim_ty>()
+                        .ok_or(Error::invalid_expression(
+                            format!("Cannot cast to primitive array: {}", $left_arr.data_type()))
+                        )?;
+                    let list_array = $right_arr.as_list_opt::<i32>()
+                        .ok_or(Error::invalid_expression(
+                            format!("Cannot cast to list array: {}", $right_arr.data_type()))
+                        )?;
+                in_list(prim_array, list_array).map(wrap_comparison_result)
+            }
+        )+
+            _ => Err(ArrowError::CastError(
+                        format!("Bad Comparison between: {:?} and {:?}",
+                            $left_arr.data_type(),
+                            $right_arr.data_type())
+                        )
+                )
+        }.map_err(Error::generic_err);
+    };
+}
+
+pub(crate) use prim_array_cmp;
+
+/// Get the indicies in `parquet_schema` of the specified columns in `requested_schema`. This
+/// returns a tuples of (mask_indicies: Vec<parquet_schema_index>, reorder_indicies:
+/// Vec<requested_index>). `mask_indicies` is used for generating the mask for reading from the
+
 fn make_arrow_error(s: String) -> Error {
     Error::Arrow(arrow_schema::ArrowError::InvalidArgumentError(s))
 }
@@ -498,6 +531,7 @@ fn get_indices(
 /// Get the indices in `parquet_schema` of the specified columns in `requested_schema`. This returns
 /// a tuple of (mask_indices: Vec<parquet_schema_index>, reorder_indices:
 /// Vec<requested_index>). `mask_indices` is used for generating the mask for reading from the
+
 /// parquet file, and simply contains an entry for each index we wish to select from the parquet
 /// file set to the index of the requested column in the parquet. `reorder_indices` is used for
 /// re-ordering. See the documentation for [`ReorderIndex`] to understand what each element in the

kernel/src/error.rs

@@ -105,7 +105,7 @@ pub enum Error {
     #[error("Invalid url: {0}")]
     InvalidUrl(#[from] url::ParseError),
 
-    /// serde enountered malformed json
+    /// serde encountered malformed json
     #[error(transparent)]
     MalformedJson(serde_json::Error),
 
@@ -145,12 +145,16 @@ pub enum Error {
     InvalidTableLocation(String),
 
     /// Precision or scale not compliant with delta specification
-    #[error("Inavlid decimal: {0}")]
+    #[error("Invalid decimal: {0}")]
     InvalidDecimal(String),
 
-    /// Incosistent data passed to struct scalar
+    /// Inconsistent data passed to struct scalar
     #[error("Invalid struct data: {0}")]
     InvalidStructData(String),
+
+    /// Expressions did not parse or evaluate correctly
+    #[error("Invalid expression evaluation: {0}")]
+    InvalidExpressionEvaluation(String),
 }
 
 // Convenience constructors for Error types that take a String argument
@@ -196,6 +200,9 @@ impl Error {
     pub fn invalid_struct_data(msg: impl ToString) -> Self {
         Self::InvalidStructData(msg.to_string())
     }
+    pub fn invalid_expression(msg: impl ToString) -> Self {
+        Self::InvalidExpressionEvaluation(msg.to_string())
+    }
 
     pub fn internal_error(msg: impl ToString) -> Self {
         Self::InternalError(msg.to_string()).with_backtrace()

kernel/src/expressions/mod.rs

@@ -5,7 +5,7 @@ use std::fmt::{Display, Formatter};
 
 use itertools::Itertools;
 
-pub use self::scalars::{Scalar, StructData};
+pub use self::scalars::{ArrayData, Scalar, StructData};
 
 mod scalars;
 
@@ -34,6 +34,10 @@ pub enum BinaryOperator {
     NotEqual,
     /// Distinct
     Distinct,
+    /// IN
+    In,
+    /// NOT IN
+    NotIn,
 }
 
 impl BinaryOperator {
@@ -50,8 +54,8 @@ impl BinaryOperator {
         }
     }
 
-    // invert an operator. Returns Some<InvertedOp> if the operator supports inversion, None if it
-    // cannot be inverted
+    /// invert an operator. Returns Some<InvertedOp> if the operator supports inversion, None if it
+    /// cannot be inverted
     pub(crate) fn invert(&self) -> Option<BinaryOperator> {
         use BinaryOperator::*;
         match self {
@@ -61,6 +65,8 @@ impl BinaryOperator {
             GreaterThanOrEqual => Some(LessThan),
             Equal => Some(NotEqual),
             NotEqual => Some(Equal),
+            In => Some(NotIn),
+            NotIn => Some(In),
             _ => None,
         }
     }
@@ -97,8 +103,10 @@ impl Display for BinaryOperator {
             Self::NotEqual => write!(f, "!="),
             // TODO(roeap): AFAIK DISTINCT does not have a commonly used operator symbol
             // so ideally this would not be used as we use Display for rendering expressions
-            // in our code we take care of this, bot thers might now ...
+            // in our code we take care of this, but theirs might not ...
             Self::Distinct => write!(f, "DISTINCT"),
+            Self::In => write!(f, "IN"),
+            Self::NotIn => write!(f, "NOT IN"),
         }
     }
 }