feat: support 'col IN (a, b, c)' type expressions

kernel/src/engine/arrow_conversion.rs

@@ -208,7 +208,7 @@ impl TryFrom<&ArrowDataType> for DataType {
             ArrowDataType::Date64 => Ok(DataType::DATE),
             ArrowDataType::Timestamp(TimeUnit::Microsecond, None) => Ok(DataType::TIMESTAMP_NTZ),
             ArrowDataType::Timestamp(TimeUnit::Microsecond, Some(tz))
-                if tz.eq_ignore_ascii_case("utc") =>
+                if tz.eq_ignore_ascii_case("utc") || tz.eq_ignore_ascii_case("+00:00") =>
             {
                 Ok(DataType::TIMESTAMP)
             }

kernel/src/engine/arrow_expression.rs

@@ -32,6 +32,7 @@ use crate::expressions::{
     BinaryExpression, BinaryOperator, Expression, Scalar, UnaryExpression, UnaryOperator,
     VariadicExpression, VariadicOperator,
 };
+use crate::predicates::PredicateEvaluatorDefaults;
 use crate::schema::{ArrayType, DataType, MapType, PrimitiveType, Schema, SchemaRef, StructField};
 use crate::{EngineData, ExpressionEvaluator, ExpressionHandler};
 
@@ -280,6 +281,84 @@ fn evaluate_expression(
                     (ArrowDataType::Decimal256(_, _), Decimal256Type)
                 }
             }
