From cd3c5f3915a5d185d92639bb1c6bc294bdd2e09d Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Mon, 25 Dec 2023 15:45:31 +0300 Subject: [PATCH 1/5] Initial commit --- .../optimizer/src/eliminate_cross_join.rs | 1 + datafusion/optimizer/src/push_down_filter.rs | 10 ++- datafusion/sqllogictest/src/test_context.rs | 63 ++++++++++++++++++- datafusion/sqllogictest/test_files/joins.slt | 21 +++++++ 4 files changed, 92 insertions(+), 3 deletions(-) diff --git a/datafusion/optimizer/src/eliminate_cross_join.rs b/datafusion/optimizer/src/eliminate_cross_join.rs index 7c866950a622e..d9e96a9f2543a 100644 --- a/datafusion/optimizer/src/eliminate_cross_join.rs +++ b/datafusion/optimizer/src/eliminate_cross_join.rs @@ -45,6 +45,7 @@ impl EliminateCrossJoin { /// 'select ... from a, b where (a.x = b.y and b.xx = 100) or (a.x = b.y and b.xx = 200);' /// 'select ... from a, b, c where (a.x = b.y and b.xx = 100 and a.z = c.z) /// or (a.x = b.y and b.xx = 200 and a.z=c.z);' +/// 'select ... from a, b where a.x > b.y' /// For above queries, the join predicate is available in filters and they are moved to /// join nodes appropriately /// This fix helps to improve the performance of TPCH Q19. issue#78 diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 4eed39a089415..de4d3fadd1359 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -965,11 +965,10 @@ impl PushDownFilter { } } -/// Convert cross join to join by pushing down filter predicate to the join condition +/// Converts cross join to the inner join with empty equality predicate and empty filter condition. fn convert_cross_join_to_inner_join(cross_join: CrossJoin) -> Result { let CrossJoin { left, right, .. } = cross_join; let join_schema = build_join_schema(left.schema(), right.schema(), &JoinType::Inner)?; - // predicate is given Ok(Join { left, right, @@ -991,6 +990,13 @@ fn convert_to_cross_join_if_beneficial(plan: LogicalPlan) -> Result .cross_join(join.right.as_ref().clone())? .build(); } + } else if let LogicalPlan::Filter(filter) = &plan { + let new_input = + convert_to_cross_join_if_beneficial(filter.input.as_ref().clone())?; + return Ok(LogicalPlan::Filter(Filter::try_new( + filter.predicate.clone(), + Arc::new(new_input), + )?)); } Ok(plan) } diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 941dcb69d2f4d..a16ef037c08d7 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -15,9 +15,11 @@ // specific language governing permissions and limitations // under the License. +use arrow::array::{ArrayRef, Int64Array}; use async_trait::async_trait; use datafusion::execution::context::SessionState; -use datafusion::logical_expr::Expr; +use datafusion::logical_expr::{create_udf, Expr, ScalarUDF, Volatility}; +use datafusion::physical_expr::functions::make_scalar_function; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion::{ @@ -33,6 +35,7 @@ use datafusion::{ datasource::{MemTable, TableProvider, TableType}, prelude::{CsvReadOptions, SessionContext}, }; +use datafusion_common::cast::{as_float64_array, as_int32_array}; use datafusion_common::DataFusionError; use log::info; use std::collections::HashMap; @@ -102,6 +105,8 @@ impl TestContext { } "joins.slt" => { info!("Registering partition table tables"); + let twice = create_twice_udf(); + test_ctx.ctx.register_udf(twice); register_partition_table(&mut test_ctx).await; } "metadata.slt" => { @@ -348,3 +353,59 @@ pub async fn register_metadata_tables(ctx: &SessionContext) { ctx.register_batch("table_with_metadata", batch).unwrap(); } + +/// Create a UDF function named "TWICE" +fn create_twice_udf() -> ScalarUDF { + // First, declare the actual implementation of the calculation + let pow = |args: &[ArrayRef]| { + // in DataFusion, all `args` and output are dynamically-typed arrays, which means that we need to: + // 1. cast the values to the type we want + // 2. perform the computation for every element in the array (using a loop or SIMD) and construct the result + + // this is guaranteed by DataFusion based on the function's signature. + assert_eq!(args.len(), 2); + + // 1. cast both arguments to f64. These casts MUST be aligned with the signature or this function panics! + let base = as_float64_array(&args[0]).expect("cast failed"); + let exponent = as_float64_array(&args[1]).expect("cast failed"); + + // this is guaranteed by DataFusion. We place it just to make it obvious. + assert_eq!(exponent.len(), base.len()); + + // 2. perform the computation + let array = base + .iter() + .zip(exponent.iter()) + .map(|(base, exponent)| { + match (base, exponent) { + // in arrow, any value can be null. + // Here we decide to make our UDF to return null when either base or exponent is null. + (Some(base), Some(exponent)) => Some(base.powf(exponent)), + _ => None, + } + }) + .collect::(); + + // `Ok` because no error occurred during the calculation (we should add one if exponent was [0, 1[ and the base < 0 because that panics!) + // `Arc` because arrays are immutable, thread-safe, trait objects. + Ok(Arc::new(array) as ArrayRef) + }; + // the function above expects an `ArrayRef`, but DataFusion may pass a scalar to a UDF. + // thus, we use `make_scalar_function` to decorare the closure so that it can handle both Arrays and Scalar values. + let pow = make_scalar_function(pow); + + // Next: + // * give it a name so that it shows nicely when the plan is printed + // * declare what input it expects + // * declare its return type + let pow = create_udf( + "pow_udf", + // expects two f64 + vec![DataType::Float64, DataType::Float64], + // returns f64 + Arc::new(DataType::Float64), + Volatility::Immutable, + pow, + ); + pow +} diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index eee213811f443..76bbedbda100a 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -3483,6 +3483,27 @@ NestedLoopJoinExec: join_type=Inner, filter=a@0 > a@1 ----CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true --CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true +# Filter conditions are not pushed down through UDFs. +query TT +EXPLAIN SELECT * +FROM annotated_data as t1, annotated_data as t2 +WHERE POW_UDF(t1.a, 2.0) > POW_UDF(t2.a, 2.0) +---- +logical_plan +Filter: pow_udf(CAST(t1.a AS Float64), Float64(2)) > pow_udf(CAST(t2.a AS Float64), Float64(2)) +--CrossJoin: +----SubqueryAlias: t1 +------TableScan: annotated_data projection=[a0, a, b, c, d] +----SubqueryAlias: t2 +------TableScan: annotated_data projection=[a0, a, b, c, d] +physical_plan +CoalesceBatchesExec: target_batch_size=2 +--FilterExec: pow_udf(CAST(a@1 AS Float64), 2) > pow_udf(CAST(a@6 AS Float64), 2) +----CrossJoinExec +------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true +------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 +--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true + #### # Config teardown #### From 7d75198df5e227af83ed5ccd5d58bf45fd727390 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Mon, 25 Dec 2023 15:46:28 +0300 Subject: [PATCH 2/5] Minor changes --- datafusion/sqllogictest/src/test_context.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index a16ef037c08d7..1d8d6892c7899 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -105,7 +105,7 @@ impl TestContext { } "joins.slt" => { info!("Registering partition table tables"); - let twice = create_twice_udf(); + let twice = create_pow_udf(); test_ctx.ctx.register_udf(twice); register_partition_table(&mut test_ctx).await; } @@ -354,8 +354,8 @@ pub async fn register_metadata_tables(ctx: &SessionContext) { ctx.register_batch("table_with_metadata", batch).unwrap(); } -/// Create a UDF function named "TWICE" -fn create_twice_udf() -> ScalarUDF { +/// Create a UDF function named "pow_udf" +fn create_pow_udf() -> ScalarUDF { // First, declare the actual implementation of the calculation let pow = |args: &[ArrayRef]| { // in DataFusion, all `args` and output are dynamically-typed arrays, which means that we need to: From 463f81684f7636e7935ffe026921ec872d2bec19 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Mon, 25 Dec 2023 15:56:44 +0300 Subject: [PATCH 3/5] Simplifications --- datafusion/sqllogictest/src/test_context.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 1d8d6892c7899..bafbb674f81dc 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, Int64Array}; +use arrow::array::ArrayRef; use async_trait::async_trait; use datafusion::execution::context::SessionState; use datafusion::logical_expr::{create_udf, Expr, ScalarUDF, Volatility}; @@ -35,7 +35,7 @@ use datafusion::{ datasource::{MemTable, TableProvider, TableType}, prelude::{CsvReadOptions, SessionContext}, }; -use datafusion_common::cast::{as_float64_array, as_int32_array}; +use datafusion_common::cast::as_float64_array; use datafusion_common::DataFusionError; use log::info; use std::collections::HashMap; @@ -105,8 +105,8 @@ impl TestContext { } "joins.slt" => { info!("Registering partition table tables"); - let twice = create_pow_udf(); - test_ctx.ctx.register_udf(twice); + let pow = create_pow_udf(); + test_ctx.ctx.register_udf(pow); register_partition_table(&mut test_ctx).await; } "metadata.slt" => { @@ -398,7 +398,7 @@ fn create_pow_udf() -> ScalarUDF { // * give it a name so that it shows nicely when the plan is printed // * declare what input it expects // * declare its return type - let pow = create_udf( + create_udf( "pow_udf", // expects two f64 vec![DataType::Float64, DataType::Float64], @@ -406,6 +406,5 @@ fn create_pow_udf() -> ScalarUDF { Arc::new(DataType::Float64), Volatility::Immutable, pow, - ); - pow + ) } From 7fb086eb4a8f305cc8d262f34709c287ed1d9fd8 Mon Sep 17 00:00:00 2001 From: Mustafa Akur Date: Mon, 25 Dec 2023 16:34:20 +0300 Subject: [PATCH 4/5] Update UDF example --- datafusion/sqllogictest/src/test_context.rs | 34 ++++++++++---------- datafusion/sqllogictest/test_files/joins.slt | 9 +++--- 2 files changed, 22 insertions(+), 21 deletions(-) diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index bafbb674f81dc..479901e4c82fd 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -105,8 +105,8 @@ impl TestContext { } "joins.slt" => { info!("Registering partition table tables"); - let pow = create_pow_udf(); - test_ctx.ctx.register_udf(pow); + let example = create_example_udf(); + test_ctx.ctx.register_udf(example); register_partition_table(&mut test_ctx).await; } "metadata.slt" => { @@ -354,10 +354,10 @@ pub async fn register_metadata_tables(ctx: &SessionContext) { ctx.register_batch("table_with_metadata", batch).unwrap(); } -/// Create a UDF function named "pow_udf" -fn create_pow_udf() -> ScalarUDF { +/// Create a UDF function named "example" +fn create_example_udf() -> ScalarUDF { // First, declare the actual implementation of the calculation - let pow = |args: &[ArrayRef]| { + let adder = |args: &[ArrayRef]| { // in DataFusion, all `args` and output are dynamically-typed arrays, which means that we need to: // 1. cast the values to the type we want // 2. perform the computation for every element in the array (using a loop or SIMD) and construct the result @@ -366,45 +366,45 @@ fn create_pow_udf() -> ScalarUDF { assert_eq!(args.len(), 2); // 1. cast both arguments to f64. These casts MUST be aligned with the signature or this function panics! - let base = as_float64_array(&args[0]).expect("cast failed"); - let exponent = as_float64_array(&args[1]).expect("cast failed"); + let lhs = as_float64_array(&args[0]).expect("cast failed"); + let rhs = as_float64_array(&args[1]).expect("cast failed"); // this is guaranteed by DataFusion. We place it just to make it obvious. - assert_eq!(exponent.len(), base.len()); + assert_eq!(lhs.len(), rhs.len()); // 2. perform the computation - let array = base + let array = lhs .iter() - .zip(exponent.iter()) - .map(|(base, exponent)| { - match (base, exponent) { + .zip(rhs.iter()) + .map(|(lhs, rhs)| { + match (lhs, rhs) { // in arrow, any value can be null. // Here we decide to make our UDF to return null when either base or exponent is null. - (Some(base), Some(exponent)) => Some(base.powf(exponent)), + (Some(lhs), Some(rhs)) => Some(lhs + rhs), _ => None, } }) .collect::(); - // `Ok` because no error occurred during the calculation (we should add one if exponent was [0, 1[ and the base < 0 because that panics!) + // `Ok` because no error occurred during the calculation // `Arc` because arrays are immutable, thread-safe, trait objects. Ok(Arc::new(array) as ArrayRef) }; // the function above expects an `ArrayRef`, but DataFusion may pass a scalar to a UDF. // thus, we use `make_scalar_function` to decorare the closure so that it can handle both Arrays and Scalar values. - let pow = make_scalar_function(pow); + let adder = make_scalar_function(adder); // Next: // * give it a name so that it shows nicely when the plan is printed // * declare what input it expects // * declare its return type create_udf( - "pow_udf", + "example", // expects two f64 vec![DataType::Float64, DataType::Float64], // returns f64 Arc::new(DataType::Float64), Volatility::Immutable, - pow, + adder, ) } diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 76bbedbda100a..9a349f6000912 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -3483,14 +3483,15 @@ NestedLoopJoinExec: join_type=Inner, filter=a@0 > a@1 ----CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true --CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true -# Filter conditions are not pushed down through UDFs. +# Currently datafusion cannot pushdown filter conditions with scalar UDF into +# cross join. query TT EXPLAIN SELECT * FROM annotated_data as t1, annotated_data as t2 -WHERE POW_UDF(t1.a, 2.0) > POW_UDF(t2.a, 2.0) +WHERE EXAMPLE(t1.a, t2.a) > 3 ---- logical_plan -Filter: pow_udf(CAST(t1.a AS Float64), Float64(2)) > pow_udf(CAST(t2.a AS Float64), Float64(2)) +Filter: example(CAST(t1.a AS Float64), CAST(t2.a AS Float64)) > Float64(3) --CrossJoin: ----SubqueryAlias: t1 ------TableScan: annotated_data projection=[a0, a, b, c, d] @@ -3498,7 +3499,7 @@ Filter: pow_udf(CAST(t1.a AS Float64), Float64(2)) > pow_udf(CAST(t2.a AS Float6 ------TableScan: annotated_data projection=[a0, a, b, c, d] physical_plan CoalesceBatchesExec: target_batch_size=2 ---FilterExec: pow_udf(CAST(a@1 AS Float64), 2) > pow_udf(CAST(a@6 AS Float64), 2) +--FilterExec: example(CAST(a@1 AS Float64), CAST(a@6 AS Float64)) > 3 ----CrossJoinExec ------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true ------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 From 0712184b9e1e91b226a3c8c7c11ed3ce56421c11 Mon Sep 17 00:00:00 2001 From: Mehmet Ozan Kabak Date: Mon, 25 Dec 2023 23:46:11 +0300 Subject: [PATCH 5/5] Address review --- datafusion/optimizer/src/push_down_filter.rs | 12 +-- datafusion/sqllogictest/src/test_context.rs | 79 +++++++------------- 2 files changed, 31 insertions(+), 60 deletions(-) diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index de4d3fadd1359..9d277d18d2f7c 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -965,7 +965,8 @@ impl PushDownFilter { } } -/// Converts cross join to the inner join with empty equality predicate and empty filter condition. +/// Converts the given cross join to an inner join with an empty equality +/// predicate and an empty filter condition. fn convert_cross_join_to_inner_join(cross_join: CrossJoin) -> Result { let CrossJoin { left, right, .. } = cross_join; let join_schema = build_join_schema(left.schema(), right.schema(), &JoinType::Inner)?; @@ -981,7 +982,8 @@ fn convert_cross_join_to_inner_join(cross_join: CrossJoin) -> Result { }) } -/// Converts the inner join with empty equality predicate and empty filter condition to the cross join +/// Converts the given inner join with an empty equality predicate and an +/// empty filter condition to a cross join. fn convert_to_cross_join_if_beneficial(plan: LogicalPlan) -> Result { if let LogicalPlan::Join(join) = &plan { // Can be converted back to cross join @@ -993,10 +995,8 @@ fn convert_to_cross_join_if_beneficial(plan: LogicalPlan) -> Result } else if let LogicalPlan::Filter(filter) = &plan { let new_input = convert_to_cross_join_if_beneficial(filter.input.as_ref().clone())?; - return Ok(LogicalPlan::Filter(Filter::try_new( - filter.predicate.clone(), - Arc::new(new_input), - )?)); + return Filter::try_new(filter.predicate.clone(), Arc::new(new_input)) + .map(LogicalPlan::Filter); } Ok(plan) } diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 479901e4c82fd..a5ce7ccb9fe08 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -15,34 +15,33 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::ArrayRef; -use async_trait::async_trait; +use std::collections::HashMap; +use std::fs::File; +use std::io::Write; +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, BinaryArray, Float64Array, Int32Array, LargeBinaryArray, LargeStringArray, + StringArray, TimestampNanosecondArray, +}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use arrow::record_batch::RecordBatch; use datafusion::execution::context::SessionState; use datafusion::logical_expr::{create_udf, Expr, ScalarUDF, Volatility}; use datafusion::physical_expr::functions::make_scalar_function; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use datafusion::{ - arrow::{ - array::{ - BinaryArray, Float64Array, Int32Array, LargeBinaryArray, LargeStringArray, - StringArray, TimestampNanosecondArray, - }, - datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}, - record_batch::RecordBatch, - }, catalog::{schema::MemorySchemaProvider, CatalogProvider, MemoryCatalogProvider}, datasource::{MemTable, TableProvider, TableType}, prelude::{CsvReadOptions, SessionContext}, }; use datafusion_common::cast::as_float64_array; use datafusion_common::DataFusionError; + +use async_trait::async_trait; use log::info; -use std::collections::HashMap; -use std::fs::File; -use std::io::Write; -use std::path::Path; -use std::sync::Arc; use tempfile::TempDir; /// Context for running tests @@ -105,8 +104,8 @@ impl TestContext { } "joins.slt" => { info!("Registering partition table tables"); - let example = create_example_udf(); - test_ctx.ctx.register_udf(example); + let example_udf = create_example_udf(); + test_ctx.ctx.register_udf(example_udf); register_partition_table(&mut test_ctx).await; } "metadata.slt" => { @@ -354,55 +353,27 @@ pub async fn register_metadata_tables(ctx: &SessionContext) { ctx.register_batch("table_with_metadata", batch).unwrap(); } -/// Create a UDF function named "example" +/// Create a UDF function named "example". See the `sample_udf.rs` example +/// file for an explanation of the API. fn create_example_udf() -> ScalarUDF { - // First, declare the actual implementation of the calculation - let adder = |args: &[ArrayRef]| { - // in DataFusion, all `args` and output are dynamically-typed arrays, which means that we need to: - // 1. cast the values to the type we want - // 2. perform the computation for every element in the array (using a loop or SIMD) and construct the result - - // this is guaranteed by DataFusion based on the function's signature. - assert_eq!(args.len(), 2); - - // 1. cast both arguments to f64. These casts MUST be aligned with the signature or this function panics! + let adder = make_scalar_function(|args: &[ArrayRef]| { let lhs = as_float64_array(&args[0]).expect("cast failed"); let rhs = as_float64_array(&args[1]).expect("cast failed"); - - // this is guaranteed by DataFusion. We place it just to make it obvious. - assert_eq!(lhs.len(), rhs.len()); - - // 2. perform the computation let array = lhs .iter() .zip(rhs.iter()) - .map(|(lhs, rhs)| { - match (lhs, rhs) { - // in arrow, any value can be null. - // Here we decide to make our UDF to return null when either base or exponent is null. - (Some(lhs), Some(rhs)) => Some(lhs + rhs), - _ => None, - } + .map(|(lhs, rhs)| match (lhs, rhs) { + (Some(lhs), Some(rhs)) => Some(lhs + rhs), + _ => None, }) .collect::(); - - // `Ok` because no error occurred during the calculation - // `Arc` because arrays are immutable, thread-safe, trait objects. Ok(Arc::new(array) as ArrayRef) - }; - // the function above expects an `ArrayRef`, but DataFusion may pass a scalar to a UDF. - // thus, we use `make_scalar_function` to decorare the closure so that it can handle both Arrays and Scalar values. - let adder = make_scalar_function(adder); - - // Next: - // * give it a name so that it shows nicely when the plan is printed - // * declare what input it expects - // * declare its return type + }); create_udf( "example", - // expects two f64 + // Expects two f64 values: vec![DataType::Float64, DataType::Float64], - // returns f64 + // Returns an f64 value: Arc::new(DataType::Float64), Volatility::Immutable, adder,