Merge branch 'main' into fix-approx-percentile-cont-with-weight

Author: sesteves

.github/workflows/docs.yaml

@@ -41,7 +41,7 @@ jobs:
           path: asf-site
 
       - name: Setup Python
-        uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548  # v6.1.0
+        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6.2.0
         with:
           python-version: "3.12"
 

.github/workflows/docs_pr.yaml

@@ -45,7 +45,7 @@ jobs:
           submodules: true
           fetch-depth: 1
       - name: Setup Python
-        uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548  # v6.1.0
+        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6.2.0
         with:
           python-version: "3.12"
       - name: Install doc dependencies

Cargo.lock

@@ -108,7 +108,7 @@ async-trait = "0.1.89"
 bigdecimal = "0.4.8"
 bytes = "1.11"
 bzip2 = "0.6.1"
-chrono = { version = "0.4.42", default-features = false }
+chrono = { version = "0.4.43", default-features = false }
 criterion = "0.8"
 ctor = "0.6.3"
 dashmap = "6.0.1"

Cargo.toml

@@ -1843,9 +1843,14 @@ impl ContextProvider for SessionContextProvider<'_> {
                 self.state.execution_props().query_execution_start_time,
             );
         let simplifier = ExprSimplifier::new(simplify_context);
+        let schema = DFSchema::empty();
         let args = args
             .into_iter()
-            .map(|arg| simplifier.simplify(arg))
+            .map(|arg| {
+                simplifier
+                    .coerce(arg, &schema)
+                    .and_then(|e| simplifier.simplify(e))
+            })
             .collect::<datafusion_common::Result<Vec<_>>>()?;
         let provider = tbl_func.create_table_provider(&args)?;
 

datafusion/core/src/execution/session_state.rs

@@ -221,6 +221,31 @@ impl TableFunctionImpl for SimpleCsvTableFunc {
     }
 }
 