+            (Column(name), Literal(Scalar::Array(ad))) => {
+                fn op<T: ArrowPrimitiveType>(
+                    values: &dyn Array,
+                    from: fn(T::Native) -> Scalar,
+                ) -> impl Iterator<Item = Option<Scalar>> + '_ {
+                    values.as_primitive::<T>().iter().map(move |v| v.map(from))
+                }
+
+                fn str_op<'a>(
+                    column: impl Iterator<Item = Option<&'a str>> + 'a,
+                ) -> impl Iterator<Item = Option<Scalar>> + 'a {
+                    column.map(|v| v.map(Scalar::from))
+                }
+
+                fn op_in(
+                    inlist: &[Scalar],
+                    values: impl Iterator<Item = Option<Scalar>>,
+                ) -> BooleanArray {
+                    // `v IN (k1, ..., kN)` is logically equivalent to `v = k1 OR ... OR v = kN`, so evaluate
+                    // it as such, ensuring correct handling of NULL inputs (including `Scalar::Null`).
+                    values
+                        .map(|v| {
+                            Some(
+                                PredicateEvaluatorDefaults::finish_eval_variadic(
+                                    VariadicOperator::Or,
+                                    inlist.iter().map(|k| v.as_ref().map(|vv| vv == k)),
+                                    false,
+                                )
+                                // None is returned when no dominant value (true) is found and there is at least one NULL
+                                // In th case of IN, this is equivalent to false
+                                .unwrap_or(false),
+                            )
+                        })
+                        .collect()
+                }
+
+                #[allow(deprecated)]
+                let inlist = ad.array_elements();
+                let column = extract_column(batch, name)?;
+                let data_type = ad
+                    .array_type()
+                    .element_type()
+                    .as_primitive_opt()
+                    .ok_or_else(|| {
+                        Error::invalid_expression(format!(
+                            "IN only supports array literals with primitive elements, got: '{:?}'",
+                            ad.array_type().element_type()
+                        ))
+                    })?;
+
+                // safety: as_* methods on arrow arrays can panic, but we checked the data type before applying.
+                let arr = match (column.data_type(), data_type) {
+                    (ArrowDataType::Utf8, PrimitiveType::String) => op_in(inlist, str_op(column.as_string::<i32>().iter())),
+                    (ArrowDataType::LargeUtf8, PrimitiveType::String) => op_in(inlist, str_op(column.as_string::<i64>().iter())),
+                    (ArrowDataType::Utf8View, PrimitiveType::String) => op_in(inlist, str_op(column.as_string_view().iter())),
+                    (ArrowDataType::Int8, PrimitiveType::Byte) => op_in(inlist,op::<Int8Type>( column.as_ref(), Scalar::from)),
+                    (ArrowDataType::Int16, PrimitiveType::Short) => op_in(inlist,op::<Int16Type>(column.as_ref(), Scalar::from)),
+                    (ArrowDataType::Int32, PrimitiveType::Integer) => op_in(inlist,op::<Int32Type>(column.as_ref(), Scalar::from)),
+                    (ArrowDataType::Int64, PrimitiveType::Long) => op_in(inlist,op::<Int64Type>(column.as_ref(), Scalar::from)),
+                    (ArrowDataType::Float32, PrimitiveType::Float) => op_in(inlist,op::<Float32Type>(column.as_ref(), Scalar::from)),
+                    (ArrowDataType::Float64, PrimitiveType::Double) => op_in(inlist,op::<Float64Type>(column.as_ref(), Scalar::from)),
+                    (ArrowDataType::Date32, PrimitiveType::Date) => op_in(inlist,op::<Date32Type>(column.as_ref(), Scalar::Date)),
+                    (
+                        ArrowDataType::Timestamp(TimeUnit::Microsecond, Some(_)),
+                        PrimitiveType::Timestamp,
+                    ) => op_in(inlist,op::<TimestampMicrosecondType>(column.as_ref(), Scalar::Timestamp)),
+                    (
+                        ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
+                        PrimitiveType::TimestampNtz,
+                    ) => op_in(inlist,op::<TimestampMicrosecondType>(column.as_ref(), Scalar::TimestampNtz)),
+                    (l, r) => {
+                        return Err(Error::invalid_expression(format!(
+                        "Cannot check if value of type '{l}' is contained in array with values of type '{r}'"
+                    )))
+                    }
+                };
+                Ok(Arc::new(arr))
+            }
             (Literal(lit), Literal(Scalar::Array(ad))) => {
                 #[allow(deprecated)]
                 let exists = ad.array_elements().contains(lit);
@@ -382,8 +461,8 @@ fn new_field_with_metadata(
 
 // A helper that is a wrapper over `transform_field_and_col`. This will take apart the passed struct
 // and use that method to transform each column and then put the struct back together. Target types
-// and names for each column should be passed in `target_types_and_names`. The number of elements in
-// the `target_types_and_names` iterator _must_ be the same as the number of columns in
+// and names for each column should be passed in `target_fields`. The number of elements in
+// the `target_fields` iterator _must_ be the same as the number of columns in
 // `struct_array`. The transformation is ordinal. That is, the order of fields in `target_fields`
 // _must_ match the order of the columns in `struct_array`.
 fn transform_struct(
@@ -692,6 +771,85 @@ mod tests {
         assert_eq!(in_result.as_ref(), &in_expected);
     }
 
+    #[test]
+    fn test_column_in_array() {
+        let values = Int32Array::from(vec![0, 1, 2, 3]);
+        let field = Arc::new(Field::new("item", DataType::Int32, true));
+        let rhs = Expression::literal(Scalar::Array(ArrayData::new(
+            ArrayType::new(PrimitiveType::Integer.into(), false),
+            [Scalar::Integer(1), Scalar::Integer(3)],
+        )));
+        let schema = Schema::new([field.clone()]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values.clone())]).unwrap();
+
+        let in_op = Expression::binary(BinaryOperator::In, column_expr!("item"), rhs.clone());
+        let in_result =
+            evaluate_expression(&in_op, &batch, Some(&crate::schema::DataType::BOOLEAN)).unwrap();
+        let in_expected = BooleanArray::from(vec![false, true, false, true]);
+        assert_eq!(in_result.as_ref(), &in_expected);
+
+        let not_in_op = Expression::binary(BinaryOperator::NotIn, column_expr!("item"), rhs);
+        let not_in_result =
+            evaluate_expression(&not_in_op, &batch, Some(&crate::schema::DataType::BOOLEAN))
+                .unwrap();
+        let not_in_expected = BooleanArray::from(vec![true, false, true, false]);
+        assert_eq!(not_in_result.as_ref(), &not_in_expected);
+
+        let in_expected = BooleanArray::from(vec![false, true, false, true]);
+
+        // Date arrays
+        let values = Date32Array::from(vec![0, 1, 2, 3]);
+        let field = Arc::new(Field::new("item", DataType::Date32, true));
+        let rhs = Expression::literal(Scalar::Array(ArrayData::new(
+            ArrayType::new(PrimitiveType::Date.into(), false),
+            [Scalar::Date(1), Scalar::Date(3)],
+        )));
+        let schema = Schema::new([field.clone()]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values.clone())]).unwrap();
+        let in_op = Expression::binary(BinaryOperator::In, column_expr!("item"), rhs.clone());
+        let in_result =
+            evaluate_expression(&in_op, &batch, Some(&crate::schema::DataType::BOOLEAN)).unwrap();
+        assert_eq!(in_result.as_ref(), &in_expected);
+
+        // Timestamp arrays
+        let values = TimestampMicrosecondArray::from(vec![0, 1, 2, 3]).with_timezone("UTC");
+        let field = Arc::new(Field::new(
+            "item",
+            (&crate::schema::DataType::TIMESTAMP).try_into().unwrap(),
+            true,
+        ));
+        let rhs = Expression::literal(Scalar::Array(ArrayData::new(
+            ArrayType::new(PrimitiveType::Timestamp.into(), false),
+            [Scalar::Timestamp(1), Scalar::Timestamp(3)],
+        )));
+        let schema = Schema::new([field.clone()]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values.clone())]).unwrap();
+        let in_op = Expression::binary(BinaryOperator::In, column_expr!("item"), rhs.clone());
+        let in_result =
+            evaluate_expression(&in_op, &batch, Some(&crate::schema::DataType::BOOLEAN)).unwrap();
+        assert_eq!(in_result.as_ref(), &in_expected);
+
+        // Timestamp NTZ arrays
+        let values = TimestampMicrosecondArray::from(vec![0, 1, 2, 3]);
+        let field = Arc::new(Field::new(
+            "item",
+            (&crate::schema::DataType::TIMESTAMP_NTZ)
+                .try_into()
+                .unwrap(),
+            true,
+        ));
+        let rhs = Expression::literal(Scalar::Array(ArrayData::new(
+            ArrayType::new(PrimitiveType::TimestampNtz.into(), false),
+            [Scalar::TimestampNtz(1), Scalar::TimestampNtz(3)],
+        )));
+        let schema = Schema::new([field.clone()]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values.clone())]).unwrap();
+        let in_op = Expression::binary(BinaryOperator::In, column_expr!("item"), rhs.clone());
+        let in_result =
+            evaluate_expression(&in_op, &batch, Some(&crate::schema::DataType::BOOLEAN)).unwrap();
+        assert_eq!(in_result.as_ref(), &in_expected);
+    }
+
     #[test]
     fn test_extract_column() {
         let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);

