comphead commented on code in PR #6937:
URL: https://github.com/apache/datafusion-comet/pull/6937#discussion_r4285216229
##########
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:
It looks like `a_sliced_batch_is_read_only_within_its_slice`, added in the
last commit, covers this test now. Its `s` is null on odd rows over values of
their own, and `s.n` is valid on those rows. So a walk that leaves out a struct
level with nulls overflows `s.n`'s 1 KiB dictionary and drops it from
`expected`. `ls` checks the same one level under a list. The `_large_list` and
`_fixed_list` columns of `reads_every_leaf_type` are covered the same way by
`ll` and `fl`, which add nulls and a slice on top.
I traced this by reading the code and have not run the mutation. If it
holds, could this test and the list half of `reads_every_leaf_type` go?
##########
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);
+ };
Review Comment:
This `else` branch never runs. `Path::visit` returns early when there are no
levels, and `visit` only recurses when `below` is not empty, so `levels` is
never empty here. Could the empty check live in one place? For example,
`Path::visit` could call `split_first` itself and pass `level` and `below` in.
##########
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:
The length `bits` and `native_bytes` return is never read. Every fixed-width
candidate has `width: Some(..)`, so `add` ignores `len` for them. It is not the
plain width either: a `Decimal128(9, 2)` value reports 16 here and is written
as a 4-byte INT32. Could the readers return the plain size instead, the
candidate's `width` for fixed-width values and `4 + len` for byte arrays? Then
`add` would not need `FirstPage::width`, and no reader would return a number
nobody reads.
--
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]
