alamb commented on code in PR #2838: URL: https://github.com/apache/arrow-rs/pull/2838#discussion_r993837739
########## arrow/src/compute/kernels/interleave.rs: ########## @@ -0,0 +1,218 @@ +// 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 arrow_array::{make_array, new_empty_array, Array, ArrayRef}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::ArrowError; + +/// +/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values. Review Comment: ❤️ for the ascii art (lol though I am biased) cc @Dandandan and @yjshen ########## arrow/src/compute/kernels/interleave.rs: ########## @@ -0,0 +1,218 @@ +// 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 arrow_array::{make_array, new_empty_array, Array, ArrayRef}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::ArrowError; + +/// +/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values. +/// +/// Each element in `indices` is a pair of `usize` with the first identifying the index +/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`] +/// +/// ```text +/// ┌─────────────────┐ ┌─────────┐ ┌─────────────────┐ +/// │ A │ │ (0, 0) │ interleave( │ A │ +/// ├─────────────────┤ ├─────────┤ [values0, values1], ├─────────────────┤ +/// │ D │ │ (1, 0) │ indices │ B │ +/// └─────────────────┘ ├─────────┤ ) ├─────────────────┤ +/// values array 0 │ (1, 1) │ ─────────────────────────▶ │ C │ +/// ├─────────┤ ├─────────────────┤ +/// │ (0, 1) │ │ D │ +/// └─────────┘ └─────────────────┘ +/// ┌─────────────────┐ indices +/// │ B │ array +/// ├─────────────────┤ result +/// │ C │ +/// ├─────────────────┤ +/// │ E │ +/// └─────────────────┘ +/// values array 1 +/// ``` +/// +/// For selecting values by index from a single array see [compute::take](crate::compute::take) +/// +/// # Panics +/// +/// Panics if the arrays do not have the same data type or `values` is empty +pub fn interleave( + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result<ArrayRef, ArrowError> { + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + let data_type = values[0].data_type(); + + if values + .iter() + .skip(1) + .any(|array| array.data_type() != data_type) + { + return Err(ArrowError::InvalidArgumentError( + "It is not possible to interleave arrays of different data types." + .to_string(), + )); + } + + if indices.is_empty() { + return Ok(new_empty_array(data_type)); + } + + // TODO: Add specialized implementations (#1523) Review Comment: It seems like the specialized implementation should be tracked by a different ticket than the initial kernel, right? ########## arrow/src/compute/kernels/interleave.rs: ########## @@ -0,0 +1,218 @@ +// 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 arrow_array::{make_array, new_empty_array, Array, ArrayRef}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::ArrowError; + +/// +/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values. +/// +/// Each element in `indices` is a pair of `usize` with the first identifying the index +/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`] +/// +/// ```text +/// ┌─────────────────┐ ┌─────────┐ ┌─────────────────┐ +/// │ A │ │ (0, 0) │ interleave( │ A │ +/// ├─────────────────┤ ├─────────┤ [values0, values1], ├─────────────────┤ +/// │ D │ │ (1, 0) │ indices │ B │ +/// └─────────────────┘ ├─────────┤ ) ├─────────────────┤ +/// values array 0 │ (1, 1) │ ─────────────────────────▶ │ C │ +/// ├─────────┤ ├─────────────────┤ +/// │ (0, 1) │ │ D │ +/// └─────────┘ └─────────────────┘ +/// ┌─────────────────┐ indices +/// │ B │ array +/// ├─────────────────┤ result +/// │ C │ +/// ├─────────────────┤ +/// │ E │ +/// └─────────────────┘ +/// values array 1 +/// ``` +/// +/// For selecting values by index from a single array see [compute::take](crate::compute::take) +/// +/// # Panics +/// +/// Panics if the arrays do not have the same data type or `values` is empty +pub fn interleave( + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result<ArrayRef, ArrowError> { + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + let data_type = values[0].data_type(); + + if values + .iter() + .skip(1) + .any(|array| array.data_type() != data_type) + { + return Err(ArrowError::InvalidArgumentError( + "It is not possible to interleave arrays of different data types." + .to_string(), + )); + } + + if indices.is_empty() { + return Ok(new_empty_array(data_type)); + } + + // TODO: Add specialized implementations (#1523) + + interleave_fallback(values, indices) +} + +/// Fallback implementation of interleave using [`MutableArrayData`] +fn interleave_fallback( + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result<ArrayRef, ArrowError> { + let arrays: Vec<_> = values.iter().map(|x| x.data()).collect(); + let mut array_data = MutableArrayData::new(arrays, false, indices.len()); + + let mut cur_array = indices[0].0; + let mut start_row_idx = indices[0].1; + let mut end_row_idx = start_row_idx + 1; + + for (array, row) in indices.iter().skip(1).copied() { + if array == cur_array && row == end_row_idx { + // subsequent row in same batch + end_row_idx += 1; + continue; + } + + // emit current batch of rows for current buffer + array_data.extend(cur_array, start_row_idx, end_row_idx); + + // start new batch of rows + cur_array = array; + start_row_idx = row; + end_row_idx = start_row_idx + 1; + } + + // emit final batch of rows + array_data.extend(cur_array, start_row_idx, end_row_idx); + Ok(make_array(array_data.freeze())) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::builder::{Int32Builder, ListBuilder}; + use arrow_array::cast::{as_primitive_array, as_string_array}; + use arrow_array::types::Int32Type; + use arrow_array::{Int32Array, ListArray, StringArray}; + + #[test] + fn test_primitive() { + let a = Int32Array::from_iter_values([1, 2, 3, 4]); + let b = Int32Array::from_iter_values([5, 6, 7]); + let c = Int32Array::from_iter_values([8, 9, 10]); + let values = + interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap(); + let v = as_primitive_array::<Int32Type>(&values); + assert_eq!(v.values(), &[4, 4, 10, 8, 6]); + } + + #[test] + fn test_primitive_nulls() { + let a = Int32Array::from_iter_values([1, 2, 3, 4]); + let b = Int32Array::from_iter([Some(1), Some(4), None]); + let values = + interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (0, 2)]).unwrap(); + let v: Vec<_> = as_primitive_array::<Int32Type>(&values) + .into_iter() + .collect(); + assert_eq!(&v, &[Some(2), None, None, Some(4), Some(3)]) + } + + #[test] + fn test_primitive_empty() { + let a = Int32Array::from_iter_values([1, 2, 3, 4]); + let v = interleave(&[&a], &[]).unwrap(); + assert!(v.is_empty()); Review Comment: ```suggestion assert!(v.is_empty()); assert!(v.data_type(), DataType::Int32); ``` ########## arrow/src/compute/kernels/interleave.rs: ########## @@ -0,0 +1,218 @@ +// 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 arrow_array::{make_array, new_empty_array, Array, ArrayRef}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::ArrowError; + +/// +/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values. +/// +/// Each element in `indices` is a pair of `usize` with the first identifying the index +/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`] +/// +/// ```text +/// ┌─────────────────┐ ┌─────────┐ ┌─────────────────┐ +/// │ A │ │ (0, 0) │ interleave( │ A │ +/// ├─────────────────┤ ├─────────┤ [values0, values1], ├─────────────────┤ +/// │ D │ │ (1, 0) │ indices │ B │ +/// └─────────────────┘ ├─────────┤ ) ├─────────────────┤ +/// values array 0 │ (1, 1) │ ─────────────────────────▶ │ C │ +/// ├─────────┤ ├─────────────────┤ +/// │ (0, 1) │ │ D │ +/// └─────────┘ └─────────────────┘ +/// ┌─────────────────┐ indices +/// │ B │ array +/// ├─────────────────┤ result +/// │ C │ +/// ├─────────────────┤ +/// │ E │ +/// └─────────────────┘ +/// values array 1 +/// ``` +/// +/// For selecting values by index from a single array see [compute::take](crate::compute::take) +/// +/// # Panics +/// +/// Panics if the arrays do not have the same data type or `values` is empty +pub fn interleave( + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result<ArrayRef, ArrowError> { + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + let data_type = values[0].data_type(); + + if values + .iter() + .skip(1) + .any(|array| array.data_type() != data_type) + { + return Err(ArrowError::InvalidArgumentError( + "It is not possible to interleave arrays of different data types." + .to_string(), + )); + } + + if indices.is_empty() { + return Ok(new_empty_array(data_type)); Review Comment: 👍 ########## arrow/src/compute/kernels/interleave.rs: ########## @@ -0,0 +1,218 @@ +// 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 arrow_array::{make_array, new_empty_array, Array, ArrayRef}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::ArrowError; + +/// +/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values. +/// +/// Each element in `indices` is a pair of `usize` with the first identifying the index +/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`] +/// +/// ```text +/// ┌─────────────────┐ ┌─────────┐ ┌─────────────────┐ +/// │ A │ │ (0, 0) │ interleave( │ A │ +/// ├─────────────────┤ ├─────────┤ [values0, values1], ├─────────────────┤ +/// │ D │ │ (1, 0) │ indices │ B │ +/// └─────────────────┘ ├─────────┤ ) ├─────────────────┤ +/// values array 0 │ (1, 1) │ ─────────────────────────▶ │ C │ +/// ├─────────┤ ├─────────────────┤ +/// │ (0, 1) │ │ D │ +/// └─────────┘ └─────────────────┘ +/// ┌─────────────────┐ indices +/// │ B │ array +/// ├─────────────────┤ result +/// │ C │ +/// ├─────────────────┤ +/// │ E │ +/// └─────────────────┘ +/// values array 1 +/// ``` +/// +/// For selecting values by index from a single array see [compute::take](crate::compute::take) +/// +/// # Panics Review Comment: The function appears to return an error (rather than `panic`) in these two cases ########## arrow/src/compute/kernels/interleave.rs: ########## @@ -0,0 +1,218 @@ +// 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 arrow_array::{make_array, new_empty_array, Array, ArrayRef}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::ArrowError; + +/// +/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values. +/// +/// Each element in `indices` is a pair of `usize` with the first identifying the index +/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`] +/// +/// ```text +/// ┌─────────────────┐ ┌─────────┐ ┌─────────────────┐ +/// │ A │ │ (0, 0) │ interleave( │ A │ +/// ├─────────────────┤ ├─────────┤ [values0, values1], ├─────────────────┤ +/// │ D │ │ (1, 0) │ indices │ B │ +/// └─────────────────┘ ├─────────┤ ) ├─────────────────┤ +/// values array 0 │ (1, 1) │ ─────────────────────────▶ │ C │ +/// ├─────────┤ ├─────────────────┤ +/// │ (0, 1) │ │ D │ +/// └─────────┘ └─────────────────┘ +/// ┌─────────────────┐ indices +/// │ B │ array +/// ├─────────────────┤ result +/// │ C │ +/// ├─────────────────┤ +/// │ E │ +/// └─────────────────┘ +/// values array 1 +/// ``` +/// +/// For selecting values by index from a single array see [compute::take](crate::compute::take) +/// +/// # Panics +/// +/// Panics if the arrays do not have the same data type or `values` is empty +pub fn interleave( + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result<ArrayRef, ArrowError> { + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + let data_type = values[0].data_type(); + + if values + .iter() + .skip(1) + .any(|array| array.data_type() != data_type) + { + return Err(ArrowError::InvalidArgumentError( + "It is not possible to interleave arrays of different data types." + .to_string(), Review Comment: ```suggestion format!("It is not possible to interleave arrays of different data types ({:?} and {:?})" array.data_type(), data_type) ``` ########## arrow/src/compute/kernels/interleave.rs: ########## @@ -0,0 +1,218 @@ +// 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 arrow_array::{make_array, new_empty_array, Array, ArrayRef}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::ArrowError; + +/// +/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values. +/// +/// Each element in `indices` is a pair of `usize` with the first identifying the index +/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`] +/// +/// ```text +/// ┌─────────────────┐ ┌─────────┐ ┌─────────────────┐ +/// │ A │ │ (0, 0) │ interleave( │ A │ +/// ├─────────────────┤ ├─────────┤ [values0, values1], ├─────────────────┤ +/// │ D │ │ (1, 0) │ indices │ B │ +/// └─────────────────┘ ├─────────┤ ) ├─────────────────┤ +/// values array 0 │ (1, 1) │ ─────────────────────────▶ │ C │ +/// ├─────────┤ ├─────────────────┤ +/// │ (0, 1) │ │ D │ +/// └─────────┘ └─────────────────┘ +/// ┌─────────────────┐ indices +/// │ B │ array +/// ├─────────────────┤ result +/// │ C │ +/// ├─────────────────┤ +/// │ E │ +/// └─────────────────┘ +/// values array 1 +/// ``` +/// +/// For selecting values by index from a single array see [compute::take](crate::compute::take) +/// +/// # Panics +/// +/// Panics if the arrays do not have the same data type or `values` is empty +pub fn interleave( + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result<ArrayRef, ArrowError> { + if values.is_empty() { + return Err(ArrowError::InvalidArgumentError( + "interleave requires input of at least one array".to_string(), + )); + } + let data_type = values[0].data_type(); + + if values + .iter() + .skip(1) + .any(|array| array.data_type() != data_type) + { + return Err(ArrowError::InvalidArgumentError( + "It is not possible to interleave arrays of different data types." + .to_string(), + )); + } + + if indices.is_empty() { + return Ok(new_empty_array(data_type)); + } + + // TODO: Add specialized implementations (#1523) + + interleave_fallback(values, indices) +} + +/// Fallback implementation of interleave using [`MutableArrayData`] +fn interleave_fallback( + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result<ArrayRef, ArrowError> { + let arrays: Vec<_> = values.iter().map(|x| x.data()).collect(); + let mut array_data = MutableArrayData::new(arrays, false, indices.len()); + + let mut cur_array = indices[0].0; + let mut start_row_idx = indices[0].1; + let mut end_row_idx = start_row_idx + 1; + + for (array, row) in indices.iter().skip(1).copied() { + if array == cur_array && row == end_row_idx { + // subsequent row in same batch + end_row_idx += 1; + continue; + } + + // emit current batch of rows for current buffer + array_data.extend(cur_array, start_row_idx, end_row_idx); + + // start new batch of rows + cur_array = array; + start_row_idx = row; + end_row_idx = start_row_idx + 1; + } + + // emit final batch of rows + array_data.extend(cur_array, start_row_idx, end_row_idx); + Ok(make_array(array_data.freeze())) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::builder::{Int32Builder, ListBuilder}; + use arrow_array::cast::{as_primitive_array, as_string_array}; + use arrow_array::types::Int32Type; + use arrow_array::{Int32Array, ListArray, StringArray}; + + #[test] + fn test_primitive() { + let a = Int32Array::from_iter_values([1, 2, 3, 4]); + let b = Int32Array::from_iter_values([5, 6, 7]); + let c = Int32Array::from_iter_values([8, 9, 10]); + let values = + interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap(); Review Comment: 👍 for checking repeated indexes `(0,3)` and `(0,3)` ########## arrow/src/compute/kernels/interleave.rs: ########## @@ -0,0 +1,218 @@ +// 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 arrow_array::{make_array, new_empty_array, Array, ArrayRef}; +use arrow_data::transform::MutableArrayData; +use arrow_schema::ArrowError; + +/// +/// Takes elements by index from a list of [`Array`], creating a new [`Array`] from those values. +/// +/// Each element in `indices` is a pair of `usize` with the first identifying the index +/// of the [`Array`] in `values`, and the second the index of the value within that [`Array`] +/// +/// ```text +/// ┌─────────────────┐ ┌─────────┐ ┌─────────────────┐ +/// │ A │ │ (0, 0) │ interleave( │ A │ +/// ├─────────────────┤ ├─────────┤ [values0, values1], ├─────────────────┤ +/// │ D │ │ (1, 0) │ indices │ B │ +/// └─────────────────┘ ├─────────┤ ) ├─────────────────┤ +/// values array 0 │ (1, 1) │ ─────────────────────────▶ │ C │ +/// ├─────────┤ ├─────────────────┤ +/// │ (0, 1) │ │ D │ +/// └─────────┘ └─────────────────┘ +/// ┌─────────────────┐ indices +/// │ B │ array +/// ├─────────────────┤ result +/// │ C │ +/// ├─────────────────┤ +/// │ E │ +/// └─────────────────┘ +/// values array 1 +/// ``` +/// +/// For selecting values by index from a single array see [compute::take](crate::compute::take) Review Comment: 👍 -- 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]
