kosiew commented on code in PR #22784:
URL: https://github.com/apache/datafusion/pull/22784#discussion_r3371382754
##########
datafusion/sqllogictest/test_files/map.slt:
##########
@@ -901,3 +901,14 @@ NULL
statement ok
drop table tt;
+
+# Regression test: map with literal keys and column-reference values
+# Previously failed with "map requires key and value lists to have the same
length"
+# because scalar keys (FixedSizeList) had length=N (number of keys) while
+# array values had length=batch_size.
+query ?
+SELECT map(['a','b'], [column1, column1 * 10]) FROM (VALUES (1), (2), (3)) t;
Review Comment:
Nice coverage for literal keys with column values. Since the implementation
expands either side when mixed, it may also be useful to add the symmetric
SQL-visible case with array keys and scalar values, such as `map([column1,
column1 * 10], ['a','b'])` or an equivalent version with unique keys.
##########
datafusion/functions-nested/src/map.rs:
##########
@@ -63,15 +63,37 @@ fn can_evaluate_to_const(args: &[ColumnarValue]) -> bool {
.all(|arg| matches!(arg, ColumnarValue::Scalar(_)))
}
-fn make_map_batch(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+fn make_map_batch(args: &[ColumnarValue], number_rows: usize) ->
Result<ColumnarValue> {
let [keys_arg, values_arg] = take_function_args("make_map", args)?;
let can_evaluate_to_const = can_evaluate_to_const(args);
- let keys = get_first_array_ref(keys_arg)?;
+ // If we have a mix of scalar and array args, expand the scalar to an array
+ // so that the rest of the logic handles uniform array inputs.
+ // This handles cases like: map(['a','b'], [col1, col2]) where keys are all
+ // literals (scalar FixedSizeList) but values contain column references.
+ let (keys_arg, values_arg) = if !can_evaluate_to_const {
Review Comment:
This scalar expansion logic is repeated for both keys and values. It might
be worth pulling it into a small helper so the invariant lives in one place and
future mixed scalar/array handling is less likely to drift.
For example:
```rust
fn expand_if_scalar(arg: &ColumnarValue, rows: usize) ->
Result<ColumnarValue> {
match arg {
ColumnarValue::Scalar(s) =>
Ok(ColumnarValue::Array(s.to_array_of_size(rows)?)),
ColumnarValue::Array(a) => Ok(ColumnarValue::Array(Arc::clone(a))),
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
