Jefffrey commented on code in PR #10994:
URL: https://github.com/apache/arrow-rs/pull/10994#discussion_r4093190522
##########
arrow-select/src/take.rs:
##########
@@ -434,7 +434,78 @@ fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
}
})
.collect::<ScalarBuffer<_>>();
- Ok(type_ids)
+ Ok((type_ids, Some(null_type_id)))
+}
+
+/// Type id of a union child that can represent a newly introduced null.
+fn union_null_type_id(fields: &UnionFields) -> Result<i8, ArrowError> {
+ fields
+ .iter()
+ .find_map(|(type_id, field)|
field_can_represent_take_null(field).then_some(type_id))
+ .ok_or_else(|| {
+ ArrowError::ComputeError(
+ "Cannot take null indices from a union with no field that can
represent nulls"
+ .into(),
+ )
+ })
+}
+
+/// Whether `field` can physically store a newly introduced null.
+///
+/// Union and RunEndEncoded have no top-level validity bitmap, so a null must
+/// be represented by a descendant that is marked nullable. A field marked
+/// nullable is not sufficient if its nested type cannot store a null.
+fn field_can_represent_take_null(field: &FieldRef) -> bool {
+ if !field.is_nullable() {
+ return false;
+ }
+ match field.data_type() {
+ DataType::RunEndEncoded(_, values) =>
field_can_represent_take_null(values),
+ DataType::Union(fields, _) => fields
+ .iter()
+ .any(|(_, child)| field_can_represent_take_null(child)),
+ _ => true,
+ }
+}
+
+/// Takes a sparse union child for `indices`.
+///
+/// Null take indices select one child that can represent a null
+/// ([`take_union_type_ids`]). Values in the other children at those positions
+/// are unspecified, so they are taken with dummy indices instead of
introducing
+/// nulls that would contradict field metadata.
+fn take_sparse_union_child<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
+ values: &dyn Array,
+ represent_nulls: bool,
+ indices: &PrimitiveArray<IndexType>,
+) -> Result<ArrayRef, ArrowError> {
+ if represent_nulls || indices.null_count() == 0 {
+ return take_impl::<_, CHECKED>(values, indices);
+ }
+
+ if values.is_empty() {
+ // Dummy index 0 is OOB on an empty child. Sparse children have the
same
+ // length as the union, so a non-null index is also OOB and already
+ // panics in `take_native` via [`take_union_type_ids`].
+ return Ok(make_array(
Review Comment:
which reminds me of this PR, it mightve been applicable here
- https://github.com/apache/arrow-rs/pull/9988
though its still a rather niche case
--
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]
