This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new e36a410a2f feat: support size() for MapType input (#4580)
e36a410a2f is described below
commit e36a410a2fc3eb7d7c64270c6f401be7d7a23cad
Author: Zaki <[email protected]>
AuthorDate: Fri Jul 17 22:19:57 2026 +0800
feat: support size() for MapType input (#4580)
---
docs/source/user-guide/latest/expressions.md | 4 +-
native/spark-expr/src/array_funcs/size.rs | 141 ++++++++++++++++-----
.../main/scala/org/apache/comet/serde/arrays.scala | 6 +-
.../sql-tests/expressions/array/posexplode.sql | 4 +-
.../resources/sql-tests/expressions/array/size.sql | 18 ++-
.../org/apache/comet/CometMapExpressionSuite.scala | 51 +++++---
6 files changed, 158 insertions(+), 66 deletions(-)
diff --git a/docs/source/user-guide/latest/expressions.md
b/docs/source/user-guide/latest/expressions.md
index 4ad4d9afe3..f45fa21444 100644
--- a/docs/source/user-guide/latest/expressions.md
+++ b/docs/source/user-guide/latest/expressions.md
@@ -186,10 +186,10 @@ The tables below list every Spark built-in expression
with its current status.
| Function | Status | Notes |
| --- | --- | --- |
| `array_size` | ✅ | |
-| `cardinality` | ✅ | MapType input falls back |
+| `cardinality` | ✅ | |
| `concat` | ✅ | Binary/array children fall back |
| `reverse` | ✅ | Binary-element arrays fall back (Incompatible)
([details](compatibility/expressions/array.md)) |
-| `size` | ✅ | MapType input falls back |
+| `size` | ✅ | |
---
diff --git a/native/spark-expr/src/array_funcs/size.rs
b/native/spark-expr/src/array_funcs/size.rs
index 9bed27b33b..5ad4f1670a 100644
--- a/native/spark-expr/src/array_funcs/size.rs
+++ b/native/spark-expr/src/array_funcs/size.rs
@@ -195,6 +195,14 @@ fn spark_size_scalar(scalar: &ScalarValue) ->
Result<ScalarValue, DataFusionErro
Ok(ScalarValue::Int32(Some(len)))
}
}
+ ScalarValue::Map(array) => {
+ if array.is_null(0) {
+ Ok(ScalarValue::Int32(Some(-1)))
+ } else {
+ let len = array.value_length(0);
+ Ok(ScalarValue::Int32(Some(len)))
+ }
+ }
ScalarValue::Null => {
Ok(ScalarValue::Int32(Some(-1))) // Spark behavior: return -1 for
null
}
@@ -273,38 +281,28 @@ mod tests {
assert_eq!(result, ScalarValue::Int32(Some(-1)));
}
- // TODO: Add map array test once Arrow MapArray API constraints are
resolved
- // Currently MapArray doesn't allow nulls in entries which makes testing
complex
- // The core size() implementation supports maps correctly
- #[ignore]
#[test]
fn test_spark_size_map_array() {
- use arrow::array::{MapArray, StringArray};
+ use arrow::array::{Int32Array, MapArray, StringArray};
- // Create a simpler test with maps:
- // [{"key1": "value1", "key2": "value2"}, {"key3": "value3"}, {}, null]
+ // Create test data: [{"key1": 1, "key2": 2}, {"key3": 3}, {}, null]
+ let keys = StringArray::from(vec![Some("key1"), Some("key2"),
Some("key3")]);
+ let values = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
- // Create keys array for all entries (no nulls)
- let keys = StringArray::from(vec!["key1", "key2", "key3"]);
-
- // Create values array for all entries (no nulls)
- let values = StringArray::from(vec!["value1", "value2", "value3"]);
-
- // Create entry offsets: [0, 2, 3, 3] representing:
+ // Create entry offsets: [0, 2, 3, 3, 3] representing:
// - Map 1: entries 0-1 (2 key-value pairs)
- // - Map 2: entries 2-2 (1 key-value pair)
+ // - Map 2: entry 2 (1 key-value pair)
// - Map 3: entries 3-2 (0 key-value pairs, empty map)
// - Map 4: null (handled by null buffer)
- let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0, 2, 3, 3,
3].into());
+ let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0i32, 2, 3,
3, 3].into());
let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
- let value_field = Arc::new(Field::new("value", DataType::Utf8,
false)); // Make values non-nullable too
+ let value_field = Arc::new(Field::new("value", DataType::Int32, true));
- // Create the entries struct array
let entries = arrow::array::StructArray::new(
arrow::datatypes::Fields::from(vec![key_field, value_field]),
vec![Arc::new(keys), Arc::new(values)],
- None, // No nulls in the entries struct array itself
+ None,
);
// Create null buffer for the map array (fourth map is null)
@@ -314,27 +312,23 @@ mod tests {
null_buffer.append(true); // Empty map - not null
null_buffer.append(false); // null map
- let map_data_type = DataType::Map(
- Arc::new(Field::new(
- "entries",
- DataType::Struct(arrow::datatypes::Fields::from(vec![
- Field::new("key", DataType::Utf8, false),
- Field::new("value", DataType::Utf8, false), // Make values
non-nullable too
- ])),
- false,
- )),
- false, // keys are not sorted
- );
-
- let map_field = Arc::new(Field::new("map", map_data_type, true));
+ let map_field = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(arrow::datatypes::Fields::from(vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("value", DataType::Int32, true),
+ ])),
+ false,
+ ));
- let map_array = MapArray::new(
+ let map_array = MapArray::try_new(
map_field,
entry_offsets,
entries,
null_buffer.finish(),
- false, // keys are not sorted
- );
+ false,
+ )
+ .unwrap();
let array_ref: ArrayRef = Arc::new(map_array);
let result = spark_size_array(&array_ref).unwrap();
@@ -347,6 +341,83 @@ mod tests {
assert_eq!(result.value(3), -1); // null map returns -1
}
+ #[test]
+ fn test_spark_size_scalar_map() {
+ use arrow::array::{Int32Array, MapArray, StringArray};
+
+ // Test non-null map with 2 entries
+ let keys = StringArray::from(vec![Some("a"), Some("b")]);
+ let values = Int32Array::from(vec![Some(1), Some(2)]);
+ let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0i32,
2].into());
+
+ let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
+ let value_field = Arc::new(Field::new("value", DataType::Int32, true));
+
+ let entries = arrow::array::StructArray::new(
+ arrow::datatypes::Fields::from(vec![key_field, value_field]),
+ vec![Arc::new(keys), Arc::new(values)],
+ None,
+ );
+
+ let map_field = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(arrow::datatypes::Fields::from(vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("value", DataType::Int32, true),
+ ])),
+ false,
+ ));
+
+ let map_array = MapArray::try_new(map_field, entry_offsets, entries,
None, false).unwrap();
+ let scalar = ScalarValue::Map(Arc::new(map_array));
+ let result = spark_size_scalar(&scalar).unwrap();
+ assert_eq!(result, ScalarValue::Int32(Some(2))); // Map with 2 entries
+ }
+
+ #[test]
+ fn test_spark_size_scalar_null_map() {
+ use arrow::array::{Int32Array, MapArray, StringArray};
+
+ // Test null map - Spark returns -1 for null map input
+ let keys = StringArray::from(vec![Some("a")]);
+ let values = Int32Array::from(vec![Some(1)]);
+ let entry_offsets = arrow::buffer::OffsetBuffer::new(vec![0i32,
1].into());
+
+ let key_field = Arc::new(Field::new("key", DataType::Utf8, false));
+ let value_field = Arc::new(Field::new("value", DataType::Int32, true));
+
+ let entries = arrow::array::StructArray::new(
+ arrow::datatypes::Fields::from(vec![key_field, value_field]),
+ vec![Arc::new(keys), Arc::new(values)],
+ None,
+ );
+
+ let map_field = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(arrow::datatypes::Fields::from(vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("value", DataType::Int32, true),
+ ])),
+ false,
+ ));
+
+ // Mark the single map entry as null
+ let mut null_buffer = NullBufferBuilder::new(1);
+ null_buffer.append(false); // null map
+
+ let map_array = MapArray::try_new(
+ map_field,
+ entry_offsets,
+ entries,
+ null_buffer.finish(),
+ false,
+ )
+ .unwrap();
+ let scalar = ScalarValue::Map(Arc::new(map_array));
+ let result = spark_size_scalar(&scalar).unwrap();
+ assert_eq!(result, ScalarValue::Int32(Some(-1))); // null map returns
-1
+ }
+
#[test]
fn test_spark_size_fixed_size_list_array() {
use arrow::array::FixedSizeListArray;
diff --git a/spark/src/main/scala/org/apache/comet/serde/arrays.scala
b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
index f92827d3b1..8bc639021c 100644
--- a/spark/src/main/scala/org/apache/comet/serde/arrays.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/arrays.scala
@@ -739,15 +739,11 @@ object CometArrayFilter extends
CometExpressionSerde[ArrayFilter] {
object CometSize extends CometExpressionSerde[Size] {
- override def getUnsupportedReasons(): Seq[String] = Seq(
- "Only supports `ArrayType` input; `MapType` input is not supported")
-
override def getSupportLevel(expr: Size): SupportLevel = {
expr.child.dataType match {
case _: ArrayType => Compatible()
- case _: MapType => Unsupported(Some("size does not support map inputs"))
+ case _: MapType => Compatible()
case other =>
- // this should be unreachable because Spark only supports map and
array inputs
Unsupported(Some(s"Unsupported child data type: $other"))
}
}
diff --git
a/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql
b/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql
index 20b0547ebe..b7ba70d45a 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/posexplode.sql
@@ -95,6 +95,6 @@ INSERT INTO test_posexplode_map VALUES
(1, map('a', 1, 'b', 2)),
(2, map('c', 3))
--- posexplode over a map falls back to Spark (Comet only supports array inputs)
-query expect_fallback(size does not support map inputs)
+-- posexplode over a map falls back to Spark (Comet only supports array
inputs, not maps)
+query expect_fallback(Comet only supports explode/explode_outer for arrays,
not maps)
SELECT id, posexplode(m) FROM test_posexplode_map
diff --git a/spark/src/test/resources/sql-tests/expressions/array/size.sql
b/spark/src/test/resources/sql-tests/expressions/array/size.sql
index b006a4da0d..fb2bedef1e 100644
--- a/spark/src/test/resources/sql-tests/expressions/array/size.sql
+++ b/spark/src/test/resources/sql-tests/expressions/array/size.sql
@@ -15,15 +15,29 @@
-- specific language governing permissions and limitations
-- under the License.
+-- ConfigMatrix: spark.sql.legacy.sizeOfNull=true,false
+
statement
CREATE TABLE test_size(arr array<int>, m map<string, int>) USING parquet
statement
INSERT INTO test_size VALUES (array(1, 2, 3), map('a', 1, 'b', 2)), (array(),
map()), (NULL, NULL)
-query spark_answer_only
+query
SELECT size(arr), size(m) FROM test_size
--- literal arguments
+-- literal array arguments
query
SELECT size(array(1, 2, 3)), size(array()), size(cast(NULL as array<int>))
+
+-- literal map via CreateMap (falls back: Comet has no CreateMap serde;
+-- cast(NULL as map) avoids CreateMap and goes through CometLiteral instead)
+query spark_answer_only
+SELECT size(map('a', 1, 'b', 2)), size(map())
+
+query
+SELECT size(cast(NULL as map<string,int>))
+
+-- cardinality is a SQL alias for size
+query
+SELECT cardinality(arr), cardinality(m) FROM test_size
diff --git
a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala
b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala
index 9904919b74..2581ef5b0d 100644
--- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala
@@ -125,34 +125,47 @@ class CometMapExpressionSuite extends CometTestBase {
}
}
- test("fallback for size with map input") {
+ test("size with map input") {
withTempDir { dir =>
withTempView("t1") {
val path = new Path(dir.toURI.toString, "test.parquet")
makeParquetFileAllPrimitiveTypes(path, dictionaryEnabled = true, 100)
spark.read.parquet(path.toString).createOrReplaceTempView("t1")
- // Use column references in maps to avoid constant folding
- checkSparkAnswerAndFallbackReason(
- sql("SELECT size(case when _2 < 0 then map(_8, _9) else map() end)
from t1"),
- "size does not support map inputs")
+ checkSparkAnswer(
+ sql("SELECT size(case when _2 < 0 then map(_8, _9) else map() end)
from t1"))
}
}
}
- // fails with "map is not supported"
- ignore("size with map input") {
- withTempDir { dir =>
- withTempView("t1") {
- val path = new Path(dir.toURI.toString, "test.parquet")
- makeParquetFileAllPrimitiveTypes(path, dictionaryEnabled = true, 100)
- spark.read.parquet(path.toString).createOrReplaceTempView("t1")
+ test("size with map input - v2 reader") {
+ withTempPath { dir =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ val df = spark
+ .range(100)
+ .select(
+ col("id"),
+ when(col("id") > 1, map(col("id"), col("id"))).alias("map1"),
+ when(col("id") > 5, map(col("id"), col("id"))).alias("map2"))
+ df.write.parquet(dir.toString())
+ }
- // Use column references in maps to avoid constant folding
- checkSparkAnswerAndOperator(
- sql("SELECT size(map(_8, _9, _10, _11)) from t1 where _8 is not
null"))
- checkSparkAnswerAndOperator(
- sql("SELECT size(case when _2 < 0 then map(_8, _9) else map() end)
from t1"))
+ Seq("", "parquet").foreach { v1List =>
+ withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> v1List) {
+ val df = spark.read.parquet(dir.toString())
+ df.createOrReplaceTempView("t1")
+ if (v1List.isEmpty) {
+ checkSparkAnswer(df.select(size(col("map1"))))
+ checkSparkAnswer(df.select(size(col("map2"))))
+ checkSparkAnswer(
+ sql("SELECT size(CASE WHEN id < 50 THEN map1 ELSE map2 END) FROM
t1"))
+ } else {
+ checkSparkAnswerAndOperator(df.select(size(col("map1"))))
+ checkSparkAnswerAndOperator(df.select(size(col("map2"))))
+ checkSparkAnswerAndOperator(
+ sql("SELECT size(CASE WHEN id < 50 THEN map1 ELSE map2 END) FROM
t1"))
+ }
+ }
}
}
}
@@ -233,13 +246,11 @@ class CometMapExpressionSuite extends CometTestBase {
}
}
- test("map_from_entries - binary type") {
+ test("map_from_entries - binary type routes through codegen dispatcher") {
val table = "t2"
withTable(table) {
sql(
s"create table $table using parquet as select cast(array() as
array<binary>) as c1 from range(10)")
- // The native path is Incompatible for binary keys/values, so Comet
routes these through
- // the codegen dispatcher and still executes natively.
checkSparkAnswerAndOperator(
sql(s"select map_from_entries(array(struct(c1, 0))) from $table"))
checkSparkAnswerAndOperator(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
