Jefffrey commented on code in PR #10994:
URL: https://github.com/apache/arrow-rs/pull/10994#discussion_r4093023697
##########
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:
theres an edge case here identified by codex where if this child is a run
array with non-nullable values, this still creates a null array for it since
run array doesnt have a null bitmap (so removing the `nulls` doesnt do anything)
but this is a pretty extreme case, im not too bothered by it (can deal in a
followup if we want)
--
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]
