This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new fa1c4ee2c6 perf: optimize take for RunEndArrays & introduce
`arrow-cmp` crate (#10325)
fa1c4ee2c6 is described below
commit fa1c4ee2c6499f27bab79232e692a86d2beddaf8
Author: RIchard Baah <[email protected]>
AuthorDate: Wed Aug 5 05:48:33 2026 -0400
perf: optimize take for RunEndArrays & introduce `arrow-cmp` crate (#10325)
## The Diff (+2,037 -1,839) makes this PR appear much bigger than it is.
this is mostly moving code around & test. most of the logic is under 100
LOC
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- Closes #7710.
- revised version of https://github.com/apache/arrow-rs/pull/9865
# Rationale for this change
for a logical representation of an ree
```
let logical_repr = [1, 1, 0, 0, 1, 1];
let ree_take_results = take(logical_repr,[0,1,4,5]);
this produces
run_ends = [2,4]
values = [1,1]
```
when the result should be
`runs: [4], values: [1]`
both answers are correct but ree's should be as compact as possible.
see #7710
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
- `take()` on `RunEndEncoded` arrays now compares values instead of
physical indices when deciding run boundaries, producing a more compact
run-end representation and fixing cases where identical values across
different runs were not merged.
- Introduces `arrow-cmp`, a minimal crate that extracts
`make_comparator`/`DynComparator` from `arrow-ord` so `arrow-select` can
use slot-wise comparison without a circular dependency. Most of the line
diff is code moving **(+1,889,-1834)**, not new logic `arrow-ord`
re-exports from` arrow-cmp` so its public API is unchanged.
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
yes, I included three test to assert the compaction behavior we expect
from a Run-end array.
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
1. `take()` on `RunEndEncoded` arrays behaves differently. the output
may have fewer runs than before. Previously runs were only merged when
they hit the same physical index; now runs with equal values are also
merged across different physical indices. Any code asserting on the
exact run-end structure of take output could break.
2. New `arrow-cmp` crate is published. Users can depend on it directly
to get `make_comparator` / `DynComparator` without pulling in all of
`arrow-ord`. `arrow-ord` still re-exports both so its API is unchanged.
---------
Co-authored-by: Jeffrey Vo <[email protected]>
---
Cargo.lock | 12 +
Cargo.toml | 2 +
{arrow-ord => arrow-cmp}/Cargo.toml | 12 +-
arrow-cmp/LICENSE.txt | 1 +
arrow-cmp/NOTICE.txt | 1 +
arrow-ord/src/ord.rs => arrow-cmp/src/lib.rs | 279 ++--
arrow-ord/Cargo.toml | 1 +
arrow-ord/src/ord.rs | 1834 +-------------------------
arrow-select/Cargo.toml | 1 +
arrow-select/src/take.rs | 134 +-
dev/release/README.md | 15 +-
11 files changed, 316 insertions(+), 1976 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 3b80c386e0..f2d3786d3a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -301,6 +301,16 @@ dependencies = [
"ryu",
]
+[[package]]
+name = "arrow-cmp"
+version = "59.2.0"
+dependencies = [
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-schema",
+ "half",
+]
+
[[package]]
name = "arrow-csv"
version = "59.2.0"
@@ -460,6 +470,7 @@ version = "59.2.0"
dependencies = [
"arrow-array",
"arrow-buffer",
+ "arrow-cmp",
"arrow-data",
"arrow-schema",
"arrow-select",
@@ -512,6 +523,7 @@ dependencies = [
"ahash",
"arrow-array",
"arrow-buffer",
+ "arrow-cmp",
"arrow-data",
"arrow-schema",
"num-traits",
diff --git a/Cargo.toml b/Cargo.toml
index 3a2bc6b6ce..7f7d98cf10 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -32,6 +32,7 @@ members = [
"arrow-integration-testing",
"arrow-ipc",
"arrow-json",
+ "arrow-cmp",
"arrow-ord",
"arrow-pyarrow",
"arrow-row",
@@ -90,6 +91,7 @@ arrow-csv = { version = "59.2.0", path = "./arrow-csv" }
arrow-data = { version = "59.2.0", path = "./arrow-data" }
arrow-ipc = { version = "59.2.0", path = "./arrow-ipc" }
arrow-json = { version = "59.2.0", path = "./arrow-json" }
+arrow-cmp = { version = "59.2.0", path = "./arrow-cmp" }
arrow-ord = { version = "59.2.0", path = "./arrow-ord" }
arrow-pyarrow = { version = "59.2.0", path = "./arrow-pyarrow" }
arrow-row = { version = "59.2.0", path = "./arrow-row" }
diff --git a/arrow-ord/Cargo.toml b/arrow-cmp/Cargo.toml
similarity index 83%
copy from arrow-ord/Cargo.toml
copy to arrow-cmp/Cargo.toml
index fa351f32d8..bbeadb22ab 100644
--- a/arrow-ord/Cargo.toml
+++ b/arrow-cmp/Cargo.toml
@@ -16,9 +16,9 @@
# under the License.
[package]
-name = "arrow-ord"
+name = "arrow-cmp"
version = { workspace = true }
-description = "Ordering kernels for arrow arrays"
+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 }
@@ -29,7 +29,7 @@ edition = { workspace = true }
rust-version = { workspace = true }
[lib]
-name = "arrow_ord"
+name = "arrow_cmp"
bench = false
[package.metadata.docs.rs]
@@ -38,13 +38,7 @@ all-features = true
[dependencies]
arrow-array = { workspace = true }
arrow-buffer = { workspace = true }
-arrow-data = { workspace = true }
arrow-schema = { workspace = true }
-arrow-select = { workspace = true }
[dev-dependencies]
half = { version = "2.1", default-features = false, features = ["num-traits"] }
-rand = { version = "0.10", default-features = false, features = ["std",
"std_rng"] }
-
-[lints]
-workspace = true
diff --git a/arrow-cmp/LICENSE.txt b/arrow-cmp/LICENSE.txt
new file mode 120000
index 0000000000..4ab43736a8
--- /dev/null
+++ b/arrow-cmp/LICENSE.txt
@@ -0,0 +1 @@
+../LICENSE.txt
\ No newline at end of file
diff --git a/arrow-cmp/NOTICE.txt b/arrow-cmp/NOTICE.txt
new file mode 120000
index 0000000000..eb9f24e040
--- /dev/null
+++ b/arrow-cmp/NOTICE.txt
@@ -0,0 +1 @@
+../NOTICE.txt
\ No newline at end of file
diff --git a/arrow-ord/src/ord.rs b/arrow-cmp/src/lib.rs
similarity index 90%
copy from arrow-ord/src/ord.rs
copy to arrow-cmp/src/lib.rs
index bcc6a217bd..c49d1ae940 100644
--- a/arrow-ord/src/ord.rs
+++ b/arrow-cmp/src/lib.rs
@@ -15,7 +15,22 @@
// specific language governing permissions and limitations
// under the License.
-//! Contains functions and function factories to compare arrays.
+//! 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::*;
@@ -46,7 +61,7 @@ fn compare_run_end_encoded<R: RunEndIndexType>(
Ok(f)
}
-/// Compare the values at two arbitrary indices in two arrays.
+/// Compare values at arbitrary indices in two arrays.
pub type DynComparator = Box<dyn Fn(usize, usize) -> Ordering + Send + Sync>;
/// If parent sort order is descending we need to invert the value of
nulls_first so that
@@ -114,7 +129,10 @@ fn compare_primitive<T: ArrowPrimitiveType>(
left: &dyn Array,
right: &dyn Array,
opts: SortOptions,
-) -> DynComparator {
+) -> DynComparator
+where
+ T::Native: ArrowNativeTypeOp,
+{
let left = left.as_primitive::<T>();
let right = right.as_primitive::<T>();
let l_values = left.values().clone();
@@ -165,10 +183,29 @@ fn compare_byte_view<T: ByteViewType>(
let l = left.clone();
let r = right.clone();
compare(left, right, opts, move |i, j| {
- crate::cmp::compare_byte_view(&l, i, &r, j)
+ compare_byte_view_values(&l, i, &r, j)
})
}
+fn compare_byte_view_values<T: ByteViewType>(
+ left: &GenericByteViewArray<T>,
+ left_idx: usize,
+ right: &GenericByteViewArray<T>,
+ 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::<T>::inline_key_fast(*l_view)
+ .cmp(&GenericByteViewArray::<T>::inline_key_fast(*r_view));
+ }
+
+ unsafe { GenericByteViewArray::compare_unchecked(left, left_idx, right,
right_idx) }
+}
+
fn compare_dict<K: ArrowDictionaryKeyType>(
left: &dyn Array,
right: &dyn Array,
@@ -370,15 +407,13 @@ fn compare_union(
if left_fields != right_fields {
return Err(ArrowError::InvalidArgumentError(format!(
- "Cannot compare UnionArrays with different fields: left={:?},
right={:?}",
- left_fields, right_fields
+ "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={:?},
right={:?}",
- left_mode, right_mode
+ "Cannot compare UnionArrays with different modes:
left={left_mode:?}, right={right_mode:?}"
)));
}
@@ -404,10 +439,8 @@ fn compare_union(
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);
@@ -423,75 +456,11 @@ fn compare_union(
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
-/// ```
+/// Returns a comparison function that compares two values at two arbitrary
indices.
///
-/// # 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<BooleanArray, ArrowError> {
-/// 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))
-/// }
-/// ````
+/// 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,
@@ -568,10 +537,11 @@ pub fn make_comparator(
#[cfg(test)]
mod tests {
use super::*;
- use arrow_array::builder::{Int32Builder, ListBuilder};
+ use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder,
StringBuilder};
use arrow_buffer::{IntervalDayTime, NullBuffer, OffsetBuffer,
ScalarBuffer, i256};
- use arrow_schema::{DataType, Field, Fields, UnionFields};
+ use arrow_schema::{ArrowError, DataType, Field, Fields, UnionFields};
use half::f16;
+ use std::cmp::Ordering;
use std::sync::Arc;
#[test]
@@ -1262,31 +1232,64 @@ mod tests {
#[test]
fn test_map() {
// Create first map array demonstrating key priority over values:
- let map1 = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
- vec![
- // high value for "a", low value for "b"
- Some(vec![("a", Some(100)), ("b", Some(1))]),
- // very high value for "b", low value for "c"
- Some(vec![("b", Some(999)), ("c", Some(1))]),
- Some(vec![]),
- Some(vec![("x", Some(1))]),
- ],
- false,
- );
+ // [{"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 map2 = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
- vec![
- // low value for "a", high value for "c"
- Some(vec![("a", Some(1)), ("c", Some(999))]),
- // low value for "b", high value for "d"
- Some(vec![("b", Some(1)), ("d", Some(999))]),
- Some(vec![("a", Some(1))]),
- None,
- ],
- false,
- );
+ 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,
@@ -1347,25 +1350,59 @@ mod tests {
#[test]
fn test_map_vs_list_consistency() {
// Create map arrays and convert them to list arrays to verify
comparison consistency
- let map1 = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
- vec![
- Some(vec![("a", Some(1)), ("b", Some(2))]),
- Some(vec![("x", Some(10))]),
- Some(vec![]),
- Some(vec![("c", Some(3))]),
- ],
- false,
- );
-
- let map2 = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
- vec![
- Some(vec![("a", Some(1)), ("b", Some(2))]),
- Some(vec![("y", Some(20))]),
- Some(vec![("d", Some(4))]),
- None,
- ],
- false,
- );
+ // 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();
diff --git a/arrow-ord/Cargo.toml b/arrow-ord/Cargo.toml
index fa351f32d8..d0e914e02a 100644
--- a/arrow-ord/Cargo.toml
+++ b/arrow-ord/Cargo.toml
@@ -38,6 +38,7 @@ all-features = true
[dependencies]
arrow-array = { workspace = true }
arrow-buffer = { workspace = true }
+arrow-cmp = { workspace = true }
arrow-data = { 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 bcc6a217bd..f0fe975c99 100644
--- a/arrow-ord/src/ord.rs
+++ b/arrow-ord/src/ord.rs
@@ -17,1836 +17,4 @@
//! Contains functions and function factories to compare arrays.
-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<R: RunEndIndexType>(
- left: &dyn Array,
- right: &dyn Array,
- opts: SortOptions,
-) -> Result<DynComparator, ArrowError> {
- let left = left.as_run::<R>();
- let right = right.as_run::<R>();
-
- 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<dyn Fn(usize, usize) -> 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<A, F>(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::<true, true, _>(l, r, cmp),
- (true, false) => compare_impl::<true, false, _>(l, r, cmp),
- (false, true) => compare_impl::<false, true, _>(l, r, cmp),
- (false, false) => compare_impl::<false, false, _>(l, r, cmp),
- }
-}
-
-fn compare_impl<const NULLS_FIRST: bool, const DESCENDING: bool, F>(
- l: Option<NullBuffer>,
- r: Option<NullBuffer>,
- 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<T: ArrowPrimitiveType>(
- left: &dyn Array,
- right: &dyn Array,
- opts: SortOptions,
-) -> DynComparator {
- let left = left.as_primitive::<T>();
- let right = right.as_primitive::<T>();
- 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<T: ByteArrayType>(
- left: &dyn Array,
- right: &dyn Array,
- opts: SortOptions,
-) -> DynComparator {
- let left = left.as_bytes::<T>();
- let right = right.as_bytes::<T>();
-
- 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<T: ByteViewType>(
- left: &dyn Array,
- right: &dyn Array,
- opts: SortOptions,
-) -> DynComparator {
- let left = left.as_byte_view::<T>();
- let right = right.as_byte_view::<T>();
-
- 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<K: ArrowDictionaryKeyType>(
- left: &dyn Array,
- right: &dyn Array,
- opts: SortOptions,
-) -> Result<DynComparator, ArrowError> {
- let left = left.as_dictionary::<K>();
- let right = right.as_dictionary::<K>();
-
- 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<O: OffsetSizeTrait>(
- left: &dyn Array,
- right: &dyn Array,
- opts: SortOptions,
-) -> Result<DynComparator, ArrowError> {
- let left = left.as_list::<O>();
- let right = right.as_list::<O>();
-
- 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<DynComparator, ArrowError> {
- 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<O: OffsetSizeTrait>(
- left: &dyn Array,
- right: &dyn Array,
- opts: SortOptions,
-) -> Result<DynComparator, ArrowError> {
- let left = left.as_list_view::<O>();
- let right = right.as_list_view::<O>();
-
- 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<DynComparator, ArrowError> {
- 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<DynComparator, ArrowError> {
- 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::<Result<Vec<_>, _>>()?;
-
- 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<DynComparator, ArrowError> {
- 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<BooleanArray, ArrowError> {
-/// 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<DynComparator, ArrowError> {
- 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::<Utf8Type>(left, right, opts)),
- (LargeUtf8, LargeUtf8) => Ok(compare_bytes::<LargeUtf8Type>(left,
right, opts)),
- (Utf8View, Utf8View) => Ok(compare_byte_view::<StringViewType>(left,
right, opts)),
- (Binary, Binary) => Ok(compare_bytes::<BinaryType>(left, right, opts)),
- (LargeBinary, LargeBinary) =>
Ok(compare_bytes::<LargeBinaryType>(left, right, opts)),
- (BinaryView, BinaryView) =>
Ok(compare_byte_view::<BinaryViewType>(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::<i32>(left, right, opts),
- (LargeList(_), LargeList(_)) => compare_list::<i64>(left, right, opts),
- (ListView(_), ListView(_)) => compare_list_view::<i32>(left, right,
opts),
- (LargeListView(_), LargeListView(_)) => compare_list_view::<i64>(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};
- 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::<Decimal32Array>()
- .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::<Decimal64Array>()
- .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::<Decimal128Array>()
- .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::<Decimal256Array>()
- .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::<DictionaryArray<Int16Type>>();
-
- 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::<DictionaryArray<Int16Type>>();
- let d2 = vec!["e", "f", "g", "a"];
- let a2 = d2.into_iter().collect::<DictionaryArray<Int16Type>>();
-
- 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<T: ByteArrayType>() {
- let offsets = OffsetBuffer::from_lengths([3, 3, 1]);
- let a = GenericByteArray::<T>::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::<Utf8Type>();
- test_bytes_impl::<LargeUtf8Type>();
- test_bytes_impl::<BinaryType>();
- test_bytes_impl::<LargeBinaryType>();
- }
-
- fn assert_cmp_cases<A: Array>(
- 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<O: OffsetSizeTrait>(
- values: Vec<i32>,
- offsets: &[usize],
- sizes: &[usize],
- valid: Option<&[bool]>,
- ) -> GenericListViewArray<O> {
- let offsets = offsets
- .iter()
- .map(|v| O::from_usize(*v).unwrap())
- .collect::<ScalarBuffer<O>>();
- let sizes = sizes
- .iter()
- .map(|v| O::from_usize(*v).unwrap())
- .collect::<ScalarBuffer<O>>();
- 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<O: OffsetSizeTrait>() {
- let array = list_view_array::<O>(
- 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::<i32>();
- }
-
- #[test]
- fn test_large_list_view() {
- test_list_view_comparisons::<i64>();
- }
-
- #[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:
- let map1 = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
- vec![
- // high value for "a", low value for "b"
- Some(vec![("a", Some(100)), ("b", Some(1))]),
- // very high value for "b", low value for "c"
- Some(vec![("b", Some(999)), ("c", Some(1))]),
- Some(vec![]),
- Some(vec![("x", Some(1))]),
- ],
- false,
- );
-
- // Create second map array:
- // [{"a": 1, "c": 999}, {"b": 1, "d": 999}, {"a": 1}, None]
- let map2 = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
- vec![
- // low value for "a", high value for "c"
- Some(vec![("a", Some(1)), ("c", Some(999))]),
- // low value for "b", high value for "d"
- Some(vec![("b", Some(1)), ("d", Some(999))]),
- Some(vec![("a", Some(1))]),
- None,
- ],
- false,
- );
-
- 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
- let map1 = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
- vec![
- Some(vec![("a", Some(1)), ("b", Some(2))]),
- Some(vec![("x", Some(10))]),
- Some(vec![]),
- Some(vec![("c", Some(3))]),
- ],
- false,
- );
-
- let map2 = MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
- vec![
- Some(vec![("a", Some(1)), ("b", Some(2))]),
- Some(vec![("y", Some(20))]),
- Some(vec![("d", Some(4))]),
- None,
- ],
- false,
- );
-
- // 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::<ScalarBuffer<i8>>();
- let offsets = [0, 0, 1, 1,
2].into_iter().collect::<ScalarBuffer<i32>>();
-
- let union_fields = [
- (0, Arc::new(Field::new("A", DataType::Int32, false))),
- (1, Arc::new(Field::new("B", DataType::Utf8, false))),
- ]
- .into_iter()
- .collect::<UnionFields>();
-
- 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::<ScalarBuffer<i8>>();
- let offsets2 = [0, 0, 1, 1].into_iter().collect::<ScalarBuffer<i32>>();
-
- 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::<ScalarBuffer<i8>>();
-
- let union_fields = [
- (0, Arc::new(Field::new("a", DataType::Int32, false))),
- (1, Arc::new(Field::new("b", DataType::Utf8, false))),
- ]
- .into_iter()
- .collect::<UnionFields>();
-
- 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::<ScalarBuffer<i8>>();
- let offsets = [0, 0, 1].into_iter().collect::<ScalarBuffer<i32>>();
-
- let union_fields = [
- (0, Arc::new(Field::new("A", DataType::Int32, false))),
- (1, Arc::new(Field::new("B", DataType::Utf8, false))),
- ]
- .into_iter()
- .collect::<UnionFields>();
-
- 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::<ScalarBuffer<i8>>();
- let offsets1 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
-
- let union_fields1 = [
- (0, Arc::new(Field::new("A", DataType::Int32, false))),
- (1, Arc::new(Field::new("B", DataType::Utf8, false))),
- ]
- .into_iter()
- .collect::<UnionFields>();
-
- 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::<ScalarBuffer<i8>>();
- let offsets2 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
-
- let union_fields2 = [
- (0, Arc::new(Field::new("A", DataType::Int32, false))),
- (1, Arc::new(Field::new("C", DataType::Float64, false))),
- ]
- .into_iter()
- .collect::<UnionFields>();
-
- 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::<ScalarBuffer<i8>>();
- let offsets1 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
-
- let union_fields1 = [
- (0, Arc::new(Field::new("A", DataType::Int32, false))),
- (1, Arc::new(Field::new("B", DataType::Utf8, false))),
- ]
- .into_iter()
- .collect::<UnionFields>();
-
- 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::<ScalarBuffer<i8>>();
-
- 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::<Int32Type>::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::<Int32Type>::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::<Int32Type>::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::<Int32Type>::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::<Int16Type>::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::<Int16Type>::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::<Int64Type>::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::<Int64Type>::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::<Int32Type>::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::<Int32Type>::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::<Int32Type>::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::<Int32Type>::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::<Int64Type>::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 d6b0a7ee81..09ed8434b9 100644
--- a/arrow-select/Cargo.toml
+++ b/arrow-select/Cargo.toml
@@ -37,6 +37,7 @@ all-features = true
[dependencies]
arrow-buffer = { workspace = true }
+arrow-cmp = { workspace = true }
arrow-data = { workspace = true }
arrow-schema = { workspace = true }
arrow-array = { workspace = true }
diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs
index 772f21951e..0c550ea0b8 100644
--- a/arrow-select/src/take.rs
+++ b/arrow-select/src/take.rs
@@ -29,8 +29,9 @@ use arrow_buffer::{
ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer,
OffsetBuffer, RunEndBuffer,
ScalarBuffer, bit_util,
};
+use arrow_cmp::make_comparator;
use arrow_data::transform::MutableArrayData;
-use arrow_schema::{ArrowError, DataType, FieldRef, UnionMode};
+use arrow_schema::{ArrowError, DataType, FieldRef, SortOptions, UnionMode};
use num_traits::Zero;
@@ -959,9 +960,19 @@ fn take_run<T: RunEndIndexType, I: ArrowPrimitiveType>(
// `unwrap` is used in this function because the unwrapped values are
bounded by the corresponding `::Native`.
let mut new_run_ends_builder = BufferBuilder::<T::Native>::new(1);
let mut take_value_indices = BufferBuilder::<I::Native>::new(1);
+
+ let values_cmp = make_comparator(
+ run_array.values().as_ref(),
+ run_array.values().as_ref(),
+ SortOptions::default(),
+ )?;
+
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());
+ let prev_idx = physical_indices[ix - 1];
+ let cur_idx = physical_indices[ix];
+ let is_new_run = cur_idx != prev_idx && values_cmp(cur_idx,
prev_idx).is_ne();
+ 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());
}
}
@@ -2612,11 +2623,16 @@ mod tests {
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]);
+ // adjacent identical values are merged: [2,2,2,2,2,1,1] -> 2 runs
+ assert_eq!(
+ take_out.run_ends().values().len(),
+ 2,
+ "expected two physical runs"
+ );
+ assert_eq!(take_out.run_ends().values(), &[5_i32, 7]);
let take_out_values = take_out.values().as_primitive::<Int32Type>();
- assert_eq!(take_out_values.values(), &[2, 2, 2, 2, 1]);
+ assert_eq!(take_out_values.values(), &[2, 1]);
}
#[test]
@@ -2634,6 +2650,14 @@ mod tests {
let result = take_run(&run_array, &take_indices).unwrap();
let result = result.downcast::<Int32Array>().unwrap();
+ // [3, 5, 5, 3, 4] -> 4 physical runs (no adjacent duplicates to merge)
+ assert_eq!(
+ result.run_ends().values().len(),
+ 4,
+ "expected four physical runs"
+ );
+ assert_eq!(result.run_ends().values(), &[1_i32, 3, 4, 5]);
+
let expected = vec![3, 5, 5, 3, 4];
let actual = result.into_iter().flatten().collect::<Vec<_>>();
@@ -2915,4 +2939,102 @@ 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() {
+ // https://github.com/apache/arrow-rs/issues/7710
+ // Indices [0,1,4,5] select from [1,1,0,0,1,1] — the 0s are skipped,
+ // so the output should be a single run of 1s, not two.
+ let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
+ builder.extend([1, 1, 0, 0, 1, 1].into_iter().map(Some));
+ let ree = builder.finish();
+
+ let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]);
+ let result = take(&ree, &indexes, None).unwrap();
+ let result = result
+ .as_run::<Int32Type>()
+ .downcast::<Int32Array>()
+ .unwrap();
+
+ // Verify physical layout: all four logical values collapse into one
run.
+ assert_eq!(
+ result.run_ends().values().len(),
+ 1,
+ "expected a single physical run"
+ );
+ assert_eq!(result.run_ends().values(), &[4_i32]);
+
+ let actual = result.into_iter().flatten().collect::<Vec<_>>();
+ assert_eq!(actual, vec![1, 1, 1, 1]);
+ }
+
+ #[test]
+ fn test_take_run_end_encoded_merges_identical_string_runs() {
+ let mut builder = StringRunBuilder::<Int32Type>::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 result = take(&ree, &indexes, None).unwrap();
+ let result = result
+ .as_run::<Int32Type>()
+ .downcast::<StringArray>()
+ .unwrap();
+
+ // Verify physical layout: all four logical values collapse into one
run.
+ assert_eq!(
+ result.run_ends().values().len(),
+ 1,
+ "expected a single physical run"
+ );
+ assert_eq!(result.run_ends().values(), &[4_i32]);
+
+ let actual = result.into_iter().flatten().collect::<Vec<_>>();
+ assert_eq!(actual, vec!["bob", "bob", "bob", "bob"]);
+ }
+
+ #[test]
+ fn test_take_run_end_encoded_mixed_runs() {
+ // Validates that runs are merged whether the same logical value comes
+ // from the same physical index (repeated indices) or distinct
physical indices.
+ let mut builder = StringRunBuilder::<Int32Type>::new();
+ builder.extend(
+ ["bob", "bob", "alice", "alice", "bob", "bob", "eve", "eve"]
+ .into_iter()
+ .map(Some),
+ );
+ let ree = builder.finish();
+
+ // [bob,bob,bob,bob,bob,alice,alice,alice,eve,eve,eve]
+ 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::<Int32Type>()
+ .downcast::<StringArray>()
+ .unwrap();
+
+ // Verify physical layout: 11 logical values across exactly 3 physical
runs.
+
+ println!("run_ends_raw: {:?}", result.run_ends());
+ println!("run_ends: {:?}", result.run_ends().values());
+ println!("values : {:?}", result.values());
+ assert_eq!(
+ result.run_ends().values().len(),
+ 3,
+ "expected three physical runs"
+ );
+ assert_eq!(result.run_ends().values(), &[5_i32, 8, 11]);
+
+ let actual = result.into_iter().flatten().collect::<Vec<_>>();
+ assert_eq!(
+ actual,
+ vec![
+ "bob", "bob", "bob", "bob", "bob", "alice", "alice", "alice",
"eve", "eve", "eve"
+ ]
+ );
+ }
}
diff --git a/dev/release/README.md b/dev/release/README.md
index 8de898863d..91ed3e813d 100644
--- a/dev/release/README.md
+++ b/dev/release/README.md
@@ -45,7 +45,7 @@ crates.io, the Rust ecosystem's package manager.
We create a `CHANGELOG.md` so our users know what has been changed between
releases.
-## Prepare CHANGELOG and version:
+## Prepare CHANGELOG and version
- Ensure [`git-cliff`](https://git-cliff.org/docs/installation/) is installed
@@ -103,7 +103,7 @@ distribution servers.
Pick numbers in sequential order, with `1` for `rc1`, `2` for `rc2`, etc.
-### Create git tag for the release:
+### Create git tag for the release
While the official release artifact is a signed tarball, we also tag the
commit it was created for convenience and code archaeology.
@@ -135,11 +135,11 @@ The `create-tarball.sh` script
apache distribution svn server
2. provide you an email template to
- send to [email protected] for release voting.
+ send to <[email protected]> for release voting.
### Vote on Release Candidate tarball
-Send an email, based on the output from the script to [email protected].
+Send an email, based on the output from the script to <[email protected]>.
See an [example of how the email should
look](https://lists.apache.org/thread/2vpxdt6n7kzo72sxpr7q8yyby4495gnk).
For the release to become "official" it needs at least three Apache Arrow PMC
members to vote +1 on it.
@@ -156,7 +156,7 @@ The `dev/release/verify-release-candidate.sh` script in
this repository can assi
If the release is not approved, fix whatever the problem is and try again with
the next RC number
-### If the release is approved,
+### If the release is approved
Then, create a new release on GitHub using the tag `<version>` (e.g. `4.1.0`).
@@ -167,7 +167,7 @@ git tag <version> <version>-<rc>
git push apache <version>
```
-Move tarball to the release location in SVN, e.g.
https://dist.apache.org/repos/dist/release/arrow/arrow-rs-4.1.0/, using the
`release-tarball.sh` script:
+Move tarball to the release location in SVN, e.g.
<https://dist.apache.org/repos/dist/release/arrow/arrow-rs-4.1.0/>, using the
`release-tarball.sh` script:
```shell
./dev/release/release-tarball.sh 4.1.0 2
@@ -179,7 +179,7 @@ Congratulations! The release is now official!
The [`release.yml`] workflow automatically creates a github release for the
tag.
Check that the release is created and contains the correct changelog here:
-https://github.com/apache/arrow-rs/releases
+<https://github.com/apache/arrow-rs/releases>
[`release.yml`]:
https://github.com/apache/arrow-rs/blob/main/.github/workflows/release.yml#L1-L0
@@ -210,6 +210,7 @@ Rust Arrow Crates:
(cd arrow-schema && cargo publish)
(cd arrow-data && cargo publish)
(cd arrow-array && cargo publish)
+(cd arrow-cmp && cargo publish)
(cd arrow-select && cargo publish)
(cd arrow-ord && cargo publish)
(cd arrow-cast && cargo publish)
