diff --git a/datafusion/core/tests/dataframe/dataframe_functions.rs b/datafusion/core/tests/dataframe/dataframe_functions.rs index 1c55c48fea40..f68ae507e35e 100644 --- a/datafusion/core/tests/dataframe/dataframe_functions.rs +++ b/datafusion/core/tests/dataframe/dataframe_functions.rs @@ -34,6 +34,7 @@ use datafusion_common::{DFSchema, ScalarValue}; use datafusion_expr::expr::Alias; use datafusion_expr::ExprSchemable; use datafusion_functions_aggregate::expr_fn::{approx_median, approx_percentile_cont}; +use datafusion_functions_array::map::map; fn test_schema() -> SchemaRef { Arc::new(Schema::new(vec![ @@ -1087,3 +1088,20 @@ async fn test_fn_array_to_string() -> Result<()> { Ok(()) } + +#[tokio::test] +async fn test_fn_map() -> Result<()> { + let expr = map(vec![lit("a"), lit("b"), lit("c")], vec![lit(1), lit(2), lit(3)]); + let expected = [ + "+---------------------------------------------------------------------------------------+", + "| map(make_array(Utf8(\"a\"),Utf8(\"b\"),Utf8(\"c\")),make_array(Int32(1),Int32(2),Int32(3))) |", + "+---------------------------------------------------------------------------------------+", + "| {a: 1, b: 2, c: 3} |", + "| {a: 1, b: 2, c: 3} |", + "| {a: 1, b: 2, c: 3} |", + "| {a: 1, b: 2, c: 3} |", + "+---------------------------------------------------------------------------------------+", + ]; + assert_fn_batches!(expr, expected); + Ok(()) +} diff --git a/datafusion/functions-array/src/lib.rs b/datafusion/functions-array/src/lib.rs index 9717d29883fd..5a31c4e0f55c 100644 --- a/datafusion/functions-array/src/lib.rs +++ b/datafusion/functions-array/src/lib.rs @@ -53,6 +53,8 @@ pub mod set_ops; pub mod sort; pub mod string; pub mod utils; +pub mod map; + use datafusion_common::Result; use datafusion_execution::FunctionRegistry; use datafusion_expr::ScalarUDF; diff --git a/datafusion/functions-array/src/planner.rs b/datafusion/functions-array/src/planner.rs index fbb541d9b151..47309999b021 100644 --- a/datafusion/functions-array/src/planner.rs +++ b/datafusion/functions-array/src/planner.rs @@ -24,6 +24,7 @@ use datafusion_expr::{ planner::{ExprPlanner, PlannerResult, RawBinaryExpr, RawFieldAccessExpr}, sqlparser, Expr, ExprSchemable, GetFieldAccess, }; +use datafusion_functions::core::map_udf; use datafusion_functions::expr_fn::get_field; use datafusion_functions_aggregate::nth_value::nth_value_udaf; @@ -112,7 +113,7 @@ impl ExprPlanner for ArrayFunctionPlanner { Ok(PlannerResult::Planned(Expr::ScalarFunction( ScalarFunction::new_udf( - datafusion_functions::core::map(), + map_udf(), vec![keys, values], ), ))) diff --git a/datafusion/functions/benches/map.rs b/datafusion/functions/benches/map.rs index 811c21a41b46..520981e3e159 100644 --- a/datafusion/functions/benches/map.rs +++ b/datafusion/functions/benches/map.rs @@ -23,7 +23,7 @@ use arrow_buffer::{OffsetBuffer, ScalarBuffer}; use criterion::{black_box, criterion_group, criterion_main, Criterion}; use datafusion_common::ScalarValue; use datafusion_expr::ColumnarValue; -use datafusion_functions::core::map; +use datafusion_functions::core::map_udf; use rand::prelude::ThreadRng; use rand::Rng; use std::sync::Arc; @@ -68,7 +68,7 @@ fn criterion_benchmark(c: &mut Criterion) { b.iter(|| { black_box( - map() + map_udf() .invoke(&[keys.clone(), values.clone()]) .expect("map should work on valid values"), ); diff --git a/datafusion/functions/src/core/map.rs b/datafusion/functions/src/core/map.rs index 2deef242f8a0..9bd021374863 100644 --- a/datafusion/functions/src/core/map.rs +++ b/datafusion/functions/src/core/map.rs @@ -23,10 +23,12 @@ use arrow::array::{Array, ArrayData, ArrayRef, MapArray, StructArray}; use arrow::datatypes::{DataType, Field, SchemaBuilder}; use arrow_buffer::{Buffer, ToByteSlice}; -use datafusion_common::Result; use datafusion_common::{exec_err, ScalarValue}; +use datafusion_common::Result; use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility}; +make_udf_function!(MapFunc, MAP, map_udf); + /// Check if we can evaluate the expr to constant directly. /// /// # Example diff --git a/datafusion/functions/src/core/mod.rs b/datafusion/functions/src/core/mod.rs index cbfaa592b012..f7a2d0d244d1 100644 --- a/datafusion/functions/src/core/mod.rs +++ b/datafusion/functions/src/core/mod.rs @@ -19,6 +19,7 @@ use datafusion_expr::ScalarUDF; use std::sync::Arc; +pub use crate::core::map::map_udf; pub mod arrow_cast; pub mod arrowtypeof; @@ -43,7 +44,6 @@ make_udf_function!(r#struct::StructFunc, STRUCT, r#struct); make_udf_function!(named_struct::NamedStructFunc, NAMED_STRUCT, named_struct); make_udf_function!(getfield::GetFieldFunc, GET_FIELD, get_field); make_udf_function!(coalesce::CoalesceFunc, COALESCE, coalesce); -make_udf_function!(map::MapFunc, MAP, map); pub mod expr_fn { use datafusion_expr::{Expr, Literal}; @@ -81,7 +81,7 @@ pub mod expr_fn { "Returns `coalesce(args...)`, which evaluates to the value of the first expr which is not NULL", args, ),( - map, + map_udf, "Returns a map created from a key list and a value list", args, )); @@ -102,6 +102,6 @@ pub fn functions() -> Vec> { named_struct(), get_field(), coalesce(), - map(), + map_udf(), ] }