andygrove commented on code in PR #6937:
URL: https://github.com/apache/datafusion-comet/pull/6937#discussion_r4274915000
##########
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:
Added both in f67eab9ef8. `batch_boundaries_do_not_change_the_choice` now
also cuts `corpus(0..20_000)` with `RecordBatch::slice` at the same three batch
sizes. The new `a_sliced_batch_is_read_only_within_its_slice` slices a batch of
a list, a large list, a fixed-size list, a map, nested structs and a list of
structs at offset 1001, with nulls at every level, and checks it against the
same rows materialized.
For the materialized copy I used `take_record_batch` rather than
`concat_batches`. Given a single batch, arrow's `concat` returns
`array.slice(0, len)`, so `concat_batches` would have handed back the same
sliced arrays and the test would have compared the slice with itself.
`take_record_batch` copies the rows into fresh arrays, which is also how the
fanout arm cuts its parts.
Every value outside the slice or under a null is distinct, so reading any of
them overflows the 1 KiB dictionary. I checked the tests against three
mutations: list offsets taken relative to the slice, and null checks that
ignore the null buffer's offset, once at an intermediate level and once at the
leaf. The new test fails on all three, and none of the existing tests caught
any of them.
##########
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:
It does pay off, so I kept it and put the numbers in its doc comment in
f67eab9ef8. I timed `choose` on the benchmark's shapes in one process,
alternating between `Bytes`, plain `&[u8]` keys, and each half of `Bytes` on
its own, for 150 rounds per shape. All of them made the same choices. With
plain slices `maps` took 27% longer, `nested` 24%, `primitives` 12%, `arrays`
11%, `primitives_1k` 9% and `structs` 3%, while `sorted`, where most values
repeat the one before, took 2% less.
Most of it is the length and last-byte check, through the comparison with
`previous` as you guessed: dropping only that cost 16% on `maps` and `nested`.
Their keys, like `k0` and `k1`, mostly have the same length as the key before
them, and slice equality calls `bcmp` whenever the lengths match. Hashing
without the length is worth up to 3%.
--
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]
