You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@auron.apache.org by sl...@apache.org on 2026/04/08 06:12:54 UTC
(auron) branch master updated: [AURON #2152] Implement native function of map_from_arrays (#2165)
This is an automated email from the ASF dual-hosted git repository.
slfan1989 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git
The following commit(s) were added to refs/heads/master by this push:
new 74844702 [AURON #2152] Implement native function of map_from_arrays (#2165)
74844702 is described below
commit 74844702d1ca9928fe1c623e38da386cfb42c34b
Author: Ming Wei <we...@gmail.com>
AuthorDate: Wed Apr 8 14:12:49 2026 +0800
[AURON #2152] Implement native function of map_from_arrays (#2165)
# Which issue does this PR close?
Closes #2152
# Rationale for this change
This PR adds native support for Spark `map_from_arrays(keys, values)`
through Auron’s extension-function path.
Spark defines this function as creating a map from the given key/value
arrays, with the requirement that all elements in `keys` must be
non-null. This change implements that behavior in Auron and adds
regression coverage.
`map_from_arrays(...)` was not supported in Auron’s native execution
path.
This PR extends native map-function coverage while keeping behavior
aligned with Spark semantics, following the existing extension-function
pattern already used for Spark-specific functions such as
`map_concat(...)`.
# What changes are included in this PR?
This PR:
- adds `MapFromArrays` conversion in `NativeConverters`
- registers `Spark_MapFromArrays` in `datafusion-ext-functions`
- implements `map_from_arrays(...)` in `spark_map.rs`
- enforces Spark-compatible null-key rejection
- handles row-level null propagation
- adds Scala and Rust regression tests
# Are there any user-facing changes?
Yes.
Queries using `map_from_arrays(keys, values)` can now run through
Auron’s native extension-function path.
# How was this patch tested?
UT.
---------
Signed-off-by: weimingdiit <we...@gmail.com>
---
native-engine/datafusion-ext-functions/src/lib.rs | 1 +
.../datafusion-ext-functions/src/spark_map.rs | 288 ++++++++++++++++++++-
.../org/apache/auron/AuronFunctionSuite.scala | 35 +++
.../apache/spark/sql/auron/NativeConverters.scala | 2 +
4 files changed, 324 insertions(+), 2 deletions(-)
diff --git a/native-engine/datafusion-ext-functions/src/lib.rs b/native-engine/datafusion-ext-functions/src/lib.rs
index 9555112d..1b8a88be 100644
--- a/native-engine/datafusion-ext-functions/src/lib.rs
+++ b/native-engine/datafusion-ext-functions/src/lib.rs
@@ -65,6 +65,7 @@ pub fn create_auron_ext_function(
"Spark_ParseJson" => Arc::new(spark_get_json_object::spark_parse_json),
"Spark_MakeArray" => Arc::new(spark_make_array::array),
"Spark_MapConcat" => Arc::new(spark_map::map_concat),
+ "Spark_MapFromArrays" => Arc::new(spark_map::map_from_arrays),
"Spark_StringSpace" => Arc::new(spark_strings::string_space),
"Spark_StringRepeat" => Arc::new(spark_strings::string_repeat),
"Spark_StringSplit" => Arc::new(spark_strings::string_split),
diff --git a/native-engine/datafusion-ext-functions/src/spark_map.rs b/native-engine/datafusion-ext-functions/src/spark_map.rs
index 5c834960..72aa36f0 100644
--- a/native-engine/datafusion-ext-functions/src/spark_map.rs
+++ b/native-engine/datafusion-ext-functions/src/spark_map.rs
@@ -16,7 +16,7 @@
use std::{collections::HashSet, sync::Arc};
use arrow::{
- array::{Array, ArrayRef, MapArray, StructArray, new_empty_array},
+ array::{Array, ArrayRef, ListArray, MapArray, StructArray, new_empty_array},
buffer::{NullBuffer, OffsetBuffer, ScalarBuffer},
datatypes::{DataType, Field},
};
@@ -153,6 +153,69 @@ fn get_arg_arrays(
.collect()
}
+fn as_list_array(array: &ArrayRef) -> Result<ListArray> {
+ array
+ .as_any()
+ .downcast_ref::<ListArray>()
+ .cloned()
+ .ok_or_else(|| {
+ datafusion::error::DataFusionError::Execution(format!(
+ "map_from_arrays args must be array, found {:?}",
+ array.data_type()
+ ))
+ })
+}
+
+fn new_null_list_field() -> Arc<Field> {
+ Arc::new(Field::new_list_field(DataType::Null, true))
+}
+
+fn new_null_list_array(list_field: Arc<Field>, len: usize) -> ListArray {
+ ListArray::new_null(list_field, len)
+}
+
+fn get_list_array_field(array: &ListArray, arg_name: &str) -> Result<Arc<Field>> {
+ match array.data_type() {
+ DataType::List(field) => Ok(field.clone()),
+ data_type => {
+ df_execution_err!("map_from_arrays {arg_name} arg must be array, found {data_type:?}")
+ }
+ }
+}
+
+fn columnar_value_to_list_array(arg: &ColumnarValue, arg_name: &str) -> Result<ListArray> {
+ match arg {
+ ColumnarValue::Array(array) if matches!(array.data_type(), DataType::Null) => {
+ Ok(new_null_list_array(new_null_list_field(), array.len()))
+ }
+ ColumnarValue::Array(array) => as_list_array(array),
+ ColumnarValue::Scalar(scalar) if scalar.is_null() => {
+ let list_field = match scalar.data_type() {
+ DataType::List(field) => field,
+ _ => new_null_list_field(),
+ };
+ Ok(new_null_list_array(list_field, 1))
+ }
+ ColumnarValue::Scalar(scalar) => {
+ let array = scalar.to_array()?;
+ if matches!(array.data_type(), DataType::Null) {
+ let list_field = match scalar.data_type() {
+ DataType::List(field) => field,
+ _ => new_null_list_field(),
+ };
+ Ok(new_null_list_array(list_field, array.len()))
+ } else {
+ as_list_array(&array)
+ }
+ }
+ }
+ .map_err(|err| {
+ datafusion::error::DataFusionError::Execution(format!(
+ "map_from_arrays {arg_name} arg must be array: {err}"
+ ))
+ })
+}
+
/// Returns the union of all given maps.
///
/// This follows Spark's default duplicate-key behavior by raising an error,
@@ -270,10 +333,127 @@ pub fn map_concat(args: &[ColumnarValue]) -> Result<ColumnarValue> {
))))
}
+/// Creates a map from the given key/value arrays.
+///
+/// This follows Spark semantics by propagating null when either input array
+/// for a row is null, rejecting null keys, and raising an error on duplicate
+/// keys under Spark's default map-key dedup policy.
+pub fn map_from_arrays(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 2 {
+ return df_execution_err!("map_from_arrays requires exactly 2 arguments");
+ }
+
+ let key_array = columnar_value_to_list_array(&args[0], "keys")?;
+ let value_array = columnar_value_to_list_array(&args[1], "values")?;
+ let key_list_field = get_list_array_field(&key_array, "keys")?;
+ let value_list_field = get_list_array_field(&value_array, "values")?;
+
+ let key_field = Arc::new(Field::new("key", key_list_field.data_type().clone(), false));
+ let value_field = Arc::new(Field::new(
+ "value",
+ value_list_field.data_type().clone(),
+ value_list_field.is_nullable(),
+ ));
+ let entries_field = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(vec![key_field.as_ref().clone(), value_field.as_ref().clone()].into()),
+ false,
+ ));
+
+ let num_rows = [key_array.len(), value_array.len()]
+ .into_iter()
+ .filter(|&len| len != 1)
+ .max()
+ .unwrap_or(1);
+
+ if [key_array.len(), value_array.len()]
+ .into_iter()
+ .any(|len| len != 1 && len != num_rows)
+ {
+ return df_execution_err!("all arrays of map_from_arrays must have the same length");
+ }
+
+ let mut all_keys = Vec::new();
+ let mut all_values = Vec::new();
+ let mut offsets = Vec::with_capacity(num_rows + 1);
+ let mut valids = Vec::with_capacity(num_rows);
+ let mut next_offset = 0i32;
+
+ offsets.push(next_offset);
+
+ for row_idx in 0..num_rows {
+ let key_idx = if key_array.len() == 1 { 0 } else { row_idx };
+ let value_idx = if value_array.len() == 1 { 0 } else { row_idx };
+
+ if key_array.is_null(key_idx) || value_array.is_null(value_idx) {
+ valids.push(false);
+ offsets.push(next_offset);
+ continue;
+ }
+
+ let keys = key_array.value(key_idx);
+ let values = value_array.value(value_idx);
+
+ if keys.len() != values.len() {
+ return df_execution_err!(
+ "map_from_arrays requires key/value arrays to have the same length"
+ );
+ }
+
+ let mut row_keys = HashSet::new();
+ valids.push(true);
+
+ for i in 0..keys.len() {
+ if keys.is_null(i) {
+ return df_execution_err!("map_from_arrays does not support null map keys");
+ }
+
+ let key = compacted_scalar_value_from_array(keys.as_ref(), i)?;
+ if !row_keys.insert(key.clone()) {
+ return df_execution_err!("map_from_arrays duplicate key found: {key}");
+ }
+
+ let value = compacted_scalar_value_from_array(values.as_ref(), i)?;
+ all_keys.push(key);
+ all_values.push(value);
+ }
+
+ next_offset += keys.len() as i32;
+ offsets.push(next_offset);
+ }
+
+ let keys = if all_keys.is_empty() {
+ new_empty_array(key_field.data_type())
+ } else {
+ ScalarValue::iter_to_array(all_keys.into_iter())?
+ };
+
+ let values = if all_values.is_empty() {
+ new_empty_array(value_field.data_type())
+ } else {
+ ScalarValue::iter_to_array(all_values.into_iter())?
+ };
+
+ let entries = StructArray::from(vec![(key_field, keys), (value_field, values)]);
+ let nulls = if valids.iter().all(|valid| *valid) {
+ None
+ } else {
+ Some(NullBuffer::from(valids))
+ };
+
+ Ok(ColumnarValue::Array(Arc::new(MapArray::new(
+ entries_field,
+ OffsetBuffer::new(ScalarBuffer::from(offsets)),
+ entries,
+ nulls,
+ false,
+ ))))
+}
+
#[cfg(test)]
mod test {
use arrow::{
- array::{Int32Array, StringArray},
+ array::{Int32Array, Int32Builder, ListBuilder, NullArray, StringArray, StringBuilder},
datatypes::Fields,
};
@@ -410,6 +590,44 @@ mod test {
)
}
+ fn build_string_list_array(rows: Vec<Option<Vec<Option<&'static str>>>>) -> ListArray {
+ let mut builder = ListBuilder::new(StringBuilder::new());
+ for row in rows {
+ match row {
+ Some(values) => {
+ for value in values {
+ match value {
+ Some(value) => builder.values().append_value(value),
+ None => builder.values().append_null(),
+ }
+ }
+ builder.append(true);
+ }
+ None => builder.append(false),
+ }
+ }
+ builder.finish()
+ }
+
+ fn build_int_list_array(rows: Vec<Option<Vec<Option<i32>>>>) -> ListArray {
+ let mut builder = ListBuilder::new(Int32Builder::new());
+ for row in rows {
+ match row {
+ Some(values) => {
+ for value in values {
+ match value {
+ Some(value) => builder.values().append_value(value),
+ None => builder.values().append_null(),
+ }
+ }
+ builder.append(true);
+ }
+ None => builder.append(false),
+ }
+ }
+ builder.finish()
+ }
+
#[test]
fn test_map_concat() -> Result<()> {
let left = build_string_int_map_array(vec![
@@ -507,4 +725,70 @@ mod test {
assert!(err.to_string().contains("same length"));
}
+
+ #[test]
+ fn test_map_from_arrays() -> Result<()> {
+ let keys = Arc::new(build_string_list_array(vec![
+ Some(vec![Some("a"), Some("b")]),
+ Some(vec![Some("x")]),
+ None,
+ Some(vec![Some("m"), Some("n")]),
+ ])) as ArrayRef;
+ let values = Arc::new(build_int_list_array(vec![
+ Some(vec![Some(1), Some(2)]),
+ Some(vec![Some(10)]),
+ Some(vec![Some(20)]),
+ Some(vec![None, Some(30)]),
+ ])) as ArrayRef;
+
+ let actual = map_from_arrays(&[ColumnarValue::Array(keys), ColumnarValue::Array(values)])?
+ .into_array(4)?;
+
+ let expected = Arc::new(build_string_int_map_array(vec![
+ Some(vec![("a", Some(1)), ("b", Some(2))]),
+ Some(vec![("x", Some(10))]),
+ None,
+ Some(vec![("m", None), ("n", Some(30))]),
+ ])) as ArrayRef;
+
+ assert_eq!(&actual, &expected);
+ Ok(())
+ }
+
+ #[test]
+ fn test_map_from_arrays_rejects_null_keys() {
+ let keys = Arc::new(build_string_list_array(vec![Some(vec![Some("a"), None])])) as ArrayRef;
+ let values = Arc::new(build_int_list_array(vec![Some(vec![Some(1), Some(2)])])) as ArrayRef;
+
+ let err = map_from_arrays(&[ColumnarValue::Array(keys), ColumnarValue::Array(values)])
+ .expect_err("map_from_arrays should fail when null keys exist");
+
+ assert!(err.to_string().contains("null map keys"));
+ }
+
+ #[test]
+ fn test_map_from_arrays_null_array_propagation() -> Result<()> {
+ let keys = Arc::new(NullArray::new(1)) as ArrayRef;
+ let values = Arc::new(build_int_list_array(vec![Some(vec![Some(1), Some(2)])])) as ArrayRef;
+
+ let actual = map_from_arrays(&[ColumnarValue::Array(keys), ColumnarValue::Array(values)])?
+ .into_array(1)?;
+
+ assert!(actual.is_null(0));
+ Ok(())
+ }
+
+ #[test]
+ fn test_map_from_arrays_null_scalar_propagation() -> Result<()> {
+ let values = Arc::new(build_int_list_array(vec![Some(vec![Some(1), Some(2)])])) as ArrayRef;
+
+ let actual = map_from_arrays(&[
+ ColumnarValue::Scalar(ScalarValue::Null),
+ ColumnarValue::Array(values),
+ ])?
+ .into_array(1)?;
+
+ assert!(actual.is_null(0));
+ Ok(())
+ }
}
diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
index 2df0d6c6..ef0983a5 100644
--- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
+++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronFunctionSuite.scala
@@ -336,6 +336,41 @@ class AuronFunctionSuite extends AuronQueryTest with BaseAuronSQLSuite {
}
}
+ test("map_from_arrays function") {
+ withTable("t1") {
+ sql("create table t1(c1 array<string>, c2 array<int>) using parquet")
+ sql("""
+ |insert into t1 values
+ | (array('a', 'b'), array(1, 2)),
+ | (array('x'), array(10)),
+ | (null, array(20)),
+ | (array('m', 'n'), array(null, 30))
+ |""".stripMargin)
+ checkSparkAnswerAndOperator("select map_from_arrays(c1, c2) from t1")
+ }
+ }
+
+ test("map_from_arrays rejects null keys") {
+ withTable("t1") {
+ sql("create table t1(c1 array<string>, c2 array<int>) using parquet")
+ sql("insert into t1 values (array('a', cast(null as string)), array(1, 2))")
+ val df = sql("select map_from_arrays(c1, c2) from t1")
+ val err = intercept[Exception] {
+ df.collect()
+ }
+ assert(err.getMessage.contains("null map keys"))
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ plan
+ .collectFirst { case op if !isNativeOrPassThrough(op) => op }
+ .foreach { op =>
+ fail(s"""
+ |Found non-native operator: ${op.nodeName}
+ |plan:
+ |${plan}""".stripMargin)
+ }
+ }
+ }
+
test("acosh null propagation") {
withTable("t1") {
sql("create table t1(c1 double) using parquet")
diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala
index c0d12d60..6d2ba759 100644
--- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala
+++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/NativeConverters.scala
@@ -940,6 +940,8 @@ object NativeConverters extends Logging {
buildExtScalarFunction("Spark_Murmur3Hash", children, IntegerType)
case XxHash64(children, 42L) =>
buildExtScalarFunction("Spark_XxHash64", children, LongType)
+ case e: MapFromArrays =>
+ buildExtScalarFunction("Spark_MapFromArrays", e.children, e.dataType)
case e: Greatest =>
buildScalarFunction(pb.ScalarFunction.Greatest, e.children, e.dataType)
case e: Pow =>