Jefffrey commented on code in PR #10812:
URL: https://github.com/apache/arrow-rs/pull/10812#discussion_r3877551754
##########
arrow-select/src/take.rs:
##########
@@ -652,76 +653,155 @@ where
OffsetType::Native: OffsetSizeTrait,
PrimitiveArray<OffsetType>: From<Vec<OffsetType::Native>>,
{
- let list_offsets = values.value_offsets();
+ let src_offsets = values.value_offsets();
let child_data = values.values().to_data();
let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
- let mut new_offsets = Vec::with_capacity(indices.len() + 1);
- new_offsets.push(OffsetType::Native::zero());
+ let mut dst_offsets = Vec::with_capacity(indices.len() + 1);
+ dst_offsets.push(OffsetType::Native::zero());
- let use_nulls = child_data.null_count() > 0;
+ let field = match values.data_type() {
+ DataType::List(fld) | DataType::LargeList(fld) => fld.clone(),
+ dtype => unreachable!("take_list called with non-list data type
{dtype}"),
+ };
+
+ let is_primitive_child = child_data.null_count() == 0 &&
child_data.data_type().is_primitive();
+
+ if is_primitive_child {
+ let values_buf = &child_data.buffers()[0];
+ let bytes_per_value =
child_data.data_type().primitive_width().unwrap_or(0);
+ let child_buf_offset = child_data.offset() * bytes_per_value;
Review Comment:
we can clean this up to be similar to fixedsizelist path
- use `value_field`
- use `let Some(_) = ` pattern to avoid need for unwrap on primitive_width
##########
arrow-select/src/take.rs:
##########
@@ -652,76 +653,155 @@ where
OffsetType::Native: OffsetSizeTrait,
PrimitiveArray<OffsetType>: From<Vec<OffsetType::Native>>,
{
- let list_offsets = values.value_offsets();
+ let src_offsets = values.value_offsets();
let child_data = values.values().to_data();
let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
- let mut new_offsets = Vec::with_capacity(indices.len() + 1);
- new_offsets.push(OffsetType::Native::zero());
+ let mut dst_offsets = Vec::with_capacity(indices.len() + 1);
+ dst_offsets.push(OffsetType::Native::zero());
- let use_nulls = child_data.null_count() > 0;
+ let field = match values.data_type() {
+ DataType::List(fld) | DataType::LargeList(fld) => fld.clone(),
+ dtype => unreachable!("take_list called with non-list data type
{dtype}"),
+ };
+
+ let is_primitive_child = child_data.null_count() == 0 &&
child_data.data_type().is_primitive();
+
+ if is_primitive_child {
+ let values_buf = &child_data.buffers()[0];
+ let bytes_per_value =
child_data.data_type().primitive_width().unwrap_or(0);
+ let child_buf_offset = child_data.offset() * bytes_per_value;
+
+ let avg_row_len = child_data
+ .len()
+ .checked_div(values.len().max(1))
+ .unwrap_or(0);
+ let mut dst_buf = MutableBuffer::new(
+ avg_row_len
+ .saturating_mul(indices.len())
+ .saturating_mul(bytes_per_value),
+ );
+
+ let mut child_len = OffsetType::Native::zero();
+
+ match nulls.as_ref().filter(|n| n.null_count() > 0) {
+ None => {
+ for &idx in indices.values() {
+ let row = idx.as_usize();
+ let start = child_buf_offset + src_offsets[row].as_usize()
* bytes_per_value;
+ let end = child_buf_offset + src_offsets[row +
1].as_usize() * bytes_per_value;
+ dst_buf.extend_from_slice(&values_buf[start..end]);
+ child_len = child_len
+ .checked_add(&(src_offsets[row + 1] -
src_offsets[row]))
+ .ok_or_else(||
ArrowError::OffsetOverflowError(child_len.as_usize()))?;
+ dst_offsets.push(child_len);
+ }
+ }
+ Some(valid) => {
+ let mut prev = 0;
+ for vidx in valid.valid_indices() {
+ // Fill offsets for null values between the two valid
indices.
+ if prev < vidx {
+ dst_offsets.extend(std::iter::repeat_n(child_len, vidx
- prev));
+ }
+ let row = if CHECKED {
+ indices.value(vidx).as_usize()
+ } else {
+ // SAFETY: `vidx` comes from validity bitmap over
`indices`, so in-bounds.
Review Comment:
safety comment here might need to be extended to explain why its not safe in
all cases, i.e. why we have a `CHECKED` branch above (otherwise it sounds like
this is always safe so its confusing why its gated by `CHECKED`)
--
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]
