neilconway commented on code in PR #25754:
URL: https://github.com/apache/datafusion/pull/25754#discussion_r4284866860
##########
datafusion/functions-nested/src/utils.rs:
##########
@@ -529,6 +579,59 @@ pub(crate) mod tests {
use arrow::datatypes::Int64Type;
use datafusion_common::utils::SingleRowListArrayBuilder;
+ /// Compare `prepare_list_filter` with a per-row reference on a sliced list
+ /// with the given row lengths and validity, returning the input and output
+ /// bitmaps so callers can check whether the input was reused.
+ fn check_list_filter<O: OffsetSizeTrait>(
+ lengths: [usize; 8],
+ valid: [bool; 8],
+ ) -> (BooleanBuffer, BooleanBuffer) {
+ let list = GenericListArray::<O>::new(
+ Arc::new(Field::new_list_field(DataType::Int32, true)),
+ OffsetBuffer::from_lengths(lengths),
+ // The helper must permit selecting NULL elements.
+ arrow::array::new_null_array(&DataType::Int32,
lengths.iter().sum()),
+ Some(NullBuffer::from(valid.to_vec())),
+ )
+ .slice(1, 7);
+ let (first, len) = offset_span(list.offsets());
+ // Use an offset that is not a multiple of the bitmap's three-bit
period.
+ let keep = BooleanBuffer::collect_bool(len + 1, |i| i % 3 !=
0).slice(1, len);
+ let mut expected = keep.iter().collect::<Vec<_>>();
+ let mut expected_offsets = vec![0];
+ for (row, window) in list.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first;
+ let end = window[1].as_usize() - first;
+ if list.is_null(row) {
+ expected[start..end].fill(false);
+ }
+ expected_offsets.push(expected[..end].iter().filter(|&&v|
v).count());
+ }
+ let (offsets, mask) = prepare_list_filter(&list, &keep);
+ assert_eq!(mask.iter().collect::<Vec<_>>(), expected);
+ assert_eq!(
+ offsets.iter().map(|o| o.as_usize()).collect::<Vec<_>>(),
+ expected_offsets
+ );
+ (keep, mask)
+ }
+
+ #[test]
+ fn list_filter_mask_handles_slices_and_hidden_elements() {
+ let valid = [true, false, true, false, false, true, true, false];
+ // After slicing off the first row: an empty NULL row, a NULL row
storing
+ // an element after a valid row, consecutive NULL rows, and a trailing
one.
+ let hidden = [2, 0, 2, 1, 2, 0, 3, 2];
Review Comment:
Good idea, thanks! Done.
--
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]
