arrow_ord/
partition.rs

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

arrow_ord/
partition.rs
