Refactor InListExpr to store arrays and support structs.

adriangb · adriangb · commit f1f3b660b602 · 2025-11-02T15:13:04.000-06:00
Changes:
- Enhance InListExpr to efficiently store homogeneous lists as arrays and avoid a conversion to Vec&lt;PhysicalExpr&gt;
    by adding an internal InListStorage enum with Array and Exprs variants
- Re-use existing hashing and comparison utilities to support Struct arrays and other complex types
- Add public function `in_list_from_array(expr, list_array, negated)` for creating InList from arrays
diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs
diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs
@@ -46,7 +46,7 @@ pub use cast_column::CastColumnExpr;
 pub use column::{col, with_new_schema, Column};
 pub use datafusion_expr::utils::format_state_name;
 pub use dynamic_filters::DynamicFilterPhysicalExpr;
-pub use in_list::{in_list, InListExpr};
+pub use in_list::{in_list, in_list_from_array, InListExpr};
 pub use is_not_null::{is_not_null, IsNotNullExpr};
 pub use is_null::{is_null, IsNullExpr};
 pub use like::{like, LikeExpr};
diff --git a/datafusion/physical-expr/src/utils/guarantee.rs b/datafusion/physical-expr/src/utils/guarantee.rs
@@ -93,12 +93,12 @@ impl LiteralGuarantee {
     /// Create a new instance of the guarantee if the provided operator is
     /// supported. Returns None otherwise. See [`LiteralGuarantee::analyze`] to
     /// create these structures from an predicate (boolean expression).
-    fn new<'a>(
+    fn new(
         column_name: impl Into<String>,
         guarantee: Guarantee,
-        literals: impl IntoIterator<Item = &'a ScalarValue>,
+        literals: impl IntoIterator<Item = ScalarValue>,
     ) -> Self {
-        let literals: HashSet<_> = literals.into_iter().cloned().collect();
+        let literals: HashSet<_> = literals.into_iter().collect();
 
         Self {
             column: Column::from_name(column_name),
@@ -130,11 +130,18 @@ impl LiteralGuarantee {
                     .downcast_ref::<crate::expressions::InListExpr>()
                 {
                     if let Some(inlist) = ColInList::try_new(inlist) {
-                        builder.aggregate_multi_conjunct(
-                            inlist.col,
-                            inlist.guarantee,
-                            inlist.list.iter().map(|lit| lit.value()),
-                        )
+                        match inlist.inlist.list() {
+                            Ok(list) => builder.aggregate_multi_conjunct(
+                                inlist.col,
+                                inlist.guarantee,
+                                list.iter().filter_map(|expr| {
+                                    expr.as_any()
+                                        .downcast_ref::<crate::expressions::Literal>()
+                                        .map(|lit| lit.value().clone())
+                                }),
+                            ),
+                            Err(_) => builder,
+                        }
                     } else {
                         builder
                     }
@@ -184,7 +191,7 @@ impl LiteralGuarantee {
                         builder.aggregate_multi_conjunct(
                             first_term.unwrap().col,
                             Guarantee::In,
-                            terms.iter().map(|term| term.lit.value()),
+                            terms.iter().map(|term| term.lit.value().clone()),
                         )
                     } else {
                         // Handle disjunctions with conjunctions like (a = 1 AND b = 2) OR (a = 2 AND b = 3)
@@ -221,19 +228,18 @@ impl LiteralGuarantee {
                             let literals: Vec<_> = termsets
                                 .iter()
                                 .filter_map(|terms| {
-                                    terms.iter().find(|term| term.col() == col).map(
-                                        |term| {
-                                            term.lits().into_iter().map(|lit| lit.value())
-                                        },
-                                    )
+                                    terms
+                                        .iter()
+                                        .find(|term| term.col() == col)
+                                        .and_then(|term| term.lit_values().ok())
                                 })
                                 .flatten()
                                 .collect();
 
                             builder = builder.aggregate_multi_conjunct(
                                 col,
                                 Guarantee::In,
-                                literals.into_iter(),
+                                literals,
                             );
                         }
 
@@ -296,7 +302,7 @@ impl<'a> GuaranteeBuilder<'a> {
         self.aggregate_multi_conjunct(
             col_op_lit.col,
             col_op_lit.guarantee,
-            [col_op_lit.lit.value()],
+            [col_op_lit.lit.value().clone()],
         )
     }
 
@@ -313,7 +319,7 @@ impl<'a> GuaranteeBuilder<'a> {
         mut self,
         col: &'a crate::expressions::Column,
         guarantee: Guarantee,
-        new_values: impl IntoIterator<Item = &'a ScalarValue>,
+        new_values: impl IntoIterator<Item = ScalarValue>,
     ) -> Self {
         let key = (col, guarantee);
         if let Some(index) = self.map.get(&key) {
@@ -336,20 +342,19 @@ impl<'a> GuaranteeBuilder<'a> {
                 // another `AND a != 6` we know that a must not be either 5 or 6
                 // for the expression to be true
                 Guarantee::NotIn => {
-                    let new_values: HashSet<_> = new_values.into_iter().collect();
-                    existing.literals.extend(new_values.into_iter().cloned());
+                    existing.literals.extend(new_values);
                 }
                 Guarantee::In => {
                     let intersection = new_values
                         .into_iter()
-                        .filter(|new_value| existing.literals.contains(*new_value))
+                        .filter(|new_value| existing.literals.contains(new_value))
                         .collect::<Vec<_>>();
                     // for an In guarantee, if the intersection is not empty,  we can extend the guarantee
                     // e.g. `a IN (1,2,3) AND a IN (2,3,4)` is `a IN (2,3)`
                     // otherwise, we invalidate the guarantee
                     // e.g. `a IN (1,2,3) AND a IN (4,5,6)` is `a IN ()`, which is invalid
                     if !intersection.is_empty() {
-                        existing.literals = intersection.into_iter().cloned().collect();
+                        existing.literals = intersection.into_iter().collect();
                     } else {
                         // at least one was not, so invalidate the guarantee
                         *entry = None;
@@ -436,7 +441,7 @@ impl<'a> ColOpLit<'a> {
 struct ColInList<'a> {
     col: &'a crate::expressions::Column,
     guarantee: Guarantee,
-    list: Vec<&'a crate::expressions::Literal>,
+    inlist: &'a crate::expressions::InListExpr,
 }
 
 impl<'a> ColInList<'a> {
@@ -452,11 +457,12 @@ impl<'a> ColInList<'a> {
             .as_any()
             .downcast_ref::<crate::expressions::Column>()?;
 
-        let literals = inlist
-            .list()
-            .iter()
-            .map(|e| e.as_any().downcast_ref::<crate::expressions::Literal>())
-            .collect::<Option<Vec<_>>>()?;
+        let list = inlist.list().ok()?;
+        // Verify all items are literals
+        for expr in &list {
+            expr.as_any()
+                .downcast_ref::<crate::expressions::Literal>()?;
+        }
 
         let guarantee = if inlist.negated() {
             Guarantee::NotIn
@@ -467,7 +473,7 @@ impl<'a> ColInList<'a> {
         Some(Self {
             col,
             guarantee,
-            list: literals,
+            inlist,
         })
     }
 }
@@ -503,10 +509,20 @@ impl<'a> ColOpLitOrInList<'a> {
         }
     }
 
-    fn lits(&self) -> Vec<&'a crate::expressions::Literal> {
+    fn lit_values(&self) -> datafusion_common::Result<Vec<ScalarValue>> {
         match self {
-            Self::ColOpLit(col_op_lit) => vec![col_op_lit.lit],
-            Self::ColInList(col_in_list) => col_in_list.list.clone(),
+            Self::ColOpLit(col_op_lit) => Ok(vec![col_op_lit.lit.value().clone()]),
+            Self::ColInList(col_in_list) => {
+                let list = col_in_list.inlist.list()?;
+                Ok(list
+                    .iter()
+                    .filter_map(|expr| {
+                        expr.as_any()
+                            .downcast_ref::<crate::expressions::Literal>()
+                            .map(|lit| lit.value().clone())
+                    })
+                    .collect())
+            }
         }
     }
 }
@@ -1088,7 +1104,7 @@ mod test {
         S: Into<ScalarValue> + 'a,
     {
         let literals: Vec<_> = literals.into_iter().map(|s| s.into()).collect();
-        LiteralGuarantee::new(column, Guarantee::In, literals.iter())
+        LiteralGuarantee::new(column, Guarantee::In, literals)
     }
 
     /// Guarantee that the expression is true if the column is NOT any of the specified values
@@ -1098,7 +1114,7 @@ mod test {
         S: Into<ScalarValue> + 'a,
     {
         let literals: Vec<_> = literals.into_iter().map(|s| s.into()).collect();
-        LiteralGuarantee::new(column, Guarantee::NotIn, literals.iter())
+        LiteralGuarantee::new(column, Guarantee::NotIn, literals)
     }
 
     // Schema for testing
diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs
@@ -315,7 +315,10 @@ pub fn serialize_physical_expr(
             expr_type: Some(protobuf::physical_expr_node::ExprType::InList(Box::new(
                 protobuf::PhysicalInListNode {
                     expr: Some(Box::new(serialize_physical_expr(expr.expr(), codec)?)),
-                    list: serialize_physical_exprs(expr.list(), codec)?,
+                    // TODO: serialize the inner ArrayRef directly to avoid materialization into literals
+                    // by extending the protobuf definition to support both representations and adding a public
+                    // accessor method to InListExpr to get the inner ArrayRef
+                    list: serialize_physical_exprs(&expr.list()?, codec)?,
                     negated: expr.negated(),
                 },
             ))),
diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs
@@ -1424,9 +1424,11 @@ fn build_predicate_expression(
         }
     }
     if let Some(in_list) = expr_any.downcast_ref::<phys_expr::InListExpr>() {
-        if !in_list.list().is_empty()
-            && in_list.list().len() <= MAX_LIST_VALUE_SIZE_REWRITE
-        {
+        if !in_list.is_empty() && in_list.len() <= MAX_LIST_VALUE_SIZE_REWRITE {
+            let list = match in_list.list() {
+                Ok(list) => list,
+                Err(_) => return unhandled_hook.handle(expr),
+            };
             let eq_op = if in_list.negated() {
                 Operator::NotEq
             } else {
@@ -1437,8 +1439,7 @@ fn build_predicate_expression(
             } else {
                 Operator::Or
             };
-            let change_expr = in_list
-                .list()
+            let change_expr = list
                 .iter()
                 .map(|e| {
                     Arc::new(phys_expr::BinaryExpr::new(
diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt
@@ -6408,7 +6408,7 @@ physical_plan
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
 05)--------ProjectionExec: expr=[]
 06)----------CoalesceBatchesExec: target_batch_size=8192
-07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
+07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 09)----------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
@@ -6437,7 +6437,7 @@ physical_plan
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
 05)--------ProjectionExec: expr=[]
 06)----------CoalesceBatchesExec: target_batch_size=8192
-07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
+07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 09)----------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
@@ -6466,7 +6466,7 @@ physical_plan
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
 05)--------ProjectionExec: expr=[]
 06)----------CoalesceBatchesExec: target_batch_size=8192
-07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
+07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 09)----------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
@@ -6495,7 +6495,7 @@ physical_plan
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
 05)--------ProjectionExec: expr=[]
 06)----------CoalesceBatchesExec: target_batch_size=8192
-07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
+07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 09)----------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
@@ -6524,7 +6524,7 @@ physical_plan
 04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))]
 05)--------ProjectionExec: expr=[]
 06)----------CoalesceBatchesExec: target_batch_size=8192
-07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
+07)------------FilterExec: substr(md5(CAST(value@0 AS Utf8View)), 1, 32) IN (SET) ([7f4b18de3cfeb9b4ac78c381ee2ad278, a, b, c])
 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 09)----------------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=8192]
 
diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q19.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q19.slt.part
@@ -69,7 +69,7 @@ physical_plan
 03)----CoalescePartitionsExec
 04)------AggregateExec: mode=Partial, gby=[], aggr=[sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]
 05)--------CoalesceBatchesExec: target_batch_size=8192