kernel/src/expressions/scalars.rs

@@ -89,7 +89,7 @@ impl StructData {
 
 /// A single value, which can be null. Used for representing literal values
 /// in [Expressions][crate::expressions::Expression].
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone)]
 pub enum Scalar {
     /// 32bit integer
     Integer(i32),
@@ -224,6 +224,48 @@ impl Display for Scalar {
     }
 }
 
+impl PartialEq<Scalar> for Scalar {
+    fn eq(&self, other: &Self) -> bool {
+        use Scalar::*;
+        // NOTE: We intentionally do two match arms for each variant to avoid a catch-all, so
+        // that new variants trigger compilation failures instead of being silently ignored.
+        match (self, other) {
+            (Integer(a), Integer(b)) => a == b,
+            (Integer(_), _) => false,
+            (Long(a), Long(b)) => a == b,
+            (Long(_), _) => false,
+            (Short(a), Short(b)) => a == b,
+            (Short(_), _) => false,
+            (Byte(a), Byte(b)) => a == b,
+            (Byte(_), _) => false,
+            (Float(a), Float(b)) => a == b,
+            (Float(_), _) => false,
+            (Double(a), Double(b)) => a == b,
+            (Double(_), _) => false,
+            (String(a), String(b)) => a == b,
+            (String(_), _) => false,
+            (Boolean(a), Boolean(b)) => a == b,
+            (Boolean(_), _) => false,
+            (Timestamp(a), Timestamp(b)) => a == b,
+            (Timestamp(_), _) => false,
+            (TimestampNtz(a), TimestampNtz(b)) => a == b,
+            (TimestampNtz(_), _) => false,
+            (Date(a), Date(b)) => a == b,
+            (Date(_), _) => false,
+            (Binary(a), Binary(b)) => a == b,
+            (Binary(_), _) => false,
+            (Decimal(a, _, _), Decimal(b, _, _)) => a == b,
+            (Decimal(_, _, _), _) => false,
+            (Struct(a), Struct(b)) => a == b,
+            (Struct(_), _) => false,
+            (Array(a), Array(b)) => a == b,
+            (Array(_), _) => false,
+            (Null(_), Null(_)) => false, // NOTE: NULL values are incomparable by definition
+            (Null(_), _) => false,
+        }
+    }
+}
+
 impl PartialOrd for Scalar {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         use Scalar::*;
@@ -585,6 +627,7 @@ mod tests {
         assert_eq!(&format!("{}", column_op), "3.1415927 IN Column(item)");
         assert_eq!(&format!("{}", column_not_op), "'Cool' NOT IN Column(item)");
     }
+
     #[test]
     fn test_timestamp_parse() {
         let assert_timestamp_eq = |scalar_string, micros| {
@@ -599,6 +642,7 @@ mod tests {
         assert_timestamp_eq("2011-01-11 13:06:07.123456", 1294751167123456);
         assert_timestamp_eq("1970-01-01 00:00:00", 0);
     }
+
     #[test]
     fn test_timestamp_ntz_parse() {
         let assert_timestamp_eq = |scalar_string, micros| {
@@ -627,4 +671,30 @@ mod tests {
         let p_type = PrimitiveType::Timestamp;
         assert_timestamp_fails(&p_type, "1971-07-22");
     }
+
+    #[test]
+    fn test_partial_cmp() {
+        let a = Scalar::Integer(1);
+        let b = Scalar::Integer(2);
+        assert_eq!(a.partial_cmp(&b), Some(Ordering::Less));
+        assert_eq!(b.partial_cmp(&a), Some(Ordering::Greater));
+        assert_eq!(a.partial_cmp(&a), Some(Ordering::Equal));
+        assert_eq!(b.partial_cmp(&b), Some(Ordering::Equal));
+
+        // assert that NULL values are incomparable
+        let null = Scalar::Null(DataType::INTEGER);
+        assert_eq!(null.partial_cmp(&null), None);
+    }
+
+    #[test]
+    fn test_partial_eq() {
+        let a = Scalar::Integer(1);
+        let b = Scalar::Integer(2);
+        assert!(!a.eq(&b));
+        assert!(a.eq(&a));
+
+        // assert that NULL values are incomparable
+        let null = Scalar::Null(DataType::INTEGER);
+        assert!(!null.eq(&null));
+    }
 }