Skip to content

Commit a0fd769

Browse files
committed
Add support for LargeBinaryArray in chunk iterator helper
1 parent a87d0a4 commit a0fd769

File tree

1 file changed

+18
-4
lines changed

1 file changed

+18
-4
lines changed

crates/store/re_chunk/src/iter.rs

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use arrow::{
44
array::{
55
Array as ArrowArray, ArrayRef as ArrowArrayRef, ArrowPrimitiveType, BinaryArray,
66
BooleanArray as ArrowBooleanArray, FixedSizeListArray as ArrowFixedSizeListArray,
7-
ListArray as ArrowListArray, PrimitiveArray as ArrowPrimitiveArray,
7+
LargeBinaryArray, ListArray as ArrowListArray, PrimitiveArray as ArrowPrimitiveArray,
88
StringArray as ArrowStringArray, StructArray as ArrowStructArray,
99
},
1010
buffer::{BooleanBuffer as ArrowBooleanBuffer, Buffer, ScalarBuffer as ArrowScalarBuffer},
@@ -533,7 +533,7 @@ where
533533
}))
534534
}
535535

536-
// We special case `&[u8]` so that it works both for `List[u8]` and `Binary` arrays.
536+
// We special case `&[u8]` so that it works both for `List[u8]` and `Binary/LargeBinary` arrays.
537537
fn slice_as_u8<'a>(
538538
component_descriptor: ComponentDescriptor,
539539
array: &'a dyn ArrowArray,
@@ -545,14 +545,28 @@ fn slice_as_u8<'a>(
545545
let lengths = offsets.lengths().collect_vec();
546546

547547
// NOTE: No need for validity checks here, `component_spans` already takes care of that.
548-
Either::Left(component_spans.map(move |span| {
548+
Either::Left(Either::Left(component_spans.map(move |span| {
549549
let offsets = &offsets[span.range()];
550550
let lengths = &lengths[span.range()];
551551
izip!(offsets, lengths)
552552
// NOTE: Not an actual clone, just a refbump of the underlying buffer.
553553
.map(|(&idx, &len)| values.clone().slice_with_length(idx as _, len))
554554
.collect_vec()
555-
}))
555+
})))
556+
} else if let Some(binary_array) = array.downcast_array_ref::<LargeBinaryArray>() {
557+
let values = binary_array.values();
558+
let offsets = binary_array.offsets();
559+
let lengths = offsets.lengths().collect_vec();
560+
561+
// NOTE: No need for validity checks here, `component_spans` already takes care of that.
562+
Either::Left(Either::Right(component_spans.map(move |span| {
563+
let offsets = &offsets[span.range()];
564+
let lengths = &lengths[span.range()];
565+
izip!(offsets, lengths)
566+
// NOTE: Not an actual clone, just a refbump of the underlying buffer.
567+
.map(|(&idx, &len)| values.clone().slice_with_length(idx as _, len))
568+
.collect_vec()
569+
})))
556570
} else {
557571
Either::Right(
558572
slice_as_buffer_native::<arrow::array::types::UInt8Type, u8>(

0 commit comments

Comments
 (0)