From f103e47aa333d0d36dd3fd027e4b8fb17ee2c910 Mon Sep 17 00:00:00 2001 From: Alfonso Subiotto Marques Date: Thu, 30 Apr 2026 14:13:11 +0200 Subject: [PATCH 1/2] perf[arrow-select]: add specialized REE interleave The specialized interleave works by preserving run ends as much as possible by coalescing groups of adjacent logical indices pointing to the same source and calling interleave on the run end values. Future work could additionally coalesce values across sources, but this requires a value equality check. Signed-off-by: Alfonso Subiotto Marques --- arrow-select/src/interleave.rs | 70 ++++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index f5904bc171ee..4dd8c3ee8629 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -108,6 +108,12 @@ pub fn interleave( DataType::Struct(fields) => interleave_struct(fields, values, indices), DataType::List(field) => interleave_list::(values, indices, field), DataType::LargeList(field) => interleave_list::(values, indices, field), + DataType::RunEndEncoded(r, _) => match r.data_type() { + DataType::Int16 => interleave_run_end::(values, indices), + DataType::Int32 => interleave_run_end::(values, indices), + DataType::Int64 => interleave_run_end::(values, indices), + t => unreachable!("illegal run-end type {t}"), + }, _ => interleave_fallback(values, indices) } } @@ -411,6 +417,70 @@ fn interleave_list( Ok(Arc::new(list_array)) } +/// Specialized [`interleave`] for [`RunArray`]. +fn interleave_run_end( + values: &[&dyn Array], + indices: &[(usize, usize)], +) -> Result { + if indices.is_empty() { + return Ok(new_empty_array(values[0].data_type())); + } + + let n = indices.len(); + R::Native::from_usize(n).ok_or_else(|| { + ArrowError::ComputeError(format!( + "interleave_run_end: output length {n} does not fit run-end type" + )) + })?; + + let runs: Vec<&RunArray> = values.iter().map(|a| a.as_run::()).collect(); + let value_arrays: Vec<&dyn Array> = runs.iter().map(|r| r.values().as_ref()).collect(); + + // Resolve each (array, logical_row) to (array, physical_row), so we can + // lookup physical indices by batch. + let mut phys_pairs: Vec<(usize, usize)> = vec![(0, 0); n]; + let mut grouped: Vec<(Vec, Vec)> = + (0..runs.len()).map(|_| (Vec::new(), Vec::new())).collect(); + for (out_pos, &(arr, row)) in indices.iter().enumerate() { + let row = R::Native::from_usize(row).ok_or_else(|| { + ArrowError::InvalidArgumentError(format!( + "interleave_run_end: row index {row} out of range" + )) + })?; + grouped[arr].0.push(row); + grouped[arr].1.push(out_pos); + } + for (arr_idx, (logical_rows, out_positions)) in grouped.into_iter().enumerate() { + let phys = runs[arr_idx].get_physical_indices(&logical_rows)?; + for (p, out_pos) in phys.iter().zip(out_positions.iter()) { + phys_pairs[*out_pos] = (arr_idx, *p); + } + } + + // Coalesce by physical-pair equality only: emit a new run when the + // (array_idx, physical_idx) pair changes between adjacent output rows. + // TODO: We could perform an equality check across sources to extend the + // output run, but we can't call make_comparator from this crate. + let mut run_ends_buf: Vec = Vec::with_capacity(n); + let mut dedup_pairs: Vec<(usize, usize)> = Vec::with_capacity(n); + dedup_pairs.push(phys_pairs[0]); + for i in 1..n { + if phys_pairs[i] != phys_pairs[i - 1] { + run_ends_buf.push(R::Native::from_usize(i).unwrap()); + dedup_pairs.push(phys_pairs[i]); + } + } + run_ends_buf.push(R::Native::from_usize(n).unwrap()); + + let taken_values = interleave(&value_arrays, &dedup_pairs)?; + let run_ends = PrimitiveArray::::from_iter_values(run_ends_buf); + + Ok(Arc::new(RunArray::::try_new( + &run_ends, + taken_values.as_ref(), + )?)) +} + /// Fallback implementation of interleave using [`MutableArrayData`] fn interleave_fallback( values: &[&dyn Array], From 94baaa3ebf45cd4b124d26f0be050f7fa5f2dd5a Mon Sep 17 00:00:00 2001 From: Richard Date: Tue, 5 May 2026 12:59:11 -0400 Subject: [PATCH 2/2] added equality check as well as caching. includes test from #9865 --- arrow-select/src/interleave.rs | 122 +++++++++++++++++++++++++++++++-- 1 file changed, 116 insertions(+), 6 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 4dd8c3ee8629..f798868ad1ae 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -24,6 +24,7 @@ use arrow_array::cast::AsArray; use arrow_array::types::*; use arrow_array::*; use arrow_buffer::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer}; +use arrow_data::ArrayData; use arrow_data::ByteView; use arrow_data::transform::MutableArrayData; use arrow_schema::{ArrowError, DataType, FieldRef, Fields}; @@ -435,6 +436,8 @@ fn interleave_run_end( let runs: Vec<&RunArray> = values.iter().map(|a| a.as_run::()).collect(); let value_arrays: Vec<&dyn Array> = runs.iter().map(|r| r.values().as_ref()).collect(); + // hoist value data + let value_data: Vec = value_arrays.iter().map(|a| a.to_data()).collect(); // Resolve each (array, logical_row) to (array, physical_row), so we can // lookup physical indices by batch. @@ -442,6 +445,12 @@ fn interleave_run_end( let mut grouped: Vec<(Vec, Vec)> = (0..runs.len()).map(|_| (Vec::new(), Vec::new())).collect(); for (out_pos, &(arr, row)) in indices.iter().enumerate() { + if row >= runs[arr].len() { + return Err(ArrowError::InvalidArgumentError(format!( + "interleave_run_end: row index {row} out of range for array {arr} with length {}", + runs[arr].len() + ))); + }; let row = R::Native::from_usize(row).ok_or_else(|| { ArrowError::InvalidArgumentError(format!( "interleave_run_end: row index {row} out of range" @@ -456,16 +465,16 @@ fn interleave_run_end( phys_pairs[*out_pos] = (arr_idx, *p); } } - - // Coalesce by physical-pair equality only: emit a new run when the - // (array_idx, physical_idx) pair changes between adjacent output rows. - // TODO: We could perform an equality check across sources to extend the - // output run, but we can't call make_comparator from this crate. let mut run_ends_buf: Vec = Vec::with_capacity(n); let mut dedup_pairs: Vec<(usize, usize)> = Vec::with_capacity(n); dedup_pairs.push(phys_pairs[0]); for i in 1..n { - if phys_pairs[i] != phys_pairs[i - 1] { + // this doesnt check for equality of values. + let (prev_arr, prev_row) = phys_pairs[i - 1]; + let (curr_arr, curr_row) = phys_pairs[i]; + let prev_value = value_data[prev_arr].slice(prev_row, 1); + let current_value = value_data[curr_arr].slice(curr_row, 1); + if phys_pairs[i] != phys_pairs[i - 1] && prev_value != current_value { run_ends_buf.push(R::Native::from_usize(i).unwrap()); dedup_pairs.push(phys_pairs[i]); } @@ -1614,4 +1623,105 @@ 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: Arc = + 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_break_intent() { + 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. + // greater than int16::max + assert!(interleave(&[&a, &b], &[(0, 32766), (0, 4), (1, 2), (1, 3), (1, 4)]).is_err()) + } + #[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"]); + } }