fix: Try to convert a static list into a set in Rust (#184)

apache · Mar 12, 2024 · 5dc865d · 5dc865d
1 parent 4fec40e
commit 5dc865d
Showing 1 changed file with 20 additions and 11 deletions.
diff --git a/core/src/execution/datafusion/planner.rs b/core/src/execution/datafusion/planner.rs
@@ -25,9 +25,14 @@ use datafusion::{
     common::DataFusionError,
     logical_expr::{BuiltinScalarFunction, Operator as DataFusionOperator},
     physical_expr::{
-        expressions::{BinaryExpr, Column, IsNotNullExpr, Literal as DataFusionLiteral},
+        execution_props::ExecutionProps,
+        expressions::{
+            in_list, BinaryExpr, CaseExpr, CastExpr, Column, Count, FirstValue, InListExpr,
+            IsNotNullExpr, IsNullExpr, LastValue, Literal as DataFusionLiteral, Max, Min,
+            NegativeExpr, NotExpr, Sum, UnKnownColumn,
+        },
         functions::create_physical_expr,
-        PhysicalExpr, PhysicalSortExpr,
+        AggregateExpr, PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr,
     },
     physical_plan::{
         aggregates::{AggregateMode as DFAggregateMode, PhysicalGroupBy},
@@ -39,14 +44,6 @@ use datafusion::{
     },
 };
 use datafusion_common::ScalarValue;
-use datafusion_physical_expr::{
-    execution_props::ExecutionProps,
-    expressions::{
-        CaseExpr, CastExpr, Count, FirstValue, InListExpr, IsNullExpr, LastValue, Max, Min,
-        NegativeExpr, NotExpr, Sum, UnKnownColumn,
-    },
-    AggregateExpr, ScalarFunctionExpr,
-};
 use itertools::Itertools;
 use jni::objects::GlobalRef;
 use num::{BigInt, ToPrimitive};
@@ -496,7 +493,19 @@ impl PhysicalPlanner {
                     .iter()
                     .map(|x| self.create_expr(x, input_schema.clone()).unwrap())
                     .collect::<Vec<_>>();
-                Ok(Arc::new(InListExpr::new(value, list, expr.negated, None)))
+
+                // if schema contains any dictionary type, we should use InListExpr instead of
+                // in_list as it doesn't handle value being dictionary type correctly
+                let contains_dict_type = input_schema
+                    .fields()
+                    .iter()
+                    .any(|f| matches!(f.data_type(), DataType::Dictionary(_, _)));
+                if contains_dict_type {
+                    // TODO: remove the fallback when https://github.com/apache/arrow-datafusion/issues/9530 is fixed
+                    Ok(Arc::new(InListExpr::new(value, list, expr.negated, None)))
+                } else {
+                    in_list(value, list, &expr.negated, input_schema.as_ref()).map_err(|e| e.into())
+                }
             }
             ExprStruct::If(expr) => {
                 let if_expr =