-06)----------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_partkey@0, p_partkey@0)], filter=p_brand@1 = Brand#12 AND p_container@3 IN ([SM CASE, SM BOX, SM PACK, SM PKG]) AND l_quantity@0 >= Some(100),15,2 AND l_quantity@0 <= Some(1100),15,2 AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND p_container@3 IN ([MED BAG, MED BOX, MED PKG, MED PACK]) AND l_quantity@0 >= Some(1000),15,2 AND l_quantity@0 <= Some(2000),15,2 AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND p_container@3 IN ([LG CASE, LG BOX, LG PACK, LG PKG]) AND l_quantity@0 >= Some(2000),15,2 AND l_quantity@0 <= Some(3000),15,2 AND p_size@2 <= 15, projection=[l_extendedprice@2, l_discount@3]
+06)----------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_partkey@0, p_partkey@0)], filter=p_brand@1 = Brand#12 AND p_container@3 IN (SET) ([SM CASE, SM BOX, SM PACK, SM PKG]) AND l_quantity@0 >= Some(100),15,2 AND l_quantity@0 <= Some(1100),15,2 AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND p_container@3 IN (SET) ([MED BAG, MED BOX, MED PKG, MED PACK]) AND l_quantity@0 >= Some(1000),15,2 AND l_quantity@0 <= Some(2000),15,2 AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND p_container@3 IN (SET) ([LG CASE, LG BOX, LG PACK, LG PKG]) AND l_quantity@0 >= Some(2000),15,2 AND l_quantity@0 <= Some(3000),15,2 AND p_size@2 <= 15, projection=[l_extendedprice@2, l_discount@3]
 07)------------CoalesceBatchesExec: target_batch_size=8192
 08)--------------RepartitionExec: partitioning=Hash([l_partkey@0], 4), input_partitions=4
 09)----------------CoalesceBatchesExec: target_batch_size=8192
