andygrove commented on code in PR #6937:
URL: https://github.com/apache/datafusion-comet/pull/6937#discussion_r4286328097
##########
native/core/src/execution/operators/iceberg_dictionary.rs:
##########
@@ -254,139 +343,234 @@ fn collect_leaves(data_type: &DataType, leaf: &mut
Leaf, leaves: &mut Vec<Leaf>)
descend(Step::Key, fields[0].data_type(), leaf, leaves);
descend(Step::Value, fields[1].data_type(), leaf, leaves);
}
- _ => leaves.push(leaf.clone()),
+ _ => leaves.push((leaf.clone(), data_type.clone())),
},
- _ => leaves.push(leaf.clone()),
+ _ => leaves.push((leaf.clone(), data_type.clone())),
}
}
-fn descend(step: Step, data_type: &DataType, leaf: &mut Leaf, leaves: &mut
Vec<Leaf>) {
+fn descend(step: Step, data_type: &DataType, leaf: &mut Leaf, leaves: &mut
Vec<(Leaf, DataType)>) {
leaf.steps.push(step);
collect_leaves(data_type, leaf, leaves);
leaf.steps.pop();
}
-/// `(row, index)` pairs, one per entry of an array that the column chunk
stores: its row of the
-/// batch, and its index into the array.
-type Entries<'a> = Box<dyn Iterator<Item = (u32, usize)> + 'a>;
-
impl Leaf {
- /// The leaf's array in `batch`, with the entries of it that lie under the
rows `top` yields and
- /// whose ancestors are all present, in the order the column chunk stores
them. Nothing is
- /// expanded ahead: each entry is produced when the page asks for it.
`None` when `batch` does
- /// not have the shape the schema promised.
- fn entries<'a>(
- &self,
- batch: &'a RecordBatch,
- top: Entries<'a>,
- ) -> Option<(&'a dyn Array, Entries<'a>)> {
+ /// The arrays of `batch` the way to the leaf passes through, or `None`
when `batch` does not
+ /// have the shape the schema promised.
+ fn path<'a>(&self, batch: &'a RecordBatch) -> Option<Path<'a>> {
let mut array = batch.columns().get(self.column)?.as_ref();
- let mut entries = top;
+ let mut levels = Vec::with_capacity(self.steps.len());
for step in &self.steps {
- let present = present(array, entries);
- (array, entries) = match (step, array.data_type()) {
+ let (children, next) = match (step, array.data_type()) {
(Step::Field(field), DataType::Struct(_)) => {
- (array.as_struct().columns().get(*field)?.as_ref(),
present)
+ (Children::Same, array.as_struct().columns().get(*field)?)
}
(Step::Element, DataType::List(_)) => {
let list = array.as_list::<i32>();
- (
- list.values().as_ref(),
- children(present, list.value_offsets()),
- )
+ (Children::Offsets(list.value_offsets()), list.values())
}
(Step::Element, DataType::LargeList(_)) => {
let list = array.as_list::<i64>();
- (
- list.values().as_ref(),
- children(present, list.value_offsets()),
- )
- }
- (Step::Element, DataType::FixedSizeList(_, size)) => {
- let size = *size as usize;
- let elements: Entries<'a> = Box::new(present.flat_map(move
|(row, index)| {
- (index * size..(index + 1) * size).map(move |element|
(row, element))
- }));
- (array.as_fixed_size_list().values().as_ref(), elements)
+ (Children::LargeOffsets(list.value_offsets()),
list.values())
}
+ (Step::Element, DataType::FixedSizeList(_, size)) => (
+ Children::Fixed(*size as usize),
+ array.as_fixed_size_list().values(),
+ ),
(Step::Key, DataType::Map(_, _)) => {
let map = array.as_map();
- (map.keys().as_ref(), children(present,
map.value_offsets()))
+ (Children::Offsets(map.value_offsets()), map.keys())
}
(Step::Value, DataType::Map(_, _)) => {
let map = array.as_map();
- (
- map.values().as_ref(),
- children(present, map.value_offsets()),
- )
+ (Children::Offsets(map.value_offsets()), map.values())
}
_ => return None,
};
+ // A null buffer with no nulls in it is as good as none, and a
struct without nulls
+ // leads every entry straight to its field's, so neither needs a
level.
+ let nulls = present_nulls(array);
+ if nulls.is_some() || !matches!(children, Children::Same) {
+ levels.push(Level { nulls, children });
+ }
+ array = next.as_ref();
}
- Some((array, entries))
+ Some(Path {
+ levels,
+ leaf: array,
+ })
}
}
-/// The entries whose value in `array` is not null.
-fn present<'a>(array: &'a dyn Array, entries: Entries<'a>) -> Entries<'a> {
- match array.nulls() {
- Some(nulls) => Box::new(entries.filter(move |&(_, index)|
nulls.is_valid(index))),
- None => entries,
- }
+/// The arrays the way from a batch's column down to one leaf column passes
through.
+struct Path<'a> {
+ /// One for each step, from the batch's column down.
+ levels: Vec<Level<'a>>,
+ /// The leaf column's array.
+ leaf: &'a dyn Array,
}
-/// The entries of the children of every list or map entry in `entries`.
-fn children<'a, O: OffsetSizeTrait>(entries: Entries<'a>, offsets: &'a [O]) ->
Entries<'a> {
- Box::new(entries.flat_map(move |(row, index)| {
- (offsets[index].as_usize()..offsets[index + 1].as_usize()).map(move
|child| (row, child))
- }))
+/// An array on the way to a leaf, and how its entries lead into the next one.
+struct Level<'a> {
+ /// A null entry has nothing under it in the column chunk, even where a
null list spans child
+ /// values.
+ nulls: Option<&'a NullBuffer>,
+ children: Children<'a>,
}
-type ValueBytes<'a> = Box<dyn Fn(usize) -> &'a [u8] + 'a>;
+/// The entries of the next array that lie under entry `i` of a level.
+#[derive(Clone, Copy)]
+enum Children<'a> {
+ /// A struct's: entry `i` of its field.
+ Same,
+ /// A list's or a map's: entries `offsets[i]` up to `offsets[i + 1]` of
its elements, keys or
+ /// values.
+ Offsets(&'a [i32]),
+ /// A large list's.
+ LargeOffsets(&'a [i64]),
+ /// A fixed-size list's: the `size` elements from `i * size`.
+ Fixed(usize),
+}
-/// The bytes that tell `array`'s values apart: a fixed-width value's native
bytes, or a string's or
-/// binary's content. `None` for a type an Iceberg schema cannot produce.
-fn value_bytes<'a>(array: &'a dyn Array) -> Option<ValueBytes<'a>> {
- let bytes: ValueBytes<'a> = downcast_primitive_array!(
- array => Box::new(move |i| array.values()[i].to_byte_slice()),
- DataType::Utf8 => {
- let array = array.as_string::<i32>();
- Box::new(move |i| array.value(i).as_bytes())
- }
- DataType::LargeUtf8 => {
- let array = array.as_string::<i64>();
- Box::new(move |i| array.value(i).as_bytes())
- }
- DataType::Utf8View => {
- let array = array.as_string_view();
- Box::new(move |i| array.value(i).as_bytes())
- }
- DataType::Binary => {
- let array = array.as_binary::<i32>();
- Box::new(move |i| array.value(i))
- }
- DataType::LargeBinary => {
- let array = array.as_binary::<i64>();
- Box::new(move |i| array.value(i))
+impl Children<'_> {
+ /// Where the entries under entry `index` start in the next array. Those
under the entries
+ /// before `index` end there too, so this also holds for one past the last
entry.
+ fn start(self, index: usize) -> usize {
+ match self {
+ Children::Same => index,
+ Children::Offsets(offsets) => offsets[index] as usize,
+ Children::LargeOffsets(offsets) => offsets[index] as usize,
+ Children::Fixed(size) => index * size,
}
- DataType::BinaryView => {
- let array = array.as_binary_view();
- Box::new(move |i| array.value(i))
+ }
+}
+
+impl Path<'_> {
+ /// How many entries of the leaf array lie under the batch's first `rows`
rows, null or not.
+ fn entries(&self, rows: usize) -> usize {
+ let (start, end) = self.levels.iter().fold((0, rows), |(start, end),
level| {
+ (level.children.start(start), level.children.start(end))
+ });
+ end - start
+ }
+
+ /// Calls `f` with the leaf index of every entry under the batch's row
`row` whose ancestors
+ /// are all present, in the order the column chunk stores them, until `f`
returns `false`.
+ /// Returns whether every entry was visited. Nothing is expanded ahead, so
a page that ends
+ /// partway through a batch reads no further.
+ fn visit(&self, row: usize, f: &mut impl FnMut(usize) -> bool) -> bool {
+ // A leaf with no levels above it has one entry per row: the row
itself.
+ if self.levels.is_empty() {
+ return f(row);
}
- _ => return None,
- );
- Some(bytes)
+ visit(&self.levels, row, f)
+ }
+}
+
+/// `array`'s nulls, or `None` when it has none.
+fn present_nulls(array: &dyn Array) -> Option<&NullBuffer> {
+ array.nulls().filter(|nulls| nulls.null_count() > 0)
+}
+
+fn visit(levels: &[Level], index: usize, f: &mut impl FnMut(usize) -> bool) ->
bool {
+ let Some((level, below)) = levels.split_first() else {
+ return f(index);
+ };
+ if level.nulls.is_some_and(|nulls| nulls.is_null(index)) {
+ return true;
+ }
+ let mut children = level.children.start(index)..level.children.start(index
+ 1);
+ // The leaf's entries are visited here rather than one call deeper each.
+ if below.is_empty() {
+ children.all(f)
+ } else {
+ children.all(|child| visit(below, child, &mut *f))
+ }
+}
+
+/// Reads fixed-width values by their bits, zero-extended to `K`, which is at
least as wide.
+fn bits<T: ArrowNativeType, K: Bits>(values: &[T]) -> impl Fn(usize) -> (K,
usize) + '_ {
+ move |index| (K::of(values[index]), size_of::<T>())
+}
+
+/// Reads fixed-width values too wide for [`bits`] by their native bytes.
+fn native_bytes<'a, T: ArrowNativeType>(values: &'a [T]) -> impl Fn(usize) ->
(Bytes<'a>, usize) {
+ move |index| Bytes::of(values[index].to_byte_slice())
Review Comment:
Done in c2f688c4a8. The readers now return the plain size: the candidate's
`width` for fixed-width values and `4 + len` for byte arrays.
`FirstPage::width` is gone and `add` takes the size directly. `Bytes::of` went
with it, since nothing else used it. `judge_primitive` takes the width from the
candidate and makes no choice if it is somehow missing.
##########
native/core/src/execution/operators/iceberg_dictionary.rs:
##########
@@ -1253,6 +1677,26 @@ mod tests {
);
}
+ /// The fields of a null struct are not part of the column chunk either.
+ #[test]
+ fn values_under_a_null_struct_are_not_counted() {
Review Comment:
I kept `values_under_a_null_struct_are_not_counted`, for the reason viirya
gives: it is small and pins the struct-null check on its own. I did drop the
large-list and fixed-size-list half of `reads_every_leaf_type`, since `ll` and
`fl` in the sliced test cover those types with nulls and a slice on top. That
is in c2f688c4a8. I read the code to check this and did not run a mutation.
--
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]
