arrow_ord/
partition.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Defines partition kernel for `ArrayRef`
6
7use std::ops::Range;
8
9use arrow_array::{Array, ArrayRef};
10use arrow_buffer::BooleanBuffer;
11use arrow_schema::{ArrowError, SortOptions};
12
13use crate::cmp::{distinct, supports_distinct};
14use crate::ord::make_comparator;
15
16/// A computed set of partitions, see [`partition`]
17#[derive(Debug, Clone)]
18pub struct Partitions(Option<BooleanBuffer>);
19
20impl Partitions {
21    /// Returns the range of each partition
22    ///
23    /// Consecutive ranges will be contiguous: i.e [`(a, b)` and `(b, c)`], and
24    /// `start = 0` and `end = self.len()` for the first and last range respectively
25    pub fn ranges(&self) -> Vec<Range<usize>> {
26        let Some(boundaries) = &self.0 else {
27            return vec![];
28        };
29
30        let mut out = vec![];
31        let mut current = 0;
32        for idx in boundaries.set_indices() {
33            let t = current;
34            current = idx + 1;
35            out.push(t..current)
36        }
37        let last = boundaries.len() + 1;
38        if current != last {
39            out.push(current..last)
40        }
41        out
42    }
43
44    /// Returns the number of partitions
45    pub fn len(&self) -> usize {
46        match &self.0 {
47            Some(b) => b.count_set_bits() + 1,
48            None => 0,
49        }
50    }
51
52    /// Returns true if this contains no partitions
53    pub fn is_empty(&self) -> bool {
54        self.0.is_none()
55    }
56}
57
58/// Given a list of lexicographically sorted columns, computes the [`Partitions`],
59/// where a partition consists of the set of consecutive rows with equal values
60///
61/// Returns an error if no columns are specified or all columns do not
62/// have the same number of rows.
63///
64/// # Example:
65///
66/// For example, given columns `x`, `y` and `z`, calling
67/// [`partition`]`(values, (x, y))` will divide the
68/// rows into ranges where the values of `(x, y)` are equal:
69///
70/// ```text
71/// ┌ ─ ┬───┬ ─ ─┌───┐─ ─ ┬───┬ ─ ─ ┐
72///     │ 1 │    │ 1 │    │ A │        Range: 0..1 (x=1, y=1)
73/// ├ ─ ┼───┼ ─ ─├───┤─ ─ ┼───┼ ─ ─ ┤
74///     │ 1 │    │ 2 │    │ B │
75/// │   ├───┤    ├───┤    ├───┤     │
76///     │ 1 │    │ 2 │    │ C │        Range: 1..4 (x=1, y=2)
77/// │   ├───┤    ├───┤    ├───┤     │
78///     │ 1 │    │ 2 │    │ D │
79/// ├ ─ ┼───┼ ─ ─├───┤─ ─ ┼───┼ ─ ─ ┤
80///     │ 2 │    │ 1 │    │ E │        Range: 4..5 (x=2, y=1)
81/// ├ ─ ┼───┼ ─ ─├───┤─ ─ ┼───┼ ─ ─ ┤
82///     │ 3 │    │ 1 │    │ F │        Range: 5..6 (x=3, y=1)
83/// └ ─ ┴───┴ ─ ─└───┘─ ─ ┴───┴ ─ ─ ┘
84///
85///       x        y        z     partition(&[x, y])
86/// ```
87///
88/// # Example Code
89///
90/// ```
91/// # use std::{sync::Arc, ops::Range};
92/// # use arrow_array::{RecordBatch, Int64Array, StringArray, ArrayRef};
93/// # use arrow_ord::sort::{SortColumn, SortOptions};
94/// # use arrow_ord::partition::partition;
95/// let batch = RecordBatch::try_from_iter(vec![
96///     ("x", Arc::new(Int64Array::from(vec![1, 1, 1, 1, 2, 3])) as ArrayRef),
97///     ("y", Arc::new(Int64Array::from(vec![1, 2, 2, 2, 1, 1])) as ArrayRef),
98///     ("z", Arc::new(StringArray::from(vec!["A", "B", "C", "D", "E", "F"])) as ArrayRef),
99/// ]).unwrap();
100///
101/// // Partition on first two columns
102/// let ranges = partition(&batch.columns()[..2]).unwrap().ranges();
103///
104/// let expected = vec![
105///     (0..1),
106///     (1..4),
107///     (4..5),
108///     (5..6),
109/// ];
110///
111/// assert_eq!(ranges, expected);
112/// ```
113pub fn partition(columns: &[ArrayRef]) -> Result<Partitions, ArrowError> {
114    if columns.is_empty() {
115        return Err(ArrowError::InvalidArgumentError(
116            "Partition requires at least one column".to_string(),
117        ));
118    }
119    let num_rows = columns[0].len();
120    if columns.iter().any(|item| item.len() != num_rows) {
121        return Err(ArrowError::InvalidArgumentError(
122            "Partition columns have different row counts".to_string(),
123        ));
124    }
125
126    match num_rows {
127        0 => return Ok(Partitions(None)),
128        1 => return Ok(Partitions(Some(BooleanBuffer::new_unset(0)))),
129        _ => {}
130    }
131
132    let acc = find_boundaries(&columns[0])?;
133    let acc = columns.iter().skip(1).try_fold(acc, |mut acc, c| {
134        find_boundaries(c.as_ref()).map(|b| {
135            // Do OrAssign to try to reuse the existing buffer if possible
136            acc |= &b;
137
138            acc
139        })
140    })?;
141
142    Ok(Partitions(Some(acc)))
143}
144
145/// Returns a mask with bits set whenever the value or nullability changes
146fn find_boundaries(v: &dyn Array) -> Result<BooleanBuffer, ArrowError> {
147    let slice_len = v.len() - 1;
148    let v1 = v.slice(0, slice_len);
149    let v2 = v.slice(1, slice_len);
150
151    if supports_distinct(v.data_type()) {
152        return Ok(distinct(&v1, &v2)?.values().clone());
153    }
154    // Given that we're only comparing values, null ordering in the input or
155    // sort options do not matter.
156    let cmp = make_comparator(&v1, &v2, SortOptions::default())?;
157    Ok((0..slice_len).map(|i| !cmp(i, i).is_eq()).collect())
158}
159
160#[cfg(test)]
161mod tests {
162    use std::sync::Arc;
163
164    use super::*;
165    use arrow_array::*;
166    use arrow_schema::DataType;
167
168    #[test]
169    fn test_partition_empty() {
170        let err = partition(&[]).unwrap_err();
171        assert_eq!(
172            err.to_string(),
173            "Invalid argument error: Partition requires at least one column"
174        );
175    }
176
177    #[test]
178    fn test_partition_unaligned_rows() {
179        let input = vec![
180            Arc::new(Int64Array::from(vec![None, Some(-1)])) as _,
181            Arc::new(StringArray::from(vec![Some("foo")])) as _,
182        ];
183        let err = partition(&input).unwrap_err();
184        assert_eq!(
185            err.to_string(),
186            "Invalid argument error: Partition columns have different row counts"
187        )
188    }
189
190    #[test]
191    fn test_partition_small() {
192        let results = partition(&[
193            Arc::new(Int32Array::new(vec![].into(), None)) as _,
194            Arc::new(Int32Array::new(vec![].into(), None)) as _,
195            Arc::new(Int32Array::new(vec![].into(), None)) as _,
196        ])
197        .unwrap();
198        assert_eq!(results.len(), 0);
199        assert!(results.is_empty());
200
201        let results = partition(&[
202            Arc::new(Int32Array::from(vec![1])) as _,
203            Arc::new(Int32Array::from(vec![1])) as _,
204        ])
205        .unwrap()
206        .ranges();
207        assert_eq!(results.len(), 1);
208        assert_eq!(results[0], 0..1);
209    }
210
211    #[test]
212    fn test_partition_single_column() {
213        let a = Int64Array::from(vec![1, 2, 2, 2, 2, 2, 2, 2, 9]);
214        let input = vec![Arc::new(a) as _];
215        assert_eq!(
216            partition(&input).unwrap().ranges(),
217            vec![(0..1), (1..8), (8..9)],
218        );
219    }
220
221    #[test]
222    fn test_partition_all_equal_values() {
223        let a = Int64Array::from_value(1, 1000);
224        let input = vec![Arc::new(a) as _];
225        assert_eq!(partition(&input).unwrap().ranges(), vec![(0..1000)]);
226    }
227
228    #[test]
229    fn test_partition_all_null_values() {
230        let input = vec![
231            new_null_array(&DataType::Int8, 1000),
232            new_null_array(&DataType::UInt16, 1000),
233        ];
234        assert_eq!(partition(&input).unwrap().ranges(), vec![(0..1000)]);
235    }
236
237    #[test]
238    fn test_partition_unique_column_1() {
239        let input = vec![
240            Arc::new(Int64Array::from(vec![None, Some(-1)])) as _,
241            Arc::new(StringArray::from(vec![Some("foo"), Some("bar")])) as _,
242        ];
243        assert_eq!(partition(&input).unwrap().ranges(), vec![(0..1), (1..2)],);
244    }
245
246    #[test]
247    fn test_partition_unique_column_2() {
248        let input = vec![
249            Arc::new(Int64Array::from(vec![None, Some(-1), Some(-1)])) as _,
250            Arc::new(StringArray::from(vec![
251                Some("foo"),
252                Some("bar"),
253                Some("apple"),
254            ])) as _,
255        ];
256        assert_eq!(
257            partition(&input).unwrap().ranges(),
258            vec![(0..1), (1..2), (2..3),],
259        );
260    }
261
262    #[test]
263    fn test_partition_non_unique_column_1() {
264        let input = vec![
265            Arc::new(Int64Array::from(vec![None, Some(-1), Some(-1), Some(1)])) as _,
266            Arc::new(StringArray::from(vec![
267                Some("foo"),
268                Some("bar"),
269                Some("bar"),
270                Some("bar"),
271            ])) as _,
272        ];
273        assert_eq!(
274            partition(&input).unwrap().ranges(),
275            vec![(0..1), (1..3), (3..4),],
276        );
277    }
278
279    #[test]
280    fn test_partition_masked_nulls() {
281        let input = vec![
282            Arc::new(Int64Array::new(vec![1; 9].into(), None)) as _,
283            Arc::new(Int64Array::new(
284                vec![1, 1, 2, 2, 2, 3, 3, 3, 3].into(),
285                Some(vec![false, true, true, true, true, false, false, true, false].into()),
286            )) as _,
287            Arc::new(Int64Array::new(
288                vec![1, 1, 2, 2, 2, 2, 2, 3, 7].into(),
289                Some(vec![true, true, true, true, false, true, true, true, false].into()),
290            )) as _,
291        ];
292
293        assert_eq!(
294            partition(&input).unwrap().ranges(),
295            vec![(0..1), (1..2), (2..4), (4..5), (5..7), (7..8), (8..9)],
296        );
297    }
298
299    #[test]
300    fn test_partition_run_end_encoded() {
301        let run_ends = Int32Array::from(vec![2, 3, 5]);
302        let values = StringArray::from(vec!["x", "y", "x"]);
303        let ree = RunArray::try_new(&run_ends, &values).unwrap();
304        // logical: ["x", "x", "y", "x", "x"]
305        let input = vec![Arc::new(ree) as _];
306        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
307    }
308
309    #[test]
310    fn test_partition_nested_run_end_encoded() {
311        // Inner REE (values of the outer): run_ends [1, 2, 3], values ["x", "y", "x"]
312        // logical length 3: ["x", "y", "x"]
313        let inner_run_ends = Int32Array::from(vec![1, 2, 3]);
314        let inner_values = StringArray::from(vec!["x", "y", "x"]);
315        let inner_ree = RunArray::try_new(&inner_run_ends, &inner_values).unwrap();
316
317        // Outer REE: run_ends [2, 3, 5], values = inner_ree (length 3)
318        // logical: rows 0,1 → inner[0]="x", row 2 → inner[1]="y", rows 3,4 → inner[2]="x"
319        // = ["x", "x", "y", "x", "x"]
320        let outer_run_ends = Int32Array::from(vec![2, 3, 5]);
321        let outer_ree = RunArray::try_new(&outer_run_ends, &inner_ree).unwrap();
322
323        let input = vec![Arc::new(outer_ree) as ArrayRef];
324        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5]);
325    }
326
327    #[test]
328    fn test_partition_ree_with_dictionary_values() {
329        // Dictionary values: keys [0, 1, 0], dict ["x", "y"] → logical ["x", "y", "x"]
330        let dict_values = StringArray::from(vec!["x", "y"]);
331        let keys = Int32Array::from(vec![0, 1, 0]);
332        let dict = DictionaryArray::try_new(keys, Arc::new(dict_values)).unwrap();
333
334        // REE wrapping dict: run_ends [2, 3, 5] → logical [dict[0], dict[0], dict[1], dict[2], dict[2]]
335        // = ["x", "x", "y", "x", "x"]
336        let run_ends = Int32Array::from(vec![2, 3, 5]);
337        let ree = RunArray::try_new(&run_ends, &dict).unwrap();
338        let input = vec![Arc::new(ree) as ArrayRef];
339        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
340    }
341
342    #[test]
343    fn test_partition_dictionary() {
344        let values = StringArray::from(vec!["x", "y"]);
345        let keys = Int32Array::from(vec![0, 0, 1, 0, 0]);
346        let dict = DictionaryArray::try_new(keys, Arc::new(values)).unwrap();
347        // logical: ["x", "x", "y", "x", "x"]
348        let input = vec![Arc::new(dict) as _];
349        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
350    }
351
352    #[test]
353    fn test_partition_nested_dictionary() {
354        let inner_values = StringArray::from(vec!["x", "y"]);
355        let inner_keys = Int32Array::from(vec![0, 1, 0]);
356        let inner_dict = DictionaryArray::try_new(inner_keys, Arc::new(inner_values)).unwrap();
357
358        // Outer dict keys index into inner dict's logical values: ["x", "y", "x"]
359        // keys [0, 0, 1, 2, 2] → logical ["x", "x", "y", "x", "x"]
360        let outer_keys = Int32Array::from(vec![0, 0, 1, 2, 2]);
361        let outer_dict = DictionaryArray::try_new(outer_keys, Arc::new(inner_dict)).unwrap();
362        let input = vec![Arc::new(outer_dict) as ArrayRef];
363        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
364    }
365
366    #[test]
367    fn test_partition_dictionary_with_ree_values() {
368        // REE values: run_ends [2, 3], values ["x", "y"] → logical ["x", "x", "y"]
369        let run_ends = Int32Array::from(vec![2, 3]);
370        let str_values = StringArray::from(vec!["x", "y"]);
371        let ree = RunArray::try_new(&run_ends, &str_values).unwrap();
372
373        // Dictionary keys index into the REE's logical values
374        // keys [0, 0, 2, 0, 0] → logical ["x", "x", "y", "x", "x"]
375        let keys = Int32Array::from(vec![0, 0, 2, 0, 0]);
376        let dict = DictionaryArray::try_new(keys, Arc::new(ree)).unwrap();
377        let input = vec![Arc::new(dict) as ArrayRef];
378        assert_eq!(partition(&input).unwrap().ranges(), vec![0..2, 2..3, 3..5],);
379    }
380
381    #[test]
382    fn test_partition_nested() {
383        let input = vec![
384            Arc::new(
385                StructArray::try_from(vec![(
386                    "f1",
387                    Arc::new(Int64Array::from(vec![
388                        None,
389                        None,
390                        Some(1),
391                        Some(2),
392                        Some(2),
393                        Some(2),
394                        Some(3),
395                        Some(4),
396                    ])) as _,
397                )])
398                .unwrap(),
399            ) as _,
400            Arc::new(Int64Array::from(vec![1, 1, 1, 2, 3, 3, 3, 4])) as _,
401        ];
402        assert_eq!(
403            partition(&input).unwrap().ranges(),
404            vec![0..2, 2..3, 3..4, 4..6, 6..7, 7..8]
405        )
406    }
407}
arrow_ord/partition.rs

arrow_ord/
partition.rs
