diff --git a/datafusion/physical-plan/src/sorts/arrow/mod.rs b/datafusion/physical-plan/src/sorts/arrow/mod.rs new file mode 100644 index 0000000000000..0e1bbfc4a09ae --- /dev/null +++ b/datafusion/physical-plan/src/sorts/arrow/mod.rs @@ -0,0 +1,3 @@ +pub(super) mod sort; +#[expect(unused)] +pub(super) mod rank; diff --git a/datafusion/physical-plan/src/sorts/arrow/rank.rs b/datafusion/physical-plan/src/sorts/arrow/rank.rs new file mode 100644 index 0000000000000..3c92d96859559 --- /dev/null +++ b/datafusion/physical-plan/src/sorts/arrow/rank.rs @@ -0,0 +1,794 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Provides `rank` function to assign a rank to each value in an array + +use arrow::array::cast::AsArray; +use arrow::array::types::*; +use arrow::array::{ + Array, ArrowNativeTypeOp, BooleanArray, GenericByteArray, GenericByteViewArray, + downcast_primitive_array, +}; +use arrow::buffer::NullBuffer; +use arrow::error::{ArrowError}; +use std::cmp::Ordering; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType}; + +/// Whether `arrow_ord::rank` can rank an array of given data type. +pub(crate) fn can_rank(data_type: &DataType) -> bool { + data_type.is_primitive() + || matches!( + data_type, + DataType::Boolean + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + | DataType::Utf8View + | DataType::BinaryView + ) +} + +/// Assigns a rank to each value in `array` based on its position in the sorted order +/// +/// Where values are equal, they will be assigned the highest of their ranks, +/// leaving gaps in the overall rank assignment +/// +/// ``` +/// # use arrow_array::StringArray; +/// # use arrow_ord::rank::rank; +/// let array = StringArray::from(vec![Some("foo"), None, Some("foo"), None, Some("bar")]); +/// let ranks = rank(&array, None).unwrap(); +/// assert_eq!(ranks, &[5, 2, 5, 2, 3]); +/// ``` +pub fn rank(array: &dyn Array, options: Option) -> Result, ArrowError> { + let options = options.unwrap_or_default(); + let ranks = downcast_primitive_array! { + array => primitive_rank(array.values(), array.nulls(), options), + DataType::Boolean => boolean_rank(array.as_boolean(), options), + DataType::Utf8 => bytes_rank(array.as_bytes::(), options), + DataType::LargeUtf8 => bytes_rank(array.as_bytes::(), options), + DataType::Binary => bytes_rank(array.as_bytes::(), options), + DataType::LargeBinary => bytes_rank(array.as_bytes::(), options), + DataType::Utf8View => byte_view_rank(array.as_string_view(), options), + DataType::BinaryView => byte_view_rank(array.as_binary_view(), options), + d => return Err(ArrowError::ComputeError(format!("{d:?} not supported in rank"))) + }; + Ok(ranks) +} + +#[inline(never)] +fn primitive_rank( + values: &[T], + nulls: Option<&NullBuffer>, + options: SortOptions, +) -> Vec { + let len: u32 = values.len().try_into().unwrap(); + let to_sort = match nulls.filter(|n| n.null_count() > 0) { + Some(n) => n + .valid_indices() + .map(|idx| (values[idx], idx as u32)) + .collect(), + None => values.iter().copied().zip(0..len).collect(), + }; + rank_impl(values.len(), to_sort, options, T::compare, T::is_eq) +} + +#[inline(never)] +fn bytes_rank(array: &GenericByteArray, options: SortOptions) -> Vec { + let to_sort: Vec<(&[u8], u32)> = match array.nulls().filter(|n| n.null_count() > 0) { + Some(n) => n + .valid_indices() + .map(|idx| (array.value(idx).as_ref(), idx as u32)) + .collect(), + None => (0..array.len()) + .map(|idx| (array.value(idx).as_ref(), idx as u32)) + .collect(), + }; + rank_impl(array.len(), to_sort, options, Ord::cmp, PartialEq::eq) +} + +#[inline(never)] +fn byte_view_rank( + array: &GenericByteViewArray, + options: SortOptions, +) -> Vec { + // An inline view already contains the complete value. Convert it once to + // a key whose integer ordering matches the byte ordering, as is done by + // `sort_byte_view`. + if array.data_buffers().is_empty() { + let to_sort: Vec<(u128, u32)> = match array.nulls().filter(|n| n.null_count() > 0) { + Some(n) => n + .valid_indices() + .map(|idx| { + // SAFETY: `valid_indices` only yields indices in the array. + let raw = unsafe { *array.views().get_unchecked(idx) }; + (GenericByteViewArray::::inline_key_fast(raw), idx as u32) + }) + .collect(), + None => array + .views() + .iter() + .enumerate() + .map(|(idx, raw)| (GenericByteViewArray::::inline_key_fast(*raw), idx as u32)) + .collect(), + }; + return rank_impl( + array.len(), + to_sort, + options, + |a, b| a.cmp(&b), + |a, b| a == b, + ); + } + + if has_high_byte_view_key_collision_rate(array) { + let to_sort: Vec<(&[u8], u32)> = match array.nulls().filter(|n| n.null_count() > 0) { + Some(n) => n + .valid_indices() + .map(|idx| (array.value(idx).as_ref(), idx as u32)) + .collect(), + None => (0..array.len()) + .map(|idx| (array.value(idx).as_ref(), idx as u32)) + .collect(), + }; + return rank_impl(array.len(), to_sort, options, Ord::cmp, PartialEq::eq); + } + + // Cache a wider prefix than the 4 bytes stored in a non-inline view. This + // pays for the backing-buffer access once per value instead of once per + // comparison, and only resolves the complete value when two keys collide. + let to_sort: Vec<(u128, u32)> = match array.nulls().filter(|n| n.null_count() > 0) { + Some(n) => n + .valid_indices() + .map(|idx| { + // SAFETY: `valid_indices` only yields indices in the array. + let value: &[u8] = unsafe { array.value_unchecked(idx).as_ref() }; + (byte_view_key(value), idx as u32) + }) + .collect(), + None => (0..array.len()) + .map(|idx| { + // SAFETY: `idx` is in `0..array.len()`. + let value: &[u8] = unsafe { array.value_unchecked(idx).as_ref() }; + (byte_view_key(value), idx as u32) + }) + .collect(), + }; + rank_impl_by( + array.len(), + to_sort, + options, + |a, b| compare_view_key(array, a, b), + |a, b| equal_view_key(array, a, b), + ) +} + +// A 16-byte prefix fits in one `u128` and is wider than the 4-byte view prefix. +// Shorter keys collide more often; longer keys need another representation. +const BYTE_VIEW_KEY_LEN: usize = 16; + +// Four valid values per window keeps sampling cheap while allowing an early +// all-collision decision. More samples improve confidence but add buffer reads. +const BYTE_VIEW_KEY_SAMPLES_PER_WINDOW: usize = 4; + +// Capacity for the two windows; it must be at least +// `2 * BYTE_VIEW_KEY_SAMPLES_PER_WINDOW`. Increasing it alone has no effect; +// decreasing it without changing the sampling count can overflow the array. +const BYTE_VIEW_KEY_SAMPLE_SIZE: usize = 8; + +// Bound entries inspected when nulls are present. A higher limit finds valid +// samples more reliably but costs reads; a lower limit is cheaper but less +// informative for null-heavy arrays. +const BYTE_VIEW_KEY_MAX_PROBES_PER_WINDOW: usize = 32; + +// `colliding_keys * 3 >= sample_len` means roughly one third of sampled keys +// are duplicates. Lower values fall back earlier; higher values risk keeping +// the key path for collision-heavy inputs. +const BYTE_VIEW_KEY_FALLBACK_COLLISION_RATIO: usize = 3; + +/// Estimates whether cached byte-view keys collide often enough to make the +/// key-based ranking path unattractive. +/// Caching a wider key usually avoids repeated backing-buffer reads for long +/// views. If sampled keys collide frequently, resolving full values plus the +/// extra key comparison can be slower than comparing slices directly, so the +/// caller falls back to that path. The bounded two-window sample keeps this +/// check inexpensive. +fn has_high_byte_view_key_collision_rate(array: &GenericByteViewArray) -> bool { + if array.len() < 2 { + return false; + } + + let mut keys = [0_u128; BYTE_VIEW_KEY_SAMPLE_SIZE]; + let mut sample_len = 0; + let midpoint = array.len() / 2; + + // Probe small local windows in both halves. Keeping each probe local avoids + // turning collision detection itself into scattered backing-buffer reads. + for (start, end) in [(0, midpoint), (midpoint, array.len())] { + let probe_end = end.min(start.saturating_add(BYTE_VIEW_KEY_MAX_PROBES_PER_WINDOW)); + let window_start = sample_len; + let mut window_samples = 0; + + for idx in start..probe_end { + if array.is_null(idx) { + continue; + } + + // SAFETY: `idx` is within a window bounded by `array.len()`. + let value: &[u8] = unsafe { array.value_unchecked(idx).as_ref() }; + keys[sample_len] = byte_view_key(value); + sample_len += 1; + window_samples += 1; + if window_samples == BYTE_VIEW_KEY_SAMPLES_PER_WINDOW { + break; + } + } + + // Four equal keys already contribute three collisions. Even if every + // sample in the other window is distinct, that satisfies the final + // one-third threshold, so avoid touching the second backing-buffer + // window in the common all-collision case. + if window_samples == BYTE_VIEW_KEY_SAMPLES_PER_WINDOW + && keys[window_start..sample_len] + .windows(2) + .all(|w| w[0] == w[1]) + { + return true; + } + } + + if sample_len < 2 { + return false; + } + + let keys = &mut keys[..sample_len]; + keys.sort_unstable(); + let unique_keys = 1 + keys.windows(2).filter(|w| w[0] != w[1]).count(); + + // If at least roughly one third of sampled keys collide, comparing the + // wider key before every full-value comparison is likely more expensive + // than sorting slices directly. + let colliding_keys = sample_len - unique_keys; + colliding_keys * BYTE_VIEW_KEY_FALLBACK_COLLISION_RATIO >= sample_len +} + +#[inline(always)] +fn byte_view_key(value: &[u8]) -> u128 { + let mut key = [0_u8; BYTE_VIEW_KEY_LEN]; + let key_len = value.len().min(key.len()); + key[..key_len].copy_from_slice(&value[..key_len]); + + // Big-endian conversion makes integer comparison equivalent to comparing + // these bytes lexicographically. Equal keys fall back to the full values, + // covering values that differ after 16 bytes and prefixes containing zero. + u128::from_be_bytes(key) +} + +#[inline(always)] +fn compare_view_key( + array: &GenericByteViewArray, + a: &(u128, u32), + b: &(u128, u32), +) -> Ordering { + match a.0.cmp(&b.0) { + Ordering::Equal => { + // SAFETY: both indices were produced from this array above. + let full_a: &[u8] = unsafe { array.value_unchecked(a.1 as usize).as_ref() }; + let full_b: &[u8] = unsafe { array.value_unchecked(b.1 as usize).as_ref() }; + full_a.cmp(full_b) + } + ordering => ordering, + } +} + +#[inline(always)] +fn equal_view_key( + array: &GenericByteViewArray, + a: &(u128, u32), + b: &(u128, u32), +) -> bool { + if a.0 != b.0 { + return false; + } + + // SAFETY: both indices were produced from this array above. + let full_a: &[u8] = unsafe { array.value_unchecked(a.1 as usize).as_ref() }; + let full_b: &[u8] = unsafe { array.value_unchecked(b.1 as usize).as_ref() }; + full_a == full_b +} + +fn rank_impl_by( + len: usize, + mut valid: Vec<(T, u32)>, + options: SortOptions, + compare: C, + eq: E, +) -> Vec +where + C: Fn(&(T, u32), &(T, u32)) -> Ordering, + E: Fn(&(T, u32), &(T, u32)) -> bool, +{ + // Same ranking and null handling as `rank_impl`, but callbacks receive + // tuple references because key collisions use the index to read the full + // value. `rank_impl` passes copied values directly through `Fn(T, T)`. + // We can use an unstable sort as we combine equal values later + valid.sort_unstable_by(compare); + if options.descending { + valid.reverse(); + } + + let (mut valid_rank, null_rank) = match options.nulls_first { + true => (len as u32, (len - valid.len()) as u32), + false => (valid.len() as u32, len as u32), + }; + + let mut out: Vec<_> = vec![null_rank; len]; + if let Some(v) = valid.last() { + out[v.1 as usize] = valid_rank; + } + + let mut count = 1; // Number of values in rank + for w in valid.windows(2).rev() { + match eq(&w[0], &w[1]) { + true => { + count += 1; + out[w[0].1 as usize] = valid_rank; + } + false => { + valid_rank -= count; + count = 1; + out[w[0].1 as usize] = valid_rank + } + } + } + + out +} + +fn rank_impl( + len: usize, + mut valid: Vec<(T, u32)>, + options: SortOptions, + compare: C, + eq: E, +) -> Vec +where + T: Copy, + C: Fn(T, T) -> Ordering, + E: Fn(T, T) -> bool, +{ + // We can use an unstable sort as we combine equal values later + valid.sort_unstable_by(|a, b| compare(a.0, b.0)); + if options.descending { + valid.reverse(); + } + + let (mut valid_rank, null_rank) = match options.nulls_first { + true => (len as u32, (len - valid.len()) as u32), + false => (valid.len() as u32, len as u32), + }; + + let mut out: Vec<_> = vec![null_rank; len]; + if let Some(v) = valid.last() { + out[v.1 as usize] = valid_rank; + } + + let mut count = 1; // Number of values in rank + for w in valid.windows(2).rev() { + match eq(w[0].0, w[1].0) { + true => { + count += 1; + out[w[0].1 as usize] = valid_rank; + } + false => { + valid_rank -= count; + count = 1; + out[w[0].1 as usize] = valid_rank + } + } + } + + out +} + +/// Return the index for the rank when ranking boolean array +/// +/// The index is calculated as follows: +/// if is_null is true, the index is 2 +/// if is_null is false and the value is true, the index is 1 +/// otherwise, the index is 0 +/// +/// false is 0 and true is 1 because these are the value when cast to number +#[inline] +fn get_boolean_rank_index(value: bool, is_null: bool) -> usize { + let is_null_num = is_null as usize; + (is_null_num << 1) | (value as usize & !is_null_num) +} + +#[inline(never)] +fn boolean_rank(array: &BooleanArray, options: SortOptions) -> Vec { + let null_count = array.null_count() as u32; + let true_count = array.true_count() as u32; + let false_count = array.len() as u32 - null_count - true_count; + + // Rank values for [false, true, null] in that order + // + // The value for a rank is last value rank + own value count + // this means that if we have the following order: `false`, `true` and then `null` + // the ranks will be: + // - false: false_count + // - true: false_count + true_count + // - null: false_count + true_count + null_count + // + // If we have the following order: `null`, `false` and then `true` + // the ranks will be: + // - false: null_count + false_count + // - true: null_count + false_count + true_count + // - null: null_count + // + // You will notice that the last rank is always the total length of the array but we don't use it for readability on how the rank is calculated + let ranks_index: [u32; 3] = match (options.descending, options.nulls_first) { + // The order is null, true, false + (true, true) => [ + null_count + true_count + false_count, + null_count + true_count, + null_count, + ], + // The order is true, false, null + (true, false) => [ + true_count + false_count, + true_count, + true_count + false_count + null_count, + ], + // The order is null, false, true + (false, true) => [ + null_count + false_count, + null_count + false_count + true_count, + null_count, + ], + // The order is false, true, null + (false, false) => [ + false_count, + false_count + true_count, + false_count + true_count + null_count, + ], + }; + + match array.nulls().filter(|n| n.null_count() > 0) { + Some(n) => array + .values() + .iter() + .zip(n.iter()) + .map(|(value, is_valid)| ranks_index[get_boolean_rank_index(value, !is_valid)]) + .collect::>(), + None => array + .values() + .iter() + .map(|value| ranks_index[value as usize]) + .collect::>(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::*; + + fn assert_same_rank(left: &dyn Array, right: &dyn Array) { + for descending in [false, true] { + for nulls_first in [false, true] { + let options = SortOptions { + descending, + nulls_first, + }; + assert_eq!( + rank(left, Some(options)).unwrap(), + rank(right, Some(options)).unwrap() + ); + } + } + } + + #[test] + fn test_primitive() { + let descending = SortOptions { + descending: true, + nulls_first: true, + }; + + let nulls_last = SortOptions { + descending: false, + nulls_first: false, + }; + + let nulls_last_descending = SortOptions { + descending: true, + nulls_first: false, + }; + + let a = Int32Array::from(vec![Some(1), Some(1), None, Some(3), Some(3), Some(4)]); + let res = rank(&a, None).unwrap(); + assert_eq!(res, &[3, 3, 1, 5, 5, 6]); + + let res = rank(&a, Some(descending)).unwrap(); + assert_eq!(res, &[6, 6, 1, 4, 4, 2]); + + let res = rank(&a, Some(nulls_last)).unwrap(); + assert_eq!(res, &[2, 2, 6, 4, 4, 5]); + + let res = rank(&a, Some(nulls_last_descending)).unwrap(); + assert_eq!(res, &[5, 5, 6, 3, 3, 1]); + + // Test with non-zero null values + let nulls = NullBuffer::from(vec![true, true, false, true, false, false]); + let a = Int32Array::new(vec![1, 4, 3, 4, 5, 5].into(), Some(nulls)); + let res = rank(&a, None).unwrap(); + assert_eq!(res, &[4, 6, 3, 6, 3, 3]); + } + + #[test] + fn test_get_boolean_rank_index() { + assert_eq!(get_boolean_rank_index(true, true), 2); + assert_eq!(get_boolean_rank_index(false, true), 2); + assert_eq!(get_boolean_rank_index(true, false), 1); + assert_eq!(get_boolean_rank_index(false, false), 0); + } + + #[test] + fn test_nullable_booleans() { + let descending = SortOptions { + descending: true, + nulls_first: true, + }; + + let nulls_last = SortOptions { + descending: false, + nulls_first: false, + }; + + let nulls_last_descending = SortOptions { + descending: true, + nulls_first: false, + }; + + let a = BooleanArray::from(vec![Some(true), Some(true), None, Some(false), Some(false)]); + let res = rank(&a, None).unwrap(); + assert_eq!(res, &[5, 5, 1, 3, 3]); + + let res = rank(&a, Some(descending)).unwrap(); + assert_eq!(res, &[3, 3, 1, 5, 5]); + + let res = rank(&a, Some(nulls_last)).unwrap(); + assert_eq!(res, &[4, 4, 5, 2, 2]); + + let res = rank(&a, Some(nulls_last_descending)).unwrap(); + assert_eq!(res, &[2, 2, 5, 4, 4]); + + // Test with non-zero null values + let nulls = NullBuffer::from(vec![true, true, false, true, true]); + let a = BooleanArray::new(vec![true, true, true, false, false].into(), Some(nulls)); + let res = rank(&a, None).unwrap(); + assert_eq!(res, &[5, 5, 1, 3, 3]); + } + + #[test] + fn test_booleans() { + let descending = SortOptions { + descending: true, + nulls_first: true, + }; + + let nulls_last = SortOptions { + descending: false, + nulls_first: false, + }; + + let nulls_last_descending = SortOptions { + descending: true, + nulls_first: false, + }; + + let a = BooleanArray::from(vec![true, false, false, false, true]); + let res = rank(&a, None).unwrap(); + assert_eq!(res, &[5, 3, 3, 3, 5]); + + let res = rank(&a, Some(descending)).unwrap(); + assert_eq!(res, &[2, 5, 5, 5, 2]); + + let res = rank(&a, Some(nulls_last)).unwrap(); + assert_eq!(res, &[5, 3, 3, 3, 5]); + + let res = rank(&a, Some(nulls_last_descending)).unwrap(); + assert_eq!(res, &[2, 5, 5, 5, 2]); + } + + #[test] + fn test_bytes() { + let v = vec!["foo", "fo", "bar", "bar"]; + let values = StringArray::from(v.clone()); + let res = rank(&values, None).unwrap(); + assert_eq!(res, &[4, 3, 2, 2]); + + let values = LargeStringArray::from(v.clone()); + let res = rank(&values, None).unwrap(); + assert_eq!(res, &[4, 3, 2, 2]); + + let values = StringViewArray::from(v); + let res = rank(&values, None).unwrap(); + assert_eq!(res, &[4, 3, 2, 2]); + + let v: Vec<&[u8]> = vec![&[1, 2], &[0], &[1, 2, 3], &[1, 2]]; + let values = LargeBinaryArray::from(v.clone()); + let res = rank(&values, None).unwrap(); + assert_eq!(res, &[3, 1, 4, 3]); + + let values = BinaryArray::from(v.clone()); + let res = rank(&values, None).unwrap(); + assert_eq!(res, &[3, 1, 4, 3]); + + let values = BinaryViewArray::from_iter_values(v); + let res = rank(&values, None).unwrap(); + assert_eq!(res, &[3, 1, 4, 3]); + } + + #[test] + fn test_inline_byte_views() { + let string_values = vec![ + Some(""), + Some("short"), + None, + Some("0123456789qa"), // exactly the 12-byte inline limit + Some("short"), + ]; + let string_view = StringViewArray::from(string_values.clone()); + let string = StringArray::from(string_values); + + assert!(string_view.data_buffers().is_empty()); + assert_same_rank(&string_view, &string); + + let binary_values: Vec> = vec![ + Some(b""), + Some(b"short"), + None, + Some(b"0123456789qa"), + Some(b"short"), + ]; + let binary_view = BinaryViewArray::from_iter(binary_values.clone()); + let binary = BinaryArray::from_opt_vec(binary_values); + + assert!(binary_view.data_buffers().is_empty()); + assert_same_rank(&binary_view, &binary); + } + + #[test] + fn test_string_view_with_nulls() { + let values = StringViewArray::from(vec![ + Some("a string longer than twelve bytes"), + Some("bar"), + None, + Some("a string longer than twelve bytes"), + ]); + let res = rank(&values, None).unwrap(); + assert_eq!(res, &[3, 4, 1, 3]); + } + + #[test] + fn test_binary_view_with_nulls() { + let long_value = b"a binary value longer than twelve bytes".as_ref(); + let values = BinaryViewArray::from_iter([ + Some(long_value), + Some(b"bar".as_ref()), + None, + Some(long_value), + ]); + let res = rank(&values, None).unwrap(); + assert_eq!(res, &[3, 4, 1, 3]); + } + + #[test] + fn test_string_view_key_collisions() { + let values = vec![ + Some("abcdefghijklmnop"), + Some("abcdefghijklmnopA"), + Some("abcdefghijklmnopB"), + Some("abcdefghijklmnopA"), + Some("abcdefghijklmno"), + Some("short"), + None, + ]; + let expected = StringArray::from(values.clone()); + let actual = StringViewArray::from(values); + + assert_same_rank(&actual, &expected); + } + + #[test] + fn test_binary_view_key_collisions() { + let zeroes_16 = [0_u8; 16]; + let zeroes_17 = [0_u8; 17]; + let mut zeroes_then_one = [0_u8; 17]; + zeroes_then_one[16] = 1; + let mut zeroes_then_two = [0_u8; 17]; + zeroes_then_two[16] = 2; + + let values: Vec> = vec![ + Some(b""), + Some(b"\0"), + Some(&zeroes_16), + Some(&zeroes_17), + Some(&zeroes_then_one), + Some(&zeroes_then_two), + Some(&zeroes_then_one), + None, + ]; + let expected = BinaryArray::from_opt_vec(values.clone()); + let actual = BinaryViewArray::from_iter(values); + + assert_same_rank(&actual, &expected); + } + + #[test] + fn test_byte_view_high_key_collision_detection() { + const SIZE: u32 = 64; + + let same_key: StringViewArray = (0..SIZE) + .map(|i| { + let suffix = i.wrapping_mul(2_654_435_761); + Some(format!("abcdefghijklmnop{suffix:08x}")) + }) + .collect(); + + assert_eq!( + byte_view_key(same_key.value(0).as_bytes()), + byte_view_key(same_key.value(1).as_bytes()) + ); + assert!(has_high_byte_view_key_collision_rate(&same_key)); + + let clustered: StringViewArray = (0..SIZE) + .map(|i| { + let suffix = i.wrapping_mul(2_654_435_761); + let value = if i < SIZE / 2 { + format!("{suffix:016x}abcdefgh") + } else { + format!("abcdefghijklmnop{suffix:08x}") + }; + Some(value) + }) + .collect(); + assert!(has_high_byte_view_key_collision_rate(&clustered)); + + let with_nulls: StringViewArray = (0..SIZE) + .map(|i| { + (i % 2 == 0).then(|| { + let suffix = i.wrapping_mul(2_654_435_761); + format!("abcdefghijklmnop{suffix:08x}") + }) + }) + .collect(); + assert!(has_high_byte_view_key_collision_rate(&with_nulls)); + + let distinct: StringViewArray = (0..SIZE) + .map(|i| { + let suffix = i.wrapping_mul(2_654_435_761); + Some(format!("{suffix:016x}abcdefgh")) + }) + .collect(); + assert!(!has_high_byte_view_key_collision_rate(&distinct)); + } +} diff --git a/datafusion/physical-plan/src/sorts/arrow/sort.rs b/datafusion/physical-plan/src/sorts/arrow/sort.rs new file mode 100644 index 0000000000000..5d1cea9ac5e92 --- /dev/null +++ b/datafusion/physical-plan/src/sorts/arrow/sort.rs @@ -0,0 +1,862 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Sort kernels over several arrays at once. +//! +//! The rows of all arrays ("blocks") are sorted together without concatenating or +//! slicing them. A row is addressed by an [`ArrayRowIndex`], `(array index, row index)`, +//! and every comparison reads the value through `get_unchecked` on the matching array. +//! +//! Adapted from `arrow_ord::sort`. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, BooleanArray, ByteView, + DynComparator, FixedSizeBinaryArray, GenericByteArray, GenericByteViewArray, + PrimitiveArray, make_comparator, +}; +use arrow::compute::{SortOptions, interleave}; +use arrow::datatypes::{ + BinaryType, BinaryViewType, ByteArrayType, ByteViewType, DataType, LargeBinaryType, + LargeUtf8Type, StringViewType, Utf8Type, +}; +use arrow::downcast_primitive_array; +use arrow::error::ArrowError; +use arrow::row::{RowConverter, Rows, SortField}; + +/// `(array index, row index in that array)`, the shape [`interleave`] takes +pub type ArrayRowIndex = (usize, usize); + +/// Views up to this many bytes are stored inline +const MAX_INLINE_VIEW_LEN: u32 = 12; + +/// One column to be used in lexicographical sort, made of one array per block +#[derive(Clone, Debug)] +pub struct SortColumn { + /// The arrays of the column, all of the same data type + pub values: Vec, + /// Sort options for this column + pub options: Option, +} + +/// Sorts the rows of `arrays` together and gathers them into new arrays with the same +/// lengths as the inputs. Nulls are ordered according to `nulls_first`, floats by IEEE +/// 754 totalOrder. Unstable: equal elements may not keep their order. +pub fn sort(arrays: &[ArrayRef], options: Option) -> Result, ArrowError> { + let indices = sort_to_indices(arrays, options, None)?; + gather(arrays, &indices) +} + +/// Like [`sort`] but keeps only the first `limit` rows of the sort order, still chunked +/// by the input lengths (the trailing chunks come out shorter or are omitted) +pub fn sort_limit( + arrays: &[ArrayRef], + options: Option, + limit: Option, +) -> Result, ArrowError> { + let indices = sort_to_indices(arrays, options, limit)?; + gather(arrays, &indices) +} + +/// Gathers `indices` out of `arrays` into new arrays chunked like `arrays`: output `i` +/// has the length of input `i` (shorter when `indices` runs out) +pub fn gather(arrays: &[ArrayRef], indices: &[ArrayRowIndex]) -> Result, ArrowError> { + let refs: Vec<&dyn Array> = arrays.iter().map(|array| array.as_ref()).collect(); + let mut out = Vec::with_capacity(arrays.len()); + let mut start = 0; + for array in arrays { + if start >= indices.len() { + break; + } + let end = (start + array.len()).min(indices.len()); + out.push(interleave(&refs, &indices[start..end])?); + start = end; + } + Ok(out) +} + +/// Sorts the rows of `arrays` together and returns their positions in sort order. +/// `limit` keeps only the first `limit` positions ([`partial_sort`]). +pub fn sort_to_indices( + arrays: &[ArrayRef], + options: Option, + limit: Option, +) -> Result, ArrowError> { + let Some(first) = arrays.first() else { + return Ok(vec![]); + }; + if let Some(other) = arrays.iter().find(|array| array.data_type() != first.data_type()) { + return Err(ArrowError::ComputeError(format!( + "sort arrays have different data types: {} and {}", + first.data_type(), + other.data_type() + ))); + } + let total: usize = arrays.iter().map(|array| array.len()).sum(); + if total == 0 || limit == Some(0) { + return Ok(vec![]); + } + + let options = options.unwrap_or_default(); + let first = first.as_ref(); + + Ok(downcast_primitive_array! { + first => sort_primitive_like(first, arrays, options, limit), + DataType::Boolean => sort_boolean(&downcast_all::(arrays), options, limit), + DataType::Utf8 => sort_bytes(&downcast_all::>(arrays), options, limit), + DataType::LargeUtf8 => sort_bytes(&downcast_all::>(arrays), options, limit), + DataType::Binary => sort_bytes(&downcast_all::>(arrays), options, limit), + DataType::LargeBinary => sort_bytes(&downcast_all::>(arrays), options, limit), + DataType::Utf8View => sort_byte_view(&downcast_all::>(arrays), options, limit), + DataType::BinaryView => sort_byte_view(&downcast_all::>(arrays), options, limit), + DataType::FixedSizeBinary(_) => sort_fixed_size_binary(&downcast_all::(arrays), options, limit), + // Dictionaries, lists, runs, structs...: compared through arrow's comparators + _ => sort_by_comparators(&[SortColumn { values: arrays.to_vec(), options: Some(options) }], limit)?, + }) +} + +/// Sorts lexicographically by every column and returns the positions in sort order +pub fn lexsort_to_indices( + columns: &[SortColumn], + limit: Option, +) -> Result, ArrowError> { + let Some(first) = columns.first() else { + return Err(ArrowError::InvalidArgumentError( + "Sort requires at least one column".to_string(), + )); + }; + for column in columns { + if column.values.len() != first.values.len() + || column + .values + .iter() + .zip(&first.values) + .any(|(array, other)| array.len() != other.len()) + { + return Err(ArrowError::ComputeError( + "lexical sort columns have different row counts".to_string(), + )); + } + } + if first.values.is_empty() { + return Ok(Vec::new()); + } + if columns.len() == 1 { + return sort_to_indices(&first.values, first.options, limit); + } + let fields = columns + .iter() + .map(|column| { + SortField::new_with_options( + column.values[0].data_type().clone(), + column.options.unwrap_or_default(), + ) + }) + .collect::>(); + if RowConverter::supports_fields(&fields) { + return sort_by_rows(columns, fields, limit); + } + sort_by_comparators(columns, limit) +} + +/// Sorts lexicographically by every column and gathers each column like [`gather`] +pub fn lexsort( + columns: &[SortColumn], + limit: Option, +) -> Result>, ArrowError> { + let indices = lexsort_to_indices(columns, limit)?; + columns + .iter() + .map(|column| gather(&column.values, &indices)) + .collect() +} + +/// The arrays typed as `A`, they were checked to share one data type +fn downcast_all(arrays: &[ArrayRef]) -> Vec<&A> { + arrays + .iter() + .map(|array| { + array + .as_any() + .downcast_ref::() + .expect("all arrays were checked to have the same data type") + }) + .collect() +} + +/// Positions of the non null rows and of the null rows, both in array order +fn partition_validity(arrays: &[&A]) -> (Vec, Vec) { + let null_count: usize = arrays.iter().map(|array| array.null_count()).sum(); + let total: usize = arrays.iter().map(|array| array.len()).sum(); + let mut valids = Vec::with_capacity(total - null_count); + let mut nulls = Vec::with_capacity(null_count); + for (array_index, array) in arrays.iter().enumerate() { + match array.nulls().filter(|nulls| nulls.null_count() > 0) { + None => valids.extend((0..array.len()).map(|row| (array_index, row))), + Some(validity) => { + for row in 0..array.len() { + if validity.is_valid(row) { + valids.push((array_index, row)); + } else { + nulls.push((array_index, row)); + } + } + } + } + } + (valids, nulls) +} + +/// `T` is inferred from the first array, which the downcast macro hands over typed +fn sort_primitive_like( + _first: &PrimitiveArray, + arrays: &[ArrayRef], + options: SortOptions, + limit: Option, +) -> Vec { + let arrays = downcast_all::>(arrays); + let (valids, nulls) = partition_validity(&arrays); + let values: Vec<&[T::Native]> = arrays.iter().map(|array| array.values().as_ref()).collect(); + let mut valids: Vec<(ArrayRowIndex, T::Native)> = valids + .into_iter() + // SAFETY: the positions come from the arrays themselves + .map(|position| (position, unsafe { *values.get_unchecked(position.0).get_unchecked(position.1) })) + .collect(); + sort_impl(options, &mut valids, &nulls, limit, |a, b| a.1.compare(b.1)) +} + +fn sort_boolean( + arrays: &[&BooleanArray], + options: SortOptions, + limit: Option, +) -> Vec { + let (valids, nulls) = partition_validity(arrays); + let mut valids: Vec<(ArrayRowIndex, bool)> = valids + .into_iter() + // SAFETY: the positions come from the arrays themselves + .map(|position| (position, unsafe { arrays.get_unchecked(position.0).value_unchecked(position.1) })) + .collect(); + sort_impl(options, &mut valids, &nulls, limit, |a, b| a.1.cmp(&b.1)) +} + +fn sort_bytes( + arrays: &[&GenericByteArray], + options: SortOptions, + limit: Option, +) -> Vec { + let (valids, nulls) = partition_validity(arrays); + // SAFETY: the positions come from the arrays themselves + let bytes = |position: ArrayRowIndex| -> &[u8] { + unsafe { arrays.get_unchecked(position.0).value_unchecked(position.1).as_ref() } + }; + // Most byte sequences differ in their first bytes: a 4 byte big endian prefix compared as + // one u32 (left padded when shorter) decides nearly every comparison without touching + // the full values + let mut valids: Vec<(ArrayRowIndex, (u32, u64))> = valids + .into_iter() + .map(|position| (position, prefix_and_len(bytes(position)))) + .collect(); + sort_impl(options, &mut valids, &nulls, limit, |a, b| { + compare_prefixed(a.1, b.1).unwrap_or_else(|| bytes(a.0).cmp(bytes(b.0))) + }) +} + +/// `(4 byte big endian prefix, len)` of `slice`, shorter slices are left padded +fn prefix_and_len(slice: &[u8]) -> (u32, u64) { + let prefix = if slice.len() >= 4 { + // SAFETY: at least 4 readable bytes + u32::from_be(unsafe { std::ptr::read_unaligned(slice.as_ptr().cast::()) }) + } else if slice.is_empty() { + 0 + } else { + let mut prefix = 0u32; + for &byte in slice { + prefix = (prefix << 8) | byte as u32; + } + // len is in [1, 3] so the shift is in [8, 24] + prefix << (8 * (4 - slice.len())) + }; + (prefix, slice.len() as u64) +} + +/// The order decided by the prefixes alone, `None` when the full values must be compared +fn compare_prefixed(a: (u32, u64), b: (u32, u64)) -> Option { + let ord = a.0.cmp(&b.0); + if ord != Ordering::Equal { + return Some(ord); + } + // padded prefixes are equal only when both are complete and equal + if a.1 < 4 || b.1 < 4 { + let ord = a.1.cmp(&b.1); + if ord != Ordering::Equal { + return Some(ord); + } + } + None +} + +fn sort_byte_view( + arrays: &[&GenericByteViewArray], + options: SortOptions, + limit: Option, +) -> Vec { + let (valids, nulls) = partition_validity(arrays); + let mut valids: Vec<(ArrayRowIndex, u128)> = valids + .into_iter() + // SAFETY: the positions come from the arrays themselves + .map(|position| (position, unsafe { *arrays.get_unchecked(position.0).views().get_unchecked(position.1) })) + .collect(); + + if arrays.iter().all(|array| array.data_buffers().is_empty()) { + // every view is inline, the key is the view itself + return sort_impl(options, &mut valids, &nulls, limit, |a, b| { + GenericByteViewArray::::inline_key_fast(a.1) + .cmp(&GenericByteViewArray::::inline_key_fast(b.1)) + }); + } + + sort_impl(options, &mut valids, &nulls, limit, |a, b| { + let (raw_a, raw_b) = (a.1, b.1); + if (raw_a as u32) <= MAX_INLINE_VIEW_LEN && (raw_b as u32) <= MAX_INLINE_VIEW_LEN { + return GenericByteViewArray::::inline_key_fast(raw_a) + .cmp(&GenericByteViewArray::::inline_key_fast(raw_b)); + } + let prefix_a = ByteView::from(raw_a).prefix.swap_bytes(); + let prefix_b = ByteView::from(raw_b).prefix.swap_bytes(); + if prefix_a != prefix_b { + return prefix_a.cmp(&prefix_b); + } + // SAFETY: the positions come from the arrays themselves + let full_a: &[u8] = unsafe { arrays.get_unchecked(a.0.0).value_unchecked(a.0.1).as_ref() }; + let full_b: &[u8] = unsafe { arrays.get_unchecked(b.0.0).value_unchecked(b.0.1).as_ref() }; + full_a.cmp(full_b) + }) +} + +fn sort_fixed_size_binary( + arrays: &[&FixedSizeBinaryArray], + options: SortOptions, + limit: Option, +) -> Vec { + let (valids, nulls) = partition_validity(arrays); + let mut valids: Vec<(ArrayRowIndex, &[u8])> = valids + .into_iter() + // SAFETY: the positions come from the arrays themselves + .map(|position| (position, unsafe { arrays.get_unchecked(position.0).value_unchecked(position.1) })) + .collect(); + sort_impl(options, &mut valids, &nulls, limit, |a, b| a.1.cmp(b.1)) +} + +/// Sorts `valids` by `cmp` (reversed when descending), then lays out nulls and valids +/// according to `nulls_first`, keeping only `limit` positions +#[inline(never)] +fn sort_impl( + options: SortOptions, + valids: &mut [(ArrayRowIndex, K)], + nulls: &[ArrayRowIndex], + limit: Option, + mut cmp: impl FnMut(&(ArrayRowIndex, K), &(ArrayRowIndex, K)) -> Ordering, +) -> Vec { + let valid_limit = match (limit, options.nulls_first) { + (Some(limit), true) => limit.saturating_sub(nulls.len()).min(valids.len()), + _ => valids.len(), + }; + if options.descending { + sort_unstable_by(valids, valid_limit, |a, b| cmp(a, b).reverse()); + } else { + sort_unstable_by(valids, valid_limit, cmp); + } + + let len = valids.len() + nulls.len(); + let limit = limit.unwrap_or(len).min(len); + let mut out = Vec::with_capacity(limit); + if options.nulls_first { + out.extend_from_slice(&nulls[..nulls.len().min(limit)]); + let remaining = limit - out.len(); + out.extend(valids.iter().map(|(position, _)| *position).take(remaining)); + } else { + out.extend(valids.iter().map(|(position, _)| *position).take(limit)); + let remaining = limit - out.len(); + out.extend_from_slice(&nulls[..remaining]); + } + out +} + +/// Sorts every row of `columns` with one arrow comparator per column and pair of arrays, +/// the fallback for types without a typed path and for multi column sorts +/// Multi column sort through the row format: every array is encoded to rows on its own +/// (nothing is concatenated) and rows compare as plain bytes, which is far cheaper than +/// a chain of dynamic comparators per column +fn sort_by_rows( + columns: &[SortColumn], + fields: Vec, + limit: Option, +) -> Result, ArrowError> { + let converter = RowConverter::new(fields)?; + let rows = (0..columns[0].values.len()) + .map(|array_index| { + let arrays = columns + .iter() + .map(|column| Arc::clone(&column.values[array_index])) + .collect::>(); + converter.convert_columns(&arrays) + }) + .collect::, _>>()?; + + // Fixed width rows are compared as one or two big endian integers carried next to + // the index, which beats a `memcmp` through two pointer chases per comparison + let fixed_width = columns + .iter() + .map(|column| match column.values[0].data_type() { + DataType::Boolean => Some(1 + 1), + data_type => data_type.primitive_width().map(|width| 1 + width), + }) + .sum::>(); + match fixed_width { + Some(width) if width <= 16 => { + return Ok(sort_by_fixed_rows(&rows, limit, |row| { + let mut key = [0u8; 16]; + key[..width].copy_from_slice(row); + u128::from_be_bytes(key) + })); + } + Some(width) if width <= 32 => { + return Ok(sort_by_fixed_rows(&rows, limit, |row| { + let mut key = [0u8; 32]; + key[..width].copy_from_slice(row); + let (high, low) = key.split_at(16); + ( + u128::from_be_bytes(high.try_into().unwrap()), + u128::from_be_bytes(low.try_into().unwrap()), + ) + })); + } + _ => {} + } + + let mut indices: Vec = rows + .iter() + .enumerate() + .flat_map(|(array_index, rows)| (0..rows.num_rows()).map(move |row| (array_index, row))) + .collect(); + let len = limit.unwrap_or(indices.len()).min(indices.len()); + sort_unstable_by(&mut indices, len, |a, b| { + // SAFETY: the indices were built from the row counts of these very `rows` + unsafe { + rows.get_unchecked(a.0) + .row_unchecked(a.1) + .cmp(&rows.get_unchecked(b.0).row_unchecked(b.1)) + } + }); + indices.truncate(len); + Ok(indices) +} + +/// Sorts rows of one fixed width by the integer key `make_key` builds from their bytes +fn sort_by_fixed_rows( + rows: &[Rows], + limit: Option, + make_key: impl Fn(&[u8]) -> K, +) -> Vec { + let make_key = &make_key; + let mut keyed: Vec<(K, ArrayRowIndex)> = rows + .iter() + .enumerate() + .flat_map(|(array_index, rows)| { + rows.iter() + .enumerate() + .map(move |(row_index, row)| (make_key(row.as_ref()), (array_index, row_index))) + }) + .collect(); + let len = limit.unwrap_or(keyed.len()).min(keyed.len()); + sort_unstable_by(&mut keyed, len, |a, b| a.0.cmp(&b.0)); + keyed.into_iter().take(len).map(|(_, index)| index).collect() +} + +fn sort_by_comparators( + columns: &[SortColumn], + limit: Option, +) -> Result, ArrowError> { + let arrays_count = columns[0].values.len(); + // comparators[column][left array][right array] + let comparators: Vec>> = columns + .iter() + .map(|column| { + let options = column.options.unwrap_or_default(); + (0..arrays_count) + .map(|left| { + (0..arrays_count) + .map(|right| { + make_comparator( + column.values[left].as_ref(), + column.values[right].as_ref(), + options, + ) + }) + .collect::, _>>() + }) + .collect::, _>>() + }) + .collect::, _>>()?; + + let mut indices: Vec = columns[0] + .values + .iter() + .enumerate() + .flat_map(|(array_index, array)| (0..array.len()).map(move |row| (array_index, row))) + .collect(); + let len = limit.unwrap_or(indices.len()).min(indices.len()); + sort_unstable_by(&mut indices, len, |a, b| { + for column in &comparators { + // SAFETY: array indices come from `columns[0].values`, every column has as many + let ord = unsafe { column.get_unchecked(a.0).get_unchecked(b.0) }(a.1, b.1); + if ord != Ordering::Equal { + return ord; + } + } + Ordering::Equal + }); + indices.truncate(len); + Ok(indices) +} + +#[inline] +fn sort_unstable_by(array: &mut [T], limit: usize, cmp: F) +where + F: FnMut(&T, &T) -> Ordering, +{ + if array.len() == limit { + array.sort_unstable_by(cmp); + } else { + partial_sort(array, limit, cmp); + } +} + +/// Sorts only the first `limit` elements into place. Unstable. +pub fn partial_sort(v: &mut [T], limit: usize, mut is_less: F) +where + F: FnMut(&T, &T) -> Ordering, +{ + if let Some(n) = limit.checked_sub(1) { + let (before, _mid, _after) = v.select_nth_unstable_by(n, &mut is_less); + before.sort_unstable_by(is_less); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayRef, BooleanArray, DictionaryArray, Int32Array, Int64Array, StringArray, + StringViewArray, + }; + use arrow::compute::{ + SortColumn as ArrowSortColumn, cast, concat, lexsort_to_indices as arrow_lexsort_to_indices, + sort_to_indices as arrow_sort_to_indices, take, + }; + use arrow::datatypes::Int32Type; + use std::sync::Arc; + + /// Deterministic values without depending on the rand API + struct Lcg(u64); + + impl Lcg { + fn next(&mut self) -> u64 { + self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + self.0 >> 33 + } + + fn below(&mut self, n: u64) -> u64 { + self.next() % n + } + + fn maybe(&mut self, value: T) -> Option { + (self.below(10) != 0).then_some(value) + } + } + + const LENGTHS: [usize; 5] = [0, 13, 1, 50, 20]; + + fn ints(rng: &mut Lcg) -> Vec { + LENGTHS + .iter() + .map(|&len| { + let values = (0..len).map(|_| { + let value = rng.below(40) as i32 - 20; + rng.maybe(value) + }); + Arc::new(Int32Array::from(values.collect::>())) as ArrayRef + }) + .collect() + } + + fn word(rng: &mut Lcg) -> String { + let len = rng.below(20) as usize; + (0..len).map(|_| (b'a' + rng.below(4) as u8) as char).collect() + } + + fn strings(rng: &mut Lcg) -> Vec { + LENGTHS + .iter() + .map(|&len| { + let values = (0..len).map(|_| { + let value = word(rng); + rng.maybe(value) + }); + Arc::new(StringArray::from(values.collect::>())) as ArrayRef + }) + .collect() + } + + fn string_views(rng: &mut Lcg) -> Vec { + LENGTHS + .iter() + .map(|&len| { + let values = (0..len).map(|_| { + let value = word(rng); + rng.maybe(value) + }); + Arc::new(StringViewArray::from(values.collect::>())) as ArrayRef + }) + .collect() + } + + fn booleans(rng: &mut Lcg) -> Vec { + LENGTHS + .iter() + .map(|&len| { + let values = (0..len).map(|_| { + let value = rng.below(2) == 0; + rng.maybe(value) + }); + Arc::new(BooleanArray::from(values.collect::>())) as ArrayRef + }) + .collect() + } + + fn dictionaries(rng: &mut Lcg) -> Vec { + LENGTHS + .iter() + .map(|&len| { + let values: Vec> = (0..len) + .map(|_| { + let value = word(rng); + rng.maybe(value) + }) + .collect(); + Arc::new( + values + .iter() + .map(|value| value.as_deref()) + .collect::>(), + ) as ArrayRef + }) + .collect() + } + + fn all_options() -> Vec> { + let mut options = vec![None]; + for descending in [false, true] { + for nulls_first in [false, true] { + options.push(Some(SortOptions { descending, nulls_first })); + } + } + options + } + + /// Values in our sort order must equal arrow's sort of the concatenated arrays + fn check_single(arrays: &[ArrayRef]) { + let refs: Vec<&dyn Array> = arrays.iter().map(|array| array.as_ref()).collect(); + let whole = concat(&refs).unwrap(); + for options in all_options() { + for limit in [None, Some(1), Some(7), Some(1000)] { + let ours = sort_to_indices(arrays, options, limit).unwrap(); + let ours = interleave(&refs, &ours).unwrap(); + let expected = arrow_sort_to_indices(&whole, options, limit).unwrap(); + let expected = take(&whole, &expected, None).unwrap(); + // dictionaries may end up with different dictionaries but the same values + let (ours, expected) = if let DataType::Dictionary(_, _) = whole.data_type() { + (cast(&ours, &DataType::Utf8).unwrap(), cast(&expected, &DataType::Utf8).unwrap()) + } else { + (ours, expected) + }; + assert_eq!(&ours, &expected, "options {options:?} limit {limit:?}"); + } + } + } + + #[test] + fn matches_arrow_on_concatenated_input() { + let mut rng = Lcg(7); + check_single(&ints(&mut rng)); + check_single(&strings(&mut rng)); + check_single(&string_views(&mut rng)); + check_single(&booleans(&mut rng)); + check_single(&dictionaries(&mut rng)); + } + + #[test] + fn lexsort_matches_arrow_on_concatenated_input() { + let mut rng = Lcg(11); + // few distinct ints so the second column decides often + let ints: Vec = LENGTHS + .iter() + .map(|&len| { + let values = (0..len).map(|_| { + let value = rng.below(3) as i32; + rng.maybe(value) + }); + Arc::new(Int32Array::from(values.collect::>())) as ArrayRef + }) + .collect(); + let views = string_views(&mut rng); + let int_refs: Vec<&dyn Array> = ints.iter().map(|array| array.as_ref()).collect(); + let view_refs: Vec<&dyn Array> = views.iter().map(|array| array.as_ref()).collect(); + let whole_ints = concat(&int_refs).unwrap(); + let whole_views = concat(&view_refs).unwrap(); + + for int_options in all_options() { + for view_options in all_options() { + for limit in [None, Some(5)] { + let columns = [ + SortColumn { values: ints.clone(), options: int_options }, + SortColumn { values: views.clone(), options: view_options }, + ]; + let ours = lexsort_to_indices(&columns, limit).unwrap(); + let expected = arrow_lexsort_to_indices( + &[ + ArrowSortColumn { values: Arc::clone(&whole_ints), options: int_options }, + ArrowSortColumn { values: Arc::clone(&whole_views), options: view_options }, + ], + limit, + ) + .unwrap(); + assert_eq!( + &interleave(&int_refs, &ours).unwrap(), + &take(&whole_ints, &expected, None).unwrap() + ); + assert_eq!( + &interleave(&view_refs, &ours).unwrap(), + &take(&whole_views, &expected, None).unwrap() + ); + + // `lexsort` gathers the same order, chunked like the inputs + let sorted = lexsort(&columns, limit).unwrap(); + let gathered: Vec<&dyn Array> = sorted[1].iter().map(|array| array.as_ref()).collect(); + assert_eq!( + &concat(&gathered).unwrap(), + &take(&whole_views, &expected, None).unwrap() + ); + } + } + } + } + + /// Fixed width columns take the integer key paths (one `u128` up to 16 bytes of row, + /// two above that), check both against arrow on the concatenated input + #[test] + fn fixed_width_lexsort_matches_arrow_on_concatenated_input() { + let mut rng = Lcg(29); + let mut column = |distinct: u64, wide: bool| -> Vec { + LENGTHS + .iter() + .map(|&len| { + if wide { + let values = (0..len).map(|_| { + let value = rng.below(distinct) as i64 - 1; + rng.maybe(value) + }); + Arc::new(Int64Array::from(values.collect::>())) as ArrayRef + } else { + let values = (0..len).map(|_| { + let value = rng.below(distinct) as i32 - 1; + rng.maybe(value) + }); + Arc::new(Int32Array::from(values.collect::>())) as ArrayRef + } + }) + .collect() + }; + // 5 + 9 = 14 bytes per row, and 5 + 9 + 9 = 23 bytes per row + let two = vec![column(3, false), column(4, true)]; + let three = vec![column(3, false), column(3, true), column(50, true)]; + for columns in [two, three] { + let whole: Vec = columns + .iter() + .map(|arrays| { + let refs: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect(); + concat(&refs).unwrap() + }) + .collect(); + for first_options in all_options() { + for rest_options in all_options() { + for limit in [None, Some(7)] { + let options = |i: usize| if i == 0 { first_options } else { rest_options }; + let sort_columns: Vec = columns + .iter() + .enumerate() + .map(|(i, values)| SortColumn { values: values.clone(), options: options(i) }) + .collect(); + let ours = lexsort_to_indices(&sort_columns, limit).unwrap(); + let arrow_columns: Vec = whole + .iter() + .enumerate() + .map(|(i, values)| ArrowSortColumn { values: Arc::clone(values), options: options(i) }) + .collect(); + let expected = arrow_lexsort_to_indices(&arrow_columns, limit).unwrap(); + for (arrays, whole) in columns.iter().zip(&whole) { + let refs: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect(); + assert_eq!( + &interleave(&refs, &ours).unwrap(), + &take(whole, &expected, None).unwrap() + ); + } + } + } + } + } + } + + #[test] + fn sort_keeps_input_chunking() { + let mut rng = Lcg(3); + let arrays = strings(&mut rng); + let sorted = sort(&arrays, None).unwrap(); + assert_eq!(sorted.iter().map(|array| array.len()).collect::>(), LENGTHS); + let limited = sort_limit(&arrays, None, Some(15)).unwrap(); + assert_eq!(limited.iter().map(|array| array.len()).collect::>(), vec![0, 13, 1, 1]); + } + + #[test] + fn rejects_mixed_types_and_row_counts() { + let ints: ArrayRef = Arc::new(Int32Array::from(vec![1])); + let strings: ArrayRef = Arc::new(StringArray::from(vec!["a"])); + assert!(sort_to_indices(&[Arc::clone(&ints), strings], None, None).is_err()); + assert!( + lexsort_to_indices( + &[ + SortColumn { values: vec![Arc::clone(&ints)], options: None }, + SortColumn { values: vec![Arc::clone(&ints), ints], options: None }, + ], + None + ) + .is_err() + ); + } +} diff --git a/datafusion/physical-plan/src/sorts/mod.rs b/datafusion/physical-plan/src/sorts/mod.rs index 4af2a3629fa9c..2a3dee537f59f 100644 --- a/datafusion/physical-plan/src/sorts/mod.rs +++ b/datafusion/physical-plan/src/sorts/mod.rs @@ -29,5 +29,6 @@ pub mod sort; pub mod sort_preserving_merge; mod stream; pub mod streaming_merge; +mod arrow; pub(crate) use stream::IncrementalSortIterator; diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 490ea7cc85776..65f633fc875f3 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -58,7 +58,7 @@ use crate::{ }; use arrow::array::{RecordBatch, RecordBatchOptions}; -use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays}; +use arrow::compute::{lexsort_to_indices, take_arrays}; use arrow::datatypes::SchemaRef; use datafusion_common::config::SpillCompression; use datafusion_common::tree_node::TreeNodeRecursion; @@ -77,6 +77,7 @@ use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; use futures::{StreamExt, TryStreamExt}; use log::{debug, trace}; +use crate::sorts::stream::IncrementalMultiBatchSortIterator; #[cfg(test)] mod spill_tests; @@ -636,13 +637,17 @@ impl ExternalSorter { // so that growth can use it. self.merge_pool.release_unused(); // Concatenate memory batches together and sort - let batch = concat_batches(&self.schema, &self.in_mem_batches)?; - self.in_mem_batches.clear(); + let cap = self.in_mem_batches.capacity(); + let batches = std::mem::replace(&mut self.in_mem_batches, Vec::with_capacity(cap)); + let size_to_resize: usize = batches + .iter() + .map(get_reserved_bytes_for_record_batch) + .sum::>()?; self.reservation - .try_resize(get_reserved_bytes_for_record_batch(&batch)?) + .try_resize(size_to_resize) .map_err(Self::err_with_oom_context)?; let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation)?; + let sorted_stream = self.sort_batches_stream(batches, reservation)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -652,22 +657,36 @@ impl ExternalSorter { // left as one run per batch: the row-format merge of many small runs // beats sorting a few large runs with the lexicographic comparator. let batches = std::mem::take(&mut self.in_mem_batches); - let runs = if coalesce_runs && self.expr.len() == 1 { - self.coalesce_in_mem_batches_into_runs(batches)? + let streams = if coalesce_runs && self.expr.len() == 1 { + let runs = self.coalesce_in_mem_batches_into_runs(batches)?; + + runs + .into_iter() + .map(|batches| { + let size_to_split: usize = batches + .iter() + .map(get_reserved_bytes_for_record_batch) + .sum::>()?; + let reservation = self + .reservation + .split(size_to_split); + let input = self.sort_batches_stream(batches, reservation)?; + Ok(spawn_buffered(input, 1)) + }) + .collect::>()? } else { batches - }; - - let streams = runs - .into_iter() - .map(|batch| { - let reservation = self + .into_iter() + .map(|batch| { + let reservation = self .reservation .split(get_reserved_bytes_for_record_batch(&batch)?); - let input = self.sort_batch_stream(batch, reservation)?; - Ok(spawn_buffered(input, 1)) - }) - .collect::>()?; + let input = self.sort_batch_stream(batch, reservation)?; + Ok(spawn_buffered(input, 1)) + }) + .collect::>()? + }; + StreamingMergeBuilder::new() .with_streams(streams) @@ -691,22 +710,22 @@ impl ExternalSorter { fn coalesce_in_mem_batches_into_runs( &mut self, batches: Vec, - ) -> Result> { + ) -> Result>> { let target = self.sort_in_place_threshold_bytes.max(1); - let mut runs: Vec = Vec::new(); + let mut runs: Vec> = Vec::new(); let mut group: Vec = Vec::new(); let mut group_bytes = 0usize; // Flush a group into a run, skipping the copy for a single-batch group. let flush = |group: &mut Vec, - runs: &mut Vec, - schema: &SchemaRef| + runs: &mut Vec>| -> Result<()> { match group.len() { 0 => {} - 1 => runs.push(group.pop().unwrap()), + 1 => runs.push(std::mem::take(group)), _ => { - runs.push(concat_batches(schema, group.iter())?); + let run_group = std::mem::take(group); + runs.push(run_group); group.clear(); } } @@ -716,17 +735,18 @@ impl ExternalSorter { for batch in batches { let bytes = get_reserved_bytes_for_record_batch(&batch)?; if !group.is_empty() && group_bytes.saturating_add(bytes) > target { - flush(&mut group, &mut runs, &self.schema)?; + flush(&mut group, &mut runs)?; group_bytes = 0; } group_bytes += bytes; group.push(batch); } - flush(&mut group, &mut runs, &self.schema)?; + flush(&mut group, &mut runs)?; // Realign the reservation: concatenation may shift the footprint slightly. let total: usize = runs .iter() + .flatten() .map(get_reserved_bytes_for_record_batch) .sum::>()?; self.reservation @@ -770,14 +790,14 @@ impl ExternalSorter { // exceed the input estimate. Borrow only already-reserved spill // workspace; any remainder still uses the original sort consumer. let total_sorted_size: usize = sorted_batches - .iter() - .map(get_record_batch_memory_size) - .sum(); + .iter() + .map(get_record_batch_memory_size) + .sum(); let mut workspace = - merge_pool.borrow(total_sorted_size.saturating_sub(reservation.size())); + merge_pool.borrow(total_sorted_size.saturating_sub(reservation.size())); reservation - .try_resize(total_sorted_size - workspace.size()) - .map_err(Self::err_with_oom_context)?; + .try_resize(total_sorted_size - workspace.size()) + .map_err(Self::err_with_oom_context)?; if workspace.size() == 0 { return Ok(Box::pin(ReservationStream::new( @@ -804,7 +824,93 @@ impl ExternalSorter { futures::stream::iter(batches), )) as SendableRecordBatchStream) }) - .try_flatten(); + .try_flatten(); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + /// Sorts a single `RecordBatch` into a single stream. + /// + /// This may output multiple batches depending on the size of the + /// sorted data and the target batch size. + /// For single-batch output cases, `reservation` will be freed immediately after sorting, + /// as the batch will be output and is expected to be reserved by the consumer of the stream. + /// For multi-batch output cases, `reservation` and any borrowed spill workspace + /// cover the sorted output, releasing its memory as each batch is output. + /// (This leads to the same behaviour, as futures are only evaluated when polled by the consumer.) + fn sort_batches_stream( + &self, + batches: Vec, + reservation: MemoryReservation, + ) -> Result { + assert_ne!(batches.len(), 0); + if batches.len() == 1 { + return self.sort_batch_stream( + batches.into_iter().next().unwrap(), + reservation, + ); + } + + let mut expected_size = 0; + for batch in &batches { + expected_size += get_reserved_bytes_for_record_batch(batch)?; + } + + assert_eq!( + expected_size, + reservation.size() + ); + + let schema = batches[0].schema(); + let expressions = self.expr.clone(); + let batch_size = self.batch_size; + let merge_pool = Arc::clone(&self.merge_pool); + + let stream = futures::stream::once(async move { + let schema = batches[0].schema(); + + // Sort the batch immediately and get all output batches + let sorted_batches = sort_batches_chunked(batches, &expressions, batch_size)?; + + // Chunked output can retain shared buffers in every batch and + // exceed the input estimate. Borrow only already-reserved spill + // workspace; any remainder still uses the original sort consumer. + let total_sorted_size: usize = sorted_batches + .iter() + .map(get_record_batch_memory_size) + .sum(); + let mut workspace = + merge_pool.borrow(total_sorted_size.saturating_sub(reservation.size())); + reservation + .try_resize(total_sorted_size - workspace.size()) + .map_err(Self::err_with_oom_context)?; + + if workspace.size() == 0 { + return Ok(Box::pin(ReservationStream::new( + Arc::clone(&schema), + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(sorted_batches.into_iter().map(Ok)), + )), + reservation, + )) as SendableRecordBatchStream); + } + + // Return borrowed workspace first so the merge's cursors can reuse + // it immediately. Both reservations also release on stream drop. + let batches = sorted_batches.into_iter().map(move |batch| { + let size = get_record_batch_memory_size(&batch); + let borrowed = size.min(workspace.size()); + workspace.shrink(borrowed); + reservation.shrink(size - borrowed); + Ok(batch) + }); + Result::<_, DataFusionError>::Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(batches), + )) as SendableRecordBatchStream) + }) + .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } @@ -956,6 +1062,16 @@ pub fn sort_batch_chunked( ) -> Result> { IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size).collect() } +/// Sort a batch and return the result as multiple batches of size `batch_size`. +/// This is useful when you want to avoid creating one large sorted batch in memory, +/// and instead want to process the sorted data in smaller chunks. +fn sort_batches_chunked( + batches: Vec, + expressions: &LexOrdering, + batch_size: usize, +) -> Result> { + IncrementalMultiBatchSortIterator::new(batches, expressions.clone(), batch_size).collect() +} /// Sort execution plan. /// @@ -2131,7 +2247,7 @@ mod tests { use crate::test::{assert_is_pending, make_partition}; use arrow::array::*; - use arrow::compute::SortOptions; + use arrow::compute::{concat_batches, SortOptions}; use arrow::datatypes::*; use datafusion_common::ScalarValue; use datafusion_common::cast::as_primitive_array; diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index bb9c00949369e..ae08ab97e3234 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -15,11 +15,12 @@ // specific language governing permissions and limitations // under the License. +use crate::sorts::arrow::sort::SortColumn; use crate::sorts::cursor::{ArrayValues, CursorArray, RowValues}; use crate::{EmptyRecordBatchStream, SendableRecordBatchStream}; use crate::{PhysicalExpr, PhysicalSortExpr}; use arrow::array::{Array, UInt32Array}; -use arrow::compute::take_record_batch; +use arrow::compute::{interleave, take_record_batch}; use arrow::datatypes::Schema; use arrow::record_batch::RecordBatch; use arrow::row::{RowConverter, Rows, SortField}; @@ -29,6 +30,7 @@ use datafusion_execution::memory_pool::MemoryReservation; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use futures::stream::{Fuse, StreamExt}; +use std::collections::VecDeque; use std::iter::FusedIterator; use std::marker::PhantomData; use std::mem; @@ -386,6 +388,168 @@ impl Iterator for IncrementalSortIterator { impl FusedIterator for IncrementalSortIterator {} +#[derive(Debug, PartialEq)] +enum IncrementalMultiBranchIteratorState { + Init { + input_batches: Vec, + }, + CalculatedIndices { + input_batches: Vec, + // x batches of interleave (batch_index, row index) + indices: VecDeque>, + }, + Done, +} + +pub(crate) struct IncrementalMultiBatchSortIterator { + state: IncrementalMultiBranchIteratorState, + total_len: usize, + cursor: usize, + expressions: LexOrdering, + batch_size: usize, +} + +impl IncrementalMultiBatchSortIterator { + pub(crate) fn new( + batches: Vec, + expressions: LexOrdering, + batch_size: usize, + ) -> Self { + let total_len = batches.iter().map(|b| b.num_rows()).sum(); + Self { + total_len, + cursor: 0, + state: if total_len > 0 { + IncrementalMultiBranchIteratorState::Init { + input_batches: batches, + } + } else { + IncrementalMultiBranchIteratorState::Done + }, + expressions, + batch_size, + } + } + + fn take_next_output_batch( + input_batches: &[RecordBatch], + indices: &[(usize, usize)], + ) -> Result { + let schema = input_batches[0].schema(); + let columns = (0..schema.fields().len()) + .map(|i| { + let column_values: Vec<&dyn Array> = input_batches + .iter() + .map(|batch| batch.column(i).as_ref()) + .collect(); + Ok(interleave(&column_values, indices)?) + }) + .collect::>>()?; + + Ok(RecordBatch::try_new(schema, columns)?) + } + + fn on_take_next( + &mut self, + input_batches: Vec, + mut indices: VecDeque>, + ) -> Result { + let next_indices = indices.pop_front().expect("must not have empty indices"); + self.cursor += next_indices.len(); + + let output = Self::take_next_output_batch(&input_batches, &next_indices); + + self.state = if indices.is_empty() { + IncrementalMultiBranchIteratorState::Done + } else { + IncrementalMultiBranchIteratorState::CalculatedIndices { + input_batches, + indices, + } + }; + + output + } + + fn next_batch(&mut self) -> Result { + match mem::replace( + &mut self.state, + IncrementalMultiBranchIteratorState::Done, + ) { + IncrementalMultiBranchIteratorState::Init { input_batches } => { + let columns = self + .expressions + .iter() + .map(|sort_expr| { + let values = input_batches + .iter() + .map(|block| { + sort_expr + .expr + .evaluate(block)? + .into_array(block.num_rows()) + }) + .collect::>>()?; + Ok(SortColumn { + values, + options: Some(sort_expr.options), + }) + }) + .collect::>>()?; + + let order = super::arrow::sort::lexsort_to_indices(&columns, None)?; + + drop(columns); + + let indices = order + .chunks(self.batch_size) + .map(|ch| ch.to_vec()) + .collect::>(); + + self.on_take_next(input_batches, indices) + } + IncrementalMultiBranchIteratorState::CalculatedIndices { + input_batches, + indices, + } => self.on_take_next(input_batches, indices), + IncrementalMultiBranchIteratorState::Done => { + unreachable!("must not be done if reached here") + } + } + } +} + +impl Iterator for IncrementalMultiBatchSortIterator { + type Item = Result; + + fn next(&mut self) -> Option { + if self.total_len <= self.cursor { + assert_eq!(self.state, IncrementalMultiBranchIteratorState::Done); + return None; + } + + match self.next_batch() { + Ok(batch) => Some(Ok(batch)), + Err(e) => { + self.cursor = self.total_len; + // Release the memory + self.state = IncrementalMultiBranchIteratorState::Done; + Some(Err(e)) + } + } + } + + fn size_hint(&self) -> (usize, Option) { + let num_rows_left = self.total_len - self.cursor; + let num_batches = num_rows_left.div_ceil(self.batch_size); + (num_batches, Some(num_batches)) + } +} + +impl ExactSizeIterator for IncrementalMultiBatchSortIterator {} + +impl FusedIterator for IncrementalMultiBatchSortIterator {} + #[cfg(test)] mod tests { use super::*;