@@ -78,6 +78,6 @@ physical_plan
 12)------------CoalesceBatchesExec: target_batch_size=8192
 13)--------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4
 14)----------------CoalesceBatchesExec: target_batch_size=8192
-15)------------------FilterExec: (p_brand@1 = Brand#12 AND p_container@3 IN ([SM CASE, SM BOX, SM PACK, SM PKG]) AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND p_container@3 IN ([MED BAG, MED BOX, MED PKG, MED PACK]) AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND p_container@3 IN ([LG CASE, LG BOX, LG PACK, LG PKG]) AND p_size@2 <= 15) AND p_size@2 >= 1
+15)------------------FilterExec: (p_brand@1 = Brand#12 AND p_container@3 IN (SET) ([SM CASE, SM BOX, SM PACK, SM PKG]) AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND p_container@3 IN (SET) ([MED BAG, MED BOX, MED PKG, MED PACK]) AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND p_container@3 IN (SET) ([LG CASE, LG BOX, LG PACK, LG PKG]) AND p_size@2 <= 15) AND p_size@2 >= 1
 16)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 17)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_brand, p_size, p_container], file_type=csv, has_header=false
diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part
@@ -91,7 +91,7 @@ physical_plan
 15)----------------------------CoalesceBatchesExec: target_batch_size=8192
 16)------------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4
 17)--------------------------------CoalesceBatchesExec: target_batch_size=8192
