From 662d0ca32108dea3824330f4a0e74ffb718fff0b Mon Sep 17 00:00:00 2001 From: Richard Date: Thu, 30 Apr 2026 14:41:43 -0400 Subject: [PATCH 1/7] reproduced take issue & updated take_run() to provide expected behavoir of merging runs --- arrow-select/src/interleave.rs | 18 +++++++ arrow-select/src/take.rs | 90 ++++++++++++++++++++++++++++++---- 2 files changed, 98 insertions(+), 10 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index f5904bc171ee..5300227bac2c 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -1544,4 +1544,22 @@ mod tests { &[3] ); } + + #[test] + fn test_interleave_run_end_encoded_merges_identical_runs() { + // TODO: repoducer for interleave + let mut builder = PrimitiveRunBuilder::::new(); + builder.extend([0, 0, 0, 1, 1, 0, 0, 1, 1, 1].into_iter().map(Some)); + let a = builder.finish(); + + let mut builder = PrimitiveRunBuilder::::new(); + builder.extend([2, 2, 1, 1, 1, 0, 1, 0, 0, 0].into_iter().map(Some)); + let b = builder.finish(); + + println!("a: {a:#?}"); + println!("b: {b:#?}"); + let result = interleave(&[&a, &b], &[(0, 3), (0, 4), (1, 2), (1, 3), (1, 4)]).unwrap(); + println!("result: {result:#?}"); + assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _))); + } } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index ee813f5353c2..10160da89e30 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -17,10 +17,6 @@ //! Defines take kernel for [Array] -use std::fmt::Display; -use std::mem::ManuallyDrop; -use std::sync::Arc; - use arrow_array::builder::{BufferBuilder, UInt32Builder}; use arrow_array::cast::AsArray; use arrow_array::types::*; @@ -31,6 +27,9 @@ use arrow_buffer::{ }; use arrow_data::ArrayDataBuilder; use arrow_schema::{ArrowError, DataType, FieldRef, UnionMode}; +use std::fmt::Display; +use std::mem::ManuallyDrop; +use std::sync::Arc; use num_traits::{One, Zero}; @@ -295,6 +294,7 @@ fn take_impl( values => Ok(Arc::new(take_dict(values, indices)?)), t => unimplemented!("Take not supported for dictionary type {:?}", t) } + // TODO: update this line DataType::RunEndEncoded(_, _) => downcast_run_array! { values => Ok(Arc::new(take_run(values, indices)?)), t => unimplemented!("Take not supported for run type {:?}", t) @@ -871,6 +871,7 @@ fn take_run( run_array: &RunArray, logical_indices: &PrimitiveArray, ) -> Result, ArrowError> { + //use arrow_ord::ord::make_comparator; // get physical indices for the input logical indices let physical_indices = run_array.get_physical_indices(logical_indices.values())?; @@ -880,9 +881,16 @@ fn take_run( let mut new_run_ends_builder = BufferBuilder::::new(1); let mut take_value_indices = BufferBuilder::::new(1); let mut new_physical_len = 1; + let ree_values = run_array.values().to_data(); + for ix in 1..physical_indices.len() { - if physical_indices[ix] != physical_indices[ix - 1] { - take_value_indices.append(I::Native::from_usize(physical_indices[ix - 1]).unwrap()); + // add a condition to check if the values these indexes represent are acutally the same. + let prev_idx = physical_indices[ix - 1]; + let prev_value = ree_values.slice(prev_idx, 1); + let cur_idx = physical_indices[ix]; + let cur_value = ree_values.slice(cur_idx, 1); + if cur_idx != prev_idx && cur_value != prev_value { + take_value_indices.append(I::Native::from_usize(prev_idx).unwrap()); new_run_ends_builder.append(T::Native::from_usize(ix).unwrap()); new_physical_len += 1; } @@ -910,7 +918,6 @@ fn take_run( .build_unchecked() .into() }; - let new_values = take(run_array.values(), &take_value_indices, None)?; let builder = ArrayDataBuilder::new(run_array.data_type().clone()) @@ -1139,6 +1146,7 @@ pub fn take_record_batch( #[cfg(test)] mod tests { use super::*; + use arrow_array::builder::StringRunBuilder; use arrow_array::builder::*; use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano}; use arrow_data::ArrayData; @@ -2541,14 +2549,17 @@ mod tests { let take_indices: PrimitiveArray = vec![7, 2, 3, 7, 11, 4, 6].into_iter().collect(); - let take_out = take_run(&run_array, &take_indices).unwrap(); + // [2,2,2,2,2,1,1] + // ree [5,7] + // values: [2,1] + let take_out = take_run(&run_array, &take_indices).unwrap(); assert_eq!(take_out.len(), 7); assert_eq!(take_out.run_ends().len(), 7); - assert_eq!(take_out.run_ends().values(), &[1_i32, 3, 4, 5, 7]); + assert_eq!(take_out.run_ends().values(), &[5, 7]); let take_out_values = take_out.values().as_primitive::(); - assert_eq!(take_out_values.values(), &[2, 2, 2, 2, 1]); + assert_eq!(take_out_values.values(), &[2, 1]); } #[test] @@ -2820,4 +2831,63 @@ mod tests { assert_eq!(run_result.run_ends().len(), 0); assert_eq!(run_result.values().len(), 0); } + #[test] + fn test_take_run_end_encoded_merges_identical_runs() { + // see https://github.com/apache/arrow-rs/issues/7710 + let mut builder = PrimitiveRunBuilder::::new(); + builder.extend([1, 1, 0, 0, 1, 1].into_iter().map(Some)); + let ree = builder.finish(); + //let ree = Arc::new(ree); + let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]); + let interleaved = take(&ree, &indexes, None).unwrap(); + let interleaved = interleaved.as_run::(); + assert_eq!(interleaved.run_ends().values(), &[4]); + assert_eq!( + interleaved.values().as_primitive::().values(), + &[1] + ); + } + #[test] + fn test_take_run_end_encoded_merges_identical_string_runs() { + // Similar to `test_take_props` but using strings (bob/alice) + let mut builder = StringRunBuilder::::new(); + builder.extend( + ["bob", "bob", "alice", "alice", "bob", "bob"] + .into_iter() + .map(Some), + ); + let ree = builder.finish(); + + let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]); + let interleaved = take(&ree, &indexes, None).unwrap(); + let interleaved = interleaved.as_run::(); + assert_eq!(interleaved.run_ends().values(), &[4]); + assert_eq!(interleaved.values().as_string::().value(0), "bob"); + } + + #[test] + fn test_take_run_end_encoded_mixed_runs() { + // similar to the test above but we use repeated runs composed of the same logical index as well as different logical indices. + // this validates that the implementation merges runs that have the same value even if they are passed in as the same/different indices + let mut builder = StringRunBuilder::::new(); + builder.extend( + ["bob", "bob", "alice", "alice", "bob", "bob", "eve", "eve"] + .into_iter() + .map(Some), + ); + // [2,4,6,8] + // [bob,alice,bob,eve] + let ree = builder.finish(); + + let indexes = Int32Array::from_iter_values(vec![0, 0, 1, 4, 5, 2, 3, 2, 6, 7, 6]); + // [bob,bob,bob,bob,bob,alice,alice,alice,eve,eve,eve] + let interleaved = take(&ree, &indexes, None).unwrap(); + let interleaved = interleaved.as_run::(); + + assert_eq!(interleaved.len(), 11); + assert_eq!(interleaved.run_ends().values(), &[5, 8, 11]); + assert_eq!(interleaved.values().as_string::().value(0), "bob"); + assert_eq!(interleaved.values().as_string::().value(1), "alice"); + assert_eq!(interleaved.values().as_string::().value(2), "eve"); + } } From 780bbe0c41609f49c2873e0599c5c1dc143ace76 Mon Sep 17 00:00:00 2001 From: Richard Date: Fri, 1 May 2026 14:02:50 -0400 Subject: [PATCH 2/7] compact REE interleave calls --- arrow-select/src/interleave.rs | 125 ++++++++++++++++++++++++++++++--- arrow-select/src/take.rs | 59 ++++++---------- 2 files changed, 140 insertions(+), 44 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 5300227bac2c..8299ced7e293 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -19,13 +19,14 @@ use crate::concat::concat; use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values}; +use crate::take::take; use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder}; use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; use arrow_buffer::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer}; -use arrow_data::ByteView; use arrow_data::transform::MutableArrayData; +use arrow_data::{ArrayDataBuilder, ByteView}; use arrow_schema::{ArrowError, DataType, FieldRef, Fields}; use std::sync::Arc; @@ -423,7 +424,6 @@ fn interleave_fallback( 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 @@ -442,7 +442,45 @@ fn interleave_fallback( // emit final batch of rows array_data.extend(cur_array, start_row_idx, end_row_idx); - Ok(make_array(array_data.freeze())) + compact_runs(make_array(array_data.freeze())) +} + +/// Merge adjacent runs of a `RunEndEncoded` array that share equal values. +/// Returns the input unchanged for non-REE arrays or arrays already compact. +fn compact_runs(array: ArrayRef) -> Result { + if !matches!(array.data_type(), DataType::RunEndEncoded(_, _)) { + return Ok(array); + } + let data = array.to_data(); + let run_ends = &data.child_data()[0]; + let values = &data.child_data()[1]; + + if values.len() <= 1 { + return Ok(array); + } + + let mut keep: Vec = Vec::with_capacity(values.len()); + for i in 0..values.len() { + let last_in_group = i + 1 == values.len() || values.slice(i, 1) != values.slice(i + 1, 1); + if last_in_group { + keep.push(i as u32); + } + } + + if keep.len() == values.len() { + return Ok(array); + } + + let idx = UInt32Array::from(keep); + let new_run_ends = take(&make_array(run_ends.clone()), &idx, None)?; + let new_values = take(&make_array(values.clone()), &idx, None)?; + + let new_data = ArrayDataBuilder::new(array.data_type().clone()) + .len(array.len()) + .add_child_data(new_run_ends.to_data()) + .add_child_data(new_values.to_data()) + .build()?; + Ok(make_array(new_data)) } /// Fallback implementation for interleaving dictionaries when it was determined @@ -1547,7 +1585,6 @@ mod tests { #[test] fn test_interleave_run_end_encoded_merges_identical_runs() { - // TODO: repoducer for interleave let mut builder = PrimitiveRunBuilder::::new(); builder.extend([0, 0, 0, 1, 1, 0, 0, 1, 1, 1].into_iter().map(Some)); let a = builder.finish(); @@ -1556,10 +1593,82 @@ mod tests { builder.extend([2, 2, 1, 1, 1, 0, 1, 0, 0, 0].into_iter().map(Some)); let b = builder.finish(); - println!("a: {a:#?}"); - println!("b: {b:#?}"); + // logical: [1, 1, 1, 1, 1] across an a→b boundary; should compact to one run. let result = interleave(&[&a, &b], &[(0, 3), (0, 4), (1, 2), (1, 3), (1, 4)]).unwrap(); - println!("result: {result:#?}"); - assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _))); + let result = result.as_run::(); + + assert_eq!(result.run_ends().values(), &[5]); + let values = result.values().as_primitive::(); + assert_eq!(values.values(), &[1]); + } + + #[test] + fn test_interleave_run_end_encoded_partial_compaction() { + let mut builder = PrimitiveRunBuilder::::new(); + builder.extend([1, 1, 2, 2].into_iter().map(Some)); + let a = builder.finish(); + + let mut builder = PrimitiveRunBuilder::::new(); + builder.extend([1, 1, 2, 2].into_iter().map(Some)); + let b = builder.finish(); + + // logical: [1, 1, 1, 2, 2, 1, 1] — fallback emits 5 raw runs; + // compaction must merge adjacent equal pairs but keep the trailing 1s + // distinct from the leading 1s (separated by 2s). + let indices = &[(0, 0), (0, 1), (1, 0), (0, 2), (1, 3), (1, 0), (1, 1)]; + let result = interleave(&[&a, &b], indices).unwrap(); + let result = result.as_run::(); + + assert_eq!(result.run_ends().values(), &[3, 5, 7]); + let values = result.values().as_primitive::(); + assert_eq!(values.values(), &[1, 2, 1]); + } + + #[test] + fn test_interleave_run_end_encoded_pulls_identical_values() { + use arrow_array::builder::StringRunBuilder; + + let mut builder = StringRunBuilder::::new(); + builder.extend( + [ + "alice", "alice", "bob", "bob", "charlie", "charlie", "david", + ] + .into_iter() + .map(Some), + ); + let a = builder.finish(); + + let mut builder = StringRunBuilder::::new(); + builder.extend( + ["alice", "bob", "charlie", "david", "david", "eve"] + .into_iter() + .map(Some), + ); + let b = builder.finish(); + + // logical: ["alice","alice","bob","bob","charlie","charlie","david","david","david","alice","alice"] + let result = interleave( + &[&a, &b], + &[ + (0, 0), + (1, 0), + (0, 2), + (1, 1), + (0, 4), + (1, 2), + (0, 6), + (1, 3), + (1, 4), + (0, 0), + (1, 0), + ], + ) + .unwrap(); + let result = result.as_run::(); + + assert_eq!(result.run_ends().values(), &[2, 4, 6, 9, 11]); + let values = result.values().as_string::(); + let values: Vec<_> = values.into_iter().map(Option::unwrap).collect(); + assert_eq!(values, vec!["alice", "bob", "charlie", "david", "alice"]); } } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 10160da89e30..22ef34a909f9 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -294,7 +294,6 @@ fn take_impl( values => Ok(Arc::new(take_dict(values, indices)?)), t => unimplemented!("Take not supported for dictionary type {:?}", t) } - // TODO: update this line DataType::RunEndEncoded(_, _) => downcast_run_array! { values => Ok(Arc::new(take_run(values, indices)?)), t => unimplemented!("Take not supported for run type {:?}", t) @@ -871,7 +870,6 @@ fn take_run( run_array: &RunArray, logical_indices: &PrimitiveArray, ) -> Result, ArrowError> { - //use arrow_ord::ord::make_comparator; // get physical indices for the input logical indices let physical_indices = run_array.get_physical_indices(logical_indices.values())?; @@ -884,12 +882,9 @@ fn take_run( let ree_values = run_array.values().to_data(); for ix in 1..physical_indices.len() { - // add a condition to check if the values these indexes represent are acutally the same. let prev_idx = physical_indices[ix - 1]; - let prev_value = ree_values.slice(prev_idx, 1); let cur_idx = physical_indices[ix]; - let cur_value = ree_values.slice(cur_idx, 1); - if cur_idx != prev_idx && cur_value != prev_value { + if cur_idx != prev_idx && ree_values.slice(cur_idx, 1) != ree_values.slice(prev_idx, 1) { take_value_indices.append(I::Native::from_usize(prev_idx).unwrap()); new_run_ends_builder.append(T::Native::from_usize(ix).unwrap()); new_physical_len += 1; @@ -918,6 +913,7 @@ fn take_run( .build_unchecked() .into() }; + let new_values = take(run_array.values(), &take_value_indices, None)?; let builder = ArrayDataBuilder::new(run_array.data_type().clone()) @@ -2549,10 +2545,6 @@ mod tests { let take_indices: PrimitiveArray = vec![7, 2, 3, 7, 11, 4, 6].into_iter().collect(); - // [2,2,2,2,2,1,1] - // ree [5,7] - // values: [2,1] - let take_out = take_run(&run_array, &take_indices).unwrap(); assert_eq!(take_out.len(), 7); assert_eq!(take_out.run_ends().len(), 7); @@ -2831,25 +2823,23 @@ mod tests { assert_eq!(run_result.run_ends().len(), 0); assert_eq!(run_result.values().len(), 0); } + #[test] fn test_take_run_end_encoded_merges_identical_runs() { // see https://github.com/apache/arrow-rs/issues/7710 let mut builder = PrimitiveRunBuilder::::new(); builder.extend([1, 1, 0, 0, 1, 1].into_iter().map(Some)); let ree = builder.finish(); - //let ree = Arc::new(ree); + let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]); - let interleaved = take(&ree, &indexes, None).unwrap(); - let interleaved = interleaved.as_run::(); - assert_eq!(interleaved.run_ends().values(), &[4]); - assert_eq!( - interleaved.values().as_primitive::().values(), - &[1] - ); + let result = take(&ree, &indexes, None).unwrap(); + let result = result.as_run::(); + assert_eq!(result.run_ends().values(), &[4]); + assert_eq!(result.values().as_primitive::().values(), &[1]); } + #[test] fn test_take_run_end_encoded_merges_identical_string_runs() { - // Similar to `test_take_props` but using strings (bob/alice) let mut builder = StringRunBuilder::::new(); builder.extend( ["bob", "bob", "alice", "alice", "bob", "bob"] @@ -2859,35 +2849,32 @@ mod tests { let ree = builder.finish(); let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]); - let interleaved = take(&ree, &indexes, None).unwrap(); - let interleaved = interleaved.as_run::(); - assert_eq!(interleaved.run_ends().values(), &[4]); - assert_eq!(interleaved.values().as_string::().value(0), "bob"); + let result = take(&ree, &indexes, None).unwrap(); + let result = result.as_run::(); + assert_eq!(result.run_ends().values(), &[4]); + assert_eq!(result.values().as_string::().value(0), "bob"); } #[test] fn test_take_run_end_encoded_mixed_runs() { - // similar to the test above but we use repeated runs composed of the same logical index as well as different logical indices. - // this validates that the implementation merges runs that have the same value even if they are passed in as the same/different indices + // Validates that runs are merged whether the same logical value comes + // from the same physical index (repeated indices) or distinct ones. let mut builder = StringRunBuilder::::new(); builder.extend( ["bob", "bob", "alice", "alice", "bob", "bob", "eve", "eve"] .into_iter() .map(Some), ); - // [2,4,6,8] - // [bob,alice,bob,eve] let ree = builder.finish(); let indexes = Int32Array::from_iter_values(vec![0, 0, 1, 4, 5, 2, 3, 2, 6, 7, 6]); - // [bob,bob,bob,bob,bob,alice,alice,alice,eve,eve,eve] - let interleaved = take(&ree, &indexes, None).unwrap(); - let interleaved = interleaved.as_run::(); - - assert_eq!(interleaved.len(), 11); - assert_eq!(interleaved.run_ends().values(), &[5, 8, 11]); - assert_eq!(interleaved.values().as_string::().value(0), "bob"); - assert_eq!(interleaved.values().as_string::().value(1), "alice"); - assert_eq!(interleaved.values().as_string::().value(2), "eve"); + let result = take(&ree, &indexes, None).unwrap(); + let result = result.as_run::(); + + assert_eq!(result.len(), 11); + assert_eq!(result.run_ends().values(), &[5, 8, 11]); + assert_eq!(result.values().as_string::().value(0), "bob"); + assert_eq!(result.values().as_string::().value(1), "alice"); + assert_eq!(result.values().as_string::().value(2), "eve"); } } From 58a62c0d48db6960d17dfc7b08c606551c129dd0 Mon Sep 17 00:00:00 2001 From: Richard Date: Tue, 5 May 2026 13:06:50 -0400 Subject: [PATCH 3/7] removed interleave optimization from PR since #9919 resolves it --- arrow-select/src/interleave.rs | 133 +-------------------------------- 1 file changed, 3 insertions(+), 130 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 8299ced7e293..f5904bc171ee 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -19,14 +19,13 @@ use crate::concat::concat; use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values}; -use crate::take::take; use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder}; use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; use arrow_buffer::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer}; +use arrow_data::ByteView; use arrow_data::transform::MutableArrayData; -use arrow_data::{ArrayDataBuilder, ByteView}; use arrow_schema::{ArrowError, DataType, FieldRef, Fields}; use std::sync::Arc; @@ -424,6 +423,7 @@ fn interleave_fallback( 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 @@ -442,45 +442,7 @@ fn interleave_fallback( // emit final batch of rows array_data.extend(cur_array, start_row_idx, end_row_idx); - compact_runs(make_array(array_data.freeze())) -} - -/// Merge adjacent runs of a `RunEndEncoded` array that share equal values. -/// Returns the input unchanged for non-REE arrays or arrays already compact. -fn compact_runs(array: ArrayRef) -> Result { - if !matches!(array.data_type(), DataType::RunEndEncoded(_, _)) { - return Ok(array); - } - let data = array.to_data(); - let run_ends = &data.child_data()[0]; - let values = &data.child_data()[1]; - - if values.len() <= 1 { - return Ok(array); - } - - let mut keep: Vec = Vec::with_capacity(values.len()); - for i in 0..values.len() { - let last_in_group = i + 1 == values.len() || values.slice(i, 1) != values.slice(i + 1, 1); - if last_in_group { - keep.push(i as u32); - } - } - - if keep.len() == values.len() { - return Ok(array); - } - - let idx = UInt32Array::from(keep); - let new_run_ends = take(&make_array(run_ends.clone()), &idx, None)?; - let new_values = take(&make_array(values.clone()), &idx, None)?; - - let new_data = ArrayDataBuilder::new(array.data_type().clone()) - .len(array.len()) - .add_child_data(new_run_ends.to_data()) - .add_child_data(new_values.to_data()) - .build()?; - Ok(make_array(new_data)) + Ok(make_array(array_data.freeze())) } /// Fallback implementation for interleaving dictionaries when it was determined @@ -1582,93 +1544,4 @@ mod tests { &[3] ); } - - #[test] - fn test_interleave_run_end_encoded_merges_identical_runs() { - let mut builder = PrimitiveRunBuilder::::new(); - builder.extend([0, 0, 0, 1, 1, 0, 0, 1, 1, 1].into_iter().map(Some)); - let a = builder.finish(); - - let mut builder = PrimitiveRunBuilder::::new(); - builder.extend([2, 2, 1, 1, 1, 0, 1, 0, 0, 0].into_iter().map(Some)); - let b = builder.finish(); - - // logical: [1, 1, 1, 1, 1] across an a→b boundary; should compact to one run. - let result = interleave(&[&a, &b], &[(0, 3), (0, 4), (1, 2), (1, 3), (1, 4)]).unwrap(); - let result = result.as_run::(); - - assert_eq!(result.run_ends().values(), &[5]); - let values = result.values().as_primitive::(); - assert_eq!(values.values(), &[1]); - } - - #[test] - fn test_interleave_run_end_encoded_partial_compaction() { - let mut builder = PrimitiveRunBuilder::::new(); - builder.extend([1, 1, 2, 2].into_iter().map(Some)); - let a = builder.finish(); - - let mut builder = PrimitiveRunBuilder::::new(); - builder.extend([1, 1, 2, 2].into_iter().map(Some)); - let b = builder.finish(); - - // logical: [1, 1, 1, 2, 2, 1, 1] — fallback emits 5 raw runs; - // compaction must merge adjacent equal pairs but keep the trailing 1s - // distinct from the leading 1s (separated by 2s). - let indices = &[(0, 0), (0, 1), (1, 0), (0, 2), (1, 3), (1, 0), (1, 1)]; - let result = interleave(&[&a, &b], indices).unwrap(); - let result = result.as_run::(); - - assert_eq!(result.run_ends().values(), &[3, 5, 7]); - let values = result.values().as_primitive::(); - assert_eq!(values.values(), &[1, 2, 1]); - } - - #[test] - fn test_interleave_run_end_encoded_pulls_identical_values() { - use arrow_array::builder::StringRunBuilder; - - let mut builder = StringRunBuilder::::new(); - builder.extend( - [ - "alice", "alice", "bob", "bob", "charlie", "charlie", "david", - ] - .into_iter() - .map(Some), - ); - let a = builder.finish(); - - let mut builder = StringRunBuilder::::new(); - builder.extend( - ["alice", "bob", "charlie", "david", "david", "eve"] - .into_iter() - .map(Some), - ); - let b = builder.finish(); - - // logical: ["alice","alice","bob","bob","charlie","charlie","david","david","david","alice","alice"] - let result = interleave( - &[&a, &b], - &[ - (0, 0), - (1, 0), - (0, 2), - (1, 1), - (0, 4), - (1, 2), - (0, 6), - (1, 3), - (1, 4), - (0, 0), - (1, 0), - ], - ) - .unwrap(); - let result = result.as_run::(); - - assert_eq!(result.run_ends().values(), &[2, 4, 6, 9, 11]); - let values = result.values().as_string::(); - let values: Vec<_> = values.into_iter().map(Option::unwrap).collect(); - assert_eq!(values, vec!["alice", "bob", "charlie", "david", "alice"]); - } } From 6729927396feda235f18720af992ebbe8e79d659 Mon Sep 17 00:00:00 2001 From: Richard Date: Wed, 6 May 2026 13:03:31 -0400 Subject: [PATCH 4/7] refactored code base a bit and updated approach to comparissins to take_run() --- arrow-array/src/comparator.rs | 520 +++++++++++++++++++++++++++++++ arrow-array/src/lib.rs | 2 + arrow-ord/src/ord.rs | 556 +--------------------------------- arrow-select/src/take.rs | 11 +- 4 files changed, 538 insertions(+), 551 deletions(-) create mode 100644 arrow-array/src/comparator.rs diff --git a/arrow-array/src/comparator.rs b/arrow-array/src/comparator.rs new file mode 100644 index 000000000000..29ee871a6422 --- /dev/null +++ b/arrow-array/src/comparator.rs @@ -0,0 +1,520 @@ +// 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. + +//! Internal comparator factories for comparing arbitrary array slots. + +use crate::cast::AsArray; +use crate::types::*; +use crate::*; +use arrow_buffer::{ArrowNativeType, NullBuffer}; +use arrow_schema::{ArrowError, DataType, SortOptions}; +use std::{cmp::Ordering, collections::HashMap}; + +fn compare_run_end_encoded( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_run::(); + let right = right.as_run::(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + + let l_run_ends = left.run_ends().clone(); + let r_run_ends = right.run_ends().clone(); + + let f = compare(left, right, opts, move |i, j| { + let l_physical = l_run_ends.get_physical_index(i); + let r_physical = r_run_ends.get_physical_index(j); + cmp(l_physical, r_physical) + }); + Ok(f) +} + +/// Compare values at arbitrary indices in two arrays. +pub type DynComparator = Box Ordering + Send + Sync>; + +/// If parent sort order is descending we need to invert the value of nulls_first so that +/// when the parent is sorted based on the produced ranks, nulls are still ordered correctly +fn child_opts(opts: SortOptions) -> SortOptions { + SortOptions { + descending: false, + nulls_first: opts.nulls_first != opts.descending, + } +} + +fn compare(l: &A, r: &A, opts: SortOptions, cmp: F) -> DynComparator +where + A: Array + Clone, + F: Fn(usize, usize) -> Ordering + Send + Sync + 'static, +{ + let l = l.logical_nulls().filter(|x| x.null_count() > 0); + let r = r.logical_nulls().filter(|x| x.null_count() > 0); + match (opts.nulls_first, opts.descending) { + (true, true) => compare_impl::(l, r, cmp), + (true, false) => compare_impl::(l, r, cmp), + (false, true) => compare_impl::(l, r, cmp), + (false, false) => compare_impl::(l, r, cmp), + } +} + +fn compare_impl( + l: Option, + r: Option, + cmp: F, +) -> DynComparator +where + F: Fn(usize, usize) -> Ordering + Send + Sync + 'static, +{ + let cmp = move |i, j| match DESCENDING { + true => cmp(i, j).reverse(), + false => cmp(i, j), + }; + + let (left_null, right_null) = match NULLS_FIRST { + true => (Ordering::Less, Ordering::Greater), + false => (Ordering::Greater, Ordering::Less), + }; + + match (l, r) { + (None, None) => Box::new(cmp), + (Some(l), None) => Box::new(move |i, j| match l.is_null(i) { + true => left_null, + false => cmp(i, j), + }), + (None, Some(r)) => Box::new(move |i, j| match r.is_null(j) { + true => right_null, + false => cmp(i, j), + }), + (Some(l), Some(r)) => Box::new(move |i, j| match (l.is_null(i), r.is_null(j)) { + (true, true) => Ordering::Equal, + (true, false) => left_null, + (false, true) => right_null, + (false, false) => cmp(i, j), + }), + } +} + +fn compare_primitive( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> DynComparator +where + T::Native: ArrowNativeTypeOp, +{ + let left = left.as_primitive::(); + let right = right.as_primitive::(); + let l_values = left.values().clone(); + let r_values = right.values().clone(); + + compare(&left, &right, opts, move |i, j| { + l_values[i].compare(r_values[j]) + }) +} + +fn compare_boolean(left: &dyn Array, right: &dyn Array, opts: SortOptions) -> DynComparator { + let left = left.as_boolean(); + let right = right.as_boolean(); + + let l_values = left.values().clone(); + let r_values = right.values().clone(); + + compare(left, right, opts, move |i, j| { + l_values.value(i).cmp(&r_values.value(j)) + }) +} + +fn compare_bytes( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> DynComparator { + let left = left.as_bytes::(); + let right = right.as_bytes::(); + + let l = left.clone(); + let r = right.clone(); + compare(left, right, opts, move |i, j| { + let l: &[u8] = l.value(i).as_ref(); + let r: &[u8] = r.value(j).as_ref(); + l.cmp(r) + }) +} + +fn compare_byte_view( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> DynComparator { + let left = left.as_byte_view::(); + let right = right.as_byte_view::(); + + let l = left.clone(); + let r = right.clone(); + compare(left, right, opts, move |i, j| { + compare_byte_view_values(&l, i, &r, j) + }) +} + +fn compare_byte_view_values( + left: &GenericByteViewArray, + left_idx: usize, + right: &GenericByteViewArray, + right_idx: usize, +) -> Ordering { + assert!(left_idx < left.len()); + assert!(right_idx < right.len()); + + if left.data_buffers().is_empty() && right.data_buffers().is_empty() { + let l_view = unsafe { left.views().get_unchecked(left_idx) }; + let r_view = unsafe { right.views().get_unchecked(right_idx) }; + return GenericByteViewArray::::inline_key_fast(*l_view) + .cmp(&GenericByteViewArray::::inline_key_fast(*r_view)); + } + + unsafe { GenericByteViewArray::compare_unchecked(left, left_idx, right, right_idx) } +} + +fn compare_dict( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_dictionary::(); + let right = right.as_dictionary::(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + let left_keys = left.keys().values().clone(); + let right_keys = right.keys().values().clone(); + + let f = compare(left, right, opts, move |i, j| { + let l = left_keys[i].as_usize(); + let r = right_keys[j].as_usize(); + cmp(l, r) + }); + Ok(f) +} + +fn compare_list( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_list::(); + let right = right.as_list::(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + + let l_o = left.offsets().clone(); + let r_o = right.offsets().clone(); + let f = compare(left, right, opts, move |i, j| { + let l_end = l_o[i + 1].as_usize(); + let l_start = l_o[i].as_usize(); + + let r_end = r_o[j + 1].as_usize(); + let r_start = r_o[j].as_usize(); + + for (i, j) in (l_start..l_end).zip(r_start..r_end) { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + (l_end - l_start).cmp(&(r_end - r_start)) + }); + Ok(f) +} + +fn compare_fixed_list( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_fixed_size_list(); + let right = right.as_fixed_size_list(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + + let l_size = left.value_length().to_usize().unwrap(); + let r_size = right.value_length().to_usize().unwrap(); + let size_cmp = l_size.cmp(&r_size); + + let f = compare(left, right, opts, move |i, j| { + let l_start = i * l_size; + let l_end = l_start + l_size; + let r_start = j * r_size; + let r_end = r_start + r_size; + for (i, j) in (l_start..l_end).zip(r_start..r_end) { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + size_cmp + }); + Ok(f) +} + +fn compare_list_view( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_list_view::(); + let right = right.as_list_view::(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + + let l_offsets = left.offsets().clone(); + let l_sizes = left.sizes().clone(); + let r_offsets = right.offsets().clone(); + let r_sizes = right.sizes().clone(); + + let f = compare(left, right, opts, move |i, j| { + let l_start = l_offsets[i].as_usize(); + let l_len = l_sizes[i].as_usize(); + let l_end = l_start + l_len; + + let r_start = r_offsets[j].as_usize(); + let r_len = r_sizes[j].as_usize(); + let r_end = r_start + r_len; + + for (i, j) in (l_start..l_end).zip(r_start..r_end) { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + l_len.cmp(&r_len) + }); + Ok(f) +} + +fn compare_map( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_map(); + let right = right.as_map(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.entries(), right.entries(), c_opts)?; + + let l_o = left.offsets().clone(); + let r_o = right.offsets().clone(); + let f = compare(left, right, opts, move |i, j| { + let l_end = l_o[i + 1].as_usize(); + let l_start = l_o[i].as_usize(); + + let r_end = r_o[j + 1].as_usize(); + let r_start = r_o[j].as_usize(); + + for (i, j) in (l_start..l_end).zip(r_start..r_end) { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + (l_end - l_start).cmp(&(r_end - r_start)) + }); + Ok(f) +} + +fn compare_struct( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_struct(); + let right = right.as_struct(); + + if left.columns().len() != right.columns().len() { + return Err(ArrowError::InvalidArgumentError( + "Cannot compare StructArray with different number of columns".to_string(), + )); + } + + let c_opts = child_opts(opts); + let columns = left.columns().iter().zip(right.columns()); + let comparators = columns + .map(|(l, r)| make_comparator(l, r, c_opts)) + .collect::, _>>()?; + + let f = compare(left, right, opts, move |i, j| { + for cmp in &comparators { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + Ordering::Equal + }); + Ok(f) +} + +fn compare_union( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_union(); + let right = right.as_union(); + + let (left_fields, left_mode) = match left.data_type() { + DataType::Union(fields, mode) => (fields, mode), + _ => unreachable!(), + }; + let (right_fields, right_mode) = match right.data_type() { + DataType::Union(fields, mode) => (fields, mode), + _ => unreachable!(), + }; + + if left_fields != right_fields { + return Err(ArrowError::InvalidArgumentError(format!( + "Cannot compare UnionArrays with different fields: left={left_fields:?}, right={right_fields:?}" + ))); + } + + if left_mode != right_mode { + return Err(ArrowError::InvalidArgumentError(format!( + "Cannot compare UnionArrays with different modes: left={left_mode:?}, right={right_mode:?}" + ))); + } + + let c_opts = child_opts(opts); + + let mut field_comparators = HashMap::with_capacity(left_fields.len()); + + for (type_id, _field) in left_fields.iter() { + let left_child = left.child(type_id); + let right_child = right.child(type_id); + let cmp = make_comparator(left_child.as_ref(), right_child.as_ref(), c_opts)?; + + field_comparators.insert(type_id, cmp); + } + + let left_type_ids = left.type_ids().clone(); + let right_type_ids = right.type_ids().clone(); + + let left_offsets = left.offsets().cloned(); + let right_offsets = right.offsets().cloned(); + + let f = compare(left, right, opts, move |i, j| { + let left_type_id = left_type_ids[i]; + let right_type_id = right_type_ids[j]; + + match left_type_id.cmp(&right_type_id) { + Ordering::Equal => { + let left_offset = left_offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i); + let right_offset = right_offsets.as_ref().map(|o| o[j] as usize).unwrap_or(j); + + let cmp = field_comparators + .get(&left_type_id) + .expect("type id not found in field_comparators"); + + cmp(left_offset, right_offset) + } + other => other, + } + }); + Ok(f) +} + +/// Returns a comparison function that compares two values at two arbitrary indices. +/// +/// If `nulls_first` is true, null values are considered less than any non-null +/// value; otherwise they are considered greater. This is primarily shared by +/// crates that need repeated slot comparisons without constructing sliced arrays. +pub fn make_comparator( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + use arrow_schema::DataType::*; + + macro_rules! primitive_helper { + ($t:ty, $left:expr, $right:expr, $nulls_first:expr) => { + Ok(compare_primitive::<$t>($left, $right, $nulls_first)) + }; + } + downcast_primitive! { + left.data_type(), right.data_type() => (primitive_helper, left, right, opts), + (Boolean, Boolean) => Ok(compare_boolean(left, right, opts)), + (Utf8, Utf8) => Ok(compare_bytes::(left, right, opts)), + (LargeUtf8, LargeUtf8) => Ok(compare_bytes::(left, right, opts)), + (Utf8View, Utf8View) => Ok(compare_byte_view::(left, right, opts)), + (Binary, Binary) => Ok(compare_bytes::(left, right, opts)), + (LargeBinary, LargeBinary) => Ok(compare_bytes::(left, right, opts)), + (BinaryView, BinaryView) => Ok(compare_byte_view::(left, right, opts)), + (FixedSizeBinary(_), FixedSizeBinary(_)) => { + let left = left.as_fixed_size_binary(); + let right = right.as_fixed_size_binary(); + + let l = left.clone(); + let r = right.clone(); + Ok(compare(left, right, opts, move |i, j| { + l.value(i).cmp(r.value(j)) + })) + }, + (List(_), List(_)) => compare_list::(left, right, opts), + (LargeList(_), LargeList(_)) => compare_list::(left, right, opts), + (ListView(_), ListView(_)) => compare_list_view::(left, right, opts), + (LargeListView(_), LargeListView(_)) => compare_list_view::(left, right, opts), + (FixedSizeList(_, _), FixedSizeList(_, _)) => compare_fixed_list(left, right, opts), + (Struct(_), Struct(_)) => compare_struct(left, right, opts), + (Dictionary(l_key, _), Dictionary(r_key, _)) => { + macro_rules! dict_helper { + ($t:ty, $left:expr, $right:expr, $opts: expr) => { + compare_dict::<$t>($left, $right, $opts) + }; + } + downcast_integer! { + l_key.as_ref(), r_key.as_ref() => (dict_helper, left, right, opts), + _ => unreachable!() + } + }, + (RunEndEncoded(l_run_ends, _), RunEndEncoded(r_run_ends, _)) => { + macro_rules! run_end_helper { + ($t:ty, $left:expr, $right:expr, $opts:expr) => { + compare_run_end_encoded::<$t>($left, $right, $opts) + }; + } + downcast_run_end_index! { + l_run_ends.data_type(), r_run_ends.data_type() => (run_end_helper, left, right, opts), + _ => Err(ArrowError::InvalidArgumentError(format!( + "Cannot compare RunEndEncoded arrays with different run ends types: left={:?}, right={:?}", + l_run_ends.data_type(), + r_run_ends.data_type() + ))) + } + }, + (Map(_, _), Map(_, _)) => compare_map(left, right, opts), + (Null, Null) => Ok(Box::new(|_, _| Ordering::Equal)), + (Union(_, _), Union(_, _)) => compare_union(left, right, opts), + (lhs, rhs) => Err(ArrowError::InvalidArgumentError(match lhs == rhs { + true => format!("The data type type {lhs:?} has no natural order"), + false => "Can't compare arrays of different types".to_string(), + })) + } +} diff --git a/arrow-array/src/lib.rs b/arrow-array/src/lib.rs index a5f9bf5e711c..ca17370d255a 100644 --- a/arrow-array/src/lib.rs +++ b/arrow-array/src/lib.rs @@ -248,6 +248,8 @@ pub use scalar::*; pub mod builder; pub mod cast; +#[doc(hidden)] +pub mod comparator; mod delta; #[cfg(feature = "ffi")] pub mod ffi; diff --git a/arrow-ord/src/ord.rs b/arrow-ord/src/ord.rs index 5951a587888a..77892d12332f 100644 --- a/arrow-ord/src/ord.rs +++ b/arrow-ord/src/ord.rs @@ -17,556 +17,16 @@ //! Contains functions and function factories to compare arrays. -use arrow_array::cast::AsArray; +pub use arrow_array::comparator::{DynComparator, make_comparator}; + +#[cfg(test)] use arrow_array::types::*; +#[cfg(test)] use arrow_array::*; -use arrow_buffer::{ArrowNativeType, NullBuffer}; -use arrow_schema::{ArrowError, DataType, SortOptions}; -use std::{cmp::Ordering, collections::HashMap}; - -fn compare_run_end_encoded( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_run::(); - let right = right.as_run::(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - - let l_run_ends = left.run_ends().clone(); - let r_run_ends = right.run_ends().clone(); - - let f = compare(left, right, opts, move |i, j| { - let l_physical = l_run_ends.get_physical_index(i); - let r_physical = r_run_ends.get_physical_index(j); - cmp(l_physical, r_physical) - }); - Ok(f) -} - -/// Compare the values at two arbitrary indices in two arrays. -pub type DynComparator = Box Ordering + Send + Sync>; - -/// If parent sort order is descending we need to invert the value of nulls_first so that -/// when the parent is sorted based on the produced ranks, nulls are still ordered correctly -fn child_opts(opts: SortOptions) -> SortOptions { - SortOptions { - descending: false, - nulls_first: opts.nulls_first != opts.descending, - } -} - -fn compare(l: &A, r: &A, opts: SortOptions, cmp: F) -> DynComparator -where - A: Array + Clone, - F: Fn(usize, usize) -> Ordering + Send + Sync + 'static, -{ - let l = l.logical_nulls().filter(|x| x.null_count() > 0); - let r = r.logical_nulls().filter(|x| x.null_count() > 0); - match (opts.nulls_first, opts.descending) { - (true, true) => compare_impl::(l, r, cmp), - (true, false) => compare_impl::(l, r, cmp), - (false, true) => compare_impl::(l, r, cmp), - (false, false) => compare_impl::(l, r, cmp), - } -} - -fn compare_impl( - l: Option, - r: Option, - cmp: F, -) -> DynComparator -where - F: Fn(usize, usize) -> Ordering + Send + Sync + 'static, -{ - let cmp = move |i, j| match DESCENDING { - true => cmp(i, j).reverse(), - false => cmp(i, j), - }; - - let (left_null, right_null) = match NULLS_FIRST { - true => (Ordering::Less, Ordering::Greater), - false => (Ordering::Greater, Ordering::Less), - }; - - match (l, r) { - (None, None) => Box::new(cmp), - (Some(l), None) => Box::new(move |i, j| match l.is_null(i) { - true => left_null, - false => cmp(i, j), - }), - (None, Some(r)) => Box::new(move |i, j| match r.is_null(j) { - true => right_null, - false => cmp(i, j), - }), - (Some(l), Some(r)) => Box::new(move |i, j| match (l.is_null(i), r.is_null(j)) { - (true, true) => Ordering::Equal, - (true, false) => left_null, - (false, true) => right_null, - (false, false) => cmp(i, j), - }), - } -} - -fn compare_primitive( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> DynComparator -where - T::Native: ArrowNativeTypeOp, -{ - let left = left.as_primitive::(); - let right = right.as_primitive::(); - let l_values = left.values().clone(); - let r_values = right.values().clone(); - - compare(&left, &right, opts, move |i, j| { - l_values[i].compare(r_values[j]) - }) -} - -fn compare_boolean(left: &dyn Array, right: &dyn Array, opts: SortOptions) -> DynComparator { - let left = left.as_boolean(); - let right = right.as_boolean(); - - let l_values = left.values().clone(); - let r_values = right.values().clone(); - - compare(left, right, opts, move |i, j| { - l_values.value(i).cmp(&r_values.value(j)) - }) -} - -fn compare_bytes( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> DynComparator { - let left = left.as_bytes::(); - let right = right.as_bytes::(); - - let l = left.clone(); - let r = right.clone(); - compare(left, right, opts, move |i, j| { - let l: &[u8] = l.value(i).as_ref(); - let r: &[u8] = r.value(j).as_ref(); - l.cmp(r) - }) -} - -fn compare_byte_view( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> DynComparator { - let left = left.as_byte_view::(); - let right = right.as_byte_view::(); - - let l = left.clone(); - let r = right.clone(); - compare(left, right, opts, move |i, j| { - crate::cmp::compare_byte_view(&l, i, &r, j) - }) -} - -fn compare_dict( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_dictionary::(); - let right = right.as_dictionary::(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - let left_keys = left.keys().values().clone(); - let right_keys = right.keys().values().clone(); - - let f = compare(left, right, opts, move |i, j| { - let l = left_keys[i].as_usize(); - let r = right_keys[j].as_usize(); - cmp(l, r) - }); - Ok(f) -} - -fn compare_list( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_list::(); - let right = right.as_list::(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - - let l_o = left.offsets().clone(); - let r_o = right.offsets().clone(); - let f = compare(left, right, opts, move |i, j| { - let l_end = l_o[i + 1].as_usize(); - let l_start = l_o[i].as_usize(); - - let r_end = r_o[j + 1].as_usize(); - let r_start = r_o[j].as_usize(); - - for (i, j) in (l_start..l_end).zip(r_start..r_end) { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - (l_end - l_start).cmp(&(r_end - r_start)) - }); - Ok(f) -} - -fn compare_fixed_list( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_fixed_size_list(); - let right = right.as_fixed_size_list(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - - let l_size = left.value_length().to_usize().unwrap(); - let r_size = right.value_length().to_usize().unwrap(); - let size_cmp = l_size.cmp(&r_size); - - let f = compare(left, right, opts, move |i, j| { - let l_start = i * l_size; - let l_end = l_start + l_size; - let r_start = j * r_size; - let r_end = r_start + r_size; - for (i, j) in (l_start..l_end).zip(r_start..r_end) { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - size_cmp - }); - Ok(f) -} - -fn compare_list_view( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_list_view::(); - let right = right.as_list_view::(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - - let l_offsets = left.offsets().clone(); - let l_sizes = left.sizes().clone(); - let r_offsets = right.offsets().clone(); - let r_sizes = right.sizes().clone(); - - let f = compare(left, right, opts, move |i, j| { - let l_start = l_offsets[i].as_usize(); - let l_len = l_sizes[i].as_usize(); - let l_end = l_start + l_len; - - let r_start = r_offsets[j].as_usize(); - let r_len = r_sizes[j].as_usize(); - let r_end = r_start + r_len; - - for (i, j) in (l_start..l_end).zip(r_start..r_end) { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - l_len.cmp(&r_len) - }); - Ok(f) -} - -fn compare_map( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_map(); - let right = right.as_map(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.entries(), right.entries(), c_opts)?; - - let l_o = left.offsets().clone(); - let r_o = right.offsets().clone(); - let f = compare(left, right, opts, move |i, j| { - let l_end = l_o[i + 1].as_usize(); - let l_start = l_o[i].as_usize(); - - let r_end = r_o[j + 1].as_usize(); - let r_start = r_o[j].as_usize(); - - for (i, j) in (l_start..l_end).zip(r_start..r_end) { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - (l_end - l_start).cmp(&(r_end - r_start)) - }); - Ok(f) -} - -fn compare_struct( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_struct(); - let right = right.as_struct(); - - if left.columns().len() != right.columns().len() { - return Err(ArrowError::InvalidArgumentError( - "Cannot compare StructArray with different number of columns".to_string(), - )); - } - - let c_opts = child_opts(opts); - let columns = left.columns().iter().zip(right.columns()); - let comparators = columns - .map(|(l, r)| make_comparator(l, r, c_opts)) - .collect::, _>>()?; - - let f = compare(left, right, opts, move |i, j| { - for cmp in &comparators { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - Ordering::Equal - }); - Ok(f) -} - -fn compare_union( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_union(); - let right = right.as_union(); - - let (left_fields, left_mode) = match left.data_type() { - DataType::Union(fields, mode) => (fields, mode), - _ => unreachable!(), - }; - let (right_fields, right_mode) = match right.data_type() { - DataType::Union(fields, mode) => (fields, mode), - _ => unreachable!(), - }; - - if left_fields != right_fields { - return Err(ArrowError::InvalidArgumentError(format!( - "Cannot compare UnionArrays with different fields: left={:?}, right={:?}", - left_fields, right_fields - ))); - } - - if left_mode != right_mode { - return Err(ArrowError::InvalidArgumentError(format!( - "Cannot compare UnionArrays with different modes: left={:?}, right={:?}", - left_mode, right_mode - ))); - } - - let c_opts = child_opts(opts); - - let mut field_comparators = HashMap::with_capacity(left_fields.len()); - - for (type_id, _field) in left_fields.iter() { - let left_child = left.child(type_id); - let right_child = right.child(type_id); - let cmp = make_comparator(left_child.as_ref(), right_child.as_ref(), c_opts)?; - - field_comparators.insert(type_id, cmp); - } - - let left_type_ids = left.type_ids().clone(); - let right_type_ids = right.type_ids().clone(); - - let left_offsets = left.offsets().cloned(); - let right_offsets = right.offsets().cloned(); - - let f = compare(left, right, opts, move |i, j| { - let left_type_id = left_type_ids[i]; - let right_type_id = right_type_ids[j]; - - // first, compare by type_id - match left_type_id.cmp(&right_type_id) { - Ordering::Equal => { - // second, compare by values - let left_offset = left_offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i); - let right_offset = right_offsets.as_ref().map(|o| o[j] as usize).unwrap_or(j); - - let cmp = field_comparators - .get(&left_type_id) - .expect("type id not found in field_comparators"); - - cmp(left_offset, right_offset) - } - other => other, - } - }); - Ok(f) -} - -/// Returns a comparison function that compares two values at two different positions -/// between the two arrays. -/// -/// For comparing arrays element-wise, see also the vectorised kernels in [`crate::cmp`]. -/// -/// If `nulls_first` is true `NULL` values will be considered less than any non-null value, -/// otherwise they will be considered greater. -/// -/// # Basic Usage -/// -/// ``` -/// # use std::cmp::Ordering; -/// # use arrow_array::Int32Array; -/// # use arrow_ord::ord::make_comparator; -/// # use arrow_schema::SortOptions; -/// # -/// let array1 = Int32Array::from(vec![1, 2]); -/// let array2 = Int32Array::from(vec![3, 4]); -/// -/// let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); -/// // 1 (index 0 of array1) is smaller than 4 (index 1 of array2) -/// assert_eq!(cmp(0, 1), Ordering::Less); -/// -/// let array1 = Int32Array::from(vec![Some(1), None]); -/// let array2 = Int32Array::from(vec![None, Some(2)]); -/// let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); -/// -/// assert_eq!(cmp(0, 1), Ordering::Less); // Some(1) vs Some(2) -/// assert_eq!(cmp(1, 1), Ordering::Less); // None vs Some(2) -/// assert_eq!(cmp(1, 0), Ordering::Equal); // None vs None -/// assert_eq!(cmp(0, 0), Ordering::Greater); // Some(1) vs None -/// ``` -/// -/// # Postgres-compatible Nested Comparison -/// -/// Whilst SQL prescribes ternary logic for nulls, that is comparing a value against a NULL yields -/// a NULL, many systems, including postgres, instead apply a total ordering to comparison of -/// nested nulls. That is nulls within nested types are either greater than any value (postgres), -/// or less than any value (Spark). -/// -/// In particular -/// -/// ```ignore -/// { a: 1, b: null } == { a: 1, b: null } => true -/// { a: 1, b: null } == { a: 1, b: 1 } => false -/// { a: 1, b: null } == null => null -/// null == null => null -/// ``` -/// -/// This could be implemented as below -/// -/// ``` -/// # use arrow_array::{Array, BooleanArray}; -/// # use arrow_buffer::NullBuffer; -/// # use arrow_ord::cmp; -/// # use arrow_ord::ord::make_comparator; -/// # use arrow_schema::{ArrowError, SortOptions}; -/// fn eq(a: &dyn Array, b: &dyn Array) -> Result { -/// if !a.data_type().is_nested() { -/// return cmp::eq(&a, &b); // Use faster vectorised kernel -/// } -/// -/// let cmp = make_comparator(a, b, SortOptions::default())?; -/// let len = a.len().min(b.len()); -/// let values = (0..len).map(|i| cmp(i, i).is_eq()).collect(); -/// let nulls = NullBuffer::union(a.nulls(), b.nulls()); -/// Ok(BooleanArray::new(values, nulls)) -/// } -/// ```` -pub fn make_comparator( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - use arrow_schema::DataType::*; - - macro_rules! primitive_helper { - ($t:ty, $left:expr, $right:expr, $nulls_first:expr) => { - Ok(compare_primitive::<$t>($left, $right, $nulls_first)) - }; - } - downcast_primitive! { - left.data_type(), right.data_type() => (primitive_helper, left, right, opts), - (Boolean, Boolean) => Ok(compare_boolean(left, right, opts)), - (Utf8, Utf8) => Ok(compare_bytes::(left, right, opts)), - (LargeUtf8, LargeUtf8) => Ok(compare_bytes::(left, right, opts)), - (Utf8View, Utf8View) => Ok(compare_byte_view::(left, right, opts)), - (Binary, Binary) => Ok(compare_bytes::(left, right, opts)), - (LargeBinary, LargeBinary) => Ok(compare_bytes::(left, right, opts)), - (BinaryView, BinaryView) => Ok(compare_byte_view::(left, right, opts)), - (FixedSizeBinary(_), FixedSizeBinary(_)) => { - let left = left.as_fixed_size_binary(); - let right = right.as_fixed_size_binary(); - - let l = left.clone(); - let r = right.clone(); - Ok(compare(left, right, opts, move |i, j| { - l.value(i).cmp(r.value(j)) - })) - }, - (List(_), List(_)) => compare_list::(left, right, opts), - (LargeList(_), LargeList(_)) => compare_list::(left, right, opts), - (ListView(_), ListView(_)) => compare_list_view::(left, right, opts), - (LargeListView(_), LargeListView(_)) => compare_list_view::(left, right, opts), - (FixedSizeList(_, _), FixedSizeList(_, _)) => compare_fixed_list(left, right, opts), - (Struct(_), Struct(_)) => compare_struct(left, right, opts), - (Dictionary(l_key, _), Dictionary(r_key, _)) => { - macro_rules! dict_helper { - ($t:ty, $left:expr, $right:expr, $opts: expr) => { - compare_dict::<$t>($left, $right, $opts) - }; - } - downcast_integer! { - l_key.as_ref(), r_key.as_ref() => (dict_helper, left, right, opts), - _ => unreachable!() - } - }, - (RunEndEncoded(l_run_ends, _), RunEndEncoded(r_run_ends, _)) => { - macro_rules! run_end_helper { - ($t:ty, $left:expr, $right:expr, $opts:expr) => { - compare_run_end_encoded::<$t>($left, $right, $opts) - }; - } - downcast_run_end_index! { - l_run_ends.data_type(), r_run_ends.data_type() => (run_end_helper, left, right, opts), - _ => Err(ArrowError::InvalidArgumentError(format!( - "Cannot compare RunEndEncoded arrays with different run ends types: left={:?}, right={:?}", - l_run_ends.data_type(), - r_run_ends.data_type() - ))) - } - }, - (Map(_, _), Map(_, _)) => compare_map(left, right, opts), - (Null, Null) => Ok(Box::new(|_, _| Ordering::Equal)), - (Union(_, _), Union(_, _)) => compare_union(left, right, opts), - (lhs, rhs) => Err(ArrowError::InvalidArgumentError(match lhs == rhs { - true => format!("The data type type {lhs:?} has no natural order"), - false => "Can't compare arrays of different types".to_string(), - })) - } -} +#[cfg(test)] +use arrow_schema::{ArrowError, SortOptions}; +#[cfg(test)] +use std::cmp::Ordering; #[cfg(test)] mod tests { diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index 22ef34a909f9..f3bfabff844b 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -19,6 +19,7 @@ use arrow_array::builder::{BufferBuilder, UInt32Builder}; use arrow_array::cast::AsArray; +use arrow_array::comparator::make_comparator; use arrow_array::types::*; use arrow_array::*; use arrow_buffer::{ @@ -26,7 +27,7 @@ use arrow_buffer::{ bit_util, }; use arrow_data::ArrayDataBuilder; -use arrow_schema::{ArrowError, DataType, FieldRef, UnionMode}; +use arrow_schema::{ArrowError, DataType, FieldRef, SortOptions, UnionMode}; use std::fmt::Display; use std::mem::ManuallyDrop; use std::sync::Arc; @@ -879,12 +880,16 @@ fn take_run( let mut new_run_ends_builder = BufferBuilder::::new(1); let mut take_value_indices = BufferBuilder::::new(1); let mut new_physical_len = 1; - let ree_values = run_array.values().to_data(); + let values_cmp = make_comparator( + run_array.values().as_ref(), + run_array.values().as_ref(), + SortOptions::default(), + )?; for ix in 1..physical_indices.len() { let prev_idx = physical_indices[ix - 1]; let cur_idx = physical_indices[ix]; - if cur_idx != prev_idx && ree_values.slice(cur_idx, 1) != ree_values.slice(prev_idx, 1) { + if cur_idx != prev_idx && values_cmp(cur_idx, prev_idx).is_ne() { take_value_indices.append(I::Native::from_usize(prev_idx).unwrap()); new_run_ends_builder.append(T::Native::from_usize(ix).unwrap()); new_physical_len += 1; From c8b102719ac3cd4dcdd5154e57379a42732fcc66 Mon Sep 17 00:00:00 2001 From: Richard Date: Fri, 8 May 2026 12:21:34 -0400 Subject: [PATCH 5/7] simplify test --- arrow-select/src/take.rs | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index f3bfabff844b..a77d032e241d 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -2838,9 +2838,10 @@ mod tests { let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]); let result = take(&ree, &indexes, None).unwrap(); - let result = result.as_run::(); - assert_eq!(result.run_ends().values(), &[4]); - assert_eq!(result.values().as_primitive::().values(), &[1]); + let result = result.as_run::().downcast::().unwrap(); + + let actual = result.into_iter().flatten().collect::>(); + assert_eq!(actual, vec![1, 1, 1, 1]); } #[test] @@ -2855,9 +2856,13 @@ mod tests { let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]); let result = take(&ree, &indexes, None).unwrap(); - let result = result.as_run::(); - assert_eq!(result.run_ends().values(), &[4]); - assert_eq!(result.values().as_string::().value(0), "bob"); + let result = result + .as_run::() + .downcast::() + .unwrap(); + + let actual = result.into_iter().flatten().collect::>(); + assert_eq!(actual, vec!["bob", "bob", "bob", "bob"]); } #[test] @@ -2874,12 +2879,15 @@ mod tests { let indexes = Int32Array::from_iter_values(vec![0, 0, 1, 4, 5, 2, 3, 2, 6, 7, 6]); let result = take(&ree, &indexes, None).unwrap(); - let result = result.as_run::(); + let result = result + .as_run::() + .downcast::() + .unwrap(); - assert_eq!(result.len(), 11); - assert_eq!(result.run_ends().values(), &[5, 8, 11]); - assert_eq!(result.values().as_string::().value(0), "bob"); - assert_eq!(result.values().as_string::().value(1), "alice"); - assert_eq!(result.values().as_string::().value(2), "eve"); + let actual = result.into_iter().flatten().collect::>(); + assert_eq!( + actual, + vec!["bob", "bob", "bob", "bob", "bob", "alice", "alice", "alice", "eve", "eve", "eve"] + ); } } From c99e33335d977fc7bc62d36063246efb18cf0259 Mon Sep 17 00:00:00 2001 From: Richard Date: Fri, 15 May 2026 00:49:20 -0400 Subject: [PATCH 6/7] move REE take comparator out of arrow-array into new arrow-ord-basic crate gated by run_end_encoded feature Co-Authored-By: Claude Opus 4.7 (1M context) --- Cargo.toml | 2 + arrow-array/src/lib.rs | 2 - arrow-ord-basic/Cargo.toml | 42 +++++++++++++++++++ .../src/lib.rs | 25 ++++++++--- arrow-ord/Cargo.toml | 1 + arrow-ord/src/ord.rs | 2 +- arrow-select/Cargo.toml | 7 ++++ arrow-select/src/take.rs | 24 +++++++++-- arrow/Cargo.toml | 6 ++- 9 files changed, 99 insertions(+), 12 deletions(-) create mode 100644 arrow-ord-basic/Cargo.toml rename arrow-array/src/comparator.rs => arrow-ord-basic/src/lib.rs (95%) diff --git a/Cargo.toml b/Cargo.toml index 4ca4b068f65a..63b92f5dc297 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ members = [ "arrow-ipc", "arrow-json", "arrow-ord", + "arrow-ord-basic", "arrow-pyarrow", "arrow-row", "arrow-schema", @@ -95,6 +96,7 @@ arrow-data = { version = "58.1.0", path = "./arrow-data" } arrow-ipc = { version = "58.1.0", path = "./arrow-ipc" } arrow-json = { version = "58.1.0", path = "./arrow-json" } arrow-ord = { version = "58.1.0", path = "./arrow-ord" } +arrow-ord-basic = { version = "58.1.0", path = "./arrow-ord-basic" } arrow-pyarrow = { version = "58.1.0", path = "./arrow-pyarrow" } arrow-row = { version = "58.1.0", path = "./arrow-row" } arrow-schema = { version = "58.1.0", path = "./arrow-schema" } diff --git a/arrow-array/src/lib.rs b/arrow-array/src/lib.rs index ca17370d255a..a5f9bf5e711c 100644 --- a/arrow-array/src/lib.rs +++ b/arrow-array/src/lib.rs @@ -248,8 +248,6 @@ pub use scalar::*; pub mod builder; pub mod cast; -#[doc(hidden)] -pub mod comparator; mod delta; #[cfg(feature = "ffi")] pub mod ffi; diff --git a/arrow-ord-basic/Cargo.toml b/arrow-ord-basic/Cargo.toml new file mode 100644 index 000000000000..a4668567bb15 --- /dev/null +++ b/arrow-ord-basic/Cargo.toml @@ -0,0 +1,42 @@ +# 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. + +[package] +name = "arrow-ord-basic" +version = { workspace = true } +description = "Basic comparator/ordering building blocks for Apache Arrow, shared by crates that need slot-wise comparison without pulling in the full arrow-ord kernels" +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +keywords = { workspace = true } +include = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } + +[lib] +name = "arrow_ord_basic" +bench = false + +[package.metadata.docs.rs] +all-features = true + +[dependencies] +arrow-array = { workspace = true } +arrow-buffer = { workspace = true } +arrow-data = { workspace = true } +arrow-schema = { workspace = true } diff --git a/arrow-array/src/comparator.rs b/arrow-ord-basic/src/lib.rs similarity index 95% rename from arrow-array/src/comparator.rs rename to arrow-ord-basic/src/lib.rs index 29ee871a6422..813a5a802db3 100644 --- a/arrow-array/src/comparator.rs +++ b/arrow-ord-basic/src/lib.rs @@ -15,11 +15,26 @@ // specific language governing permissions and limitations // under the License. -//! Internal comparator factories for comparing arbitrary array slots. - -use crate::cast::AsArray; -use crate::types::*; -use crate::*; +//! Basic comparator factories shared by Arrow crates that need to compare +//! arbitrary array slots without pulling in the full [`arrow-ord`] crate. +//! +//! The only public surface is [`make_comparator`] (with [`DynComparator`] as the +//! returned function type). `arrow-ord` re-exports both from here, so its +//! public API is unchanged. +//! +//! This crate exists so that crates such as `arrow-select` can use slot-wise +//! comparison (e.g. for the run-end-encoded `take` fast path) without taking on +//! the full ordering kernel suite — which would either create a circular +//! dependency (`arrow-ord` already depends on `arrow-select`) or force every +//! downstream user of `arrow-array` to compile the comparator machinery whether +//! they need it or not. + +#![deny(rustdoc::broken_intra_doc_links)] +#![warn(missing_docs)] + +use arrow_array::cast::AsArray; +use arrow_array::types::*; +use arrow_array::*; use arrow_buffer::{ArrowNativeType, NullBuffer}; use arrow_schema::{ArrowError, DataType, SortOptions}; use std::{cmp::Ordering, collections::HashMap}; diff --git a/arrow-ord/Cargo.toml b/arrow-ord/Cargo.toml index ae76841bda39..00b8bfd6373f 100644 --- a/arrow-ord/Cargo.toml +++ b/arrow-ord/Cargo.toml @@ -39,6 +39,7 @@ all-features = true arrow-array = { workspace = true } arrow-buffer = { workspace = true } arrow-data = { workspace = true } +arrow-ord-basic = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true } diff --git a/arrow-ord/src/ord.rs b/arrow-ord/src/ord.rs index 77892d12332f..53838039e68c 100644 --- a/arrow-ord/src/ord.rs +++ b/arrow-ord/src/ord.rs @@ -17,7 +17,7 @@ //! Contains functions and function factories to compare arrays. -pub use arrow_array::comparator::{DynComparator, make_comparator}; +pub use arrow_ord_basic::{DynComparator, make_comparator}; #[cfg(test)] use arrow_array::types::*; diff --git a/arrow-select/Cargo.toml b/arrow-select/Cargo.toml index 443094e6c986..cd43c6c2a699 100644 --- a/arrow-select/Cargo.toml +++ b/arrow-select/Cargo.toml @@ -40,8 +40,15 @@ arrow-buffer = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true } arrow-array = { workspace = true } +arrow-ord-basic = { workspace = true, optional = true } num-traits = { version = "0.2.19", default-features = false, features = ["std"] } ahash = { version = "0.8", default-features = false} +[features] +# Enables the run-end-encoded `take` fast path that merges adjacent +# physical indices whose underlying values compare equal. Pulls in +# `arrow-ord-basic` for the slot-wise comparator. +run_end_encoded = ["dep:arrow-ord-basic"] + [dev-dependencies] rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index a77d032e241d..bb22eeba96e8 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -19,7 +19,6 @@ use arrow_array::builder::{BufferBuilder, UInt32Builder}; use arrow_array::cast::AsArray; -use arrow_array::comparator::make_comparator; use arrow_array::types::*; use arrow_array::*; use arrow_buffer::{ @@ -27,7 +26,11 @@ use arrow_buffer::{ bit_util, }; use arrow_data::ArrayDataBuilder; -use arrow_schema::{ArrowError, DataType, FieldRef, SortOptions, UnionMode}; +#[cfg(feature = "run_end_encoded")] +use arrow_ord_basic::make_comparator; +use arrow_schema::{ArrowError, DataType, FieldRef, UnionMode}; +#[cfg(feature = "run_end_encoded")] +use arrow_schema::SortOptions; use std::fmt::Display; use std::mem::ManuallyDrop; use std::sync::Arc; @@ -880,6 +883,13 @@ fn take_run( let mut new_run_ends_builder = BufferBuilder::::new(1); let mut take_value_indices = BufferBuilder::::new(1); let mut new_physical_len = 1; + + // With the `run_end_encoded` feature, also collapse adjacent physical runs + // whose underlying values compare equal. Without it, we still drop + // duplicate physical indices but leave logically-equal-yet-distinct runs + // separated, which keeps `arrow-select` free of the `arrow-ord-basic` + // comparator dependency for downstream users that don't need it. + #[cfg(feature = "run_end_encoded")] let values_cmp = make_comparator( run_array.values().as_ref(), run_array.values().as_ref(), @@ -889,7 +899,11 @@ fn take_run( for ix in 1..physical_indices.len() { let prev_idx = physical_indices[ix - 1]; let cur_idx = physical_indices[ix]; - if cur_idx != prev_idx && values_cmp(cur_idx, prev_idx).is_ne() { + #[cfg(feature = "run_end_encoded")] + let is_new_run = cur_idx != prev_idx && values_cmp(cur_idx, prev_idx).is_ne(); + #[cfg(not(feature = "run_end_encoded"))] + let is_new_run = cur_idx != prev_idx; + if is_new_run { take_value_indices.append(I::Native::from_usize(prev_idx).unwrap()); new_run_ends_builder.append(T::Native::from_usize(ix).unwrap()); new_physical_len += 1; @@ -2539,6 +2553,7 @@ mod tests { assert_eq!(null_buf.as_slice(), &[0b11111111]); } + #[cfg(feature = "run_end_encoded")] #[test] fn test_take_runs() { let logical_array: Vec = vec![1_i32, 1, 2, 2, 1, 1, 1, 2, 2, 1, 1, 2, 2]; @@ -2829,6 +2844,7 @@ mod tests { assert_eq!(run_result.values().len(), 0); } + #[cfg(feature = "run_end_encoded")] #[test] fn test_take_run_end_encoded_merges_identical_runs() { // see https://github.com/apache/arrow-rs/issues/7710 @@ -2844,6 +2860,7 @@ mod tests { assert_eq!(actual, vec![1, 1, 1, 1]); } + #[cfg(feature = "run_end_encoded")] #[test] fn test_take_run_end_encoded_merges_identical_string_runs() { let mut builder = StringRunBuilder::::new(); @@ -2865,6 +2882,7 @@ mod tests { assert_eq!(actual, vec!["bob", "bob", "bob", "bob"]); } + #[cfg(feature = "run_end_encoded")] #[test] fn test_take_run_end_encoded_mixed_runs() { // Validates that runs are merged whether the same logical value comes diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml index 8e56457ff0a5..a8304cc25b99 100644 --- a/arrow/Cargo.toml +++ b/arrow/Cargo.toml @@ -61,13 +61,17 @@ half = { version = "2.1", default-features = false, features = ["rand_distr"], o all-features = true [features] -default = ["csv", "ipc", "json"] +default = ["csv", "ipc", "json", "run_end_encoded"] async = ["arrow-array/async"] ipc_compression = ["ipc", "arrow-ipc/lz4", "arrow-ipc/zstd"] csv = ["arrow-csv"] ipc = ["arrow-ipc"] json = ["arrow-json"] prettyprint = ["arrow-cast/prettyprint"] +# Enables the run-end-encoded `take` fast path in `arrow-select` that merges +# adjacent physical runs whose underlying values compare equal. Pulls in the +# `arrow-ord-basic` slot-wise comparator. +run_end_encoded = ["arrow-select/run_end_encoded"] # The test utils feature enables code used in benchmarks and tests but # not the core arrow code itself. Be aware that `rand` must be kept as # an optional dependency for supporting compile to wasm32-unknown-unknown From eb8ef670d5a45bfbbf3a041c9d6360a311f4bc33 Mon Sep 17 00:00:00 2001 From: Richard Date: Fri, 10 Jul 2026 16:45:04 -0400 Subject: [PATCH 7/7] rename arrow-ord-basic to arrow-cmp, move comparator tests, remove run_end_encoded feature gate --- Cargo.toml | 4 +- {arrow-ord-basic => arrow-cmp}/Cargo.toml | 7 +- arrow-cmp/src/lib.rs | 1889 +++++++++++++++++++++ arrow-ord-basic/src/lib.rs | 535 ------ arrow-ord/Cargo.toml | 2 +- arrow-ord/src/ord.rs | 1364 +-------------- arrow-select/Cargo.toml | 8 +- arrow-select/src/take.rs | 20 +- arrow/Cargo.toml | 6 +- 9 files changed, 1902 insertions(+), 1933 deletions(-) rename {arrow-ord-basic => arrow-cmp}/Cargo.toml (91%) create mode 100644 arrow-cmp/src/lib.rs delete mode 100644 arrow-ord-basic/src/lib.rs diff --git a/Cargo.toml b/Cargo.toml index 63b92f5dc297..4b6c612a5312 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,8 +32,8 @@ members = [ "arrow-integration-testing", "arrow-ipc", "arrow-json", + "arrow-cmp", "arrow-ord", - "arrow-ord-basic", "arrow-pyarrow", "arrow-row", "arrow-schema", @@ -96,7 +96,7 @@ arrow-data = { version = "58.1.0", path = "./arrow-data" } arrow-ipc = { version = "58.1.0", path = "./arrow-ipc" } arrow-json = { version = "58.1.0", path = "./arrow-json" } arrow-ord = { version = "58.1.0", path = "./arrow-ord" } -arrow-ord-basic = { version = "58.1.0", path = "./arrow-ord-basic" } +arrow-cmp = { version = "58.1.0", path = "./arrow-cmp" } arrow-pyarrow = { version = "58.1.0", path = "./arrow-pyarrow" } arrow-row = { version = "58.1.0", path = "./arrow-row" } arrow-schema = { version = "58.1.0", path = "./arrow-schema" } diff --git a/arrow-ord-basic/Cargo.toml b/arrow-cmp/Cargo.toml similarity index 91% rename from arrow-ord-basic/Cargo.toml rename to arrow-cmp/Cargo.toml index a4668567bb15..ebbca584a719 100644 --- a/arrow-ord-basic/Cargo.toml +++ b/arrow-cmp/Cargo.toml @@ -16,7 +16,7 @@ # under the License. [package] -name = "arrow-ord-basic" +name = "arrow-cmp" version = { workspace = true } description = "Basic comparator/ordering building blocks for Apache Arrow, shared by crates that need slot-wise comparison without pulling in the full arrow-ord kernels" homepage = { workspace = true } @@ -29,7 +29,7 @@ edition = { workspace = true } rust-version = { workspace = true } [lib] -name = "arrow_ord_basic" +name = "arrow_cmp" bench = false [package.metadata.docs.rs] @@ -40,3 +40,6 @@ arrow-array = { workspace = true } arrow-buffer = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true } + +[dev-dependencies] +half = { version = "2.1", default-features = false, features = ["num-traits"] } diff --git a/arrow-cmp/src/lib.rs b/arrow-cmp/src/lib.rs new file mode 100644 index 000000000000..932941d3fbc0 --- /dev/null +++ b/arrow-cmp/src/lib.rs @@ -0,0 +1,1889 @@ +// 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. + +//! Basic comparator factories shared by Arrow crates that need to compare +//! arbitrary array slots without pulling in the full [`arrow-ord`] crate. +//! +//! The only public surface is [`make_comparator`] (with [`DynComparator`] as the +//! returned function type). `arrow-ord` re-exports both from here, so its +//! public API is unchanged. +//! +//! This crate exists so that crates such as `arrow-select` can use slot-wise +//! comparison (e.g. for the run-end-encoded `take` fast path) without taking on +//! the full ordering kernel suite — which would either create a circular +//! dependency (`arrow-ord` already depends on `arrow-select`) or force every +//! downstream user of `arrow-array` to compile the comparator machinery whether +//! they need it or not. + +#![deny(rustdoc::broken_intra_doc_links)] +#![warn(missing_docs)] + +use arrow_array::cast::AsArray; +use arrow_array::types::*; +use arrow_array::*; +use arrow_buffer::{ArrowNativeType, NullBuffer}; +use arrow_schema::{ArrowError, DataType, SortOptions}; +use std::{cmp::Ordering, collections::HashMap}; + +fn compare_run_end_encoded( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_run::(); + let right = right.as_run::(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + + let l_run_ends = left.run_ends().clone(); + let r_run_ends = right.run_ends().clone(); + + let f = compare(left, right, opts, move |i, j| { + let l_physical = l_run_ends.get_physical_index(i); + let r_physical = r_run_ends.get_physical_index(j); + cmp(l_physical, r_physical) + }); + Ok(f) +} + +/// Compare values at arbitrary indices in two arrays. +pub type DynComparator = Box Ordering + Send + Sync>; + +/// If parent sort order is descending we need to invert the value of nulls_first so that +/// when the parent is sorted based on the produced ranks, nulls are still ordered correctly +fn child_opts(opts: SortOptions) -> SortOptions { + SortOptions { + descending: false, + nulls_first: opts.nulls_first != opts.descending, + } +} + +fn compare(l: &A, r: &A, opts: SortOptions, cmp: F) -> DynComparator +where + A: Array + Clone, + F: Fn(usize, usize) -> Ordering + Send + Sync + 'static, +{ + let l = l.logical_nulls().filter(|x| x.null_count() > 0); + let r = r.logical_nulls().filter(|x| x.null_count() > 0); + match (opts.nulls_first, opts.descending) { + (true, true) => compare_impl::(l, r, cmp), + (true, false) => compare_impl::(l, r, cmp), + (false, true) => compare_impl::(l, r, cmp), + (false, false) => compare_impl::(l, r, cmp), + } +} + +fn compare_impl( + l: Option, + r: Option, + cmp: F, +) -> DynComparator +where + F: Fn(usize, usize) -> Ordering + Send + Sync + 'static, +{ + let cmp = move |i, j| match DESCENDING { + true => cmp(i, j).reverse(), + false => cmp(i, j), + }; + + let (left_null, right_null) = match NULLS_FIRST { + true => (Ordering::Less, Ordering::Greater), + false => (Ordering::Greater, Ordering::Less), + }; + + match (l, r) { + (None, None) => Box::new(cmp), + (Some(l), None) => Box::new(move |i, j| match l.is_null(i) { + true => left_null, + false => cmp(i, j), + }), + (None, Some(r)) => Box::new(move |i, j| match r.is_null(j) { + true => right_null, + false => cmp(i, j), + }), + (Some(l), Some(r)) => Box::new(move |i, j| match (l.is_null(i), r.is_null(j)) { + (true, true) => Ordering::Equal, + (true, false) => left_null, + (false, true) => right_null, + (false, false) => cmp(i, j), + }), + } +} + +fn compare_primitive( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> DynComparator +where + T::Native: ArrowNativeTypeOp, +{ + let left = left.as_primitive::(); + let right = right.as_primitive::(); + let l_values = left.values().clone(); + let r_values = right.values().clone(); + + compare(&left, &right, opts, move |i, j| { + l_values[i].compare(r_values[j]) + }) +} + +fn compare_boolean(left: &dyn Array, right: &dyn Array, opts: SortOptions) -> DynComparator { + let left = left.as_boolean(); + let right = right.as_boolean(); + + let l_values = left.values().clone(); + let r_values = right.values().clone(); + + compare(left, right, opts, move |i, j| { + l_values.value(i).cmp(&r_values.value(j)) + }) +} + +fn compare_bytes( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> DynComparator { + let left = left.as_bytes::(); + let right = right.as_bytes::(); + + let l = left.clone(); + let r = right.clone(); + compare(left, right, opts, move |i, j| { + let l: &[u8] = l.value(i).as_ref(); + let r: &[u8] = r.value(j).as_ref(); + l.cmp(r) + }) +} + +fn compare_byte_view( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> DynComparator { + let left = left.as_byte_view::(); + let right = right.as_byte_view::(); + + let l = left.clone(); + let r = right.clone(); + compare(left, right, opts, move |i, j| { + compare_byte_view_values(&l, i, &r, j) + }) +} + +fn compare_byte_view_values( + left: &GenericByteViewArray, + left_idx: usize, + right: &GenericByteViewArray, + right_idx: usize, +) -> Ordering { + assert!(left_idx < left.len()); + assert!(right_idx < right.len()); + + if left.data_buffers().is_empty() && right.data_buffers().is_empty() { + let l_view = unsafe { left.views().get_unchecked(left_idx) }; + let r_view = unsafe { right.views().get_unchecked(right_idx) }; + return GenericByteViewArray::::inline_key_fast(*l_view) + .cmp(&GenericByteViewArray::::inline_key_fast(*r_view)); + } + + unsafe { GenericByteViewArray::compare_unchecked(left, left_idx, right, right_idx) } +} + +fn compare_dict( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_dictionary::(); + let right = right.as_dictionary::(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + let left_keys = left.keys().values().clone(); + let right_keys = right.keys().values().clone(); + + let f = compare(left, right, opts, move |i, j| { + let l = left_keys[i].as_usize(); + let r = right_keys[j].as_usize(); + cmp(l, r) + }); + Ok(f) +} + +fn compare_list( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_list::(); + let right = right.as_list::(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + + let l_o = left.offsets().clone(); + let r_o = right.offsets().clone(); + let f = compare(left, right, opts, move |i, j| { + let l_end = l_o[i + 1].as_usize(); + let l_start = l_o[i].as_usize(); + + let r_end = r_o[j + 1].as_usize(); + let r_start = r_o[j].as_usize(); + + for (i, j) in (l_start..l_end).zip(r_start..r_end) { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + (l_end - l_start).cmp(&(r_end - r_start)) + }); + Ok(f) +} + +fn compare_fixed_list( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_fixed_size_list(); + let right = right.as_fixed_size_list(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + + let l_size = left.value_length().to_usize().unwrap(); + let r_size = right.value_length().to_usize().unwrap(); + let size_cmp = l_size.cmp(&r_size); + + let f = compare(left, right, opts, move |i, j| { + let l_start = i * l_size; + let l_end = l_start + l_size; + let r_start = j * r_size; + let r_end = r_start + r_size; + for (i, j) in (l_start..l_end).zip(r_start..r_end) { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + size_cmp + }); + Ok(f) +} + +fn compare_list_view( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_list_view::(); + let right = right.as_list_view::(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; + + let l_offsets = left.offsets().clone(); + let l_sizes = left.sizes().clone(); + let r_offsets = right.offsets().clone(); + let r_sizes = right.sizes().clone(); + + let f = compare(left, right, opts, move |i, j| { + let l_start = l_offsets[i].as_usize(); + let l_len = l_sizes[i].as_usize(); + let l_end = l_start + l_len; + + let r_start = r_offsets[j].as_usize(); + let r_len = r_sizes[j].as_usize(); + let r_end = r_start + r_len; + + for (i, j) in (l_start..l_end).zip(r_start..r_end) { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + l_len.cmp(&r_len) + }); + Ok(f) +} + +fn compare_map( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_map(); + let right = right.as_map(); + + let c_opts = child_opts(opts); + let cmp = make_comparator(left.entries(), right.entries(), c_opts)?; + + let l_o = left.offsets().clone(); + let r_o = right.offsets().clone(); + let f = compare(left, right, opts, move |i, j| { + let l_end = l_o[i + 1].as_usize(); + let l_start = l_o[i].as_usize(); + + let r_end = r_o[j + 1].as_usize(); + let r_start = r_o[j].as_usize(); + + for (i, j) in (l_start..l_end).zip(r_start..r_end) { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + (l_end - l_start).cmp(&(r_end - r_start)) + }); + Ok(f) +} + +fn compare_struct( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_struct(); + let right = right.as_struct(); + + if left.columns().len() != right.columns().len() { + return Err(ArrowError::InvalidArgumentError( + "Cannot compare StructArray with different number of columns".to_string(), + )); + } + + let c_opts = child_opts(opts); + let columns = left.columns().iter().zip(right.columns()); + let comparators = columns + .map(|(l, r)| make_comparator(l, r, c_opts)) + .collect::, _>>()?; + + let f = compare(left, right, opts, move |i, j| { + for cmp in &comparators { + match cmp(i, j) { + Ordering::Equal => continue, + r => return r, + } + } + Ordering::Equal + }); + Ok(f) +} + +fn compare_union( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + let left = left.as_union(); + let right = right.as_union(); + + let (left_fields, left_mode) = match left.data_type() { + DataType::Union(fields, mode) => (fields, mode), + _ => unreachable!(), + }; + let (right_fields, right_mode) = match right.data_type() { + DataType::Union(fields, mode) => (fields, mode), + _ => unreachable!(), + }; + + if left_fields != right_fields { + return Err(ArrowError::InvalidArgumentError(format!( + "Cannot compare UnionArrays with different fields: left={left_fields:?}, right={right_fields:?}" + ))); + } + + if left_mode != right_mode { + return Err(ArrowError::InvalidArgumentError(format!( + "Cannot compare UnionArrays with different modes: left={left_mode:?}, right={right_mode:?}" + ))); + } + + let c_opts = child_opts(opts); + + let mut field_comparators = HashMap::with_capacity(left_fields.len()); + + for (type_id, _field) in left_fields.iter() { + let left_child = left.child(type_id); + let right_child = right.child(type_id); + let cmp = make_comparator(left_child.as_ref(), right_child.as_ref(), c_opts)?; + + field_comparators.insert(type_id, cmp); + } + + let left_type_ids = left.type_ids().clone(); + let right_type_ids = right.type_ids().clone(); + + let left_offsets = left.offsets().cloned(); + let right_offsets = right.offsets().cloned(); + + let f = compare(left, right, opts, move |i, j| { + let left_type_id = left_type_ids[i]; + let right_type_id = right_type_ids[j]; + + match left_type_id.cmp(&right_type_id) { + Ordering::Equal => { + let left_offset = left_offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i); + let right_offset = right_offsets.as_ref().map(|o| o[j] as usize).unwrap_or(j); + + let cmp = field_comparators + .get(&left_type_id) + .expect("type id not found in field_comparators"); + + cmp(left_offset, right_offset) + } + other => other, + } + }); + Ok(f) +} + +/// Returns a comparison function that compares two values at two arbitrary indices. +/// +/// If `nulls_first` is true, null values are considered less than any non-null +/// value; otherwise they are considered greater. This is primarily shared by +/// crates that need repeated slot comparisons without constructing sliced arrays. +pub fn make_comparator( + left: &dyn Array, + right: &dyn Array, + opts: SortOptions, +) -> Result { + use arrow_schema::DataType::*; + + macro_rules! primitive_helper { + ($t:ty, $left:expr, $right:expr, $nulls_first:expr) => { + Ok(compare_primitive::<$t>($left, $right, $nulls_first)) + }; + } + downcast_primitive! { + left.data_type(), right.data_type() => (primitive_helper, left, right, opts), + (Boolean, Boolean) => Ok(compare_boolean(left, right, opts)), + (Utf8, Utf8) => Ok(compare_bytes::(left, right, opts)), + (LargeUtf8, LargeUtf8) => Ok(compare_bytes::(left, right, opts)), + (Utf8View, Utf8View) => Ok(compare_byte_view::(left, right, opts)), + (Binary, Binary) => Ok(compare_bytes::(left, right, opts)), + (LargeBinary, LargeBinary) => Ok(compare_bytes::(left, right, opts)), + (BinaryView, BinaryView) => Ok(compare_byte_view::(left, right, opts)), + (FixedSizeBinary(_), FixedSizeBinary(_)) => { + let left = left.as_fixed_size_binary(); + let right = right.as_fixed_size_binary(); + + let l = left.clone(); + let r = right.clone(); + Ok(compare(left, right, opts, move |i, j| { + l.value(i).cmp(r.value(j)) + })) + }, + (List(_), List(_)) => compare_list::(left, right, opts), + (LargeList(_), LargeList(_)) => compare_list::(left, right, opts), + (ListView(_), ListView(_)) => compare_list_view::(left, right, opts), + (LargeListView(_), LargeListView(_)) => compare_list_view::(left, right, opts), + (FixedSizeList(_, _), FixedSizeList(_, _)) => compare_fixed_list(left, right, opts), + (Struct(_), Struct(_)) => compare_struct(left, right, opts), + (Dictionary(l_key, _), Dictionary(r_key, _)) => { + macro_rules! dict_helper { + ($t:ty, $left:expr, $right:expr, $opts: expr) => { + compare_dict::<$t>($left, $right, $opts) + }; + } + downcast_integer! { + l_key.as_ref(), r_key.as_ref() => (dict_helper, left, right, opts), + _ => unreachable!() + } + }, + (RunEndEncoded(l_run_ends, _), RunEndEncoded(r_run_ends, _)) => { + macro_rules! run_end_helper { + ($t:ty, $left:expr, $right:expr, $opts:expr) => { + compare_run_end_encoded::<$t>($left, $right, $opts) + }; + } + downcast_run_end_index! { + l_run_ends.data_type(), r_run_ends.data_type() => (run_end_helper, left, right, opts), + _ => Err(ArrowError::InvalidArgumentError(format!( + "Cannot compare RunEndEncoded arrays with different run ends types: left={:?}, right={:?}", + l_run_ends.data_type(), + r_run_ends.data_type() + ))) + } + }, + (Map(_, _), Map(_, _)) => compare_map(left, right, opts), + (Null, Null) => Ok(Box::new(|_, _| Ordering::Equal)), + (Union(_, _), Union(_, _)) => compare_union(left, right, opts), + (lhs, rhs) => Err(ArrowError::InvalidArgumentError(match lhs == rhs { + true => format!("The data type type {lhs:?} has no natural order"), + false => "Can't compare arrays of different types".to_string(), + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder}; + use arrow_buffer::{IntervalDayTime, NullBuffer, OffsetBuffer, ScalarBuffer, i256}; + use arrow_schema::{ArrowError, DataType, Field, Fields, UnionFields}; + use half::f16; + use std::cmp::Ordering; + use std::sync::Arc; + + #[test] + fn test_fixed_size_binary() { + let items = vec![vec![1u8], vec![2u8]]; + let array = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap(); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + } + + #[test] + fn test_fixed_size_binary_fixed_size_binary() { + let items = vec![vec![1u8]]; + let array1 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap(); + let items = vec![vec![2u8]]; + let array2 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap(); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + } + + #[test] + fn test_i32() { + let array = Int32Array::from(vec![1, 2]); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, (cmp)(0, 1)); + } + + #[test] + fn test_i32_i32() { + let array1 = Int32Array::from(vec![1]); + let array2 = Int32Array::from(vec![2]); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + } + + #[test] + fn test_f16() { + let array = Float16Array::from(vec![f16::from_f32(1.0), f16::from_f32(2.0)]); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + } + + #[test] + fn test_f64() { + let array = Float64Array::from(vec![1.0, 2.0]); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + } + + #[test] + fn test_f64_nan() { + let array = Float64Array::from(vec![1.0, f64::NAN]); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + assert_eq!(Ordering::Equal, cmp(1, 1)); + } + + #[test] + fn test_f64_zeros() { + let array = Float64Array::from(vec![-0.0, 0.0]); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + assert_eq!(Ordering::Greater, cmp(1, 0)); + } + + #[test] + fn test_interval_day_time() { + let array = IntervalDayTimeArray::from(vec![ + // 0 days, 1 second + IntervalDayTimeType::make_value(0, 1000), + // 1 day, 2 milliseconds + IntervalDayTimeType::make_value(1, 2), + // 90M milliseconds (which is more than is in 1 day) + IntervalDayTimeType::make_value(0, 90_000_000), + ]); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + assert_eq!(Ordering::Greater, cmp(1, 0)); + + // somewhat confusingly, while 90M milliseconds is more than 1 day, + // it will compare less as the comparison is done on the underlying + // values not field by field + assert_eq!(Ordering::Greater, cmp(1, 2)); + assert_eq!(Ordering::Less, cmp(2, 1)); + } + + #[test] + fn test_interval_year_month() { + let array = IntervalYearMonthArray::from(vec![ + // 1 year, 0 months + IntervalYearMonthType::make_value(1, 0), + // 0 years, 13 months + IntervalYearMonthType::make_value(0, 13), + // 1 year, 1 month + IntervalYearMonthType::make_value(1, 1), + ]); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + assert_eq!(Ordering::Greater, cmp(1, 0)); + + // the underlying representation is months, so both quantities are the same + assert_eq!(Ordering::Equal, cmp(1, 2)); + assert_eq!(Ordering::Equal, cmp(2, 1)); + } + + #[test] + fn test_interval_month_day_nano() { + let array = IntervalMonthDayNanoArray::from(vec![ + // 100 days + IntervalMonthDayNanoType::make_value(0, 100, 0), + // 1 month + IntervalMonthDayNanoType::make_value(1, 0, 0), + // 100 day, 1 nanoseconds + IntervalMonthDayNanoType::make_value(0, 100, 2), + ]); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + assert_eq!(Ordering::Greater, cmp(1, 0)); + + // somewhat confusingly, while 100 days is more than 1 month in all cases + // it will compare less as the comparison is done on the underlying + // values not field by field + assert_eq!(Ordering::Greater, cmp(1, 2)); + assert_eq!(Ordering::Less, cmp(2, 1)); + } + + #[test] + fn test_decimali32() { + let array = vec![Some(5_i32), Some(2_i32), Some(3_i32)] + .into_iter() + .collect::() + .with_precision_and_scale(8, 6) + .unwrap(); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + assert_eq!(Ordering::Less, cmp(1, 0)); + assert_eq!(Ordering::Greater, cmp(0, 2)); + } + + #[test] + fn test_decimali64() { + let array = vec![Some(5_i64), Some(2_i64), Some(3_i64)] + .into_iter() + .collect::() + .with_precision_and_scale(16, 6) + .unwrap(); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + assert_eq!(Ordering::Less, cmp(1, 0)); + assert_eq!(Ordering::Greater, cmp(0, 2)); + } + + #[test] + fn test_decimali128() { + let array = vec![Some(5_i128), Some(2_i128), Some(3_i128)] + .into_iter() + .collect::() + .with_precision_and_scale(23, 6) + .unwrap(); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + assert_eq!(Ordering::Less, cmp(1, 0)); + assert_eq!(Ordering::Greater, cmp(0, 2)); + } + + #[test] + fn test_decimali256() { + let array = vec![ + Some(i256::from_i128(5_i128)), + Some(i256::from_i128(2_i128)), + Some(i256::from_i128(3_i128)), + ] + .into_iter() + .collect::() + .with_precision_and_scale(53, 6) + .unwrap(); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + assert_eq!(Ordering::Less, cmp(1, 0)); + assert_eq!(Ordering::Greater, cmp(0, 2)); + } + + #[test] + fn test_dict() { + let data = vec!["a", "b", "c", "a", "a", "c", "c"]; + let array = data.into_iter().collect::>(); + + let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + assert_eq!(Ordering::Equal, cmp(3, 4)); + assert_eq!(Ordering::Greater, cmp(2, 3)); + } + + #[test] + fn test_multiple_dict() { + let d1 = vec!["a", "b", "c", "d"]; + let a1 = d1.into_iter().collect::>(); + let d2 = vec!["e", "f", "g", "a"]; + let a2 = d2.into_iter().collect::>(); + + let cmp = make_comparator(&a1, &a2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + assert_eq!(Ordering::Equal, cmp(0, 3)); + assert_eq!(Ordering::Greater, cmp(1, 3)); + } + + #[test] + fn test_primitive_dict() { + let values = Int32Array::from(vec![1_i32, 0, 2, 5]); + let keys = Int8Array::from_iter_values([0, 0, 1, 3]); + let array1 = DictionaryArray::new(keys, Arc::new(values)); + + let values = Int32Array::from(vec![2_i32, 3, 4, 5]); + let keys = Int8Array::from_iter_values([0, 1, 1, 3]); + let array2 = DictionaryArray::new(keys, Arc::new(values)); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + assert_eq!(Ordering::Less, cmp(0, 3)); + assert_eq!(Ordering::Equal, cmp(3, 3)); + assert_eq!(Ordering::Greater, cmp(3, 1)); + assert_eq!(Ordering::Greater, cmp(3, 2)); + } + + #[test] + fn test_float_dict() { + let values = Float32Array::from(vec![1.0, 0.5, 2.1, 5.5]); + let keys = Int8Array::from_iter_values([0, 0, 1, 3]); + let array1 = DictionaryArray::try_new(keys, Arc::new(values)).unwrap(); + + let values = Float32Array::from(vec![1.2, 3.2, 4.0, 5.5]); + let keys = Int8Array::from_iter_values([0, 1, 1, 3]); + let array2 = DictionaryArray::new(keys, Arc::new(values)); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + assert_eq!(Ordering::Less, cmp(0, 3)); + assert_eq!(Ordering::Equal, cmp(3, 3)); + assert_eq!(Ordering::Greater, cmp(3, 1)); + assert_eq!(Ordering::Greater, cmp(3, 2)); + } + + #[test] + fn test_timestamp_dict() { + let values = TimestampSecondArray::from(vec![1, 0, 2, 5]); + let keys = Int8Array::from_iter_values([0, 0, 1, 3]); + let array1 = DictionaryArray::new(keys, Arc::new(values)); + + let values = TimestampSecondArray::from(vec![2, 3, 4, 5]); + let keys = Int8Array::from_iter_values([0, 1, 1, 3]); + let array2 = DictionaryArray::new(keys, Arc::new(values)); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + assert_eq!(Ordering::Less, cmp(0, 3)); + assert_eq!(Ordering::Equal, cmp(3, 3)); + assert_eq!(Ordering::Greater, cmp(3, 1)); + assert_eq!(Ordering::Greater, cmp(3, 2)); + } + + #[test] + fn test_interval_dict() { + let v1 = IntervalDayTime::new(0, 1); + let v2 = IntervalDayTime::new(0, 2); + let v3 = IntervalDayTime::new(12, 2); + + let values = IntervalDayTimeArray::from(vec![Some(v1), Some(v2), None, Some(v3)]); + let keys = Int8Array::from_iter_values([0, 0, 1, 3]); + let array1 = DictionaryArray::new(keys, Arc::new(values)); + + let values = IntervalDayTimeArray::from(vec![Some(v3), Some(v2), None, Some(v1)]); + let keys = Int8Array::from_iter_values([0, 1, 1, 3]); + let array2 = DictionaryArray::new(keys, Arc::new(values)); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); // v1 vs v3 + assert_eq!(Ordering::Equal, cmp(0, 3)); // v1 vs v1 + assert_eq!(Ordering::Greater, cmp(3, 3)); // v3 vs v1 + assert_eq!(Ordering::Greater, cmp(3, 1)); // v3 vs v2 + assert_eq!(Ordering::Greater, cmp(3, 2)); // v3 vs v2 + } + + #[test] + fn test_duration_dict() { + let values = DurationSecondArray::from(vec![1, 0, 2, 5]); + let keys = Int8Array::from_iter_values([0, 0, 1, 3]); + let array1 = DictionaryArray::new(keys, Arc::new(values)); + + let values = DurationSecondArray::from(vec![2, 3, 4, 5]); + let keys = Int8Array::from_iter_values([0, 1, 1, 3]); + let array2 = DictionaryArray::new(keys, Arc::new(values)); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + assert_eq!(Ordering::Less, cmp(0, 3)); + assert_eq!(Ordering::Equal, cmp(3, 3)); + assert_eq!(Ordering::Greater, cmp(3, 1)); + assert_eq!(Ordering::Greater, cmp(3, 2)); + } + + #[test] + fn test_decimal_dict() { + let values = Decimal128Array::from(vec![1, 0, 2, 5]); + let keys = Int8Array::from_iter_values([0, 0, 1, 3]); + let array1 = DictionaryArray::new(keys, Arc::new(values)); + + let values = Decimal128Array::from(vec![2, 3, 4, 5]); + let keys = Int8Array::from_iter_values([0, 1, 1, 3]); + let array2 = DictionaryArray::new(keys, Arc::new(values)); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + assert_eq!(Ordering::Less, cmp(0, 3)); + assert_eq!(Ordering::Equal, cmp(3, 3)); + assert_eq!(Ordering::Greater, cmp(3, 1)); + assert_eq!(Ordering::Greater, cmp(3, 2)); + } + + #[test] + fn test_decimal256_dict() { + let values = Decimal256Array::from(vec![ + i256::from_i128(1), + i256::from_i128(0), + i256::from_i128(2), + i256::from_i128(5), + ]); + let keys = Int8Array::from_iter_values([0, 0, 1, 3]); + let array1 = DictionaryArray::new(keys, Arc::new(values)); + + let values = Decimal256Array::from(vec![ + i256::from_i128(2), + i256::from_i128(3), + i256::from_i128(4), + i256::from_i128(5), + ]); + let keys = Int8Array::from_iter_values([0, 1, 1, 3]); + let array2 = DictionaryArray::new(keys, Arc::new(values)); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 0)); + assert_eq!(Ordering::Less, cmp(0, 3)); + assert_eq!(Ordering::Equal, cmp(3, 3)); + assert_eq!(Ordering::Greater, cmp(3, 1)); + assert_eq!(Ordering::Greater, cmp(3, 2)); + } + + fn test_bytes_impl() { + let offsets = OffsetBuffer::from_lengths([3, 3, 1]); + let a = GenericByteArray::::new(offsets, b"abcdefa".into(), None); + let cmp = make_comparator(&a, &a, SortOptions::default()).unwrap(); + + assert_eq!(Ordering::Less, cmp(0, 1)); + assert_eq!(Ordering::Greater, cmp(0, 2)); + assert_eq!(Ordering::Equal, cmp(1, 1)); + } + + #[test] + fn test_bytes() { + test_bytes_impl::(); + test_bytes_impl::(); + test_bytes_impl::(); + test_bytes_impl::(); + } + + fn assert_cmp_cases( + array1: &A, + array2: &A, + opts: SortOptions, + cases: &[(usize, usize, Ordering)], + ) { + let cmp = make_comparator(array1, array2, opts).unwrap(); + for (left, right, expected) in cases { + assert_eq!(cmp(*left, *right), *expected); + } + } + + #[test] + fn test_lists() { + let mut a = ListBuilder::new(ListBuilder::new(Int32Builder::new())); + a.extend([ + Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]), + Some(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![Some(1)]), + ]), + Some(vec![]), + ]); + let a = a.finish(); + let mut b = ListBuilder::new(ListBuilder::new(Int32Builder::new())); + b.extend([ + Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]), + Some(vec![ + Some(vec![Some(1), Some(2), None]), + Some(vec![Some(1)]), + ]), + Some(vec![ + Some(vec![Some(1), Some(2), Some(3), Some(4)]), + Some(vec![Some(1)]), + ]), + None, + ]); + let b = b.finish(); + + // Ascending with nulls first. + assert_cmp_cases( + &a, + &b, + SortOptions { + descending: false, + nulls_first: true, + }, + &[ + (0, 0, Ordering::Equal), + (0, 1, Ordering::Less), + (0, 2, Ordering::Less), + (1, 2, Ordering::Less), + (1, 3, Ordering::Greater), + (2, 0, Ordering::Less), + ], + ); + + // Descending with nulls first. + assert_cmp_cases( + &a, + &b, + SortOptions { + descending: true, + nulls_first: true, + }, + &[ + (0, 0, Ordering::Equal), + (0, 1, Ordering::Less), + (0, 2, Ordering::Less), + (1, 2, Ordering::Greater), + (1, 3, Ordering::Greater), + (2, 0, Ordering::Greater), + ], + ); + + // Descending with nulls last. + assert_cmp_cases( + &a, + &b, + SortOptions { + descending: true, + nulls_first: false, + }, + &[ + (0, 0, Ordering::Equal), + (0, 1, Ordering::Greater), + (0, 2, Ordering::Greater), + (1, 2, Ordering::Greater), + (1, 3, Ordering::Less), + (2, 0, Ordering::Greater), + ], + ); + + // Ascending with nulls last. + assert_cmp_cases( + &a, + &b, + SortOptions { + descending: false, + nulls_first: false, + }, + &[ + (0, 0, Ordering::Equal), + (0, 1, Ordering::Greater), + (0, 2, Ordering::Greater), + (1, 2, Ordering::Less), + (1, 3, Ordering::Less), + (2, 0, Ordering::Less), + ], + ); + } + + fn list_view_array( + values: Vec, + offsets: &[usize], + sizes: &[usize], + valid: Option<&[bool]>, + ) -> GenericListViewArray { + let offsets = offsets + .iter() + .map(|v| O::from_usize(*v).unwrap()) + .collect::>(); + let sizes = sizes + .iter() + .map(|v| O::from_usize(*v).unwrap()) + .collect::>(); + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let values = Int32Array::from(values); + let nulls = valid.map(NullBuffer::from); + GenericListViewArray::new(field, offsets, sizes, Arc::new(values), nulls) + } + + fn test_list_view_comparisons() { + let array = list_view_array::( + vec![1, 2, 3, 4, 5], + &[0, 2, 1, 0, 3], + &[2, 2, 2, 0, 2], + Some(&[true, true, true, true, false]), + ); + + // Ascending with nulls first (non-monotonic offsets and empty list). + assert_cmp_cases( + &array, + &array, + SortOptions { + descending: false, + nulls_first: true, + }, + &[ + (0, 2, Ordering::Less), // [1,2] < [2,3] + (1, 2, Ordering::Greater), // [3,4] > [2,3] + (3, 0, Ordering::Less), // [] < [1,2] + (4, 0, Ordering::Less), // null < [1,2] + ], + ); + + // Ascending with nulls last. + assert_cmp_cases( + &array, + &array, + SortOptions { + descending: false, + nulls_first: false, + }, + &[ + (0, 2, Ordering::Less), + (1, 2, Ordering::Greater), + (3, 0, Ordering::Less), + (4, 0, Ordering::Greater), // null last + ], + ); + + // Descending with nulls first. + assert_cmp_cases( + &array, + &array, + SortOptions { + descending: true, + nulls_first: true, + }, + &[ + (0, 2, Ordering::Greater), + (1, 2, Ordering::Less), + (3, 0, Ordering::Greater), + (4, 0, Ordering::Less), + ], + ); + + // Descending with nulls last. + assert_cmp_cases( + &array, + &array, + SortOptions { + descending: true, + nulls_first: false, + }, + &[ + (0, 2, Ordering::Greater), + (1, 2, Ordering::Less), + (3, 0, Ordering::Greater), + (4, 0, Ordering::Greater), + ], + ); + } + + #[test] + fn test_list_view() { + test_list_view_comparisons::(); + } + + #[test] + fn test_large_list_view() { + test_list_view_comparisons::(); + } + + #[test] + fn test_struct() { + let fields = Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new_list("b", Field::new_list_field(DataType::Int32, true), true), + ]); + + let a = Int32Array::from(vec![Some(1), Some(2), None, None]); + let mut b = ListBuilder::new(Int32Builder::new()); + b.extend([Some(vec![Some(1), Some(2)]), Some(vec![None]), None, None]); + let b = b.finish(); + + let nulls = Some(NullBuffer::from_iter([true, true, true, false])); + let values = vec![Arc::new(a) as _, Arc::new(b) as _]; + let s1 = StructArray::new(fields.clone(), values, nulls); + + let a = Int32Array::from(vec![None, Some(2), None]); + let mut b = ListBuilder::new(Int32Builder::new()); + b.extend([None, None, Some(vec![])]); + let b = b.finish(); + + let values = vec![Arc::new(a) as _, Arc::new(b) as _]; + let s2 = StructArray::new(fields.clone(), values, None); + + let opts = SortOptions { + descending: false, + nulls_first: true, + }; + let cmp = make_comparator(&s1, &s2, opts).unwrap(); + assert_eq!(cmp(0, 1), Ordering::Less); // (1, [1, 2]) cmp (2, None) + assert_eq!(cmp(0, 0), Ordering::Greater); // (1, [1, 2]) cmp (None, None) + assert_eq!(cmp(1, 1), Ordering::Greater); // (2, [None]) cmp (2, None) + assert_eq!(cmp(2, 2), Ordering::Less); // (None, None) cmp (None, []) + assert_eq!(cmp(3, 0), Ordering::Less); // None cmp (None, []) + assert_eq!(cmp(2, 0), Ordering::Equal); // (None, None) cmp (None, None) + assert_eq!(cmp(3, 0), Ordering::Less); // None cmp (None, None) + + let opts = SortOptions { + descending: true, + nulls_first: true, + }; + let cmp = make_comparator(&s1, &s2, opts).unwrap(); + assert_eq!(cmp(0, 1), Ordering::Greater); // (1, [1, 2]) cmp (2, None) + assert_eq!(cmp(0, 0), Ordering::Greater); // (1, [1, 2]) cmp (None, None) + assert_eq!(cmp(1, 1), Ordering::Greater); // (2, [None]) cmp (2, None) + assert_eq!(cmp(2, 2), Ordering::Less); // (None, None) cmp (None, []) + assert_eq!(cmp(3, 0), Ordering::Less); // None cmp (None, []) + assert_eq!(cmp(2, 0), Ordering::Equal); // (None, None) cmp (None, None) + assert_eq!(cmp(3, 0), Ordering::Less); // None cmp (None, None) + + let opts = SortOptions { + descending: true, + nulls_first: false, + }; + let cmp = make_comparator(&s1, &s2, opts).unwrap(); + assert_eq!(cmp(0, 1), Ordering::Greater); // (1, [1, 2]) cmp (2, None) + assert_eq!(cmp(0, 0), Ordering::Less); // (1, [1, 2]) cmp (None, None) + assert_eq!(cmp(1, 1), Ordering::Less); // (2, [None]) cmp (2, None) + assert_eq!(cmp(2, 2), Ordering::Greater); // (None, None) cmp (None, []) + assert_eq!(cmp(3, 0), Ordering::Greater); // None cmp (None, []) + assert_eq!(cmp(2, 0), Ordering::Equal); // (None, None) cmp (None, None) + assert_eq!(cmp(3, 0), Ordering::Greater); // None cmp (None, None) + + let opts = SortOptions { + descending: false, + nulls_first: false, + }; + let cmp = make_comparator(&s1, &s2, opts).unwrap(); + assert_eq!(cmp(0, 1), Ordering::Less); // (1, [1, 2]) cmp (2, None) + assert_eq!(cmp(0, 0), Ordering::Less); // (1, [1, 2]) cmp (None, None) + assert_eq!(cmp(1, 1), Ordering::Less); // (2, [None]) cmp (2, None) + assert_eq!(cmp(2, 2), Ordering::Greater); // (None, None) cmp (None, []) + assert_eq!(cmp(3, 0), Ordering::Greater); // None cmp (None, []) + assert_eq!(cmp(2, 0), Ordering::Equal); // (None, None) cmp (None, None) + assert_eq!(cmp(3, 0), Ordering::Greater); // None cmp (None, None) + } + + #[test] + fn test_map() { + // Create first map array demonstrating key priority over values: + // [{"a": 100, "b": 1}, {"b": 999, "c": 1}, {}, {"x": 1}] + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map1_builder = MapBuilder::new(None, string_builder, int_builder); + + // {"a": 100, "b": 1} - high value for "a", low value for "b" + map1_builder.keys().append_value("a"); + map1_builder.values().append_value(100); + map1_builder.keys().append_value("b"); + map1_builder.values().append_value(1); + map1_builder.append(true).unwrap(); + + // {"b": 999, "c": 1} - very high value for "b", low value for "c" + map1_builder.keys().append_value("b"); + map1_builder.values().append_value(999); + map1_builder.keys().append_value("c"); + map1_builder.values().append_value(1); + map1_builder.append(true).unwrap(); + + // {} + map1_builder.append(true).unwrap(); + + // {"x": 1} + map1_builder.keys().append_value("x"); + map1_builder.values().append_value(1); + map1_builder.append(true).unwrap(); + + let map1 = map1_builder.finish(); + + // Create second map array: + // [{"a": 1, "c": 999}, {"b": 1, "d": 999}, {"a": 1}, None] + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map2_builder = MapBuilder::new(None, string_builder, int_builder); + + // {"a": 1, "c": 999} - low value for "a", high value for "c" + map2_builder.keys().append_value("a"); + map2_builder.values().append_value(1); + map2_builder.keys().append_value("c"); + map2_builder.values().append_value(999); + map2_builder.append(true).unwrap(); + + // {"b": 1, "d": 999} - low value for "b", high value for "d" + map2_builder.keys().append_value("b"); + map2_builder.values().append_value(1); + map2_builder.keys().append_value("d"); + map2_builder.values().append_value(999); + map2_builder.append(true).unwrap(); + + // {"a": 1} + map2_builder.keys().append_value("a"); + map2_builder.values().append_value(1); + map2_builder.append(true).unwrap(); + + // None + map2_builder.append(false).unwrap(); + + let map2 = map2_builder.finish(); + + let opts = SortOptions { + descending: false, + nulls_first: true, + }; + let cmp = make_comparator(&map1, &map2, opts).unwrap(); + + // Test that keys have priority over values: + // {"a": 100, "b": 1} vs {"a": 1, "c": 999} + // First entries match (a:100 vs a:1), but 100 > 1, so Greater + assert_eq!(cmp(0, 0), Ordering::Greater); + + // {"b": 999, "c": 1} vs {"b": 1, "d": 999} + // First entries match (b:999 vs b:1), but 999 > 1, so Greater + assert_eq!(cmp(1, 1), Ordering::Greater); + + // Key comparison: "a" < "b", so {"a": 100, "b": 1} < {"b": 999, "c": 1} + assert_eq!(cmp(0, 1), Ordering::Less); + + // Empty map vs non-empty + assert_eq!(cmp(2, 2), Ordering::Less); // {} < {"a": 1} + + // Non-null vs null + assert_eq!(cmp(3, 3), Ordering::Greater); // {"x": 1} > None + + // Key priority test: "x" > "a", regardless of values + assert_eq!(cmp(3, 0), Ordering::Greater); // {"x": 1} > {"a": 1, "c": 999} + + // Empty vs non-empty + assert_eq!(cmp(2, 0), Ordering::Less); // {} < {"a": 1, "c": 999} + + let opts = SortOptions { + descending: true, + nulls_first: true, + }; + let cmp = make_comparator(&map1, &map2, opts).unwrap(); + + // With descending=true, value comparison is reversed + assert_eq!(cmp(0, 0), Ordering::Less); // {"a": 100, "b": 1} vs {"a": 1, "c": 999} (reversed) + assert_eq!(cmp(1, 1), Ordering::Less); // {"b": 999, "c": 1} vs {"b": 1, "d": 999} (reversed) + assert_eq!(cmp(0, 1), Ordering::Greater); // {"a": 100, "b": 1} vs {"b": 999, "c": 1} (key order reversed) + assert_eq!(cmp(3, 3), Ordering::Greater); // {"x": 1} > None + assert_eq!(cmp(2, 2), Ordering::Greater); // {} > {"a": 1} (reversed) + + let opts = SortOptions { + descending: false, + nulls_first: false, + }; + let cmp = make_comparator(&map1, &map2, opts).unwrap(); + + // Same key priority behavior with nulls_first=false + assert_eq!(cmp(0, 0), Ordering::Greater); // {"a": 100, "b": 1} vs {"a": 1, "c": 999} + assert_eq!(cmp(1, 1), Ordering::Greater); // {"b": 999, "c": 1} vs {"b": 1, "d": 999} + assert_eq!(cmp(3, 3), Ordering::Less); // {"x": 1} < None (nulls last) + assert_eq!(cmp(2, 2), Ordering::Less); // {} < {"a": 1} + } + + #[test] + fn test_map_vs_list_consistency() { + // Create map arrays and convert them to list arrays to verify comparison consistency + // Map arrays: [{"a": 1, "b": 2}, {"x": 10}, {}, {"c": 3}] + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map1_builder = MapBuilder::new(None, string_builder, int_builder); + + // {"a": 1, "b": 2} + map1_builder.keys().append_value("a"); + map1_builder.values().append_value(1); + map1_builder.keys().append_value("b"); + map1_builder.values().append_value(2); + map1_builder.append(true).unwrap(); + + // {"x": 10} + map1_builder.keys().append_value("x"); + map1_builder.values().append_value(10); + map1_builder.append(true).unwrap(); + + // {} + map1_builder.append(true).unwrap(); + + // {"c": 3} + map1_builder.keys().append_value("c"); + map1_builder.values().append_value(3); + map1_builder.append(true).unwrap(); + + let map1 = map1_builder.finish(); + + // Second map array: [{"a": 1, "b": 2}, {"y": 20}, {"d": 4}, None] + let string_builder = StringBuilder::new(); + let int_builder = Int32Builder::new(); + let mut map2_builder = MapBuilder::new(None, string_builder, int_builder); + + // {"a": 1, "b": 2} + map2_builder.keys().append_value("a"); + map2_builder.values().append_value(1); + map2_builder.keys().append_value("b"); + map2_builder.values().append_value(2); + map2_builder.append(true).unwrap(); + + // {"y": 20} + map2_builder.keys().append_value("y"); + map2_builder.values().append_value(20); + map2_builder.append(true).unwrap(); + + // {"d": 4} + map2_builder.keys().append_value("d"); + map2_builder.values().append_value(4); + map2_builder.append(true).unwrap(); + + // None + map2_builder.append(false).unwrap(); + + let map2 = map2_builder.finish(); + + // Convert map arrays to list arrays (Map entries are struct arrays with key-value pairs) + let list1: ListArray = map1.clone().into(); + let list2: ListArray = map2.clone().into(); + + let test_cases = [ + SortOptions { + descending: false, + nulls_first: true, + }, + SortOptions { + descending: true, + nulls_first: true, + }, + SortOptions { + descending: false, + nulls_first: false, + }, + SortOptions { + descending: true, + nulls_first: false, + }, + ]; + + for opts in test_cases { + let map_cmp = make_comparator(&map1, &map2, opts).unwrap(); + let list_cmp = make_comparator(&list1, &list2, opts).unwrap(); + + // Test all possible index combinations + for i in 0..map1.len() { + for j in 0..map2.len() { + let map_result = map_cmp(i, j); + let list_result = list_cmp(i, j); + assert_eq!( + map_result, list_result, + "Map comparison and List comparison should be equal for indices ({i}, {j}) with opts {opts:?}. Map: {map_result:?}, List: {list_result:?}" + ); + } + } + } + } + + #[test] + fn test_dense_union() { + // create a dense union array with Int32 (type_id = 0) and Utf8 (type_id=1) + // the values are: [1, "b", 2, "a", 3] + // type_ids are: [0, 1, 0, 1, 0] + // offsets are: [0, 0, 1, 1, 2] from [1, 2, 3] and ["b", "a"] + let int_array = Int32Array::from(vec![1, 2, 3]); + let str_array = StringArray::from(vec!["b", "a"]); + + let type_ids = [0, 1, 0, 1, 0].into_iter().collect::>(); + let offsets = [0, 0, 1, 1, 2].into_iter().collect::>(); + + let union_fields = [ + (0, Arc::new(Field::new("A", DataType::Int32, false))), + (1, Arc::new(Field::new("B", DataType::Utf8, false))), + ] + .into_iter() + .collect::(); + + let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)]; + + let array1 = + UnionArray::try_new(union_fields.clone(), type_ids, Some(offsets), children).unwrap(); + + // create a second array: [2, "a", 1, "c"] + // type ids are: [0, 1, 0, 1] + // offsets are: [0, 0, 1, 1] from [2, 1] and ["a", "c"] + let int_array2 = Int32Array::from(vec![2, 1]); + let str_array2 = StringArray::from(vec!["a", "c"]); + let type_ids2 = [0, 1, 0, 1].into_iter().collect::>(); + let offsets2 = [0, 0, 1, 1].into_iter().collect::>(); + + let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)]; + + let array2 = + UnionArray::try_new(union_fields, type_ids2, Some(offsets2), children2).unwrap(); + + let opts = SortOptions { + descending: false, + nulls_first: true, + }; + + // comparing + // [1, "b", 2, "a", 3] + // [2, "a", 1, "c"] + let cmp = make_comparator(&array1, &array2, opts).unwrap(); + + // array1[0] = (type_id=0, value=1) + // array2[0] = (type_id=0, value=2) + assert_eq!(cmp(0, 0), Ordering::Less); // 1 < 2 + + // array1[0] = (type_id=0, value=1) + // array2[1] = (type_id=1, value="a") + assert_eq!(cmp(0, 1), Ordering::Less); // type_id 0 < 1 + + // array1[1] = (type_id=1, value="b") + // array2[1] = (type_id=1, value="a") + assert_eq!(cmp(1, 1), Ordering::Greater); // "b" > "a" + + // array1[2] = (type_id=0, value=2) + // array2[0] = (type_id=0, value=2) + assert_eq!(cmp(2, 0), Ordering::Equal); // 2 == 2 + + // array1[3] = (type_id=1, value="a") + // array2[1] = (type_id=1, value="a") + assert_eq!(cmp(3, 1), Ordering::Equal); // "a" == "a" + + // array1[1] = (type_id=1, value="b") + // array2[3] = (type_id=1, value="c") + assert_eq!(cmp(1, 3), Ordering::Less); // "b" < "c" + + let opts_desc = SortOptions { + descending: true, + nulls_first: true, + }; + let cmp_desc = make_comparator(&array1, &array2, opts_desc).unwrap(); + + assert_eq!(cmp_desc(0, 0), Ordering::Greater); // 1 > 2 (reversed) + assert_eq!(cmp_desc(0, 1), Ordering::Greater); // type_id 0 < 1, reversed to Greater + assert_eq!(cmp_desc(1, 1), Ordering::Less); // "b" < "a" (reversed) + } + + #[test] + fn test_sparse_union() { + // create a sparse union array with Int32 (type_id=0) and Utf8 (type_id=1) + // values: [1, "b", 3] + // note, in sparse unions, child arrays have the same length as the union + let int_array = Int32Array::from(vec![Some(1), None, Some(3)]); + let str_array = StringArray::from(vec![None, Some("b"), None]); + let type_ids = [0, 1, 0].into_iter().collect::>(); + + let union_fields = [ + (0, Arc::new(Field::new("a", DataType::Int32, false))), + (1, Arc::new(Field::new("b", DataType::Utf8, false))), + ] + .into_iter() + .collect::(); + + let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)]; + + let array = UnionArray::try_new(union_fields, type_ids, None, children).unwrap(); + + let opts = SortOptions::default(); + let cmp = make_comparator(&array, &array, opts).unwrap(); + + // array[0] = (type_id=0, value=1), array[2] = (type_id=0, value=3) + assert_eq!(cmp(0, 2), Ordering::Less); // 1 < 3 + // array[0] = (type_id=0, value=1), array[1] = (type_id=1, value="b") + assert_eq!(cmp(0, 1), Ordering::Less); // type_id 0 < 1 + } + + #[test] + #[should_panic(expected = "index out of bounds")] + fn test_union_out_of_bounds() { + // create a dense union array with 3 elements + let int_array = Int32Array::from(vec![1, 2]); + let str_array = StringArray::from(vec!["a"]); + + let type_ids = [0, 1, 0].into_iter().collect::>(); + let offsets = [0, 0, 1].into_iter().collect::>(); + + let union_fields = [ + (0, Arc::new(Field::new("A", DataType::Int32, false))), + (1, Arc::new(Field::new("B", DataType::Utf8, false))), + ] + .into_iter() + .collect::(); + + let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)]; + + let array = UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap(); + + let opts = SortOptions::default(); + let cmp = make_comparator(&array, &array, opts).unwrap(); + + // oob + cmp(0, 3); + } + + #[test] + fn test_union_incompatible_fields() { + // create first union with Int32 and Utf8 + let int_array1 = Int32Array::from(vec![1, 2]); + let str_array1 = StringArray::from(vec!["a", "b"]); + + let type_ids1 = [0, 1].into_iter().collect::>(); + let offsets1 = [0, 0].into_iter().collect::>(); + + let union_fields1 = [ + (0, Arc::new(Field::new("A", DataType::Int32, false))), + (1, Arc::new(Field::new("B", DataType::Utf8, false))), + ] + .into_iter() + .collect::(); + + let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)]; + + let array1 = + UnionArray::try_new(union_fields1, type_ids1, Some(offsets1), children1).unwrap(); + + // create second union with Int32 and Float64 (incompatible with first) + let int_array2 = Int32Array::from(vec![3, 4]); + let float_array2 = Float64Array::from(vec![1.0, 2.0]); + + let type_ids2 = [0, 1].into_iter().collect::>(); + let offsets2 = [0, 0].into_iter().collect::>(); + + let union_fields2 = [ + (0, Arc::new(Field::new("A", DataType::Int32, false))), + (1, Arc::new(Field::new("C", DataType::Float64, false))), + ] + .into_iter() + .collect::(); + + let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(float_array2)]; + + let array2 = + UnionArray::try_new(union_fields2, type_ids2, Some(offsets2), children2).unwrap(); + + let opts = SortOptions::default(); + + let Result::Err(ArrowError::InvalidArgumentError(out)) = + make_comparator(&array1, &array2, opts) + else { + panic!("expected error when making comparator of incompatible union arrays"); + }; + + assert_eq!( + &out, + "Cannot compare UnionArrays with different fields: left=[(0, Field { name: \"A\", data_type: Int32 }), (1, Field { name: \"B\", data_type: Utf8 })], right=[(0, Field { name: \"A\", data_type: Int32 }), (1, Field { name: \"C\", data_type: Float64 })]" + ); + } + + #[test] + fn test_union_incompatible_modes() { + // create first union as Dense with Int32 and Utf8 + let int_array1 = Int32Array::from(vec![1, 2]); + let str_array1 = StringArray::from(vec!["a", "b"]); + + let type_ids1 = [0, 1].into_iter().collect::>(); + let offsets1 = [0, 0].into_iter().collect::>(); + + let union_fields1 = [ + (0, Arc::new(Field::new("A", DataType::Int32, false))), + (1, Arc::new(Field::new("B", DataType::Utf8, false))), + ] + .into_iter() + .collect::(); + + let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)]; + + let array1 = + UnionArray::try_new(union_fields1.clone(), type_ids1, Some(offsets1), children1) + .unwrap(); + + // create second union as Sparse with same fields (Int32 and Utf8) + let int_array2 = Int32Array::from(vec![Some(3), None]); + let str_array2 = StringArray::from(vec![None, Some("c")]); + + let type_ids2 = [0, 1].into_iter().collect::>(); + + let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)]; + + let array2 = UnionArray::try_new(union_fields1, type_ids2, None, children2).unwrap(); + + let opts = SortOptions::default(); + + let Result::Err(ArrowError::InvalidArgumentError(out)) = + make_comparator(&array1, &array2, opts) + else { + panic!("expected error when making comparator of union arrays with different modes"); + }; + + assert_eq!( + &out, + "Cannot compare UnionArrays with different modes: left=Dense, right=Sparse" + ); + } + + #[test] + fn test_null_array_cmp() { + let a = NullArray::new(3); + let b = NullArray::new(3); + let cmp = make_comparator(&a, &b, SortOptions::default()).unwrap(); + + assert_eq!(cmp(0, 0), Ordering::Equal); + assert_eq!(cmp(0, 1), Ordering::Equal); + assert_eq!(cmp(2, 0), Ordering::Equal); + } + + #[test] + fn test_run_end_encoded_int32() { + // Create RunEndEncoded arrays: + // array1: [1, 1, 2, 2, 2, 3] + // run_ends1: [2, 5, 6], values1: [1, 2, 3] + let run_ends1 = Int32Array::from(vec![2, 5, 6]); + let values1 = Int32Array::from(vec![1, 2, 3]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + // array2: [1, 2, 2, 3, 3, 3] + // run_ends2: [1, 3, 6], values2: [1, 2, 3] + let run_ends2 = Int32Array::from(vec![1, 3, 6]); + let values2 = Int32Array::from(vec![1, 2, 3]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + // array1[0] = 1, array2[0] = 1 + assert_eq!(cmp(0, 0), Ordering::Equal); + // array1[0] = 1, array2[1] = 2 + assert_eq!(cmp(0, 1), Ordering::Less); + // array1[2] = 2, array2[1] = 2 + assert_eq!(cmp(2, 1), Ordering::Equal); + // array1[5] = 3, array2[5] = 3 + assert_eq!(cmp(5, 5), Ordering::Equal); + // array1[1] = 1, array2[2] = 2 + assert_eq!(cmp(1, 2), Ordering::Less); + // array1[4] = 2, array2[4] = 3 + assert_eq!(cmp(4, 4), Ordering::Less); + } + + #[test] + fn test_run_end_encoded_with_nulls() { + // Create RunEndEncoded arrays with nulls: + // array1: [1, 1, null, null, 2] + // run_ends1: [2, 4, 5], values1: [1, null, 2] + let run_ends1 = Int32Array::from(vec![2, 4, 5]); + let values1 = Int32Array::from(vec![Some(1), None, Some(2)]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + // array2: [null, 1, 1, 2, null] + // run_ends2: [1, 3, 4, 5], values2: [null, 1, 2, null] + let run_ends2 = Int32Array::from(vec![1, 3, 4, 5]); + let values2 = Int32Array::from(vec![None, Some(1), Some(2), None]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + let opts = SortOptions::default(); + let cmp = make_comparator(&array1, &array2, opts).unwrap(); + + // array1[0] = 1, array2[1] = 1 + assert_eq!(cmp(0, 1), Ordering::Equal); + // array1[2] = null, array2[0] = null + assert_eq!(cmp(2, 0), Ordering::Equal); + // array1[0] = 1, array2[0] = null (nulls first by default) + assert_eq!(cmp(0, 0), Ordering::Greater); + // array1[2] = null, array2[1] = 1 + assert_eq!(cmp(2, 1), Ordering::Less); + } + + #[test] + fn test_run_end_encoded_int16() { + // Test with Int16 run ends + let run_ends1 = Int16Array::from(vec![3_i16, 5, 6]); + let values1 = StringArray::from(vec!["a", "b", "c"]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + let run_ends2 = Int16Array::from(vec![2_i16, 4, 6]); + let values2 = StringArray::from(vec!["a", "b", "c"]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + // array1: [a, a, a, b, b, c] + // array2: [a, a, b, b, c, c] + assert_eq!(cmp(0, 0), Ordering::Equal); // a vs a + assert_eq!(cmp(2, 2), Ordering::Less); // a vs b + assert_eq!(cmp(3, 2), Ordering::Equal); // b vs b + assert_eq!(cmp(5, 4), Ordering::Equal); // c vs c + } + + #[test] + fn test_run_end_encoded_int64() { + // Test with Int64 run ends + let run_ends1 = Int64Array::from(vec![2_i64, 4, 6]); + let values1 = Int64Array::from(vec![10_i64, 20, 30]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + let run_ends2 = Int64Array::from(vec![3_i64, 5, 6]); + let values2 = Int64Array::from(vec![10_i64, 20, 30]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); + + // array1: [10, 10, 20, 20, 30, 30] + // array2: [10, 10, 10, 20, 20, 30] + assert_eq!(cmp(0, 0), Ordering::Equal); // 10 vs 10 + assert_eq!(cmp(1, 2), Ordering::Equal); // 10 vs 10 + assert_eq!(cmp(2, 3), Ordering::Equal); // 20 vs 20 + assert_eq!(cmp(4, 4), Ordering::Greater); // 30 vs 20 + } + + #[test] + fn test_run_end_encoded_sliced() { + // Create a RunEndEncoded array and slice it: + // original: [1, 1, 2, 2, 2, 3, 3, 4] + // run_ends: [2, 5, 7, 8], values: [1, 2, 3, 4] + let run_ends = Int32Array::from(vec![2, 5, 7, 8]); + let values = Int32Array::from(vec![1, 2, 3, 4]); + let array = RunArray::::try_new(&run_ends, &values).unwrap(); + + // slice1 = array[1..5] => [1, 2, 2, 2] + let slice1 = array.slice(1, 4); + // slice2 = array[3..7] => [2, 2, 3, 3] + let slice2 = array.slice(3, 4); + + let cmp = make_comparator(&slice1, &slice2, SortOptions::default()).unwrap(); + + // slice1[0]=1, slice2[0]=2 + assert_eq!(cmp(0, 0), Ordering::Less); + // slice1[1]=2, slice2[0]=2 + assert_eq!(cmp(1, 0), Ordering::Equal); + // slice1[3]=2, slice2[2]=3 + assert_eq!(cmp(3, 2), Ordering::Less); + // slice1[1]=2, slice2[3]=3 + assert_eq!(cmp(1, 3), Ordering::Less); + + // Compare a sliced array with an unsliced array + let run_ends2 = Int32Array::from(vec![2, 4]); + let values2 = Int32Array::from(vec![1, 2]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + let cmp = make_comparator(&slice1, &array2, SortOptions::default()).unwrap(); + + // slice1[0]=1, array2[0]=1 + assert_eq!(cmp(0, 0), Ordering::Equal); + // slice1[1]=2, array2[1]=1 + assert_eq!(cmp(1, 1), Ordering::Greater); + // slice1[3]=2, array2[3]=2 + assert_eq!(cmp(3, 3), Ordering::Equal); + } + + #[test] + fn test_run_end_encoded_sliced_with_nulls() { + // Create a RunEndEncoded array with nulls: + // original: [1, 1, null, null, 2, 2, null, 3] + // run_ends: [2, 4, 6, 7, 8], values: [1, null, 2, null, 3] + let run_ends = Int32Array::from(vec![2, 4, 6, 7, 8]); + let values = Int32Array::from(vec![Some(1), None, Some(2), None, Some(3)]); + let array = RunArray::::try_new(&run_ends, &values).unwrap(); + + // slice1 = array[1..6] => [1, null, null, 2, 2] + let slice1 = array.slice(1, 5); + // slice2 = array[3..8] => [null, 2, 2, null, 3] + let slice2 = array.slice(3, 5); + + let opts = SortOptions::default(); // nulls_first=true, descending=false + let cmp = make_comparator(&slice1, &slice2, opts).unwrap(); + + // slice1[0]=1, slice2[0]=null + assert_eq!(cmp(0, 0), Ordering::Greater); + // slice1[1]=null, slice2[0]=null + assert_eq!(cmp(1, 0), Ordering::Equal); + // slice1[1]=null, slice2[1]=2 + assert_eq!(cmp(1, 1), Ordering::Less); + // slice1[3]=2, slice2[1]=2 + assert_eq!(cmp(3, 1), Ordering::Equal); + // slice1[4]=2, slice2[4]=3 + assert_eq!(cmp(4, 4), Ordering::Less); + // slice1[3]=2, slice2[3]=null + assert_eq!(cmp(3, 3), Ordering::Greater); + } + + #[test] + fn test_run_end_encoded_different_types() { + // Test with different run end types - should fail + let run_ends1 = Int32Array::from(vec![2, 4, 6]); + let values1 = Int32Array::from(vec![1, 2, 3]); + let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); + + let run_ends2 = Int64Array::from(vec![2_i64, 4, 6]); + let values2 = Int64Array::from(vec![1_i64, 2, 3]); + let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); + + let result = make_comparator(&array1, &array2, SortOptions::default()); + assert!(result.is_err()); + let err = match result { + Err(e) => e.to_string(), + Ok(_) => panic!("Expected error"), + }; + assert!(err.contains("Cannot compare RunEndEncoded arrays")); + } +} diff --git a/arrow-ord-basic/src/lib.rs b/arrow-ord-basic/src/lib.rs deleted file mode 100644 index 813a5a802db3..000000000000 --- a/arrow-ord-basic/src/lib.rs +++ /dev/null @@ -1,535 +0,0 @@ -// 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. - -//! Basic comparator factories shared by Arrow crates that need to compare -//! arbitrary array slots without pulling in the full [`arrow-ord`] crate. -//! -//! The only public surface is [`make_comparator`] (with [`DynComparator`] as the -//! returned function type). `arrow-ord` re-exports both from here, so its -//! public API is unchanged. -//! -//! This crate exists so that crates such as `arrow-select` can use slot-wise -//! comparison (e.g. for the run-end-encoded `take` fast path) without taking on -//! the full ordering kernel suite — which would either create a circular -//! dependency (`arrow-ord` already depends on `arrow-select`) or force every -//! downstream user of `arrow-array` to compile the comparator machinery whether -//! they need it or not. - -#![deny(rustdoc::broken_intra_doc_links)] -#![warn(missing_docs)] - -use arrow_array::cast::AsArray; -use arrow_array::types::*; -use arrow_array::*; -use arrow_buffer::{ArrowNativeType, NullBuffer}; -use arrow_schema::{ArrowError, DataType, SortOptions}; -use std::{cmp::Ordering, collections::HashMap}; - -fn compare_run_end_encoded( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_run::(); - let right = right.as_run::(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - - let l_run_ends = left.run_ends().clone(); - let r_run_ends = right.run_ends().clone(); - - let f = compare(left, right, opts, move |i, j| { - let l_physical = l_run_ends.get_physical_index(i); - let r_physical = r_run_ends.get_physical_index(j); - cmp(l_physical, r_physical) - }); - Ok(f) -} - -/// Compare values at arbitrary indices in two arrays. -pub type DynComparator = Box Ordering + Send + Sync>; - -/// If parent sort order is descending we need to invert the value of nulls_first so that -/// when the parent is sorted based on the produced ranks, nulls are still ordered correctly -fn child_opts(opts: SortOptions) -> SortOptions { - SortOptions { - descending: false, - nulls_first: opts.nulls_first != opts.descending, - } -} - -fn compare(l: &A, r: &A, opts: SortOptions, cmp: F) -> DynComparator -where - A: Array + Clone, - F: Fn(usize, usize) -> Ordering + Send + Sync + 'static, -{ - let l = l.logical_nulls().filter(|x| x.null_count() > 0); - let r = r.logical_nulls().filter(|x| x.null_count() > 0); - match (opts.nulls_first, opts.descending) { - (true, true) => compare_impl::(l, r, cmp), - (true, false) => compare_impl::(l, r, cmp), - (false, true) => compare_impl::(l, r, cmp), - (false, false) => compare_impl::(l, r, cmp), - } -} - -fn compare_impl( - l: Option, - r: Option, - cmp: F, -) -> DynComparator -where - F: Fn(usize, usize) -> Ordering + Send + Sync + 'static, -{ - let cmp = move |i, j| match DESCENDING { - true => cmp(i, j).reverse(), - false => cmp(i, j), - }; - - let (left_null, right_null) = match NULLS_FIRST { - true => (Ordering::Less, Ordering::Greater), - false => (Ordering::Greater, Ordering::Less), - }; - - match (l, r) { - (None, None) => Box::new(cmp), - (Some(l), None) => Box::new(move |i, j| match l.is_null(i) { - true => left_null, - false => cmp(i, j), - }), - (None, Some(r)) => Box::new(move |i, j| match r.is_null(j) { - true => right_null, - false => cmp(i, j), - }), - (Some(l), Some(r)) => Box::new(move |i, j| match (l.is_null(i), r.is_null(j)) { - (true, true) => Ordering::Equal, - (true, false) => left_null, - (false, true) => right_null, - (false, false) => cmp(i, j), - }), - } -} - -fn compare_primitive( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> DynComparator -where - T::Native: ArrowNativeTypeOp, -{ - let left = left.as_primitive::(); - let right = right.as_primitive::(); - let l_values = left.values().clone(); - let r_values = right.values().clone(); - - compare(&left, &right, opts, move |i, j| { - l_values[i].compare(r_values[j]) - }) -} - -fn compare_boolean(left: &dyn Array, right: &dyn Array, opts: SortOptions) -> DynComparator { - let left = left.as_boolean(); - let right = right.as_boolean(); - - let l_values = left.values().clone(); - let r_values = right.values().clone(); - - compare(left, right, opts, move |i, j| { - l_values.value(i).cmp(&r_values.value(j)) - }) -} - -fn compare_bytes( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> DynComparator { - let left = left.as_bytes::(); - let right = right.as_bytes::(); - - let l = left.clone(); - let r = right.clone(); - compare(left, right, opts, move |i, j| { - let l: &[u8] = l.value(i).as_ref(); - let r: &[u8] = r.value(j).as_ref(); - l.cmp(r) - }) -} - -fn compare_byte_view( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> DynComparator { - let left = left.as_byte_view::(); - let right = right.as_byte_view::(); - - let l = left.clone(); - let r = right.clone(); - compare(left, right, opts, move |i, j| { - compare_byte_view_values(&l, i, &r, j) - }) -} - -fn compare_byte_view_values( - left: &GenericByteViewArray, - left_idx: usize, - right: &GenericByteViewArray, - right_idx: usize, -) -> Ordering { - assert!(left_idx < left.len()); - assert!(right_idx < right.len()); - - if left.data_buffers().is_empty() && right.data_buffers().is_empty() { - let l_view = unsafe { left.views().get_unchecked(left_idx) }; - let r_view = unsafe { right.views().get_unchecked(right_idx) }; - return GenericByteViewArray::::inline_key_fast(*l_view) - .cmp(&GenericByteViewArray::::inline_key_fast(*r_view)); - } - - unsafe { GenericByteViewArray::compare_unchecked(left, left_idx, right, right_idx) } -} - -fn compare_dict( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_dictionary::(); - let right = right.as_dictionary::(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - let left_keys = left.keys().values().clone(); - let right_keys = right.keys().values().clone(); - - let f = compare(left, right, opts, move |i, j| { - let l = left_keys[i].as_usize(); - let r = right_keys[j].as_usize(); - cmp(l, r) - }); - Ok(f) -} - -fn compare_list( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_list::(); - let right = right.as_list::(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - - let l_o = left.offsets().clone(); - let r_o = right.offsets().clone(); - let f = compare(left, right, opts, move |i, j| { - let l_end = l_o[i + 1].as_usize(); - let l_start = l_o[i].as_usize(); - - let r_end = r_o[j + 1].as_usize(); - let r_start = r_o[j].as_usize(); - - for (i, j) in (l_start..l_end).zip(r_start..r_end) { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - (l_end - l_start).cmp(&(r_end - r_start)) - }); - Ok(f) -} - -fn compare_fixed_list( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_fixed_size_list(); - let right = right.as_fixed_size_list(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - - let l_size = left.value_length().to_usize().unwrap(); - let r_size = right.value_length().to_usize().unwrap(); - let size_cmp = l_size.cmp(&r_size); - - let f = compare(left, right, opts, move |i, j| { - let l_start = i * l_size; - let l_end = l_start + l_size; - let r_start = j * r_size; - let r_end = r_start + r_size; - for (i, j) in (l_start..l_end).zip(r_start..r_end) { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - size_cmp - }); - Ok(f) -} - -fn compare_list_view( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_list_view::(); - let right = right.as_list_view::(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?; - - let l_offsets = left.offsets().clone(); - let l_sizes = left.sizes().clone(); - let r_offsets = right.offsets().clone(); - let r_sizes = right.sizes().clone(); - - let f = compare(left, right, opts, move |i, j| { - let l_start = l_offsets[i].as_usize(); - let l_len = l_sizes[i].as_usize(); - let l_end = l_start + l_len; - - let r_start = r_offsets[j].as_usize(); - let r_len = r_sizes[j].as_usize(); - let r_end = r_start + r_len; - - for (i, j) in (l_start..l_end).zip(r_start..r_end) { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - l_len.cmp(&r_len) - }); - Ok(f) -} - -fn compare_map( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_map(); - let right = right.as_map(); - - let c_opts = child_opts(opts); - let cmp = make_comparator(left.entries(), right.entries(), c_opts)?; - - let l_o = left.offsets().clone(); - let r_o = right.offsets().clone(); - let f = compare(left, right, opts, move |i, j| { - let l_end = l_o[i + 1].as_usize(); - let l_start = l_o[i].as_usize(); - - let r_end = r_o[j + 1].as_usize(); - let r_start = r_o[j].as_usize(); - - for (i, j) in (l_start..l_end).zip(r_start..r_end) { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - (l_end - l_start).cmp(&(r_end - r_start)) - }); - Ok(f) -} - -fn compare_struct( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_struct(); - let right = right.as_struct(); - - if left.columns().len() != right.columns().len() { - return Err(ArrowError::InvalidArgumentError( - "Cannot compare StructArray with different number of columns".to_string(), - )); - } - - let c_opts = child_opts(opts); - let columns = left.columns().iter().zip(right.columns()); - let comparators = columns - .map(|(l, r)| make_comparator(l, r, c_opts)) - .collect::, _>>()?; - - let f = compare(left, right, opts, move |i, j| { - for cmp in &comparators { - match cmp(i, j) { - Ordering::Equal => continue, - r => return r, - } - } - Ordering::Equal - }); - Ok(f) -} - -fn compare_union( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - let left = left.as_union(); - let right = right.as_union(); - - let (left_fields, left_mode) = match left.data_type() { - DataType::Union(fields, mode) => (fields, mode), - _ => unreachable!(), - }; - let (right_fields, right_mode) = match right.data_type() { - DataType::Union(fields, mode) => (fields, mode), - _ => unreachable!(), - }; - - if left_fields != right_fields { - return Err(ArrowError::InvalidArgumentError(format!( - "Cannot compare UnionArrays with different fields: left={left_fields:?}, right={right_fields:?}" - ))); - } - - if left_mode != right_mode { - return Err(ArrowError::InvalidArgumentError(format!( - "Cannot compare UnionArrays with different modes: left={left_mode:?}, right={right_mode:?}" - ))); - } - - let c_opts = child_opts(opts); - - let mut field_comparators = HashMap::with_capacity(left_fields.len()); - - for (type_id, _field) in left_fields.iter() { - let left_child = left.child(type_id); - let right_child = right.child(type_id); - let cmp = make_comparator(left_child.as_ref(), right_child.as_ref(), c_opts)?; - - field_comparators.insert(type_id, cmp); - } - - let left_type_ids = left.type_ids().clone(); - let right_type_ids = right.type_ids().clone(); - - let left_offsets = left.offsets().cloned(); - let right_offsets = right.offsets().cloned(); - - let f = compare(left, right, opts, move |i, j| { - let left_type_id = left_type_ids[i]; - let right_type_id = right_type_ids[j]; - - match left_type_id.cmp(&right_type_id) { - Ordering::Equal => { - let left_offset = left_offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i); - let right_offset = right_offsets.as_ref().map(|o| o[j] as usize).unwrap_or(j); - - let cmp = field_comparators - .get(&left_type_id) - .expect("type id not found in field_comparators"); - - cmp(left_offset, right_offset) - } - other => other, - } - }); - Ok(f) -} - -/// Returns a comparison function that compares two values at two arbitrary indices. -/// -/// If `nulls_first` is true, null values are considered less than any non-null -/// value; otherwise they are considered greater. This is primarily shared by -/// crates that need repeated slot comparisons without constructing sliced arrays. -pub fn make_comparator( - left: &dyn Array, - right: &dyn Array, - opts: SortOptions, -) -> Result { - use arrow_schema::DataType::*; - - macro_rules! primitive_helper { - ($t:ty, $left:expr, $right:expr, $nulls_first:expr) => { - Ok(compare_primitive::<$t>($left, $right, $nulls_first)) - }; - } - downcast_primitive! { - left.data_type(), right.data_type() => (primitive_helper, left, right, opts), - (Boolean, Boolean) => Ok(compare_boolean(left, right, opts)), - (Utf8, Utf8) => Ok(compare_bytes::(left, right, opts)), - (LargeUtf8, LargeUtf8) => Ok(compare_bytes::(left, right, opts)), - (Utf8View, Utf8View) => Ok(compare_byte_view::(left, right, opts)), - (Binary, Binary) => Ok(compare_bytes::(left, right, opts)), - (LargeBinary, LargeBinary) => Ok(compare_bytes::(left, right, opts)), - (BinaryView, BinaryView) => Ok(compare_byte_view::(left, right, opts)), - (FixedSizeBinary(_), FixedSizeBinary(_)) => { - let left = left.as_fixed_size_binary(); - let right = right.as_fixed_size_binary(); - - let l = left.clone(); - let r = right.clone(); - Ok(compare(left, right, opts, move |i, j| { - l.value(i).cmp(r.value(j)) - })) - }, - (List(_), List(_)) => compare_list::(left, right, opts), - (LargeList(_), LargeList(_)) => compare_list::(left, right, opts), - (ListView(_), ListView(_)) => compare_list_view::(left, right, opts), - (LargeListView(_), LargeListView(_)) => compare_list_view::(left, right, opts), - (FixedSizeList(_, _), FixedSizeList(_, _)) => compare_fixed_list(left, right, opts), - (Struct(_), Struct(_)) => compare_struct(left, right, opts), - (Dictionary(l_key, _), Dictionary(r_key, _)) => { - macro_rules! dict_helper { - ($t:ty, $left:expr, $right:expr, $opts: expr) => { - compare_dict::<$t>($left, $right, $opts) - }; - } - downcast_integer! { - l_key.as_ref(), r_key.as_ref() => (dict_helper, left, right, opts), - _ => unreachable!() - } - }, - (RunEndEncoded(l_run_ends, _), RunEndEncoded(r_run_ends, _)) => { - macro_rules! run_end_helper { - ($t:ty, $left:expr, $right:expr, $opts:expr) => { - compare_run_end_encoded::<$t>($left, $right, $opts) - }; - } - downcast_run_end_index! { - l_run_ends.data_type(), r_run_ends.data_type() => (run_end_helper, left, right, opts), - _ => Err(ArrowError::InvalidArgumentError(format!( - "Cannot compare RunEndEncoded arrays with different run ends types: left={:?}, right={:?}", - l_run_ends.data_type(), - r_run_ends.data_type() - ))) - } - }, - (Map(_, _), Map(_, _)) => compare_map(left, right, opts), - (Null, Null) => Ok(Box::new(|_, _| Ordering::Equal)), - (Union(_, _), Union(_, _)) => compare_union(left, right, opts), - (lhs, rhs) => Err(ArrowError::InvalidArgumentError(match lhs == rhs { - true => format!("The data type type {lhs:?} has no natural order"), - false => "Can't compare arrays of different types".to_string(), - })) - } -} diff --git a/arrow-ord/Cargo.toml b/arrow-ord/Cargo.toml index 00b8bfd6373f..ade7f12e3a2f 100644 --- a/arrow-ord/Cargo.toml +++ b/arrow-ord/Cargo.toml @@ -39,7 +39,7 @@ all-features = true arrow-array = { workspace = true } arrow-buffer = { workspace = true } arrow-data = { workspace = true } -arrow-ord-basic = { workspace = true } +arrow-cmp = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true } diff --git a/arrow-ord/src/ord.rs b/arrow-ord/src/ord.rs index 53838039e68c..f0fe975c996b 100644 --- a/arrow-ord/src/ord.rs +++ b/arrow-ord/src/ord.rs @@ -17,1366 +17,4 @@ //! Contains functions and function factories to compare arrays. -pub use arrow_ord_basic::{DynComparator, make_comparator}; - -#[cfg(test)] -use arrow_array::types::*; -#[cfg(test)] -use arrow_array::*; -#[cfg(test)] -use arrow_schema::{ArrowError, SortOptions}; -#[cfg(test)] -use std::cmp::Ordering; - -#[cfg(test)] -mod tests { - use super::*; - use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder}; - use arrow_buffer::{IntervalDayTime, NullBuffer, OffsetBuffer, ScalarBuffer, i256}; - use arrow_schema::{DataType, Field, Fields, UnionFields}; - use half::f16; - use std::sync::Arc; - - #[test] - fn test_fixed_size_binary() { - let items = vec![vec![1u8], vec![2u8]]; - let array = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap(); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - } - - #[test] - fn test_fixed_size_binary_fixed_size_binary() { - let items = vec![vec![1u8]]; - let array1 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap(); - let items = vec![vec![2u8]]; - let array2 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap(); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - } - - #[test] - fn test_i32() { - let array = Int32Array::from(vec![1, 2]); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, (cmp)(0, 1)); - } - - #[test] - fn test_i32_i32() { - let array1 = Int32Array::from(vec![1]); - let array2 = Int32Array::from(vec![2]); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - } - - #[test] - fn test_f16() { - let array = Float16Array::from(vec![f16::from_f32(1.0), f16::from_f32(2.0)]); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - } - - #[test] - fn test_f64() { - let array = Float64Array::from(vec![1.0, 2.0]); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - } - - #[test] - fn test_f64_nan() { - let array = Float64Array::from(vec![1.0, f64::NAN]); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - assert_eq!(Ordering::Equal, cmp(1, 1)); - } - - #[test] - fn test_f64_zeros() { - let array = Float64Array::from(vec![-0.0, 0.0]); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - assert_eq!(Ordering::Greater, cmp(1, 0)); - } - - #[test] - fn test_interval_day_time() { - let array = IntervalDayTimeArray::from(vec![ - // 0 days, 1 second - IntervalDayTimeType::make_value(0, 1000), - // 1 day, 2 milliseconds - IntervalDayTimeType::make_value(1, 2), - // 90M milliseconds (which is more than is in 1 day) - IntervalDayTimeType::make_value(0, 90_000_000), - ]); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - assert_eq!(Ordering::Greater, cmp(1, 0)); - - // somewhat confusingly, while 90M milliseconds is more than 1 day, - // it will compare less as the comparison is done on the underlying - // values not field by field - assert_eq!(Ordering::Greater, cmp(1, 2)); - assert_eq!(Ordering::Less, cmp(2, 1)); - } - - #[test] - fn test_interval_year_month() { - let array = IntervalYearMonthArray::from(vec![ - // 1 year, 0 months - IntervalYearMonthType::make_value(1, 0), - // 0 years, 13 months - IntervalYearMonthType::make_value(0, 13), - // 1 year, 1 month - IntervalYearMonthType::make_value(1, 1), - ]); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - assert_eq!(Ordering::Greater, cmp(1, 0)); - - // the underlying representation is months, so both quantities are the same - assert_eq!(Ordering::Equal, cmp(1, 2)); - assert_eq!(Ordering::Equal, cmp(2, 1)); - } - - #[test] - fn test_interval_month_day_nano() { - let array = IntervalMonthDayNanoArray::from(vec![ - // 100 days - IntervalMonthDayNanoType::make_value(0, 100, 0), - // 1 month - IntervalMonthDayNanoType::make_value(1, 0, 0), - // 100 day, 1 nanoseconds - IntervalMonthDayNanoType::make_value(0, 100, 2), - ]); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - assert_eq!(Ordering::Greater, cmp(1, 0)); - - // somewhat confusingly, while 100 days is more than 1 month in all cases - // it will compare less as the comparison is done on the underlying - // values not field by field - assert_eq!(Ordering::Greater, cmp(1, 2)); - assert_eq!(Ordering::Less, cmp(2, 1)); - } - - #[test] - fn test_decimali32() { - let array = vec![Some(5_i32), Some(2_i32), Some(3_i32)] - .into_iter() - .collect::() - .with_precision_and_scale(8, 6) - .unwrap(); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - assert_eq!(Ordering::Less, cmp(1, 0)); - assert_eq!(Ordering::Greater, cmp(0, 2)); - } - - #[test] - fn test_decimali64() { - let array = vec![Some(5_i64), Some(2_i64), Some(3_i64)] - .into_iter() - .collect::() - .with_precision_and_scale(16, 6) - .unwrap(); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - assert_eq!(Ordering::Less, cmp(1, 0)); - assert_eq!(Ordering::Greater, cmp(0, 2)); - } - - #[test] - fn test_decimali128() { - let array = vec![Some(5_i128), Some(2_i128), Some(3_i128)] - .into_iter() - .collect::() - .with_precision_and_scale(23, 6) - .unwrap(); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - assert_eq!(Ordering::Less, cmp(1, 0)); - assert_eq!(Ordering::Greater, cmp(0, 2)); - } - - #[test] - fn test_decimali256() { - let array = vec![ - Some(i256::from_i128(5_i128)), - Some(i256::from_i128(2_i128)), - Some(i256::from_i128(3_i128)), - ] - .into_iter() - .collect::() - .with_precision_and_scale(53, 6) - .unwrap(); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - assert_eq!(Ordering::Less, cmp(1, 0)); - assert_eq!(Ordering::Greater, cmp(0, 2)); - } - - #[test] - fn test_dict() { - let data = vec!["a", "b", "c", "a", "a", "c", "c"]; - let array = data.into_iter().collect::>(); - - let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - assert_eq!(Ordering::Equal, cmp(3, 4)); - assert_eq!(Ordering::Greater, cmp(2, 3)); - } - - #[test] - fn test_multiple_dict() { - let d1 = vec!["a", "b", "c", "d"]; - let a1 = d1.into_iter().collect::>(); - let d2 = vec!["e", "f", "g", "a"]; - let a2 = d2.into_iter().collect::>(); - - let cmp = make_comparator(&a1, &a2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - assert_eq!(Ordering::Equal, cmp(0, 3)); - assert_eq!(Ordering::Greater, cmp(1, 3)); - } - - #[test] - fn test_primitive_dict() { - let values = Int32Array::from(vec![1_i32, 0, 2, 5]); - let keys = Int8Array::from_iter_values([0, 0, 1, 3]); - let array1 = DictionaryArray::new(keys, Arc::new(values)); - - let values = Int32Array::from(vec![2_i32, 3, 4, 5]); - let keys = Int8Array::from_iter_values([0, 1, 1, 3]); - let array2 = DictionaryArray::new(keys, Arc::new(values)); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - assert_eq!(Ordering::Less, cmp(0, 3)); - assert_eq!(Ordering::Equal, cmp(3, 3)); - assert_eq!(Ordering::Greater, cmp(3, 1)); - assert_eq!(Ordering::Greater, cmp(3, 2)); - } - - #[test] - fn test_float_dict() { - let values = Float32Array::from(vec![1.0, 0.5, 2.1, 5.5]); - let keys = Int8Array::from_iter_values([0, 0, 1, 3]); - let array1 = DictionaryArray::try_new(keys, Arc::new(values)).unwrap(); - - let values = Float32Array::from(vec![1.2, 3.2, 4.0, 5.5]); - let keys = Int8Array::from_iter_values([0, 1, 1, 3]); - let array2 = DictionaryArray::new(keys, Arc::new(values)); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - assert_eq!(Ordering::Less, cmp(0, 3)); - assert_eq!(Ordering::Equal, cmp(3, 3)); - assert_eq!(Ordering::Greater, cmp(3, 1)); - assert_eq!(Ordering::Greater, cmp(3, 2)); - } - - #[test] - fn test_timestamp_dict() { - let values = TimestampSecondArray::from(vec![1, 0, 2, 5]); - let keys = Int8Array::from_iter_values([0, 0, 1, 3]); - let array1 = DictionaryArray::new(keys, Arc::new(values)); - - let values = TimestampSecondArray::from(vec![2, 3, 4, 5]); - let keys = Int8Array::from_iter_values([0, 1, 1, 3]); - let array2 = DictionaryArray::new(keys, Arc::new(values)); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - assert_eq!(Ordering::Less, cmp(0, 3)); - assert_eq!(Ordering::Equal, cmp(3, 3)); - assert_eq!(Ordering::Greater, cmp(3, 1)); - assert_eq!(Ordering::Greater, cmp(3, 2)); - } - - #[test] - fn test_interval_dict() { - let v1 = IntervalDayTime::new(0, 1); - let v2 = IntervalDayTime::new(0, 2); - let v3 = IntervalDayTime::new(12, 2); - - let values = IntervalDayTimeArray::from(vec![Some(v1), Some(v2), None, Some(v3)]); - let keys = Int8Array::from_iter_values([0, 0, 1, 3]); - let array1 = DictionaryArray::new(keys, Arc::new(values)); - - let values = IntervalDayTimeArray::from(vec![Some(v3), Some(v2), None, Some(v1)]); - let keys = Int8Array::from_iter_values([0, 1, 1, 3]); - let array2 = DictionaryArray::new(keys, Arc::new(values)); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); // v1 vs v3 - assert_eq!(Ordering::Equal, cmp(0, 3)); // v1 vs v1 - assert_eq!(Ordering::Greater, cmp(3, 3)); // v3 vs v1 - assert_eq!(Ordering::Greater, cmp(3, 1)); // v3 vs v2 - assert_eq!(Ordering::Greater, cmp(3, 2)); // v3 vs v2 - } - - #[test] - fn test_duration_dict() { - let values = DurationSecondArray::from(vec![1, 0, 2, 5]); - let keys = Int8Array::from_iter_values([0, 0, 1, 3]); - let array1 = DictionaryArray::new(keys, Arc::new(values)); - - let values = DurationSecondArray::from(vec![2, 3, 4, 5]); - let keys = Int8Array::from_iter_values([0, 1, 1, 3]); - let array2 = DictionaryArray::new(keys, Arc::new(values)); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - assert_eq!(Ordering::Less, cmp(0, 3)); - assert_eq!(Ordering::Equal, cmp(3, 3)); - assert_eq!(Ordering::Greater, cmp(3, 1)); - assert_eq!(Ordering::Greater, cmp(3, 2)); - } - - #[test] - fn test_decimal_dict() { - let values = Decimal128Array::from(vec![1, 0, 2, 5]); - let keys = Int8Array::from_iter_values([0, 0, 1, 3]); - let array1 = DictionaryArray::new(keys, Arc::new(values)); - - let values = Decimal128Array::from(vec![2, 3, 4, 5]); - let keys = Int8Array::from_iter_values([0, 1, 1, 3]); - let array2 = DictionaryArray::new(keys, Arc::new(values)); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - assert_eq!(Ordering::Less, cmp(0, 3)); - assert_eq!(Ordering::Equal, cmp(3, 3)); - assert_eq!(Ordering::Greater, cmp(3, 1)); - assert_eq!(Ordering::Greater, cmp(3, 2)); - } - - #[test] - fn test_decimal256_dict() { - let values = Decimal256Array::from(vec![ - i256::from_i128(1), - i256::from_i128(0), - i256::from_i128(2), - i256::from_i128(5), - ]); - let keys = Int8Array::from_iter_values([0, 0, 1, 3]); - let array1 = DictionaryArray::new(keys, Arc::new(values)); - - let values = Decimal256Array::from(vec![ - i256::from_i128(2), - i256::from_i128(3), - i256::from_i128(4), - i256::from_i128(5), - ]); - let keys = Int8Array::from_iter_values([0, 1, 1, 3]); - let array2 = DictionaryArray::new(keys, Arc::new(values)); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 0)); - assert_eq!(Ordering::Less, cmp(0, 3)); - assert_eq!(Ordering::Equal, cmp(3, 3)); - assert_eq!(Ordering::Greater, cmp(3, 1)); - assert_eq!(Ordering::Greater, cmp(3, 2)); - } - - fn test_bytes_impl() { - let offsets = OffsetBuffer::from_lengths([3, 3, 1]); - let a = GenericByteArray::::new(offsets, b"abcdefa".into(), None); - let cmp = make_comparator(&a, &a, SortOptions::default()).unwrap(); - - assert_eq!(Ordering::Less, cmp(0, 1)); - assert_eq!(Ordering::Greater, cmp(0, 2)); - assert_eq!(Ordering::Equal, cmp(1, 1)); - } - - #[test] - fn test_bytes() { - test_bytes_impl::(); - test_bytes_impl::(); - test_bytes_impl::(); - test_bytes_impl::(); - } - - fn assert_cmp_cases( - array1: &A, - array2: &A, - opts: SortOptions, - cases: &[(usize, usize, Ordering)], - ) { - let cmp = make_comparator(array1, array2, opts).unwrap(); - for (left, right, expected) in cases { - assert_eq!(cmp(*left, *right), *expected); - } - } - - #[test] - fn test_lists() { - let mut a = ListBuilder::new(ListBuilder::new(Int32Builder::new())); - a.extend([ - Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]), - Some(vec![ - Some(vec![Some(1), Some(2), Some(3)]), - Some(vec![Some(1)]), - ]), - Some(vec![]), - ]); - let a = a.finish(); - let mut b = ListBuilder::new(ListBuilder::new(Int32Builder::new())); - b.extend([ - Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]), - Some(vec![ - Some(vec![Some(1), Some(2), None]), - Some(vec![Some(1)]), - ]), - Some(vec![ - Some(vec![Some(1), Some(2), Some(3), Some(4)]), - Some(vec![Some(1)]), - ]), - None, - ]); - let b = b.finish(); - - // Ascending with nulls first. - assert_cmp_cases( - &a, - &b, - SortOptions { - descending: false, - nulls_first: true, - }, - &[ - (0, 0, Ordering::Equal), - (0, 1, Ordering::Less), - (0, 2, Ordering::Less), - (1, 2, Ordering::Less), - (1, 3, Ordering::Greater), - (2, 0, Ordering::Less), - ], - ); - - // Descending with nulls first. - assert_cmp_cases( - &a, - &b, - SortOptions { - descending: true, - nulls_first: true, - }, - &[ - (0, 0, Ordering::Equal), - (0, 1, Ordering::Less), - (0, 2, Ordering::Less), - (1, 2, Ordering::Greater), - (1, 3, Ordering::Greater), - (2, 0, Ordering::Greater), - ], - ); - - // Descending with nulls last. - assert_cmp_cases( - &a, - &b, - SortOptions { - descending: true, - nulls_first: false, - }, - &[ - (0, 0, Ordering::Equal), - (0, 1, Ordering::Greater), - (0, 2, Ordering::Greater), - (1, 2, Ordering::Greater), - (1, 3, Ordering::Less), - (2, 0, Ordering::Greater), - ], - ); - - // Ascending with nulls last. - assert_cmp_cases( - &a, - &b, - SortOptions { - descending: false, - nulls_first: false, - }, - &[ - (0, 0, Ordering::Equal), - (0, 1, Ordering::Greater), - (0, 2, Ordering::Greater), - (1, 2, Ordering::Less), - (1, 3, Ordering::Less), - (2, 0, Ordering::Less), - ], - ); - } - - fn list_view_array( - values: Vec, - offsets: &[usize], - sizes: &[usize], - valid: Option<&[bool]>, - ) -> GenericListViewArray { - let offsets = offsets - .iter() - .map(|v| O::from_usize(*v).unwrap()) - .collect::>(); - let sizes = sizes - .iter() - .map(|v| O::from_usize(*v).unwrap()) - .collect::>(); - let field = Arc::new(Field::new_list_field(DataType::Int32, true)); - let values = Int32Array::from(values); - let nulls = valid.map(NullBuffer::from); - GenericListViewArray::new(field, offsets, sizes, Arc::new(values), nulls) - } - - fn test_list_view_comparisons() { - let array = list_view_array::( - vec![1, 2, 3, 4, 5], - &[0, 2, 1, 0, 3], - &[2, 2, 2, 0, 2], - Some(&[true, true, true, true, false]), - ); - - // Ascending with nulls first (non-monotonic offsets and empty list). - assert_cmp_cases( - &array, - &array, - SortOptions { - descending: false, - nulls_first: true, - }, - &[ - (0, 2, Ordering::Less), // [1,2] < [2,3] - (1, 2, Ordering::Greater), // [3,4] > [2,3] - (3, 0, Ordering::Less), // [] < [1,2] - (4, 0, Ordering::Less), // null < [1,2] - ], - ); - - // Ascending with nulls last. - assert_cmp_cases( - &array, - &array, - SortOptions { - descending: false, - nulls_first: false, - }, - &[ - (0, 2, Ordering::Less), - (1, 2, Ordering::Greater), - (3, 0, Ordering::Less), - (4, 0, Ordering::Greater), // null last - ], - ); - - // Descending with nulls first. - assert_cmp_cases( - &array, - &array, - SortOptions { - descending: true, - nulls_first: true, - }, - &[ - (0, 2, Ordering::Greater), - (1, 2, Ordering::Less), - (3, 0, Ordering::Greater), - (4, 0, Ordering::Less), - ], - ); - - // Descending with nulls last. - assert_cmp_cases( - &array, - &array, - SortOptions { - descending: true, - nulls_first: false, - }, - &[ - (0, 2, Ordering::Greater), - (1, 2, Ordering::Less), - (3, 0, Ordering::Greater), - (4, 0, Ordering::Greater), - ], - ); - } - - #[test] - fn test_list_view() { - test_list_view_comparisons::(); - } - - #[test] - fn test_large_list_view() { - test_list_view_comparisons::(); - } - - #[test] - fn test_struct() { - let fields = Fields::from(vec![ - Field::new("a", DataType::Int32, true), - Field::new_list("b", Field::new_list_field(DataType::Int32, true), true), - ]); - - let a = Int32Array::from(vec![Some(1), Some(2), None, None]); - let mut b = ListBuilder::new(Int32Builder::new()); - b.extend([Some(vec![Some(1), Some(2)]), Some(vec![None]), None, None]); - let b = b.finish(); - - let nulls = Some(NullBuffer::from_iter([true, true, true, false])); - let values = vec![Arc::new(a) as _, Arc::new(b) as _]; - let s1 = StructArray::new(fields.clone(), values, nulls); - - let a = Int32Array::from(vec![None, Some(2), None]); - let mut b = ListBuilder::new(Int32Builder::new()); - b.extend([None, None, Some(vec![])]); - let b = b.finish(); - - let values = vec![Arc::new(a) as _, Arc::new(b) as _]; - let s2 = StructArray::new(fields.clone(), values, None); - - let opts = SortOptions { - descending: false, - nulls_first: true, - }; - let cmp = make_comparator(&s1, &s2, opts).unwrap(); - assert_eq!(cmp(0, 1), Ordering::Less); // (1, [1, 2]) cmp (2, None) - assert_eq!(cmp(0, 0), Ordering::Greater); // (1, [1, 2]) cmp (None, None) - assert_eq!(cmp(1, 1), Ordering::Greater); // (2, [None]) cmp (2, None) - assert_eq!(cmp(2, 2), Ordering::Less); // (None, None) cmp (None, []) - assert_eq!(cmp(3, 0), Ordering::Less); // None cmp (None, []) - assert_eq!(cmp(2, 0), Ordering::Equal); // (None, None) cmp (None, None) - assert_eq!(cmp(3, 0), Ordering::Less); // None cmp (None, None) - - let opts = SortOptions { - descending: true, - nulls_first: true, - }; - let cmp = make_comparator(&s1, &s2, opts).unwrap(); - assert_eq!(cmp(0, 1), Ordering::Greater); // (1, [1, 2]) cmp (2, None) - assert_eq!(cmp(0, 0), Ordering::Greater); // (1, [1, 2]) cmp (None, None) - assert_eq!(cmp(1, 1), Ordering::Greater); // (2, [None]) cmp (2, None) - assert_eq!(cmp(2, 2), Ordering::Less); // (None, None) cmp (None, []) - assert_eq!(cmp(3, 0), Ordering::Less); // None cmp (None, []) - assert_eq!(cmp(2, 0), Ordering::Equal); // (None, None) cmp (None, None) - assert_eq!(cmp(3, 0), Ordering::Less); // None cmp (None, None) - - let opts = SortOptions { - descending: true, - nulls_first: false, - }; - let cmp = make_comparator(&s1, &s2, opts).unwrap(); - assert_eq!(cmp(0, 1), Ordering::Greater); // (1, [1, 2]) cmp (2, None) - assert_eq!(cmp(0, 0), Ordering::Less); // (1, [1, 2]) cmp (None, None) - assert_eq!(cmp(1, 1), Ordering::Less); // (2, [None]) cmp (2, None) - assert_eq!(cmp(2, 2), Ordering::Greater); // (None, None) cmp (None, []) - assert_eq!(cmp(3, 0), Ordering::Greater); // None cmp (None, []) - assert_eq!(cmp(2, 0), Ordering::Equal); // (None, None) cmp (None, None) - assert_eq!(cmp(3, 0), Ordering::Greater); // None cmp (None, None) - - let opts = SortOptions { - descending: false, - nulls_first: false, - }; - let cmp = make_comparator(&s1, &s2, opts).unwrap(); - assert_eq!(cmp(0, 1), Ordering::Less); // (1, [1, 2]) cmp (2, None) - assert_eq!(cmp(0, 0), Ordering::Less); // (1, [1, 2]) cmp (None, None) - assert_eq!(cmp(1, 1), Ordering::Less); // (2, [None]) cmp (2, None) - assert_eq!(cmp(2, 2), Ordering::Greater); // (None, None) cmp (None, []) - assert_eq!(cmp(3, 0), Ordering::Greater); // None cmp (None, []) - assert_eq!(cmp(2, 0), Ordering::Equal); // (None, None) cmp (None, None) - assert_eq!(cmp(3, 0), Ordering::Greater); // None cmp (None, None) - } - - #[test] - fn test_map() { - // Create first map array demonstrating key priority over values: - // [{"a": 100, "b": 1}, {"b": 999, "c": 1}, {}, {"x": 1}] - let string_builder = StringBuilder::new(); - let int_builder = Int32Builder::new(); - let mut map1_builder = MapBuilder::new(None, string_builder, int_builder); - - // {"a": 100, "b": 1} - high value for "a", low value for "b" - map1_builder.keys().append_value("a"); - map1_builder.values().append_value(100); - map1_builder.keys().append_value("b"); - map1_builder.values().append_value(1); - map1_builder.append(true).unwrap(); - - // {"b": 999, "c": 1} - very high value for "b", low value for "c" - map1_builder.keys().append_value("b"); - map1_builder.values().append_value(999); - map1_builder.keys().append_value("c"); - map1_builder.values().append_value(1); - map1_builder.append(true).unwrap(); - - // {} - map1_builder.append(true).unwrap(); - - // {"x": 1} - map1_builder.keys().append_value("x"); - map1_builder.values().append_value(1); - map1_builder.append(true).unwrap(); - - let map1 = map1_builder.finish(); - - // Create second map array: - // [{"a": 1, "c": 999}, {"b": 1, "d": 999}, {"a": 1}, None] - let string_builder = StringBuilder::new(); - let int_builder = Int32Builder::new(); - let mut map2_builder = MapBuilder::new(None, string_builder, int_builder); - - // {"a": 1, "c": 999} - low value for "a", high value for "c" - map2_builder.keys().append_value("a"); - map2_builder.values().append_value(1); - map2_builder.keys().append_value("c"); - map2_builder.values().append_value(999); - map2_builder.append(true).unwrap(); - - // {"b": 1, "d": 999} - low value for "b", high value for "d" - map2_builder.keys().append_value("b"); - map2_builder.values().append_value(1); - map2_builder.keys().append_value("d"); - map2_builder.values().append_value(999); - map2_builder.append(true).unwrap(); - - // {"a": 1} - map2_builder.keys().append_value("a"); - map2_builder.values().append_value(1); - map2_builder.append(true).unwrap(); - - // None - map2_builder.append(false).unwrap(); - - let map2 = map2_builder.finish(); - - let opts = SortOptions { - descending: false, - nulls_first: true, - }; - let cmp = make_comparator(&map1, &map2, opts).unwrap(); - - // Test that keys have priority over values: - // {"a": 100, "b": 1} vs {"a": 1, "c": 999} - // First entries match (a:100 vs a:1), but 100 > 1, so Greater - assert_eq!(cmp(0, 0), Ordering::Greater); - - // {"b": 999, "c": 1} vs {"b": 1, "d": 999} - // First entries match (b:999 vs b:1), but 999 > 1, so Greater - assert_eq!(cmp(1, 1), Ordering::Greater); - - // Key comparison: "a" < "b", so {"a": 100, "b": 1} < {"b": 999, "c": 1} - assert_eq!(cmp(0, 1), Ordering::Less); - - // Empty map vs non-empty - assert_eq!(cmp(2, 2), Ordering::Less); // {} < {"a": 1} - - // Non-null vs null - assert_eq!(cmp(3, 3), Ordering::Greater); // {"x": 1} > None - - // Key priority test: "x" > "a", regardless of values - assert_eq!(cmp(3, 0), Ordering::Greater); // {"x": 1} > {"a": 1, "c": 999} - - // Empty vs non-empty - assert_eq!(cmp(2, 0), Ordering::Less); // {} < {"a": 1, "c": 999} - - let opts = SortOptions { - descending: true, - nulls_first: true, - }; - let cmp = make_comparator(&map1, &map2, opts).unwrap(); - - // With descending=true, value comparison is reversed - assert_eq!(cmp(0, 0), Ordering::Less); // {"a": 100, "b": 1} vs {"a": 1, "c": 999} (reversed) - assert_eq!(cmp(1, 1), Ordering::Less); // {"b": 999, "c": 1} vs {"b": 1, "d": 999} (reversed) - assert_eq!(cmp(0, 1), Ordering::Greater); // {"a": 100, "b": 1} vs {"b": 999, "c": 1} (key order reversed) - assert_eq!(cmp(3, 3), Ordering::Greater); // {"x": 1} > None - assert_eq!(cmp(2, 2), Ordering::Greater); // {} > {"a": 1} (reversed) - - let opts = SortOptions { - descending: false, - nulls_first: false, - }; - let cmp = make_comparator(&map1, &map2, opts).unwrap(); - - // Same key priority behavior with nulls_first=false - assert_eq!(cmp(0, 0), Ordering::Greater); // {"a": 100, "b": 1} vs {"a": 1, "c": 999} - assert_eq!(cmp(1, 1), Ordering::Greater); // {"b": 999, "c": 1} vs {"b": 1, "d": 999} - assert_eq!(cmp(3, 3), Ordering::Less); // {"x": 1} < None (nulls last) - assert_eq!(cmp(2, 2), Ordering::Less); // {} < {"a": 1} - } - - #[test] - fn test_map_vs_list_consistency() { - // Create map arrays and convert them to list arrays to verify comparison consistency - // Map arrays: [{"a": 1, "b": 2}, {"x": 10}, {}, {"c": 3}] - let string_builder = StringBuilder::new(); - let int_builder = Int32Builder::new(); - let mut map1_builder = MapBuilder::new(None, string_builder, int_builder); - - // {"a": 1, "b": 2} - map1_builder.keys().append_value("a"); - map1_builder.values().append_value(1); - map1_builder.keys().append_value("b"); - map1_builder.values().append_value(2); - map1_builder.append(true).unwrap(); - - // {"x": 10} - map1_builder.keys().append_value("x"); - map1_builder.values().append_value(10); - map1_builder.append(true).unwrap(); - - // {} - map1_builder.append(true).unwrap(); - - // {"c": 3} - map1_builder.keys().append_value("c"); - map1_builder.values().append_value(3); - map1_builder.append(true).unwrap(); - - let map1 = map1_builder.finish(); - - // Second map array: [{"a": 1, "b": 2}, {"y": 20}, {"d": 4}, None] - let string_builder = StringBuilder::new(); - let int_builder = Int32Builder::new(); - let mut map2_builder = MapBuilder::new(None, string_builder, int_builder); - - // {"a": 1, "b": 2} - map2_builder.keys().append_value("a"); - map2_builder.values().append_value(1); - map2_builder.keys().append_value("b"); - map2_builder.values().append_value(2); - map2_builder.append(true).unwrap(); - - // {"y": 20} - map2_builder.keys().append_value("y"); - map2_builder.values().append_value(20); - map2_builder.append(true).unwrap(); - - // {"d": 4} - map2_builder.keys().append_value("d"); - map2_builder.values().append_value(4); - map2_builder.append(true).unwrap(); - - // None - map2_builder.append(false).unwrap(); - - let map2 = map2_builder.finish(); - - // Convert map arrays to list arrays (Map entries are struct arrays with key-value pairs) - let list1: ListArray = map1.clone().into(); - let list2: ListArray = map2.clone().into(); - - let test_cases = [ - SortOptions { - descending: false, - nulls_first: true, - }, - SortOptions { - descending: true, - nulls_first: true, - }, - SortOptions { - descending: false, - nulls_first: false, - }, - SortOptions { - descending: true, - nulls_first: false, - }, - ]; - - for opts in test_cases { - let map_cmp = make_comparator(&map1, &map2, opts).unwrap(); - let list_cmp = make_comparator(&list1, &list2, opts).unwrap(); - - // Test all possible index combinations - for i in 0..map1.len() { - for j in 0..map2.len() { - let map_result = map_cmp(i, j); - let list_result = list_cmp(i, j); - assert_eq!( - map_result, list_result, - "Map comparison and List comparison should be equal for indices ({i}, {j}) with opts {opts:?}. Map: {map_result:?}, List: {list_result:?}" - ); - } - } - } - } - - #[test] - fn test_dense_union() { - // create a dense union array with Int32 (type_id = 0) and Utf8 (type_id=1) - // the values are: [1, "b", 2, "a", 3] - // type_ids are: [0, 1, 0, 1, 0] - // offsets are: [0, 0, 1, 1, 2] from [1, 2, 3] and ["b", "a"] - let int_array = Int32Array::from(vec![1, 2, 3]); - let str_array = StringArray::from(vec!["b", "a"]); - - let type_ids = [0, 1, 0, 1, 0].into_iter().collect::>(); - let offsets = [0, 0, 1, 1, 2].into_iter().collect::>(); - - let union_fields = [ - (0, Arc::new(Field::new("A", DataType::Int32, false))), - (1, Arc::new(Field::new("B", DataType::Utf8, false))), - ] - .into_iter() - .collect::(); - - let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)]; - - let array1 = - UnionArray::try_new(union_fields.clone(), type_ids, Some(offsets), children).unwrap(); - - // create a second array: [2, "a", 1, "c"] - // type ids are: [0, 1, 0, 1] - // offsets are: [0, 0, 1, 1] from [2, 1] and ["a", "c"] - let int_array2 = Int32Array::from(vec![2, 1]); - let str_array2 = StringArray::from(vec!["a", "c"]); - let type_ids2 = [0, 1, 0, 1].into_iter().collect::>(); - let offsets2 = [0, 0, 1, 1].into_iter().collect::>(); - - let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)]; - - let array2 = - UnionArray::try_new(union_fields, type_ids2, Some(offsets2), children2).unwrap(); - - let opts = SortOptions { - descending: false, - nulls_first: true, - }; - - // comparing - // [1, "b", 2, "a", 3] - // [2, "a", 1, "c"] - let cmp = make_comparator(&array1, &array2, opts).unwrap(); - - // array1[0] = (type_id=0, value=1) - // array2[0] = (type_id=0, value=2) - assert_eq!(cmp(0, 0), Ordering::Less); // 1 < 2 - - // array1[0] = (type_id=0, value=1) - // array2[1] = (type_id=1, value="a") - assert_eq!(cmp(0, 1), Ordering::Less); // type_id 0 < 1 - - // array1[1] = (type_id=1, value="b") - // array2[1] = (type_id=1, value="a") - assert_eq!(cmp(1, 1), Ordering::Greater); // "b" > "a" - - // array1[2] = (type_id=0, value=2) - // array2[0] = (type_id=0, value=2) - assert_eq!(cmp(2, 0), Ordering::Equal); // 2 == 2 - - // array1[3] = (type_id=1, value="a") - // array2[1] = (type_id=1, value="a") - assert_eq!(cmp(3, 1), Ordering::Equal); // "a" == "a" - - // array1[1] = (type_id=1, value="b") - // array2[3] = (type_id=1, value="c") - assert_eq!(cmp(1, 3), Ordering::Less); // "b" < "c" - - let opts_desc = SortOptions { - descending: true, - nulls_first: true, - }; - let cmp_desc = make_comparator(&array1, &array2, opts_desc).unwrap(); - - assert_eq!(cmp_desc(0, 0), Ordering::Greater); // 1 > 2 (reversed) - assert_eq!(cmp_desc(0, 1), Ordering::Greater); // type_id 0 < 1, reversed to Greater - assert_eq!(cmp_desc(1, 1), Ordering::Less); // "b" < "a" (reversed) - } - - #[test] - fn test_sparse_union() { - // create a sparse union array with Int32 (type_id=0) and Utf8 (type_id=1) - // values: [1, "b", 3] - // note, in sparse unions, child arrays have the same length as the union - let int_array = Int32Array::from(vec![Some(1), None, Some(3)]); - let str_array = StringArray::from(vec![None, Some("b"), None]); - let type_ids = [0, 1, 0].into_iter().collect::>(); - - let union_fields = [ - (0, Arc::new(Field::new("a", DataType::Int32, false))), - (1, Arc::new(Field::new("b", DataType::Utf8, false))), - ] - .into_iter() - .collect::(); - - let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)]; - - let array = UnionArray::try_new(union_fields, type_ids, None, children).unwrap(); - - let opts = SortOptions::default(); - let cmp = make_comparator(&array, &array, opts).unwrap(); - - // array[0] = (type_id=0, value=1), array[2] = (type_id=0, value=3) - assert_eq!(cmp(0, 2), Ordering::Less); // 1 < 3 - // array[0] = (type_id=0, value=1), array[1] = (type_id=1, value="b") - assert_eq!(cmp(0, 1), Ordering::Less); // type_id 0 < 1 - } - - #[test] - #[should_panic(expected = "index out of bounds")] - fn test_union_out_of_bounds() { - // create a dense union array with 3 elements - let int_array = Int32Array::from(vec![1, 2]); - let str_array = StringArray::from(vec!["a"]); - - let type_ids = [0, 1, 0].into_iter().collect::>(); - let offsets = [0, 0, 1].into_iter().collect::>(); - - let union_fields = [ - (0, Arc::new(Field::new("A", DataType::Int32, false))), - (1, Arc::new(Field::new("B", DataType::Utf8, false))), - ] - .into_iter() - .collect::(); - - let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)]; - - let array = UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap(); - - let opts = SortOptions::default(); - let cmp = make_comparator(&array, &array, opts).unwrap(); - - // oob - cmp(0, 3); - } - - #[test] - fn test_union_incompatible_fields() { - // create first union with Int32 and Utf8 - let int_array1 = Int32Array::from(vec![1, 2]); - let str_array1 = StringArray::from(vec!["a", "b"]); - - let type_ids1 = [0, 1].into_iter().collect::>(); - let offsets1 = [0, 0].into_iter().collect::>(); - - let union_fields1 = [ - (0, Arc::new(Field::new("A", DataType::Int32, false))), - (1, Arc::new(Field::new("B", DataType::Utf8, false))), - ] - .into_iter() - .collect::(); - - let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)]; - - let array1 = - UnionArray::try_new(union_fields1, type_ids1, Some(offsets1), children1).unwrap(); - - // create second union with Int32 and Float64 (incompatible with first) - let int_array2 = Int32Array::from(vec![3, 4]); - let float_array2 = Float64Array::from(vec![1.0, 2.0]); - - let type_ids2 = [0, 1].into_iter().collect::>(); - let offsets2 = [0, 0].into_iter().collect::>(); - - let union_fields2 = [ - (0, Arc::new(Field::new("A", DataType::Int32, false))), - (1, Arc::new(Field::new("C", DataType::Float64, false))), - ] - .into_iter() - .collect::(); - - let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(float_array2)]; - - let array2 = - UnionArray::try_new(union_fields2, type_ids2, Some(offsets2), children2).unwrap(); - - let opts = SortOptions::default(); - - let Result::Err(ArrowError::InvalidArgumentError(out)) = - make_comparator(&array1, &array2, opts) - else { - panic!("expected error when making comparator of incompatible union arrays"); - }; - - assert_eq!( - &out, - "Cannot compare UnionArrays with different fields: left=[(0, Field { name: \"A\", data_type: Int32 }), (1, Field { name: \"B\", data_type: Utf8 })], right=[(0, Field { name: \"A\", data_type: Int32 }), (1, Field { name: \"C\", data_type: Float64 })]" - ); - } - - #[test] - fn test_union_incompatible_modes() { - // create first union as Dense with Int32 and Utf8 - let int_array1 = Int32Array::from(vec![1, 2]); - let str_array1 = StringArray::from(vec!["a", "b"]); - - let type_ids1 = [0, 1].into_iter().collect::>(); - let offsets1 = [0, 0].into_iter().collect::>(); - - let union_fields1 = [ - (0, Arc::new(Field::new("A", DataType::Int32, false))), - (1, Arc::new(Field::new("B", DataType::Utf8, false))), - ] - .into_iter() - .collect::(); - - let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)]; - - let array1 = - UnionArray::try_new(union_fields1.clone(), type_ids1, Some(offsets1), children1) - .unwrap(); - - // create second union as Sparse with same fields (Int32 and Utf8) - let int_array2 = Int32Array::from(vec![Some(3), None]); - let str_array2 = StringArray::from(vec![None, Some("c")]); - - let type_ids2 = [0, 1].into_iter().collect::>(); - - let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)]; - - let array2 = UnionArray::try_new(union_fields1, type_ids2, None, children2).unwrap(); - - let opts = SortOptions::default(); - - let Result::Err(ArrowError::InvalidArgumentError(out)) = - make_comparator(&array1, &array2, opts) - else { - panic!("expected error when making comparator of union arrays with different modes"); - }; - - assert_eq!( - &out, - "Cannot compare UnionArrays with different modes: left=Dense, right=Sparse" - ); - } - - #[test] - fn test_null_array_cmp() { - let a = NullArray::new(3); - let b = NullArray::new(3); - let cmp = make_comparator(&a, &b, SortOptions::default()).unwrap(); - - assert_eq!(cmp(0, 0), Ordering::Equal); - assert_eq!(cmp(0, 1), Ordering::Equal); - assert_eq!(cmp(2, 0), Ordering::Equal); - } - - #[test] - fn test_run_end_encoded_int32() { - // Create RunEndEncoded arrays: - // array1: [1, 1, 2, 2, 2, 3] - // run_ends1: [2, 5, 6], values1: [1, 2, 3] - let run_ends1 = Int32Array::from(vec![2, 5, 6]); - let values1 = Int32Array::from(vec![1, 2, 3]); - let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); - - // array2: [1, 2, 2, 3, 3, 3] - // run_ends2: [1, 3, 6], values2: [1, 2, 3] - let run_ends2 = Int32Array::from(vec![1, 3, 6]); - let values2 = Int32Array::from(vec![1, 2, 3]); - let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - // array1[0] = 1, array2[0] = 1 - assert_eq!(cmp(0, 0), Ordering::Equal); - // array1[0] = 1, array2[1] = 2 - assert_eq!(cmp(0, 1), Ordering::Less); - // array1[2] = 2, array2[1] = 2 - assert_eq!(cmp(2, 1), Ordering::Equal); - // array1[5] = 3, array2[5] = 3 - assert_eq!(cmp(5, 5), Ordering::Equal); - // array1[1] = 1, array2[2] = 2 - assert_eq!(cmp(1, 2), Ordering::Less); - // array1[4] = 2, array2[4] = 3 - assert_eq!(cmp(4, 4), Ordering::Less); - } - - #[test] - fn test_run_end_encoded_with_nulls() { - // Create RunEndEncoded arrays with nulls: - // array1: [1, 1, null, null, 2] - // run_ends1: [2, 4, 5], values1: [1, null, 2] - let run_ends1 = Int32Array::from(vec![2, 4, 5]); - let values1 = Int32Array::from(vec![Some(1), None, Some(2)]); - let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); - - // array2: [null, 1, 1, 2, null] - // run_ends2: [1, 3, 4, 5], values2: [null, 1, 2, null] - let run_ends2 = Int32Array::from(vec![1, 3, 4, 5]); - let values2 = Int32Array::from(vec![None, Some(1), Some(2), None]); - let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); - - let opts = SortOptions::default(); - let cmp = make_comparator(&array1, &array2, opts).unwrap(); - - // array1[0] = 1, array2[1] = 1 - assert_eq!(cmp(0, 1), Ordering::Equal); - // array1[2] = null, array2[0] = null - assert_eq!(cmp(2, 0), Ordering::Equal); - // array1[0] = 1, array2[0] = null (nulls first by default) - assert_eq!(cmp(0, 0), Ordering::Greater); - // array1[2] = null, array2[1] = 1 - assert_eq!(cmp(2, 1), Ordering::Less); - } - - #[test] - fn test_run_end_encoded_int16() { - // Test with Int16 run ends - let run_ends1 = Int16Array::from(vec![3_i16, 5, 6]); - let values1 = StringArray::from(vec!["a", "b", "c"]); - let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); - - let run_ends2 = Int16Array::from(vec![2_i16, 4, 6]); - let values2 = StringArray::from(vec!["a", "b", "c"]); - let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - // array1: [a, a, a, b, b, c] - // array2: [a, a, b, b, c, c] - assert_eq!(cmp(0, 0), Ordering::Equal); // a vs a - assert_eq!(cmp(2, 2), Ordering::Less); // a vs b - assert_eq!(cmp(3, 2), Ordering::Equal); // b vs b - assert_eq!(cmp(5, 4), Ordering::Equal); // c vs c - } - - #[test] - fn test_run_end_encoded_int64() { - // Test with Int64 run ends - let run_ends1 = Int64Array::from(vec![2_i64, 4, 6]); - let values1 = Int64Array::from(vec![10_i64, 20, 30]); - let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); - - let run_ends2 = Int64Array::from(vec![3_i64, 5, 6]); - let values2 = Int64Array::from(vec![10_i64, 20, 30]); - let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); - - let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap(); - - // array1: [10, 10, 20, 20, 30, 30] - // array2: [10, 10, 10, 20, 20, 30] - assert_eq!(cmp(0, 0), Ordering::Equal); // 10 vs 10 - assert_eq!(cmp(1, 2), Ordering::Equal); // 10 vs 10 - assert_eq!(cmp(2, 3), Ordering::Equal); // 20 vs 20 - assert_eq!(cmp(4, 4), Ordering::Greater); // 30 vs 20 - } - - #[test] - fn test_run_end_encoded_sliced() { - // Create a RunEndEncoded array and slice it: - // original: [1, 1, 2, 2, 2, 3, 3, 4] - // run_ends: [2, 5, 7, 8], values: [1, 2, 3, 4] - let run_ends = Int32Array::from(vec![2, 5, 7, 8]); - let values = Int32Array::from(vec![1, 2, 3, 4]); - let array = RunArray::::try_new(&run_ends, &values).unwrap(); - - // slice1 = array[1..5] => [1, 2, 2, 2] - let slice1 = array.slice(1, 4); - // slice2 = array[3..7] => [2, 2, 3, 3] - let slice2 = array.slice(3, 4); - - let cmp = make_comparator(&slice1, &slice2, SortOptions::default()).unwrap(); - - // slice1[0]=1, slice2[0]=2 - assert_eq!(cmp(0, 0), Ordering::Less); - // slice1[1]=2, slice2[0]=2 - assert_eq!(cmp(1, 0), Ordering::Equal); - // slice1[3]=2, slice2[2]=3 - assert_eq!(cmp(3, 2), Ordering::Less); - // slice1[1]=2, slice2[3]=3 - assert_eq!(cmp(1, 3), Ordering::Less); - - // Compare a sliced array with an unsliced array - let run_ends2 = Int32Array::from(vec![2, 4]); - let values2 = Int32Array::from(vec![1, 2]); - let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); - - let cmp = make_comparator(&slice1, &array2, SortOptions::default()).unwrap(); - - // slice1[0]=1, array2[0]=1 - assert_eq!(cmp(0, 0), Ordering::Equal); - // slice1[1]=2, array2[1]=1 - assert_eq!(cmp(1, 1), Ordering::Greater); - // slice1[3]=2, array2[3]=2 - assert_eq!(cmp(3, 3), Ordering::Equal); - } - - #[test] - fn test_run_end_encoded_sliced_with_nulls() { - // Create a RunEndEncoded array with nulls: - // original: [1, 1, null, null, 2, 2, null, 3] - // run_ends: [2, 4, 6, 7, 8], values: [1, null, 2, null, 3] - let run_ends = Int32Array::from(vec![2, 4, 6, 7, 8]); - let values = Int32Array::from(vec![Some(1), None, Some(2), None, Some(3)]); - let array = RunArray::::try_new(&run_ends, &values).unwrap(); - - // slice1 = array[1..6] => [1, null, null, 2, 2] - let slice1 = array.slice(1, 5); - // slice2 = array[3..8] => [null, 2, 2, null, 3] - let slice2 = array.slice(3, 5); - - let opts = SortOptions::default(); // nulls_first=true, descending=false - let cmp = make_comparator(&slice1, &slice2, opts).unwrap(); - - // slice1[0]=1, slice2[0]=null - assert_eq!(cmp(0, 0), Ordering::Greater); - // slice1[1]=null, slice2[0]=null - assert_eq!(cmp(1, 0), Ordering::Equal); - // slice1[1]=null, slice2[1]=2 - assert_eq!(cmp(1, 1), Ordering::Less); - // slice1[3]=2, slice2[1]=2 - assert_eq!(cmp(3, 1), Ordering::Equal); - // slice1[4]=2, slice2[4]=3 - assert_eq!(cmp(4, 4), Ordering::Less); - // slice1[3]=2, slice2[3]=null - assert_eq!(cmp(3, 3), Ordering::Greater); - } - - #[test] - fn test_run_end_encoded_different_types() { - // Test with different run end types - should fail - let run_ends1 = Int32Array::from(vec![2, 4, 6]); - let values1 = Int32Array::from(vec![1, 2, 3]); - let array1 = RunArray::::try_new(&run_ends1, &values1).unwrap(); - - let run_ends2 = Int64Array::from(vec![2_i64, 4, 6]); - let values2 = Int64Array::from(vec![1_i64, 2, 3]); - let array2 = RunArray::::try_new(&run_ends2, &values2).unwrap(); - - let result = make_comparator(&array1, &array2, SortOptions::default()); - assert!(result.is_err()); - let err = match result { - Err(e) => e.to_string(), - Ok(_) => panic!("Expected error"), - }; - assert!(err.contains("Cannot compare RunEndEncoded arrays")); - } -} +pub use arrow_cmp::{DynComparator, make_comparator}; diff --git a/arrow-select/Cargo.toml b/arrow-select/Cargo.toml index cd43c6c2a699..3e38dc9f80ad 100644 --- a/arrow-select/Cargo.toml +++ b/arrow-select/Cargo.toml @@ -40,15 +40,9 @@ arrow-buffer = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true } arrow-array = { workspace = true } -arrow-ord-basic = { workspace = true, optional = true } +arrow-cmp = { workspace = true } num-traits = { version = "0.2.19", default-features = false, features = ["std"] } ahash = { version = "0.8", default-features = false} -[features] -# Enables the run-end-encoded `take` fast path that merges adjacent -# physical indices whose underlying values compare equal. Pulls in -# `arrow-ord-basic` for the slot-wise comparator. -run_end_encoded = ["dep:arrow-ord-basic"] - [dev-dependencies] rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] } diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index bb22eeba96e8..971f8b29441f 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -26,11 +26,8 @@ use arrow_buffer::{ bit_util, }; use arrow_data::ArrayDataBuilder; -#[cfg(feature = "run_end_encoded")] -use arrow_ord_basic::make_comparator; -use arrow_schema::{ArrowError, DataType, FieldRef, UnionMode}; -#[cfg(feature = "run_end_encoded")] -use arrow_schema::SortOptions; +use arrow_cmp::make_comparator; +use arrow_schema::{ArrowError, DataType, FieldRef, SortOptions, UnionMode}; use std::fmt::Display; use std::mem::ManuallyDrop; use std::sync::Arc; @@ -884,12 +881,6 @@ fn take_run( let mut take_value_indices = BufferBuilder::::new(1); let mut new_physical_len = 1; - // With the `run_end_encoded` feature, also collapse adjacent physical runs - // whose underlying values compare equal. Without it, we still drop - // duplicate physical indices but leave logically-equal-yet-distinct runs - // separated, which keeps `arrow-select` free of the `arrow-ord-basic` - // comparator dependency for downstream users that don't need it. - #[cfg(feature = "run_end_encoded")] let values_cmp = make_comparator( run_array.values().as_ref(), run_array.values().as_ref(), @@ -899,10 +890,7 @@ fn take_run( for ix in 1..physical_indices.len() { let prev_idx = physical_indices[ix - 1]; let cur_idx = physical_indices[ix]; - #[cfg(feature = "run_end_encoded")] let is_new_run = cur_idx != prev_idx && values_cmp(cur_idx, prev_idx).is_ne(); - #[cfg(not(feature = "run_end_encoded"))] - let is_new_run = cur_idx != prev_idx; if is_new_run { take_value_indices.append(I::Native::from_usize(prev_idx).unwrap()); new_run_ends_builder.append(T::Native::from_usize(ix).unwrap()); @@ -2553,7 +2541,6 @@ mod tests { assert_eq!(null_buf.as_slice(), &[0b11111111]); } - #[cfg(feature = "run_end_encoded")] #[test] fn test_take_runs() { let logical_array: Vec = vec![1_i32, 1, 2, 2, 1, 1, 1, 2, 2, 1, 1, 2, 2]; @@ -2844,7 +2831,6 @@ mod tests { assert_eq!(run_result.values().len(), 0); } - #[cfg(feature = "run_end_encoded")] #[test] fn test_take_run_end_encoded_merges_identical_runs() { // see https://github.com/apache/arrow-rs/issues/7710 @@ -2860,7 +2846,6 @@ mod tests { assert_eq!(actual, vec![1, 1, 1, 1]); } - #[cfg(feature = "run_end_encoded")] #[test] fn test_take_run_end_encoded_merges_identical_string_runs() { let mut builder = StringRunBuilder::::new(); @@ -2882,7 +2867,6 @@ mod tests { assert_eq!(actual, vec!["bob", "bob", "bob", "bob"]); } - #[cfg(feature = "run_end_encoded")] #[test] fn test_take_run_end_encoded_mixed_runs() { // Validates that runs are merged whether the same logical value comes diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml index a8304cc25b99..8e56457ff0a5 100644 --- a/arrow/Cargo.toml +++ b/arrow/Cargo.toml @@ -61,17 +61,13 @@ half = { version = "2.1", default-features = false, features = ["rand_distr"], o all-features = true [features] -default = ["csv", "ipc", "json", "run_end_encoded"] +default = ["csv", "ipc", "json"] async = ["arrow-array/async"] ipc_compression = ["ipc", "arrow-ipc/lz4", "arrow-ipc/zstd"] csv = ["arrow-csv"] ipc = ["arrow-ipc"] json = ["arrow-json"] prettyprint = ["arrow-cast/prettyprint"] -# Enables the run-end-encoded `take` fast path in `arrow-select` that merges -# adjacent physical runs whose underlying values compare equal. Pulls in the -# `arrow-ord-basic` slot-wise comparator. -run_end_encoded = ["arrow-select/run_end_encoded"] # The test utils feature enables code used in benchmarks and tests but # not the core arrow code itself. Be aware that `rand` must be kept as # an optional dependency for supporting compile to wasm32-unknown-unknown