alamb commented on code in PR #9544: URL: https://github.com/apache/arrow-rs/pull/9544#discussion_r3024773899
########## 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 { Review Comment: perhaps DataType::is_nested would be ok? Also should it include Dictionary types? ########## 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: Review Comment: I think these should be doc comments (`///` rather than `//`) Also, can you add an explantion of the rationale for this heuristic -- basically that it is trying to figure out a mapping of source union field to target type that is "most likely" (?) I do think this heuristic makes sense ########## 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: technically this will not work if there are repeated field names in the Union (as resolve_variant will return a string and `union_extract` will return the first field with that name) You could potentially avoid this by returning the `type_id` / `index` instead We can also handle it in a follow on PR ########## arrow-cast/src/cast/union.rs: ########## @@ -0,0 +1,470 @@ +// 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, new_null_array}; +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) + ) +} Review Comment: I think the notion of "same logical type" gets a bit subtle -- for example LargeUtf8 and Utf8View might seem like the same logical type, but there are some values that LargeUTF8 can store that Utf8View can not (e.g. individual strings over 4GB) ########## 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. Review Comment: ```suggestion /// Extracts the best-matching variant from a [`UnionArray`] for a given target type, /// and casts it to that type. ``` ########## arrow-cast/src/cast/union.rs: ########## @@ -0,0 +1,470 @@ +// 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, new_null_array}; +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) + ) +} + +// 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. Review Comment: I think the user can have direct control if they want and can use [`union_extract`](https://docs.rs/arrow/latest/arrow/compute/kernels/union_extract/fn.union_extract.html) and `cast` -- 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]
