Move MAKE_MAP to ExprPlanner (apache#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

Author: goldmedal

datafusion-cli/Cargo.lock

@@ -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

datafusion/expr/src/planner.rs

@@ -53,10 +53,15 @@ 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"] }
 
 [[bench]]
 harness = false
 name = "array_expression"
+
+[[bench]]
+harness = false
+name = "map"

datafusion/functions-array/Cargo.toml

@@ -0,0 +1,69 @@
+// 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.
+
+extern crate criterion;
+
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use rand::prelude::ThreadRng;
+use rand::Rng;
+
+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![];
+    for _ in 0..1000 {
+        keys.push(rng.gen_range(0..9999).to_string());
+    }
+    keys
+}
+
+fn values(rng: &mut ThreadRng) -> Vec<i32> {
+    let mut values = vec![];
+    for _ in 0..1000 {
+        values.push(rng.gen_range(0..9999));
+    }
+    values
+}
+
+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(Expr::Literal(ScalarValue::Utf8(Some(keys[i].clone()))));
+            buffer.push(Expr::Literal(ScalarValue::Int32(Some(values[i]))));
+        }
+
+        let planner = ArrayFunctionPlanner {};
+
+        b.iter(|| {
+            black_box(
+                planner
+                    .plan_make_map(buffer.clone())
+                    .expect("map should work on valid values"),
+            );
+        });
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

datafusion/functions-array/benches/map.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;

datafusion/functions-array/src/planner.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));

datafusion/functions/benches/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,