Jefffrey commented on code in PR #18531:
URL: https://github.com/apache/datafusion/pull/18531#discussion_r2522417394
##########
datafusion/functions-nested/src/map.rs:
##########
@@ -393,17 +503,249 @@ fn make_map_array_internal<O: OffsetSizeTrait>(
.add_child_data(flattened_values.to_data())
.build()?;
- let map_data = ArrayData::builder(DataType::Map(
+ let mut map_data_builder = ArrayData::builder(DataType::Map(
Arc::new(Field::new(
"entries",
struct_data.data_type().clone(),
false,
)),
false,
))
- .len(keys.len())
+ .len(original_len) // Use the original number of rows, not the filtered
count
.add_child_data(struct_data)
- .add_buffer(Buffer::from_slice_ref(offset_buffer.as_slice()))
- .build()?;
+ .add_buffer(Buffer::from_slice_ref(offset_buffer.as_slice()));
+
+ // Add the nulls bitmap if present (to preserve NULL map values)
+ if let Some(nulls) = nulls_bitmap {
+ map_data_builder = map_data_builder.nulls(Some(nulls));
+ }
+
+ let map_data = map_data_builder.build()?;
Ok(ColumnarValue::Array(Arc::new(MapArray::from(map_data))))
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ #[test]
+ fn test_make_map_with_null_maps() {
+ // Test that NULL map values (entire map is NULL) are correctly handled
+ // This test directly calls make_map_batch with a List containing NULL
elements
+ //
+ // Background: On main branch, the code would fail with "map key
cannot be null"
+ // because it couldn't distinguish between:
+ // - NULL map (entire map is NULL) - should be allowed
+ // - null key within a map - should be rejected
+
+ // Build keys array: [['a'], NULL, ['b']]
+ // The middle NULL represents an entire NULL map, not a null key
+ let mut key_builder =
+ arrow::array::ListBuilder::new(arrow::array::StringBuilder::new());
+
+ // First map: ['a']
+ key_builder.values().append_value("a");
+ key_builder.append(true);
+
+ // Second map: NULL (entire map is NULL)
+ key_builder.append(false);
+
+ // Third map: ['b']
+ key_builder.values().append_value("b");
+ key_builder.append(true);
+
+ let keys_array = Arc::new(key_builder.finish());
+
+ // Build values array: [[1], [2], [3]]
+ let mut value_builder =
+ arrow::array::ListBuilder::new(arrow::array::Int32Builder::new());
+
+ value_builder.values().append_value(1);
+ value_builder.append(true);
+
+ value_builder.values().append_value(2);
+ value_builder.append(true);
+
+ value_builder.values().append_value(3);
+ value_builder.append(true);
+
+ let values_array = Arc::new(value_builder.finish());
+
+ // Call make_map_batch - should succeed
+ let result = make_map_batch(&[
+ ColumnarValue::Array(keys_array),
+ ColumnarValue::Array(values_array),
+ ]);
+
+ assert!(result.is_ok(), "Should handle NULL maps correctly");
+
+ // Verify the result
+ let map_array = match result.unwrap() {
+ ColumnarValue::Array(arr) => arr,
+ _ => panic!("Expected Array result"),
+ };
+
+ assert_eq!(map_array.len(), 3, "Should have 3 maps");
+ assert!(!map_array.is_null(0), "First map should not be NULL");
+ assert!(map_array.is_null(1), "Second map should be NULL");
+ assert!(!map_array.is_null(2), "Third map should not be NULL");
+ }
+
+ #[test]
+ fn test_make_map_with_null_key_within_map_should_fail() {
+ // Test that null keys WITHIN a map are properly rejected
+ // This ensures the fix doesn't accidentally allow invalid null keys
+
+ // Build keys array: [['a', NULL, 'b']]
+ // The NULL here is a null key within the map, which is invalid
+ let mut key_builder =
+ arrow::array::ListBuilder::new(arrow::array::StringBuilder::new());
+
+ key_builder.values().append_value("a");
+ key_builder.values().append_null(); // Invalid: null key
+ key_builder.values().append_value("b");
+ key_builder.append(true);
+
+ let keys_array = Arc::new(key_builder.finish());
+
+ // Build values array: [[1, 2, 3]]
+ let mut value_builder =
+ arrow::array::ListBuilder::new(arrow::array::Int32Builder::new());
+
+ value_builder.values().append_value(1);
+ value_builder.values().append_value(2);
+ value_builder.values().append_value(3);
+ value_builder.append(true);
+
+ let values_array = Arc::new(value_builder.finish());
+
+ // Call make_map_batch - should fail
+ let result = make_map_batch(&[
+ ColumnarValue::Array(keys_array),
+ ColumnarValue::Array(values_array),
+ ]);
+
+ assert!(result.is_err(), "Should reject null keys within maps");
+
+ let err_msg = result.unwrap_err().to_string();
+ assert!(
+ err_msg.contains("cannot be null"),
+ "Error should mention null keys, got: {err_msg}"
+ );
+ }
+
+ #[test]
+ fn test_make_map_with_large_list() {
Review Comment:
Could we also put the largelist & fixedsizelist tests as SLTs instead?
--
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]
