yongster commented on code in PR #10994:
URL: https://github.com/apache/arrow-rs/pull/10994#discussion_r3985620393
##########
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:
Addressed in e61862695692f7a75cfc18772b0cb50142aa644c: `union_null_type_id`
now uses `find_map`.
##########
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:
A non-null out-of-bounds index is not swallowed here.
Sparse children have the same length as the union, so `values.is_empty()`
means the union itself is empty. A non-null index is then OOB and already
panics in `take_native` via `take_union_type_ids` (or errors earlier if
`check_bounds` is set). This branch is only reached for all-null indices, where
dummy index 0 would itself be OOB.
Documented that invariant next to the empty-child path in e61862695.
##########
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:
Addressed in e61862695692f7a75cfc18772b0cb50142aa644c:
`indices_without_nulls` now uses `Native::ZERO` and `from_iter_values`.
##########
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:
Addressed in e61862695692f7a75cfc18772b0cb50142aa644c: inlined the empty
unused-child construction so the helper and its doc comment are gone.
--
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]