+/// Test that expressions passed to UDTFs are properly type-coerced
+/// This is a regression test for https://github.com/apache/datafusion/issues/19914
+#[tokio::test]
+async fn test_udtf_type_coercion() -> Result<()> {
+    use datafusion::datasource::MemTable;
+
+    #[derive(Debug)]
+    struct NoOpTableFunc;
+
+    impl TableFunctionImpl for NoOpTableFunc {
+        fn call(&self, _: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+            let schema = Arc::new(arrow::datatypes::Schema::empty());
+            Ok(Arc::new(MemTable::try_new(schema, vec![vec![]])?))
+        }
+    }
+
+    let ctx = SessionContext::new();
+    ctx.register_udtf("f", Arc::new(NoOpTableFunc));
+
+    // This should not panic - the array elements should be coerced to Float64
+    let _ = ctx.sql("SELECT * FROM f(ARRAY[0.1, 1, 2])").await?;
+
+    Ok(())
+}
+
 fn read_csv_batches(csv_path: impl AsRef<Path>) -> Result<(SchemaRef, Vec<RecordBatch>)> {
     let mut file = File::open(csv_path)?;
     let (schema, _) = Format::default()

datafusion/core/tests/user_defined/user_defined_table_functions.rs

@@ -2779,12 +2779,12 @@ mod tests {
 
         assert_snapshot!(plan, @r"
         Union
-          Cross Join: 
+          Cross Join:
             SubqueryAlias: left
               Values: (Int32(1))
             SubqueryAlias: right
               Values: (Int32(1))
-          Cross Join: 
+          Cross Join:
             SubqueryAlias: left
               Values: (Int32(1))
             SubqueryAlias: right

datafusion/expr/src/logical_plan/builder.rs

@@ -1973,13 +1973,16 @@ impl LogicalPlan {
                         };
                         match join_constraint {
                             JoinConstraint::On => {
-                                write!(
-                                    f,
-                                    "{} Join: {}{}",
-                                    join_type,
-                                    join_expr.join(", "),
-                                    filter_expr
-                                )
+                                write!(f, "{join_type} Join:",)?;
+                                if !join_expr.is_empty() || !filter_expr.is_empty() {
+                                    write!(
+                                        f,
+                                        " {}{}",
+                                        join_expr.join(", "),
+                                        filter_expr
+                                    )?;
+                                }
+                                Ok(())
                             }
                             JoinConstraint::Using => {
                                 write!(

datafusion/expr/src/logical_plan/plan.rs

@@ -23,6 +23,7 @@ use arrow::{
     util::bench_util::create_primitive_array,
 };
 use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion_common::ScalarValue;
 use datafusion_common::config::ConfigOptions;
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
 use datafusion_functions::math::iszero;
@@ -31,6 +32,8 @@ use std::sync::Arc;
 
 fn criterion_benchmark(c: &mut Criterion) {
     let iszero = iszero();
+    let config_options = Arc::new(ConfigOptions::default());
+
     for size in [1024, 4096, 8192] {
         let f32_array = Arc::new(create_primitive_array::<Float32Type>(size, 0.2));
         let batch_len = f32_array.len();
@@ -43,7 +46,6 @@ fn criterion_benchmark(c: &mut Criterion) {
             })
             .collect::<Vec<_>>();
         let return_field = Arc::new(Field::new("f", DataType::Boolean, true));
-        let config_options = Arc::new(ConfigOptions::default());
 
         c.bench_function(&format!("iszero f32 array: {size}"), |b| {
             b.iter(|| {
@@ -60,6 +62,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                 )
             })
         });
+
         let f64_array = Arc::new(create_primitive_array::<Float64Type>(size, 0.2));
         let batch_len = f64_array.len();
         let f64_args = vec![ColumnarValue::Array(f64_array)];
@@ -88,6 +91,46 @@ fn criterion_benchmark(c: &mut Criterion) {
             })
         });
     }
+
+    // Scalar benchmarks - run once since size doesn't affect scalar performance
+    let scalar_f32_args = vec![ColumnarValue::Scalar(ScalarValue::Float32(Some(1.0)))];
+    let scalar_f32_arg_fields = vec![Field::new("a", DataType::Float32, false).into()];
+    let return_field_scalar = Arc::new(Field::new("f", DataType::Boolean, false));
+
+    c.bench_function("iszero f32 scalar", |b| {
+        b.iter(|| {
+            black_box(
+                iszero
+                    .invoke_with_args(ScalarFunctionArgs {
+                        args: scalar_f32_args.clone(),
+                        arg_fields: scalar_f32_arg_fields.clone(),
+                        number_rows: 1,
+                        return_field: Arc::clone(&return_field_scalar),
+                        config_options: Arc::clone(&config_options),
+                    })
+                    .unwrap(),
+            )
+        })
+    });
+
+    let scalar_f64_args = vec![ColumnarValue::Scalar(ScalarValue::Float64(Some(1.0)))];
+    let scalar_f64_arg_fields = vec![Field::new("a", DataType::Float64, false).into()];
+
+    c.bench_function("iszero f64 scalar", |b| {
+        b.iter(|| {
+            black_box(
+                iszero
+                    .invoke_with_args(ScalarFunctionArgs {
+                        args: scalar_f64_args.clone(),
+                        arg_fields: scalar_f64_arg_fields.clone(),
+                        number_rows: 1,
+                        return_field: Arc::clone(&return_field_scalar),
+                        config_options: Arc::clone(&config_options),
+                    })
+                    .unwrap(),
+            )
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

datafusion/functions/benches/iszero.rs

@@ -18,21 +18,20 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use arrow::array::{ArrayRef, ArrowNativeTypeOp, AsArray, BooleanArray};
+use arrow::array::{ArrowNativeTypeOp, AsArray, BooleanArray};
 use arrow::datatypes::DataType::{Boolean, Float16, Float32, Float64};
 use arrow::datatypes::{DataType, Float16Type, Float32Type, Float64Type};
 
 use datafusion_common::types::NativeType;
-use datafusion_common::{Result, ScalarValue, exec_err};
+use datafusion_common::utils::take_function_args;
+use datafusion_common::{Result, ScalarValue, internal_err};
 use datafusion_expr::{Coercion, TypeSignatureClass};
 use datafusion_expr::{
     ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
     Volatility,
 };
 use datafusion_macros::user_doc;
 
-use crate::utils::make_scalar_function;
-
 #[user_doc(
     doc_section(label = "Math Functions"),
     description = "Returns true if a given number is +0.0 or -0.0 otherwise returns false.",
@@ -90,79 +89,53 @@ impl ScalarUDFImpl for IsZeroFunc {
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
-        // Handle NULL input
-        if args.args[0].data_type().is_null() {
-            return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
+        let [arg] = take_function_args(self.name(), args.args)?;
+
+        match arg {
+            ColumnarValue::Scalar(scalar) => {
+                if scalar.is_null() {
+                    return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
+                }
+
+                match scalar {
+                    ScalarValue::Float64(Some(v)) => {
+                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0.0))))
+                    }
+                    ScalarValue::Float32(Some(v)) => {
+                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0.0))))
+                    }
+                    ScalarValue::Float16(Some(v)) => Ok(ColumnarValue::Scalar(
+                        ScalarValue::Boolean(Some(v.is_zero())),
+                    )),
+                    _ => {
+                        internal_err!(
+                            "Unexpected scalar type for iszero: {:?}",
+                            scalar.data_type()
+                        )
+                    }
+                }
+            }
+            ColumnarValue::Array(array) => match array.data_type() {
+                Float64 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
+                    array.as_primitive::<Float64Type>(),
+                    |x| x == 0.0,
+                )))),
+                Float32 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
+                    array.as_primitive::<Float32Type>(),
+                    |x| x == 0.0,
+                )))),
+                Float16 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
+                    array.as_primitive::<Float16Type>(),
+                    |x| x.is_zero(),
+                )))),
+                other => {
+                    internal_err!("Unexpected data type {other:?} for function iszero")
+                }
+            },
         }
