Jefffrey commented on code in PR #10994:
URL: https://github.com/apache/arrow-rs/pull/10994#discussion_r3957910497
##########
arrow-select/src/take.rs:
##########
@@ -437,6 +441,58 @@ fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
Ok(type_ids)
}
+/// Takes a sparse union child for `indices`.
+///
+/// Null take indices are represented by selecting a nullable child
+/// ([`take_union_type_ids`]). Values in the other children at those positions
are
+/// unspecified, so a non-nullable child is taken with dummy indices instead of
+/// introducing nulls that would contradict its field metadata.
+fn take_sparse_union_child<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
+ values: &dyn Array,
+ nullable: bool,
+ indices: &PrimitiveArray<IndexType>,
+) -> Result<ArrayRef, ArrowError> {
+ if nullable || indices.null_count() == 0 {
+ return take_impl::<_, CHECKED>(values, indices);
+ }
+
+ if values.is_empty() {
+ // Null indices would otherwise take dummy index 0, which is OOB.
+ return non_null_unspecified_values(values.data_type(), indices.len());
+ }
Review Comment:
what if `indices` has a non-null index beyond bounds? or is that not a
concern since it should already error for that case earlier?
##########
arrow-select/src/take.rs:
##########
@@ -415,11 +419,11 @@ fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
let null_type_id = fields
.iter()
- .next()
+ .find(|(_, field)| field.is_nullable())
.map(|(type_id, _)| type_id)
.ok_or_else(|| {
Review Comment:
could use `find_map` here
##########
arrow-select/src/take.rs:
##########
@@ -437,6 +441,58 @@ fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
Ok(type_ids)
}
+/// Takes a sparse union child for `indices`.
+///
+/// Null take indices are represented by selecting a nullable child
+/// ([`take_union_type_ids`]). Values in the other children at those positions
are
+/// unspecified, so a non-nullable child is taken with dummy indices instead of
+/// introducing nulls that would contradict its field metadata.
+fn take_sparse_union_child<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
+ values: &dyn Array,
+ nullable: bool,
+ indices: &PrimitiveArray<IndexType>,
+) -> Result<ArrayRef, ArrowError> {
+ if nullable || indices.null_count() == 0 {
+ return take_impl::<_, CHECKED>(values, indices);
+ }
+
+ if values.is_empty() {
+ // Null indices would otherwise take dummy index 0, which is OOB.
+ return non_null_unspecified_values(values.data_type(), indices.len());
+ }
+
+ take_impl::<_, CHECKED>(values, &indices_without_nulls(indices))
+}
+
+/// Replaces null take indices with `0` and drops the null bitmap.
+fn indices_without_nulls<IndexType: ArrowPrimitiveType>(
+ indices: &PrimitiveArray<IndexType>,
+) -> PrimitiveArray<IndexType> {
+ let dummy = IndexType::Native::from_usize(0).unwrap();
+ let mut values = indices.values().to_vec();
+ if let Some(nulls) = indices.nulls() {
+ for (idx, value) in values.iter_mut().enumerate() {
+ if nulls.is_null(idx) {
+ *value = dummy;
+ }
+ }
+ }
+ PrimitiveArray::new(ScalarBuffer::from(values), None)
+}
+
+/// Builds an all-valid array of `len` whose values are unspecified.
+///
+/// Used for unused sparse-union child slots when the source child is empty.
+fn non_null_unspecified_values(data_type: &DataType, len: usize) ->
Result<ArrayRef, ArrowError> {
Review Comment:
personally id just inline this; that way we dont need a doc comment to
explain for what use case this is for
##########
arrow-select/src/take.rs:
##########
@@ -437,6 +441,58 @@ fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
Ok(type_ids)
}
+/// Takes a sparse union child for `indices`.
+///
+/// Null take indices are represented by selecting a nullable child
+/// ([`take_union_type_ids`]). Values in the other children at those positions
are
+/// unspecified, so a non-nullable child is taken with dummy indices instead of
+/// introducing nulls that would contradict its field metadata.
+fn take_sparse_union_child<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
+ values: &dyn Array,
+ nullable: bool,
+ indices: &PrimitiveArray<IndexType>,
+) -> Result<ArrayRef, ArrowError> {
+ if nullable || indices.null_count() == 0 {
+ return take_impl::<_, CHECKED>(values, indices);
+ }
+
+ if values.is_empty() {
+ // Null indices would otherwise take dummy index 0, which is OOB.
+ return non_null_unspecified_values(values.data_type(), indices.len());
+ }
+
+ take_impl::<_, CHECKED>(values, &indices_without_nulls(indices))
+}
+
+/// Replaces null take indices with `0` and drops the null bitmap.
+fn indices_without_nulls<IndexType: ArrowPrimitiveType>(
+ indices: &PrimitiveArray<IndexType>,
+) -> PrimitiveArray<IndexType> {
+ let dummy = IndexType::Native::from_usize(0).unwrap();
+ let mut values = indices.values().to_vec();
+ if let Some(nulls) = indices.nulls() {
+ for (idx, value) in values.iter_mut().enumerate() {
+ if nulls.is_null(idx) {
+ *value = dummy;
+ }
+ }
+ }
+ PrimitiveArray::new(ScalarBuffer::from(values), None)
Review Comment:
```suggestion
let dummy = IndexType::Native::ZERO;
let normalized = indices.iter().map(|idx| idx.unwrap_or(dummy));
PrimitiveArray::from_iter_values(normalized)
```
--
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]
