Add constant expression evaluator to physical expression simplifier (#19130)

This improves handling of constant expressions during pruning by trying
to evaluate them in the simplifier and the pruning machinery. This is
somewhat redundant with #19129 in the simple case of our Parquet
implementation but since there may be edge cases where one is hit and
not the other, or where users are using them independently I thought it
best to implement both approaches.

Author: adriangb

datafusion/core/tests/schema_adapter/schema_adapter_integration_tests.rs

@@ -564,29 +564,36 @@ async fn test_parquet_missing_column() -> Result<()> {
         .push_down_filters(false)
         .execute()
         .await?;
-    // There will be data: the filter is (null) is not null or a = 24.
-    // Statistics pruning doesn't handle `null is not null` so it resolves to `true or a = 24` -> `true` so no row groups are pruned
+    // There should be zero batches
+    assert_eq!(batches.len(), 0);
+    // Check another filter: `b = 'foo' and a = 24` should also prune data with only statistics-based pushdown
+    let filter = col("b").eq(lit("foo")).and(col("a").eq(lit(24)));
+    let batches = test_case
+        .clone()
+        .with_predicate(filter)
+        .push_down_filters(false)
+        .execute()
+        .await?;
+    // There should be zero batches
+    assert_eq!(batches.len(), 0);
+    // On the other hand `b is null and a = 2` should prune only the second row group with stats only pruning
+    let filter = col("b").is_null().and(col("a").eq(lit(2)));
+    let batches = test_case
+        .clone()
+        .with_predicate(filter)
+        .push_down_filters(false)
+        .execute()
+        .await?;
     #[rustfmt::skip]
     let expected = [
         "+---+---+-----+",
         "| a | b | c   |",
         "+---+---+-----+",
         "| 1 |   | 1.1 |",
         "| 2 |   | 2.2 |",
-        "| 3 |   | 3.3 |",
         "+---+---+-----+",
     ];
     assert_batches_eq!(expected, &batches);
-    // On the other hand the filter `b = 'foo' and a = 24` should prune all data even with only statistics-based pushdown
-    let filter = col("b").eq(lit("foo")).and(col("a").eq(lit(24)));
-    let batches = test_case
-        .clone()
-        .with_predicate(filter)
-        .push_down_filters(false)
-        .execute()
-        .await?;
-    // There should be zero batches
-    assert_eq!(batches.len(), 0);
 
     Ok(())
 }

datafusion/datasource-parquet/src/opener.rs

@@ -1235,7 +1235,7 @@ mod test {
                 force_filter_selections: false,
                 enable_page_index: false,
                 enable_bloom_filter: false,
-                enable_row_group_stats_pruning: true,
+                enable_row_group_stats_pruning: false, // note that this is false!
                 coerce_int96: None,
                 #[cfg(feature = "parquet_encryption")]
                 file_decryption_properties: None,

datafusion/physical-expr/src/simplifier/const_evaluator.rs

@@ -0,0 +1,103 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Constant expression evaluation for the physical expression simplifier
+
+use std::sync::Arc;
+
+use arrow::array::new_null_array;
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+use datafusion_common::{Result, ScalarValue};
+use datafusion_expr_common::columnar_value::ColumnarValue;
+use datafusion_physical_expr_common::physical_expr::is_volatile;
+
+use crate::expressions::{Column, Literal};
+use crate::PhysicalExpr;
+
+/// Simplify expressions that consist only of literals by evaluating them.
+///
+/// This function checks if all children of the given expression are literals.
+/// If so, it evaluates the expression against a dummy RecordBatch and returns
+/// the result as a new Literal.
+///
+/// # Example transformations
+/// - `1 + 2` -> `3`
+/// - `(1 + 2) * 3` -> `9` (with bottom-up traversal)
+/// - `'hello' || ' world'` -> `'hello world'`
+pub fn simplify_const_expr(
+    expr: &Arc<dyn PhysicalExpr>,
+) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+    if is_volatile(expr) || has_column_references(expr) {
+        return Ok(Transformed::no(Arc::clone(expr)));
+    }
+
+    // Create a 1-row dummy batch for evaluation
+    let batch = create_dummy_batch()?;
+
+    // Evaluate the expression
+    match expr.evaluate(&batch) {
+        Ok(ColumnarValue::Scalar(scalar)) => {
+            Ok(Transformed::yes(Arc::new(Literal::new(scalar))))
+        }
+        Ok(ColumnarValue::Array(arr)) if arr.len() == 1 => {
+            // Some operations return an array even for scalar inputs
+            let scalar = ScalarValue::try_from_array(&arr, 0)?;
+            Ok(Transformed::yes(Arc::new(Literal::new(scalar))))
+        }
+        Ok(_) => {
+            // Unexpected result - keep original expression
+            Ok(Transformed::no(Arc::clone(expr)))
+        }
+        Err(_) => {
+            // On error, keep original expression
+            // The expression might succeed at runtime due to short-circuit evaluation
+            // or other runtime conditions
+            Ok(Transformed::no(Arc::clone(expr)))
+        }
+    }
+}
+
+/// Create a 1-row dummy RecordBatch for evaluating constant expressions.
+///
+/// The batch is never actually accessed for data - it's just needed because
+/// the PhysicalExpr::evaluate API requires a RecordBatch. For expressions
+/// that only contain literals, the batch content is irrelevant.
+///
+/// This is the same approach used in the logical expression `ConstEvaluator`.
+fn create_dummy_batch() -> Result<RecordBatch> {
+    // RecordBatch requires at least one column
+    let dummy_schema = Arc::new(Schema::new(vec![Field::new("_", DataType::Null, true)]));
+    let col = new_null_array(&DataType::Null, 1);
+    Ok(RecordBatch::try_new(dummy_schema, vec![col])?)
+}
+
+/// Check if this expression has any column references.
+pub fn has_column_references(expr: &Arc<dyn PhysicalExpr>) -> bool {
+    let mut has_columns = false;
+    expr.apply(|expr| {
+        if expr.as_any().downcast_ref::<Column>().is_some() {
+            has_columns = true;
+            Ok(TreeNodeRecursion::Stop)
+        } else {
+            Ok(TreeNodeRecursion::Continue)
+        }
+    })
+    .expect("apply should not fail");
+    has_columns
+}

datafusion/physical-expr/src/simplifier/mod.rs

@@ -26,6 +26,7 @@ use std::sync::Arc;
 
 use crate::{simplifier::not::simplify_not_expr, PhysicalExpr};
 
+pub mod const_evaluator;
 pub mod not;
 pub mod unwrap_cast;
 
@@ -72,11 +73,13 @@ impl<'a> TreeNodeRewriter for PhysicalExprSimplifier<'a> {
         #[cfg(test)]
         let original_type = node.data_type(self.schema).unwrap();
 
-        // Apply NOT expression simplification first, then unwrap cast optimization
-        let rewritten =
-            simplify_not_expr(&node, self.schema)?.transform_data(|node| {
+        // Apply NOT expression simplification first, then unwrap cast optimization,
+        // then constant expression evaluation
+        let rewritten = simplify_not_expr(&node, self.schema)?
+            .transform_data(|node| {
                 unwrap_cast::unwrap_cast_in_comparison(node, self.schema)
-            })?;
+            })?
+            .transform_data(|node| const_evaluator::simplify_const_expr(&node))?;
 
         #[cfg(test)]
         assert_eq!(
@@ -491,4 +494,131 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_simplify_literal_binary_expr() {
+        let schema = Schema::empty();
+        let mut simplifier = PhysicalExprSimplifier::new(&schema);
+
+        // 1 + 2 -> 3
+        let expr: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(lit(1i32), Operator::Plus, lit(2i32)));
+        let result = simplifier.simplify(expr).unwrap();
+        let literal = as_literal(&result);
+        assert_eq!(literal.value(), &ScalarValue::Int32(Some(3)));
+    }
+
+    #[test]
+    fn test_simplify_literal_comparison() {
+        let schema = Schema::empty();
+        let mut simplifier = PhysicalExprSimplifier::new(&schema);
+
+        // 5 > 3 -> true
+        let expr: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(lit(5i32), Operator::Gt, lit(3i32)));
+        let result = simplifier.simplify(expr).unwrap();
+        let literal = as_literal(&result);
+        assert_eq!(literal.value(), &ScalarValue::Boolean(Some(true)));
+
+        // 2 > 3 -> false
+        let expr: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(lit(2i32), Operator::Gt, lit(3i32)));
+        let result = simplifier.simplify(expr).unwrap();
+        let literal = as_literal(&result);
+        assert_eq!(literal.value(), &ScalarValue::Boolean(Some(false)));
+    }
+
+    #[test]
+    fn test_simplify_nested_literal_expr() {
+        let schema = Schema::empty();
+        let mut simplifier = PhysicalExprSimplifier::new(&schema);
+
+        // (1 + 2) * 3 -> 9
+        let inner: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(lit(1i32), Operator::Plus, lit(2i32)));
+        let expr: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(inner, Operator::Multiply, lit(3i32)));
+        let result = simplifier.simplify(expr).unwrap();
+        let literal = as_literal(&result);
+        assert_eq!(literal.value(), &ScalarValue::Int32(Some(9)));
+    }
+
+    #[test]
+    fn test_simplify_deeply_nested_literals() {
+        let schema = Schema::empty();
+        let mut simplifier = PhysicalExprSimplifier::new(&schema);
+
+        // ((1 + 2) * 3) + ((4 - 1) * 2) -> 9 + 6 -> 15
+        let left: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+            Arc::new(BinaryExpr::new(lit(1i32), Operator::Plus, lit(2i32))),
+            Operator::Multiply,
+            lit(3i32),
+        ));
+        let right: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+            Arc::new(BinaryExpr::new(lit(4i32), Operator::Minus, lit(1i32))),
+            Operator::Multiply,
+            lit(2i32),
+        ));
+        let expr: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(left, Operator::Plus, right));
+        let result = simplifier.simplify(expr).unwrap();
+        let literal = as_literal(&result);
+        assert_eq!(literal.value(), &ScalarValue::Int32(Some(15)));
+    }
+
+    #[test]
+    fn test_no_simplify_with_column() {
+        let schema = test_schema();
+        let mut simplifier = PhysicalExprSimplifier::new(&schema);
+
+        // c1 + 2 should NOT be simplified (has column reference)
+        let expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+            col("c1", &schema).unwrap(),
+            Operator::Plus,
+            lit(2i32),
+        ));
+        let result = simplifier.simplify(expr).unwrap();
+        // Should remain a BinaryExpr, not become a Literal
+        assert!(result.as_any().downcast_ref::<BinaryExpr>().is_some());
+    }
+
+    #[test]
+    fn test_partial_simplify_with_column() {
+        let schema = test_schema();
+        let mut simplifier = PhysicalExprSimplifier::new(&schema);
+
+        // (1 + 2) + c1 should simplify the literal part: 3 + c1
+        let literal_part: Arc<dyn PhysicalExpr> =
+            Arc::new(BinaryExpr::new(lit(1i32), Operator::Plus, lit(2i32)));
+        let expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+            literal_part,
+            Operator::Plus,
+            col("c1", &schema).unwrap(),
+        ));
+        let result = simplifier.simplify(expr).unwrap();
+
+        // Should be a BinaryExpr with a Literal(3) on the left
+        let binary = as_binary(&result);
+        let left_literal = as_literal(binary.left());
+        assert_eq!(left_literal.value(), &ScalarValue::Int32(Some(3)));
+    }
+
+    #[test]
+    fn test_simplify_literal_string_concat() {
+        let schema = Schema::empty();
+        let mut simplifier = PhysicalExprSimplifier::new(&schema);
+
+        // 'hello' || ' world' -> 'hello world'
+        let expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
+            lit("hello"),
+            Operator::StringConcat,
+            lit(" world"),
+        ));
+        let result = simplifier.simplify(expr).unwrap();
+        let literal = as_literal(&result);
+        assert_eq!(
+            literal.value(),
+            &ScalarValue::Utf8(Some("hello world".to_string()))
+        );
+    }
 }