-18)----------------------------------FilterExec: substr(c_phone@1, 1, 2) IN ([13, 31, 23, 29, 30, 18, 17])
+18)----------------------------------FilterExec: substr(c_phone@1, 1, 2) IN (SET) ([13, 31, 23, 29, 30, 18, 17])
 19)------------------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 20)--------------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_phone, c_acctbal], file_type=csv, has_header=false
 21)----------------------------CoalesceBatchesExec: target_batch_size=8192
@@ -101,6 +101,6 @@ physical_plan
 25)----------------------CoalescePartitionsExec
 26)------------------------AggregateExec: mode=Partial, gby=[], aggr=[avg(customer.c_acctbal)]
 27)--------------------------CoalesceBatchesExec: target_batch_size=8192
-28)----------------------------FilterExec: c_acctbal@1 > Some(0),15,2 AND substr(c_phone@0, 1, 2) IN ([13, 31, 23, 29, 30, 18, 17]), projection=[c_acctbal@1]
+28)----------------------------FilterExec: c_acctbal@1 > Some(0),15,2 AND substr(c_phone@0, 1, 2) IN (SET) ([13, 31, 23, 29, 30, 18, 17]), projection=[c_acctbal@1]
 29)------------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
 30)--------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_phone, c_acctbal], file_type=csv, has_header=false
