friendlymatthew commented on code in PR #9544: URL: https://github.com/apache/arrow-rs/pull/9544#discussion_r3039735363
########## arrow-cast/src/cast/union.rs: ########## @@ -0,0 +1,495 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Cast support for union arrays. + +use crate::cast::can_cast_types; +use crate::cast_with_options; +use arrow_array::{Array, ArrayRef, UnionArray}; +use arrow_schema::{ArrowError, DataType, FieldRef, UnionFields}; +use arrow_select::union_extract::union_extract; + +use super::CastOptions; + +// this is used during variant selection to prefer a "close" type over a distant cast +// for example: when targeting Utf8View, a Utf8 variant is preferred over Int32 despite both being castable +fn same_type_family(a: &DataType, b: &DataType) -> bool { + use DataType::*; + matches!( + (a, b), + (Utf8 | LargeUtf8 | Utf8View, Utf8 | LargeUtf8 | Utf8View) + | ( + Binary | LargeBinary | BinaryView, + Binary | LargeBinary | BinaryView + ) + | (Int8 | Int16 | Int32 | Int64, Int8 | Int16 | Int32 | Int64) + | ( + UInt8 | UInt16 | UInt32 | UInt64, + UInt8 | UInt16 | UInt32 | UInt64 + ) + | (Float16 | Float32 | Float64, Float16 | Float32 | Float64) + ) +} + +fn is_complex_container(dt: &DataType) -> bool { + use DataType::*; + matches!( + dt, + List(_) + | LargeList(_) + | ListView(_) + | LargeListView(_) + | FixedSizeList(_, _) + | Struct(_) + | Map(_, _) + ) +} + +// variant selection heuristic — 3 passes with decreasing specificity: +// +// first pass: field type == target type +// second pass: field and target are in the same equivalence class +// (e.g., Utf8 and Utf8View are both strings) +// third pass: field can be cast to target +// note: this is the most permissive and may lose information +// also, the matching logic is greedy so it will pick the first 'castable' variant +// +// each pass picks the first matching variant by type_id order. +pub(crate) fn resolve_variant<'a>( + fields: &'a UnionFields, + target_type: &DataType, +) -> Option<&'a FieldRef> { + fields + .iter() + .find(|(_, f)| f.data_type() == target_type) + .or_else(|| { + fields + .iter() + .find(|(_, f)| same_type_family(f.data_type(), target_type)) + }) + .or_else(|| { + // skip complex container types in pass 3 — union extraction introduces nulls, + // and casting nullable arrays to containers like List/Struct/Map can fail when + // inner fields are non-nullable. + if is_complex_container(target_type) { + return None; + } + fields + .iter() + .find(|(_, f)| can_cast_types(f.data_type(), target_type)) + }) + .map(|(_, f)| f) +} + +/// Extracts the best-matching variant from a union array for a given target type, +/// and casts it to that type. +/// +/// Rows where a different variant is active become NULL. +/// If no variant matches, returns a null array. +/// +/// # Example +/// +/// ``` +/// # use std::sync::Arc; +/// # use arrow_schema::{DataType, Field, UnionFields}; +/// # use arrow_array::{UnionArray, StringArray, Int32Array, Array}; +/// # use arrow_cast::cast::union_extract_by_type; +/// # use arrow_cast::CastOptions; +/// let fields = UnionFields::try_new( +/// [0, 1], +/// [ +/// Field::new("int", DataType::Int32, true), +/// Field::new("str", DataType::Utf8, true), +/// ], +/// ).unwrap(); +/// +/// let union = UnionArray::try_new( +/// fields, +/// vec![0, 1, 0].into(), +/// None, +/// vec![ +/// Arc::new(Int32Array::from(vec![Some(42), None, Some(99)])), +/// Arc::new(StringArray::from(vec![None, Some("hello"), None])), +/// ], +/// ) +/// .unwrap(); +/// +/// // extract the Utf8 variant and cast to Utf8View +/// let result = union_extract_by_type(&union, &DataType::Utf8View, &CastOptions::default()).unwrap(); +/// assert_eq!(result.data_type(), &DataType::Utf8View); +/// assert!(result.is_null(0)); // Int32 row -> NULL +/// assert!(!result.is_null(1)); // Utf8 row -> "hello" +/// assert!(result.is_null(2)); // Int32 row -> NULL +/// ``` +pub fn union_extract_by_type( + union_array: &UnionArray, + target_type: &DataType, + cast_options: &CastOptions, +) -> Result<ArrayRef, ArrowError> { + let fields = match union_array.data_type() { + DataType::Union(fields, _) => fields, + _ => unreachable!("union_extract_by_type called on non-union array"), + }; + + let Some(field) = resolve_variant(fields, target_type) else { + return Err(ArrowError::CastError(format!( + "cannot cast Union with fields {} to {}", + fields + .iter() + .map(|(_, f)| f.data_type().to_string()) + .collect::<Vec<_>>() + .join(", "), + target_type + ))); + }; + + let extracted = union_extract(union_array, field.name())?; Review Comment: Filed: https://github.com/apache/arrow-rs/issues/9664 -- 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]
