RyanMarcus commented on code in PR #6752: URL: https://github.com/apache/arrow-rs/pull/6752#discussion_r1850602203
########## arrow-cast/src/cast/runend.rs: ########## @@ -0,0 +1,330 @@ +// 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. + +use std::sync::Arc; + +use arrow_array::{ + builder::ArrayBuilder, make_array, types::RunEndIndexType, Array, ArrayRef, ArrowPrimitiveType, + Float16Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, + PrimitiveArray, RunArray, TypedRunArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array, +}; +use arrow_buffer::{ArrowNativeType, NullBufferBuilder}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::{ArrowError, DataType}; + +use crate::cast_with_options; + +use super::CastOptions; + +/// Attempt to cast a run-encoded array into a new type. +/// +/// `K` is the *current* run end index type +pub(crate) fn run_end_cast<K: RunEndIndexType>( + array: &dyn Array, + to_type: &DataType, + cast_options: &CastOptions, +) -> Result<ArrayRef, ArrowError> { + let ree_array = array + .as_any() + .downcast_ref::<RunArray<K>>() + .ok_or_else(|| { + ArrowError::ComputeError( + "Internal Error: Cannot cast run end array to RunArray of the expected type" + .to_string(), + ) + })?; + + let curr_value_type = ree_array.values().data_type(); + + match (curr_value_type, to_type) { + // Potentially convert to a new value or run end type + (_, DataType::RunEndEncoded(re_t, dt)) => { + let values = cast_with_options(ree_array.values(), dt.data_type(), cast_options)?; + let re = PrimitiveArray::<K>::new(ree_array.run_ends().inner().clone(), None); + let re = cast_with_options(&re, re_t.data_type(), cast_options)?; + + // TODO: we shouldn't need to validate the new run length array + // since we can assume we are converting from a valid one, but + // there's no "unchecked" variant yet + let result: Arc<dyn Array> = match re.data_type() { + DataType::Int16 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int16Array>().unwrap(), + &values, + )?), + DataType::Int32 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int32Array>().unwrap(), + &values, + )?), + DataType::Int64 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int64Array>().unwrap(), + &values, + )?), + _ => Err(ArrowError::ComputeError(format!( + "Invalid run end type requested during cast: {:?}", + re.data_type() + )))?, + }; + + Ok(result.slice(ree_array.run_ends().offset(), ree_array.run_ends().len())) + } + // We match against the native types in order to use an optimized kernel + (DataType::Int8, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int8Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int64Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt8, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt8Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt64Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float64Array>().unwrap()), + to_type, + cast_options, + ), + // For all other types, we use an interpretation-based approach + _ => { + // TODO this could be somewhat inefficent, since the run encoded Review Comment: Or compilation, I guess. ########## arrow-cast/src/cast/runend.rs: ########## @@ -0,0 +1,330 @@ +// 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. + +use std::sync::Arc; + +use arrow_array::{ + builder::ArrayBuilder, make_array, types::RunEndIndexType, Array, ArrayRef, ArrowPrimitiveType, + Float16Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, + PrimitiveArray, RunArray, TypedRunArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array, +}; +use arrow_buffer::{ArrowNativeType, NullBufferBuilder}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::{ArrowError, DataType}; + +use crate::cast_with_options; + +use super::CastOptions; + +/// Attempt to cast a run-encoded array into a new type. +/// +/// `K` is the *current* run end index type +pub(crate) fn run_end_cast<K: RunEndIndexType>( + array: &dyn Array, + to_type: &DataType, + cast_options: &CastOptions, +) -> Result<ArrayRef, ArrowError> { + let ree_array = array + .as_any() + .downcast_ref::<RunArray<K>>() + .ok_or_else(|| { + ArrowError::ComputeError( + "Internal Error: Cannot cast run end array to RunArray of the expected type" + .to_string(), + ) + })?; + + let curr_value_type = ree_array.values().data_type(); + + match (curr_value_type, to_type) { + // Potentially convert to a new value or run end type + (_, DataType::RunEndEncoded(re_t, dt)) => { + let values = cast_with_options(ree_array.values(), dt.data_type(), cast_options)?; + let re = PrimitiveArray::<K>::new(ree_array.run_ends().inner().clone(), None); + let re = cast_with_options(&re, re_t.data_type(), cast_options)?; + + // TODO: we shouldn't need to validate the new run length array + // since we can assume we are converting from a valid one, but + // there's no "unchecked" variant yet + let result: Arc<dyn Array> = match re.data_type() { + DataType::Int16 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int16Array>().unwrap(), + &values, + )?), + DataType::Int32 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int32Array>().unwrap(), + &values, + )?), + DataType::Int64 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int64Array>().unwrap(), + &values, + )?), + _ => Err(ArrowError::ComputeError(format!( + "Invalid run end type requested during cast: {:?}", + re.data_type() + )))?, + }; + + Ok(result.slice(ree_array.run_ends().offset(), ree_array.run_ends().len())) + } + // We match against the native types in order to use an optimized kernel + (DataType::Int8, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int8Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int64Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt8, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt8Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt64Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float64Array>().unwrap()), + to_type, + cast_options, + ), + // For all other types, we use an interpretation-based approach + _ => { + // TODO this could be somewhat inefficent, since the run encoded + // array is initially transformed into a flat array of the same + // type, then casted to the (potentially) new type. For example, + // casting a run encoded array of Float32 to Float64 will first + // create a primitive array of Float32s, then convert that primitive + // array to Float64. + cast_with_options(&run_array_to_flat(ree_array)?, to_type, cast_options) + } + } +} + +/// "Unroll" a run-end encoded array of primitive values into a primitive array. +/// This function should be efficient for long run lenghts due to the use of +/// Builder's `append_value_n`. Uses `PrimitiveBuilder`, so does not do any +/// interpretation. +fn typed_run_array_to_primitive<R: RunEndIndexType, T: ArrowPrimitiveType>( + arr: TypedRunArray<R, PrimitiveArray<T>>, +) -> ArrayRef { + let mut builder = PrimitiveArray::<T>::builder( Review Comment: Ah, good idea -- this indeed reduced latency by a few ms! ########## arrow-cast/src/cast/runend.rs: ########## @@ -0,0 +1,330 @@ +// 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. + +use std::sync::Arc; + +use arrow_array::{ + builder::ArrayBuilder, make_array, types::RunEndIndexType, Array, ArrayRef, ArrowPrimitiveType, + Float16Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, + PrimitiveArray, RunArray, TypedRunArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array, +}; +use arrow_buffer::{ArrowNativeType, NullBufferBuilder}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::{ArrowError, DataType}; + +use crate::cast_with_options; + +use super::CastOptions; + +/// Attempt to cast a run-encoded array into a new type. +/// +/// `K` is the *current* run end index type +pub(crate) fn run_end_cast<K: RunEndIndexType>( + array: &dyn Array, + to_type: &DataType, + cast_options: &CastOptions, +) -> Result<ArrayRef, ArrowError> { + let ree_array = array + .as_any() + .downcast_ref::<RunArray<K>>() + .ok_or_else(|| { + ArrowError::ComputeError( + "Internal Error: Cannot cast run end array to RunArray of the expected type" + .to_string(), + ) + })?; + + let curr_value_type = ree_array.values().data_type(); + + match (curr_value_type, to_type) { + // Potentially convert to a new value or run end type + (_, DataType::RunEndEncoded(re_t, dt)) => { + let values = cast_with_options(ree_array.values(), dt.data_type(), cast_options)?; + let re = PrimitiveArray::<K>::new(ree_array.run_ends().inner().clone(), None); + let re = cast_with_options(&re, re_t.data_type(), cast_options)?; + + // TODO: we shouldn't need to validate the new run length array + // since we can assume we are converting from a valid one, but + // there's no "unchecked" variant yet + let result: Arc<dyn Array> = match re.data_type() { + DataType::Int16 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int16Array>().unwrap(), + &values, + )?), + DataType::Int32 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int32Array>().unwrap(), + &values, + )?), + DataType::Int64 => Arc::new(RunArray::try_new( + re.as_any().downcast_ref::<Int64Array>().unwrap(), + &values, + )?), + _ => Err(ArrowError::ComputeError(format!( + "Invalid run end type requested during cast: {:?}", + re.data_type() + )))?, + }; + + Ok(result.slice(ree_array.run_ends().offset(), ree_array.run_ends().len())) + } + // We match against the native types in order to use an optimized kernel + (DataType::Int8, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int8Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Int64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Int64Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt8, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt8Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::UInt64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<UInt64Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float16, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float16Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float32, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float32Array>().unwrap()), + to_type, + cast_options, + ), + (DataType::Float64, _) => cast_with_options( + &typed_run_array_to_primitive(ree_array.downcast::<Float64Array>().unwrap()), + to_type, + cast_options, + ), + // For all other types, we use an interpretation-based approach + _ => { + // TODO this could be somewhat inefficent, since the run encoded + // array is initially transformed into a flat array of the same + // type, then casted to the (potentially) new type. For example, + // casting a run encoded array of Float32 to Float64 will first + // create a primitive array of Float32s, then convert that primitive + // array to Float64. + cast_with_options(&run_array_to_flat(ree_array)?, to_type, cast_options) + } + } +} + +/// "Unroll" a run-end encoded array of primitive values into a primitive array. +/// This function should be efficient for long run lenghts due to the use of +/// Builder's `append_value_n`. Uses `PrimitiveBuilder`, so does not do any +/// interpretation. +fn typed_run_array_to_primitive<R: RunEndIndexType, T: ArrowPrimitiveType>( + arr: TypedRunArray<R, PrimitiveArray<T>>, +) -> ArrayRef { + let mut builder = PrimitiveArray::<T>::builder( + arr.run_ends() + .values() + .last() + .map(|end| end.as_usize()) + .unwrap_or(0), + ); + + // copy all values into the builder + let mut last = 0; + for (run_end, val) in arr + .run_ends() + .values() + .iter() + .zip(arr.values().values().iter().copied()) + { + let run_end = run_end.as_usize(); + let run_length = run_end - last; + builder.append_value_n(val, run_length); + last = run_end; + } + let mut result = builder.finish(); + + // if we have null values, decode them as well + if let Some(null_buffer) = arr.values().nulls() { + let mut nbb = NullBufferBuilder::new(builder.len()); + + let mut last = 0; + for (run_end, val) in arr.run_ends().values().iter().zip(null_buffer.iter()) { + let run_end = run_end.as_usize(); + let run_length = run_end - last; + if val { + nbb.append_n_non_nulls(run_length); + } else { + nbb.append_n_nulls(run_length); + } + + last = run_end; + } + + let nb = nbb.finish(); + result = PrimitiveArray::<T>::new(result.values().clone(), nb); + } + + // TODO: this slice could be optimized by only copying the relevant parts of Review Comment: I was having a hard time getting the logic right for a case like this: run ends: [5, 10] slice offset: 3 len: 5 ... because you have to take 2 of the first value and then 3 of the second value. I can give it a shot if you want but I think it'll just be a source of bugs. -- 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]
