viirya commented on code in PR #6937:
URL: https://github.com/apache/datafusion-comet/pull/6937#discussion_r4264100481
##########
native/core/src/execution/operators/iceberg_dictionary.rs:
##########
@@ -254,139 +343,230 @@ 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())
+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::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))
- }
- 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())
+}
+
+/// Reads a string or binary array's values by their content.
+fn bytes<'a, A, V>(array: A) -> impl Fn(usize) -> (Bytes<'a>, usize)
+where
+ A: ArrayAccessor<Item = &'a V>,
+ V: AsRef<[u8]> + ?Sized + 'a,
+{
+ move |index| Bytes::of(array.value(index).as_ref())
+}
+
+/// The bits of a fixed-width value, zero-extended to fill this type.
+trait Bits: Copy + Eq + Hash {
+ fn of<T: ArrowNativeType>(value: T) -> Self;
+}
+
+impl Bits for u64 {
+ fn of<T: ArrowNativeType>(value: T) -> Self {
+ let mut bits = [0; 8];
+ let bytes = value.to_byte_slice();
+ bits[..bytes.len()].copy_from_slice(bytes);
+ u64::from_ne_bytes(bits)
+ }
+}
+
+impl Bits for u128 {
+ fn of<T: ArrowNativeType>(value: T) -> Self {
+ let mut bits = [0; 16];
+ let bytes = value.to_byte_slice();
+ bits[..bytes.len()].copy_from_slice(bytes);
+ u128::from_ne_bytes(bits)
+ }
+}
+
+/// A value's bytes. Two of the same length whose last bytes differ are told
apart without
+/// comparing the rest, which spares comparing most distinct values with the
one before them.
+#[derive(Clone, Copy)]
+struct Bytes<'a>(&'a [u8]);
+
+impl<'a> Bytes<'a> {
+ /// `bytes` as a value, and how many there are.
+ fn of(bytes: &'a [u8]) -> (Self, usize) {
+ (Self(bytes), bytes.len())
+ }
+}
+
+impl PartialEq for Bytes<'_> {
+ fn eq(&self, other: &Self) -> bool {
+ self.0.len() == other.0.len() && self.0.last() == other.0.last() &&
self.0 == other.0
Review Comment:
Does this check of the length and last byte, together with the length-free
`Hash` below, measurably help on its own? Slice equality already compares
lengths first. The set only calls `eq` when the tags match, so most of the gain
would have to come from the comparison with `previous`.
If keying strings and binaries by plain `&[u8]` costs little or nothing on
the `primitives` and `sorted` shapes, I'd drop the newtype. Then there's one
less custom `Eq` and `Hash` pair to keep consistent. If it does pay off, a
sentence with the numbers in the doc comment would help the next reader.
##########
native/core/src/execution/operators/iceberg_dictionary.rs:
##########
@@ -254,139 +343,230 @@ 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())
+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 {
Review Comment:
Sliced batches reach `choose` on a real path. When `BatchSlicer::slice` in
`iceberg_write.rs` isn't gathering, it hands the clustered arm
`batch.slice(offset, len)`. This rewrite redoes all of the index arithmetic:
`Children::start`, `Path::entries`, and the sliced null buffers through
`present_nulls`. None of the in-tree chooser tests use a batch with a non-zero
offset, though. `batch_boundaries_do_not_change_the_choice` builds fresh
batches, and the differential run that did cover slicing isn't checked in.
Could `batch_boundaries_do_not_change_the_choice` also cut
`corpus(0..20_000)` with `RecordBatch::slice` at a few offsets? It would also
help to add one nested case with null lists and null structs sliced at an odd
offset, compared against the same rows materialized with `concat_batches`. I
tried both locally and they pass. They would keep this path pinned going
forward.
--
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]
