Int64Decimal in list operator support

Author: paveltiunov

datafusion/src/physical_plan/expressions/in_list.rs

@@ -24,7 +24,8 @@ use arrow::array::GenericStringArray;
 use arrow::array::{
     ArrayRef, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array,
     Int64Array, Int8Array, StringOffsetSizeTrait, UInt16Array, UInt32Array, UInt64Array,
-    UInt8Array,
+    UInt8Array, Int64Decimal0Array, Int64Decimal1Array, Int64Decimal2Array, Int64Decimal3Array, Int64Decimal4Array,
+    Int64Decimal5Array, Int64Decimal10Array
 };
 use arrow::{datatypes::DataType, record_batch::RecordBatch};
 
@@ -43,6 +44,59 @@ pub struct InListExpr {
 }
 
 macro_rules! make_contains {
+($ARRAY:expr, $LIST_VALUES:expr, $NEGATED:expr, Int64Decimal, $ARRAY_TYPE:ident, $SCALE:expr) => {{
+        let array = $ARRAY.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
+
+        let mut contains_null = false;
+        let values = $LIST_VALUES
+            .iter()
+            .flat_map(|expr| match expr {
+                ColumnarValue::Scalar(s) => match s {
+                    ScalarValue::Int64Decimal(Some(v), $SCALE) => Some(*v),
+                    ScalarValue::Int64Decimal(None, $SCALE) => {
+                        contains_null = true;
+                        None
+                    }
+                    ScalarValue::Utf8(None) => {
+                        contains_null = true;
+                        None
+                    }
+                    datatype => unimplemented!("Unexpected type {} for InList", datatype),
+                },
+                ColumnarValue::Array(_) => {
+                    unimplemented!("InList does not yet support nested columns.")
+                }
+            })
+            .collect::<Vec<_>>();
+
+        Ok(ColumnarValue::Array(Arc::new(
+            array
+                .iter()
+                .map(|x| {
+                    let contains = x.map(|x| values.contains(&x));
+                    match contains {
+                        Some(true) => {
+                            if $NEGATED {
+                                Some(false)
+                            } else {
+                                Some(true)
+                            }
+                        }
+                        Some(false) => {
+                            if contains_null {
+                                None
+                            } else if $NEGATED {
+                                Some(true)
+                            } else {
+                                Some(false)
+                            }
+                        }
+                        None => None,
+                    }
+                })
+                .collect::<BooleanArray>(),
+        )))
+    }};
     ($ARRAY:expr, $LIST_VALUES:expr, $NEGATED:expr, $SCALAR_VALUE:ident, $ARRAY_TYPE:ident) => {{
         let array = $ARRAY.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
 
@@ -262,6 +316,27 @@ impl PhysicalExpr for InListExpr {
             DataType::UInt8 => {
                 make_contains!(array, list_values, self.negated, UInt8, UInt8Array)
             }
+            DataType::Int64Decimal(0) => {
+                make_contains!(array, list_values, self.negated, Int64Decimal, Int64Decimal0Array, 0)
+            }
+            DataType::Int64Decimal(1) => {
+                make_contains!(array, list_values, self.negated, Int64Decimal, Int64Decimal1Array, 1)
+            }
+            DataType::Int64Decimal(2) => {
+                make_contains!(array, list_values, self.negated, Int64Decimal, Int64Decimal2Array, 2)
+            }
+            DataType::Int64Decimal(3) => {
+                make_contains!(array, list_values, self.negated, Int64Decimal, Int64Decimal3Array, 3)
+            }
+            DataType::Int64Decimal(4) => {
+                make_contains!(array, list_values, self.negated, Int64Decimal, Int64Decimal4Array, 4)
+            }
+            DataType::Int64Decimal(5) => {
+                make_contains!(array, list_values, self.negated, Int64Decimal, Int64Decimal5Array, 5)
+            }
+            DataType::Int64Decimal(10) => {
+                make_contains!(array, list_values, self.negated, Int64Decimal, Int64Decimal10Array, 10)
+            }
             DataType::Boolean => {
                 make_contains!(array, list_values, self.negated, Boolean, BooleanArray)
             }