-        make_scalar_function(iszero, vec![])(&args.args)
     }
 
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
 }
-
-/// Iszero SQL function
-fn iszero(args: &[ArrayRef]) -> Result<ArrayRef> {
-    match args[0].data_type() {
-        Float64 => Ok(Arc::new(BooleanArray::from_unary(
-            args[0].as_primitive::<Float64Type>(),
-            |x| x == 0.0,
-        )) as ArrayRef),
-
-        Float32 => Ok(Arc::new(BooleanArray::from_unary(
-            args[0].as_primitive::<Float32Type>(),
-            |x| x == 0.0,
-        )) as ArrayRef),
-
-        Float16 => Ok(Arc::new(BooleanArray::from_unary(
-            args[0].as_primitive::<Float16Type>(),
-            |x| x.is_zero(),
-        )) as ArrayRef),
-
-        other => exec_err!("Unsupported data type {other:?} for function iszero"),
-    }
-}
-
-#[cfg(test)]
-mod test {
-    use std::sync::Arc;
-
-    use arrow::array::{ArrayRef, Float32Array, Float64Array};
-
-    use datafusion_common::cast::as_boolean_array;
-
-    use crate::math::iszero::iszero;
-
-    #[test]
-    fn test_iszero_f64() {
-        let args: Vec<ArrayRef> =
-            vec![Arc::new(Float64Array::from(vec![1.0, 0.0, 3.0, -0.0]))];
-
-        let result = iszero(&args).expect("failed to initialize function iszero");
-        let booleans =
-            as_boolean_array(&result).expect("failed to initialize function iszero");
-
-        assert_eq!(booleans.len(), 4);
-        assert!(!booleans.value(0));
-        assert!(booleans.value(1));
-        assert!(!booleans.value(2));
-        assert!(booleans.value(3));
-    }
-
-    #[test]
-    fn test_iszero_f32() {
-        let args: Vec<ArrayRef> =
-            vec![Arc::new(Float32Array::from(vec![1.0, 0.0, 3.0, -0.0]))];
-
-        let result = iszero(&args).expect("failed to initialize function iszero");
-        let booleans =
-            as_boolean_array(&result).expect("failed to initialize function iszero");
-
-        assert_eq!(booleans.len(), 4);
-        assert!(!booleans.value(0));
-        assert!(booleans.value(1));
-        assert!(!booleans.value(2));
-        assert!(booleans.value(3));
-    }
-}