buraksenn commented on code in PR #24315:
URL: https://github.com/apache/datafusion/pull/24315#discussion_r3802033730
##########
datafusion/datasource-parquet/src/nested_schema_pruning.rs:
##########
@@ -763,6 +870,142 @@ mod tests {
assert_eq!(x.value(2), 3);
}
+ /// Pins the arrow-rs behavior this module relies on: selecting a subset
+ /// of leaves under a `List<Struct>` column with `ProjectionMask::leaves`
+ /// makes the reader emit exactly the type predicted by [`clip_for_cast`],
+ /// and null list rows / null struct elements survive (their validity is
+ /// reconstructed from the surviving leaves' definition levels).
+ #[test]
+ fn arrow_reader_emits_clipped_type_for_masked_list_struct() {
+ use arrow::array::ListArray;
+ use arrow::buffer::OffsetBuffer;
+
+ let (item_field, values) = pin_item_values();
+ // Rows: [e0, e1], NULL, [e2].
+ let events = ListArray::new(
+ Arc::clone(&item_field),
+ OffsetBuffer::from_lengths([2, 0, 1]),
+ Arc::new(values),
+ Some(NullBuffer::from(vec![true, false, true])),
+ );
+ let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new(
+ "events",
+ DataType::List(item_field),
+ true,
+ )]));
+ let batch = RecordBatch::try_new(schema,
vec![Arc::new(events)]).unwrap();
+
+ let target = list_of(struct_of(vec![int64("x"), utf8("y")]));
+ let out = roundtrip_with_clipped_mask(&batch, &target, &[0, 1]);
+
+ let events =
out.column(0).as_any().downcast_ref::<ListArray>().unwrap();
+ assert!(events.is_valid(0));
+ assert!(events.is_null(1));
+ assert!(events.is_valid(2));
+ let structs = events
+ .values()
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ assert_pin_item_values(structs);
+ }
+
+ /// Pins the reader contract for `ListView<Struct>`.
+ #[test]
+ fn arrow_reader_emits_clipped_type_for_masked_list_view_struct() {
+ use arrow::array::ListViewArray;
+
+ let (item_field, values) = pin_item_values();
+ // Rows: [e0, e1], NULL, [e2].
+ let events = ListViewArray::new(
+ Arc::clone(&item_field),
+ vec![0_i32, 0, 2].into(),
+ vec![2_i32, 0, 1].into(),
+ Arc::new(values),
+ Some(NullBuffer::from(vec![true, false, true])),
Review Comment:
Thanks for the review, I've added largelistview case as well
--
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]
