This is an automated email from the ASF dual-hosted git repository.
jayzhan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 5f0dfbb8e7 Move `MAKE_MAP` to ExprPlanner (#11452)
5f0dfbb8e7 is described below
commit 5f0dfbb8e7424964303b00f4781f8df4f445d928
Author: Jax Liu <[email protected]>
AuthorDate: Fri Jul 19 17:32:24 2024 +0800
Move `MAKE_MAP` to ExprPlanner (#11452)
* move make_map to ExprPlanner
* add benchmark for make_map
* remove todo comment
* update lock
* refactor plan_make_map
* implement make_array_strict for type checking strictly
* fix planner provider
* roll back to `make_array`
* update lock
---
datafusion-cli/Cargo.lock | 9 +-
datafusion/expr/src/planner.rs | 7 +
datafusion/functions-array/Cargo.toml | 5 +
.../{functions => functions-array}/benches/map.rs | 52 ++-----
datafusion/functions-array/src/planner.rs | 21 ++-
datafusion/functions/benches/map.rs | 23 +---
datafusion/functions/src/core/map.rs | 149 +--------------------
datafusion/functions/src/core/mod.rs | 6 -
datafusion/sql/src/expr/function.rs | 12 ++
datafusion/sqllogictest/test_files/map.slt | 14 +-
10 files changed, 72 insertions(+), 226 deletions(-)
diff --git a/datafusion-cli/Cargo.lock b/datafusion-cli/Cargo.lock
index cdf0e7f573..61d9c72b89 100644
--- a/datafusion-cli/Cargo.lock
+++ b/datafusion-cli/Cargo.lock
@@ -1331,6 +1331,7 @@ dependencies = [
"itertools",
"log",
"paste",
+ "rand",
]
[[package]]
@@ -3593,18 +3594,18 @@ checksum =
"23d434d3f8967a09480fb04132ebe0a3e088c173e6d0ee7897abbdf4eab0f8b9"
[[package]]
name = "thiserror"
-version = "1.0.62"
+version = "1.0.63"
source = "registry+https://github.com/rust-lang/crates.io-index";
-checksum = "f2675633b1499176c2dff06b0856a27976a8f9d436737b4cf4f312d4d91d8bbb"
+checksum = "c0342370b38b6a11b6cc11d6a805569958d54cfa061a29969c3b5ce2ea405724"
dependencies = [
"thiserror-impl",
]
[[package]]
name = "thiserror-impl"
-version = "1.0.62"
+version = "1.0.63"
source = "registry+https://github.com/rust-lang/crates.io-index";
-checksum = "d20468752b09f49e909e55a5d338caa8bedf615594e9d80bc4c565d30faf798c"
+checksum = "a4558b58466b9ad7ca0f102865eccc95938dca1a74a856f2b57b6629050da261"
dependencies = [
"proc-macro2",
"quote",
diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs
index 009f3512c5..415af1bf94 100644
--- a/datafusion/expr/src/planner.rs
+++ b/datafusion/expr/src/planner.rs
@@ -173,6 +173,13 @@ pub trait ExprPlanner: Send + Sync {
fn plan_overlay(&self, args: Vec<Expr>) ->
Result<PlannerResult<Vec<Expr>>> {
Ok(PlannerResult::Original(args))
}
+
+ /// Plan a make_map expression, e.g., `make_map(key1, value1, key2,
value2, ...)`
+ ///
+ /// Returns origin expression arguments if not possible
+ fn plan_make_map(&self, args: Vec<Expr>) ->
Result<PlannerResult<Vec<Expr>>> {
+ Ok(PlannerResult::Original(args))
+ }
}
/// An operator with two arguments to plan
diff --git a/datafusion/functions-array/Cargo.toml
b/datafusion/functions-array/Cargo.toml
index 73c5b9114a..de424b2596 100644
--- a/datafusion/functions-array/Cargo.toml
+++ b/datafusion/functions-array/Cargo.toml
@@ -53,6 +53,7 @@ datafusion-functions-aggregate = { workspace = true }
itertools = { version = "0.12", features = ["use_std"] }
log = { workspace = true }
paste = "1.0.14"
+rand = "0.8.5"
[dev-dependencies]
criterion = { version = "0.5", features = ["async_tokio"] }
@@ -60,3 +61,7 @@ criterion = { version = "0.5", features = ["async_tokio"] }
[[bench]]
harness = false
name = "array_expression"
+
+[[bench]]
+harness = false
+name = "map"
diff --git a/datafusion/functions/benches/map.rs
b/datafusion/functions-array/benches/map.rs
similarity index 52%
copy from datafusion/functions/benches/map.rs
copy to datafusion/functions-array/benches/map.rs
index cd863d0e33..2e9b45266a 100644
--- a/datafusion/functions/benches/map.rs
+++ b/datafusion/functions-array/benches/map.rs
@@ -17,16 +17,14 @@
extern crate criterion;
-use arrow::array::{Int32Array, ListArray, StringArray};
-use arrow::datatypes::{DataType, Field};
-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::{make_map, map};
use rand::prelude::ThreadRng;
use rand::Rng;
-use std::sync::Arc;
+
+use datafusion_common::ScalarValue;
+use datafusion_expr::planner::ExprPlanner;
+use datafusion_expr::Expr;
+use datafusion_functions_array::planner::ArrayFunctionPlanner;
fn keys(rng: &mut ThreadRng) -> Vec<String> {
let mut keys = vec![];
@@ -51,46 +49,16 @@ fn criterion_benchmark(c: &mut Criterion) {
let values = values(&mut rng);
let mut buffer = Vec::new();
for i in 0..1000 {
- buffer.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
- keys[i].clone(),
- ))));
-
buffer.push(ColumnarValue::Scalar(ScalarValue::Int32(Some(values[i]))));
+
buffer.push(Expr::Literal(ScalarValue::Utf8(Some(keys[i].clone()))));
+ buffer.push(Expr::Literal(ScalarValue::Int32(Some(values[i]))));
}
- b.iter(|| {
- black_box(
- make_map()
- .invoke(&buffer)
- .expect("map should work on valid values"),
- );
- });
- });
-
- c.bench_function("map_1000", |b| {
- let mut rng = rand::thread_rng();
- let field = Arc::new(Field::new("item", DataType::Utf8, true));
- let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 1000]));
- let key_list = ListArray::new(
- field,
- offsets,
- Arc::new(StringArray::from(keys(&mut rng))),
- None,
- );
- let field = Arc::new(Field::new("item", DataType::Int32, true));
- let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 1000]));
- let value_list = ListArray::new(
- field,
- offsets,
- Arc::new(Int32Array::from(values(&mut rng))),
- None,
- );
- let keys =
ColumnarValue::Scalar(ScalarValue::List(Arc::new(key_list)));
- let values =
ColumnarValue::Scalar(ScalarValue::List(Arc::new(value_list)));
+ let planner = ArrayFunctionPlanner {};
b.iter(|| {
black_box(
- map()
- .invoke(&[keys.clone(), values.clone()])
+ planner
+ .plan_make_map(buffer.clone())
.expect("map should work on valid values"),
);
});
diff --git a/datafusion/functions-array/src/planner.rs
b/datafusion/functions-array/src/planner.rs
index dfb620f84f..fbb541d9b1 100644
--- a/datafusion/functions-array/src/planner.rs
+++ b/datafusion/functions-array/src/planner.rs
@@ -17,7 +17,8 @@
//! SQL planning extensions like [`ArrayFunctionPlanner`] and
[`FieldAccessPlanner`]
-use datafusion_common::{utils::list_ndims, DFSchema, Result};
+use datafusion_common::{exec_err, utils::list_ndims, DFSchema, Result};
+use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::{
expr::AggregateFunctionDefinition,
planner::{ExprPlanner, PlannerResult, RawBinaryExpr, RawFieldAccessExpr},
@@ -98,6 +99,24 @@ impl ExprPlanner for ArrayFunctionPlanner {
) -> Result<PlannerResult<Vec<Expr>>> {
Ok(PlannerResult::Planned(make_array(exprs)))
}
+
+ fn plan_make_map(&self, args: Vec<Expr>) ->
Result<PlannerResult<Vec<Expr>>> {
+ if args.len() % 2 != 0 {
+ return exec_err!("make_map requires an even number of arguments");
+ }
+
+ let (keys, values): (Vec<_>, Vec<_>) =
+ args.into_iter().enumerate().partition(|(i, _)| i % 2 == 0);
+ let keys = make_array(keys.into_iter().map(|(_, e)| e).collect());
+ let values = make_array(values.into_iter().map(|(_, e)| e).collect());
+
+ Ok(PlannerResult::Planned(Expr::ScalarFunction(
+ ScalarFunction::new_udf(
+ datafusion_functions::core::map(),
+ vec![keys, values],
+ ),
+ )))
+ }
}
pub struct FieldAccessPlanner;
diff --git a/datafusion/functions/benches/map.rs
b/datafusion/functions/benches/map.rs
index cd863d0e33..811c21a41b 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::{make_map, map};
+use datafusion_functions::core::map;
use rand::prelude::ThreadRng;
use rand::Rng;
use std::sync::Arc;
@@ -45,27 +45,6 @@ fn values(rng: &mut ThreadRng) -> Vec<i32> {
}
fn criterion_benchmark(c: &mut Criterion) {
- c.bench_function("make_map_1000", |b| {
- let mut rng = rand::thread_rng();
- let keys = keys(&mut rng);
- let values = values(&mut rng);
- let mut buffer = Vec::new();
- for i in 0..1000 {
- buffer.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
- keys[i].clone(),
- ))));
-
buffer.push(ColumnarValue::Scalar(ScalarValue::Int32(Some(values[i]))));
- }
-
- b.iter(|| {
- black_box(
- make_map()
- .invoke(&buffer)
- .expect("map should work on valid values"),
- );
- });
- });
-
c.bench_function("map_1000", |b| {
let mut rng = rand::thread_rng();
let field = Arc::new(Field::new("item", DataType::Utf8, true));
diff --git a/datafusion/functions/src/core/map.rs
b/datafusion/functions/src/core/map.rs
index 1834c7ac60..2deef242f8 100644
--- a/datafusion/functions/src/core/map.rs
+++ b/datafusion/functions/src/core/map.rs
@@ -20,12 +20,11 @@ use std::collections::VecDeque;
use std::sync::Arc;
use arrow::array::{Array, ArrayData, ArrayRef, MapArray, StructArray};
-use arrow::compute::concat;
use arrow::datatypes::{DataType, Field, SchemaBuilder};
use arrow_buffer::{Buffer, ToByteSlice};
-use datafusion_common::{exec_err, internal_err, ScalarValue};
-use datafusion_common::{not_impl_err, Result};
+use datafusion_common::Result;
+use datafusion_common::{exec_err, ScalarValue};
use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
/// Check if we can evaluate the expr to constant directly.
@@ -40,41 +39,6 @@ fn can_evaluate_to_const(args: &[ColumnarValue]) -> bool {
.all(|arg| matches!(arg, ColumnarValue::Scalar(_)))
}
-fn make_map(args: &[ColumnarValue]) -> Result<ColumnarValue> {
- let can_evaluate_to_const = can_evaluate_to_const(args);
-
- let (key, value): (Vec<_>, Vec<_>) = args
- .chunks_exact(2)
- .map(|chunk| {
- if let ColumnarValue::Array(_) = chunk[0] {
- return not_impl_err!("make_map does not support array keys");
- }
- if let ColumnarValue::Array(_) = chunk[1] {
- return not_impl_err!("make_map does not support array values");
- }
- Ok((chunk[0].clone(), chunk[1].clone()))
- })
- .collect::<Result<Vec<_>>>()?
- .into_iter()
- .unzip();
-
- let keys = ColumnarValue::values_to_arrays(&key)?;
- let values = ColumnarValue::values_to_arrays(&value)?;
-
- let keys: Vec<_> = keys.iter().map(|k| k.as_ref()).collect();
- let values: Vec<_> = values.iter().map(|v| v.as_ref()).collect();
-
- let key = match concat(&keys) {
- Ok(key) => key,
- Err(e) => return internal_err!("Error concatenating keys: {}", e),
- };
- let value = match concat(&values) {
- Ok(value) => value,
- Err(e) => return internal_err!("Error concatenating values: {}", e),
- };
- make_map_batch_internal(key, value, can_evaluate_to_const)
-}
-
fn make_map_batch(args: &[ColumnarValue]) -> Result<ColumnarValue> {
if args.len() != 2 {
return exec_err!(
@@ -154,115 +118,6 @@ fn make_map_batch_internal(
})
}
-#[derive(Debug)]
-pub struct MakeMap {
- signature: Signature,
-}
-
-impl Default for MakeMap {
- fn default() -> Self {
- Self::new()
- }
-}
-
-impl MakeMap {
- pub fn new() -> Self {
- Self {
- signature: Signature::user_defined(Volatility::Immutable),
- }
- }
-}
-
-impl ScalarUDFImpl for MakeMap {
- fn as_any(&self) -> &dyn Any {
- self
- }
-
- fn name(&self) -> &str {
- "make_map"
- }
-
- fn signature(&self) -> &Signature {
- &self.signature
- }
-
- fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
- if arg_types.is_empty() {
- return exec_err!(
- "make_map requires at least one pair of arguments, got 0
instead"
- );
- }
- if arg_types.len() % 2 != 0 {
- return exec_err!(
- "make_map requires an even number of arguments, got {}
instead",
- arg_types.len()
- );
- }
-
- let key_type = &arg_types[0];
- let mut value_type = &arg_types[1];
-
- for (i, chunk) in arg_types.chunks_exact(2).enumerate() {
- if chunk[0].is_null() {
- return exec_err!("make_map key cannot be null at position {}",
i);
- }
- if &chunk[0] != key_type {
- return exec_err!(
- "make_map requires all keys to have the same type {}, got
{} instead at position {}",
- key_type,
- chunk[0],
- i
- );
- }
-
- if !chunk[1].is_null() {
- if value_type.is_null() {
- value_type = &chunk[1];
- } else if &chunk[1] != value_type {
- return exec_err!(
- "map requires all values to have the same type {}, got
{} instead at position {}",
- value_type,
- &chunk[1],
- i
- );
- }
- }
- }
-
- let mut result = Vec::new();
- for _ in 0..arg_types.len() / 2 {
- result.push(key_type.clone());
- result.push(value_type.clone());
- }
-
- Ok(result)
- }
-
- fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
- let key_type = &arg_types[0];
- let mut value_type = &arg_types[1];
-
- for chunk in arg_types.chunks_exact(2) {
- if !chunk[1].is_null() && value_type.is_null() {
- value_type = &chunk[1];
- }
- }
-
- let mut builder = SchemaBuilder::new();
- builder.push(Field::new("key", key_type.clone(), false));
- builder.push(Field::new("value", value_type.clone(), true));
- let fields = builder.finish().fields;
- Ok(DataType::Map(
- Arc::new(Field::new("entries", DataType::Struct(fields), false)),
- false,
- ))
- }
-
- fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
- make_map(args)
- }
-}
-
#[derive(Debug)]
pub struct MapFunc {
signature: Signature,
diff --git a/datafusion/functions/src/core/mod.rs
b/datafusion/functions/src/core/mod.rs
index 31bce04bee..cbfaa592b0 100644
--- a/datafusion/functions/src/core/mod.rs
+++ b/datafusion/functions/src/core/mod.rs
@@ -43,7 +43,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::MakeMap, MAKE_MAP, make_map);
make_udf_function!(map::MapFunc, MAP, map);
pub mod expr_fn {
@@ -81,10 +80,6 @@ pub mod expr_fn {
coalesce,
"Returns `coalesce(args...)`, which evaluates to the value of the
first expr which is not NULL",
args,
- ),(
- make_map,
- "Returns a map created from the given keys and values pairs. This
function isn't efficient for large maps. Use the `map` function instead.",
- args,
),(
map,
"Returns a map created from a key list and a value list",
@@ -107,7 +102,6 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
named_struct(),
get_field(),
coalesce(),
- make_map(),
map(),
]
}
diff --git a/datafusion/sql/src/expr/function.rs
b/datafusion/sql/src/expr/function.rs
index dab328cc49..4a4b16b804 100644
--- a/datafusion/sql/src/expr/function.rs
+++ b/datafusion/sql/src/expr/function.rs
@@ -21,6 +21,7 @@ use datafusion_common::{
internal_datafusion_err, not_impl_err, plan_datafusion_err, plan_err,
DFSchema,
Dependency, Result,
};
+use datafusion_expr::planner::PlannerResult;
use datafusion_expr::window_frame::{check_window_frame,
regularize_window_order_by};
use datafusion_expr::{
expr, AggregateFunction, Expr, ExprSchemable, WindowFrame,
WindowFunctionDefinition,
@@ -227,6 +228,17 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
crate::utils::normalize_ident(name.0[0].clone())
};
+ if name.eq("make_map") {
+ let mut fn_args =
+ self.function_args_to_expr(args.clone(), schema,
planner_context)?;
+ for planner in self.context_provider.get_expr_planners().iter() {
+ match planner.plan_make_map(fn_args)? {
+ PlannerResult::Planned(expr) => return Ok(expr),
+ PlannerResult::Original(args) => fn_args = args,
+ }
+ }
+ }
+
// user-defined function (UDF) should have precedence
if let Some(fm) = self.context_provider.get_function_meta(&name) {
let args = self.function_args_to_expr(args, schema,
planner_context)?;
diff --git a/datafusion/sqllogictest/test_files/map.slt
b/datafusion/sqllogictest/test_files/map.slt
index fb8917a5f4..26bfb4a592 100644
--- a/datafusion/sqllogictest/test_files/map.slt
+++ b/datafusion/sqllogictest/test_files/map.slt
@@ -131,17 +131,23 @@ SELECT MAKE_MAP([1,2], ['a', 'b'], [3,4], ['b']);
----
{[1, 2]: [a, b], [3, 4]: [b]}
-query error
+query ?
SELECT MAKE_MAP('POST', 41, 'HEAD', 'ab', 'PATCH', 30);
+----
+{POST: 41, HEAD: ab, PATCH: 30}
query error
SELECT MAKE_MAP('POST', 41, 'HEAD', 33, null, 30);
-query error
-SELECT MAKE_MAP('POST', 41, 123, 33,'PATCH', 30);
+query ?
+SELECT MAKE_MAP('POST', 41, 'HEAD', 'ab', 'PATCH', 30);
+----
+{POST: 41, HEAD: ab, PATCH: 30}
-query error
+query ?
SELECT MAKE_MAP()
+----
+{}
query error
SELECT MAKE_MAP('POST', 41, 'HEAD');
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
