alopex_core/dataframe/
partition_scan.rs

1//! Partition-aware DataFrame scan primitives.
2
3use std::collections::HashSet;
4use std::ops::Range;
5use std::sync::Arc;
6
7use crate::columnar::encoding::Column;
8use crate::columnar::segment_v2::{RecordBatch, SegmentReaderV2};
9use crate::{Error, Result};
10
11/// A typed value used to compare partition keys.
12#[derive(Clone, Debug, PartialEq, Eq, Hash)]
13pub enum PartitionValue {
14    /// SQL-style NULL key component.
15    Null,
16    /// Signed 64-bit integer key component.
17    Int64(i64),
18    /// 32-bit floating-point key component, stored by raw bits for total equality.
19    Float32(u32),
20    /// 64-bit floating-point key component, stored by raw bits for total equality.
21    Float64(u64),
22    /// Boolean key component.
23    Bool(bool),
24    /// Variable-length binary key component.
25    Binary(Vec<u8>),
26    /// Fixed-length binary key component.
27    Fixed(Vec<u8>),
28}
29
30/// A composite partition key.
31#[derive(Clone, Debug, PartialEq, Eq, Hash)]
32pub struct PartitionKey {
33    values: Vec<PartitionValue>,
34}
35
36impl PartitionKey {
37    /// Create a partition key from key components.
38    pub fn new(values: Vec<PartitionValue>) -> Self {
39        Self { values }
40    }
41
42    /// Return the key components in partition-column order.
43    pub fn values(&self) -> &[PartitionValue] {
44        &self.values
45    }
46
47    fn from_batch_row(batch: &RecordBatch, key_columns: &[usize], row_idx: usize) -> Result<Self> {
48        let mut values = Vec::with_capacity(key_columns.len());
49        for &column_idx in key_columns {
50            let column = batch
51                .columns
52                .get(column_idx)
53                .ok_or_else(|| Error::InvalidParameter {
54                    param: "key_columns".into(),
55                    reason: format!("column index {column_idx} is out of bounds"),
56                })?;
57            let bitmap = batch
58                .null_bitmaps
59                .get(column_idx)
60                .and_then(|bitmap| bitmap.as_ref());
61            if let Some(bitmap) = bitmap {
62                if !bitmap.get(row_idx) {
63                    values.push(PartitionValue::Null);
64                    continue;
65                }
66            }
67            values.push(partition_value(column, row_idx)?);
68        }
69        Ok(Self::new(values))
70    }
71}
72
73/// A contiguous chunk of rows belonging to one logical partition.
74#[derive(Clone, Debug)]
75pub struct PartitionChunk {
76    /// Partition key for this chunk.
77    pub key: PartitionKey,
78    /// Shared record batch containing the rows.
79    pub batch: Arc<RecordBatch>,
80    /// Row range in `batch` belonging to this chunk.
81    pub rows: Range<usize>,
82    /// Zero-based ordinal of the logical partition in input order.
83    pub partition_ordinal: usize,
84    /// True when this chunk starts a logical partition.
85    pub is_partition_start: bool,
86    /// True when this chunk ends a logical partition.
87    pub is_partition_end: bool,
88}
89
90impl PartitionChunk {
91    /// Return the number of rows in this chunk.
92    pub fn row_count(&self) -> usize {
93        self.rows.end.saturating_sub(self.rows.start)
94    }
95}
96
97#[derive(Clone)]
98struct RawChunk {
99    key: PartitionKey,
100    batch: Arc<RecordBatch>,
101    rows: Range<usize>,
102}
103
104/// Source of projected columnar batches for partition scanning.
105pub trait BatchSource {
106    /// Return the next record batch, or `None` when exhausted.
107    fn next_batch(&mut self) -> Result<Option<RecordBatch>>;
108}
109
110/// In-memory batch source useful for callers that already own batches.
111#[derive(Clone, Debug)]
112pub struct VecBatchSource {
113    batches: std::vec::IntoIter<RecordBatch>,
114}
115
116impl VecBatchSource {
117    /// Create a source from pre-built batches.
118    pub fn new(batches: Vec<RecordBatch>) -> Self {
119        Self {
120            batches: batches.into_iter(),
121        }
122    }
123}
124
125impl BatchSource for VecBatchSource {
126    fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
127        Ok(self.batches.next())
128    }
129}
130
131/// Row-group source backed by [`SegmentReaderV2`].
132pub struct SegmentBatchSource<'a> {
133    reader: &'a SegmentReaderV2,
134    columns: Vec<usize>,
135    row_group_idx: usize,
136}
137
138impl<'a> SegmentBatchSource<'a> {
139    /// Create a row-group source that reads only `columns` from `reader`.
140    pub fn new(reader: &'a SegmentReaderV2, columns: Vec<usize>) -> Self {
141        Self {
142            reader,
143            columns,
144            row_group_idx: 0,
145        }
146    }
147}
148
149impl BatchSource for SegmentBatchSource<'_> {
150    fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
151        if self.row_group_idx >= self.reader.row_group_count() {
152            return Ok(None);
153        }
154        let batch = self
155            .reader
156            .read_row_group_by_index(&self.columns, self.row_group_idx)?;
157        self.row_group_idx += 1;
158        Ok(Some(batch))
159    }
160}
161
162/// Streaming scanner that splits sorted input batches into partition chunks.
163pub struct PartitionScanner<S> {
164    source: S,
165    key_columns: Vec<usize>,
166    current_batch: Option<Arc<RecordBatch>>,
167    row_idx: usize,
168    lookahead: Option<RawChunk>,
169    previous_key: Option<PartitionKey>,
170    completed_partition_keys: HashSet<PartitionKey>,
171    active_partition_ordinal: Option<usize>,
172    next_partition_ordinal: usize,
173}
174
175impl<S: BatchSource> PartitionScanner<S> {
176    /// Create a scanner over batches that are already ordered by partition key.
177    ///
178    /// `key_columns` are indexes in each projected [`RecordBatch`]. An empty
179    /// key list treats the full input stream as a single partition.
180    pub fn new(source: S, key_columns: Vec<usize>) -> Result<Self> {
181        Ok(Self {
182            source,
183            key_columns,
184            current_batch: None,
185            row_idx: 0,
186            lookahead: None,
187            previous_key: None,
188            completed_partition_keys: HashSet::new(),
189            active_partition_ordinal: None,
190            next_partition_ordinal: 0,
191        })
192    }
193
194    /// Return the next partition chunk with start/end boundary markers.
195    pub fn next_chunk(&mut self) -> Result<Option<PartitionChunk>> {
196        let current = match self.take_next_raw_chunk()? {
197            Some(chunk) => chunk,
198            None => return Ok(None),
199        };
200        let next = self.read_raw_chunk()?;
201
202        let is_partition_start = self.previous_key.as_ref() != Some(&current.key);
203        let is_partition_end = next.as_ref().is_none_or(|chunk| chunk.key != current.key);
204        if is_partition_start && self.completed_partition_keys.contains(&current.key) {
205            return Err(Error::InvalidFormat(
206                "partition key reappeared after its partition ended; input must be sorted by partition key"
207                    .into(),
208            ));
209        }
210        let partition_ordinal = if is_partition_start {
211            self.next_partition_ordinal
212        } else {
213            self.active_partition_ordinal
214                .expect("continued partition has an active ordinal")
215        };
216
217        self.previous_key = Some(current.key.clone());
218        if is_partition_end {
219            // Policy A: the scanner is streaming and requires partition-key-sorted input.
220            self.completed_partition_keys.insert(current.key.clone());
221            self.next_partition_ordinal = partition_ordinal.saturating_add(1);
222            self.active_partition_ordinal = None;
223        } else {
224            self.active_partition_ordinal = Some(partition_ordinal);
225        }
226        self.lookahead = next;
227
228        Ok(Some(PartitionChunk {
229            key: current.key,
230            batch: current.batch,
231            rows: current.rows,
232            partition_ordinal,
233            is_partition_start,
234            is_partition_end,
235        }))
236    }
237
238    /// Invoke `visit` once per logical partition.
239    ///
240    /// The partition view owns chunk descriptors. Row data remains shared with
241    /// the underlying batches, so callers can evaluate window or rolling logic
242    /// without copying whole partitions.
243    pub fn for_each_partition<F>(&mut self, mut visit: F) -> Result<()>
244    where
245        F: FnMut(Partition<'_>) -> Result<()>,
246    {
247        let mut chunks = Vec::new();
248        while let Some(chunk) = self.next_chunk()? {
249            chunks.push(chunk);
250            if chunks.last().is_some_and(|chunk| chunk.is_partition_end) {
251                let key = chunks
252                    .first()
253                    .expect("partition has at least one chunk")
254                    .key
255                    .clone();
256                let partition_ordinal = chunks
257                    .first()
258                    .expect("partition has at least one chunk")
259                    .partition_ordinal;
260                visit(Partition {
261                    key,
262                    partition_ordinal,
263                    chunks: &chunks,
264                })?;
265                chunks.clear();
266            }
267        }
268        Ok(())
269    }
270
271    fn take_next_raw_chunk(&mut self) -> Result<Option<RawChunk>> {
272        if self.lookahead.is_some() {
273            return Ok(self.lookahead.take());
274        }
275        self.read_raw_chunk()
276    }
277
278    fn read_raw_chunk(&mut self) -> Result<Option<RawChunk>> {
279        loop {
280            self.ensure_current_batch()?;
281            let batch = match &self.current_batch {
282                Some(batch) => Arc::clone(batch),
283                None => return Ok(None),
284            };
285            let row_count = batch.num_rows();
286            if self.row_idx >= row_count {
287                self.current_batch = None;
288                self.row_idx = 0;
289                continue;
290            }
291
292            validate_batch(&batch, &self.key_columns)?;
293            let start = self.row_idx;
294            let key = PartitionKey::from_batch_row(&batch, &self.key_columns, start)?;
295            self.row_idx += 1;
296            while self.row_idx < row_count
297                && PartitionKey::from_batch_row(&batch, &self.key_columns, self.row_idx)? == key
298            {
299                self.row_idx += 1;
300            }
301            return Ok(Some(RawChunk {
302                key,
303                batch,
304                rows: start..self.row_idx,
305            }));
306        }
307    }
308
309    fn ensure_current_batch(&mut self) -> Result<()> {
310        while self.current_batch.is_none() {
311            let Some(batch) = self.source.next_batch()? else {
312                return Ok(());
313            };
314            validate_batch(&batch, &self.key_columns)?;
315            if batch.num_rows() == 0 {
316                continue;
317            }
318            self.current_batch = Some(Arc::new(batch));
319            self.row_idx = 0;
320        }
321        Ok(())
322    }
323}
324
325/// Borrowed view of one logical partition.
326pub struct Partition<'a> {
327    key: PartitionKey,
328    partition_ordinal: usize,
329    chunks: &'a [PartitionChunk],
330}
331
332impl<'a> Partition<'a> {
333    /// Return the partition key.
334    pub fn key(&self) -> &PartitionKey {
335        &self.key
336    }
337
338    /// Return the partition ordinal in input order.
339    pub fn partition_ordinal(&self) -> usize {
340        self.partition_ordinal
341    }
342
343    /// Return chunks belonging to the partition.
344    pub fn chunks(&self) -> &'a [PartitionChunk] {
345        self.chunks
346    }
347}
348
349fn validate_batch(batch: &RecordBatch, key_columns: &[usize]) -> Result<()> {
350    let row_count = batch.num_rows();
351    for (column_idx, column) in batch.columns.iter().enumerate() {
352        let len = column_len(column);
353        if len != row_count {
354            return Err(Error::InvalidFormat(format!(
355                "record batch column {column_idx} has {len} rows, expected {row_count}"
356            )));
357        }
358    }
359    for &column_idx in key_columns {
360        if column_idx >= batch.columns.len() {
361            return Err(Error::InvalidParameter {
362                param: "key_columns".into(),
363                reason: format!("column index {column_idx} is out of bounds"),
364            });
365        }
366    }
367    Ok(())
368}
369
370fn partition_value(column: &Column, row_idx: usize) -> Result<PartitionValue> {
371    match column {
372        Column::Int64(values) => values
373            .get(row_idx)
374            .copied()
375            .map(PartitionValue::Int64)
376            .ok_or_else(row_out_of_bounds),
377        Column::Float32(values) => values
378            .get(row_idx)
379            .map(|value| PartitionValue::Float32(value.to_bits()))
380            .ok_or_else(row_out_of_bounds),
381        Column::Float64(values) => values
382            .get(row_idx)
383            .map(|value| PartitionValue::Float64(value.to_bits()))
384            .ok_or_else(row_out_of_bounds),
385        Column::Bool(values) => values
386            .get(row_idx)
387            .copied()
388            .map(PartitionValue::Bool)
389            .ok_or_else(row_out_of_bounds),
390        Column::Binary(values) => values
391            .get(row_idx)
392            .cloned()
393            .map(PartitionValue::Binary)
394            .ok_or_else(row_out_of_bounds),
395        Column::Fixed { values, .. } => values
396            .get(row_idx)
397            .cloned()
398            .map(PartitionValue::Fixed)
399            .ok_or_else(row_out_of_bounds),
400    }
401}
402
403fn row_out_of_bounds() -> Error {
404    Error::InvalidFormat("row index out of bounds".into())
405}
406
407fn column_len(column: &Column) -> usize {
408    match column {
409        Column::Int64(values) => values.len(),
410        Column::Float32(values) => values.len(),
411        Column::Float64(values) => values.len(),
412        Column::Bool(values) => values.len(),
413        Column::Binary(values) => values.len(),
414        Column::Fixed { values, .. } => values.len(),
415    }
416}
417
418#[cfg(test)]
419mod tests {
420    use super::*;
421    use crate::columnar::encoding::{Column, LogicalType};
422    use crate::columnar::segment_v2::{
423        ColumnSchema, InMemorySegmentSource, RecordBatch, Schema, SegmentConfigV2, SegmentReaderV2,
424        SegmentWriterV2,
425    };
426
427    fn schema() -> Schema {
428        Schema {
429            columns: vec![
430                ColumnSchema {
431                    name: "partition_key".into(),
432                    logical_type: LogicalType::Int64,
433                    nullable: false,
434                    fixed_len: None,
435                },
436                ColumnSchema {
437                    name: "value".into(),
438                    logical_type: LogicalType::Int64,
439                    nullable: false,
440                    fixed_len: None,
441                },
442            ],
443        }
444    }
445
446    fn batch(keys: Vec<i64>, values: Vec<i64>) -> RecordBatch {
447        RecordBatch::new(
448            schema(),
449            vec![Column::Int64(keys), Column::Int64(values)],
450            vec![None, None],
451        )
452    }
453
454    fn chunk_values(chunk: &PartitionChunk) -> Vec<i64> {
455        let Column::Int64(values) = &chunk.batch.columns[1] else {
456            panic!("expected int64 values");
457        };
458        values[chunk.rows.clone()].to_vec()
459    }
460
461    fn key_i64(chunk: &PartitionChunk) -> i64 {
462        match &chunk.key.values()[0] {
463            PartitionValue::Int64(value) => *value,
464            other => panic!("expected int64 key, got {other:?}"),
465        }
466    }
467
468    #[test]
469    fn preserves_partition_boundaries_and_order_across_batches() {
470        let source = VecBatchSource::new(vec![
471            batch(vec![1, 1, 2], vec![10, 11, 20]),
472            batch(vec![2, 3], vec![21, 30]),
473        ]);
474        let mut scan = PartitionScanner::new(source, vec![0]).unwrap();
475
476        let mut chunks = Vec::new();
477        while let Some(chunk) = scan.next_chunk().unwrap() {
478            chunks.push(chunk);
479        }
480
481        assert_eq!(chunks.len(), 4);
482        assert_eq!(
483            chunks
484                .iter()
485                .map(|chunk| {
486                    (
487                        key_i64(chunk),
488                        chunk.rows.clone(),
489                        chunk.is_partition_start,
490                        chunk.is_partition_end,
491                        chunk.partition_ordinal,
492                        chunk_values(chunk),
493                    )
494                })
495                .collect::<Vec<_>>(),
496            vec![
497                (1, 0..2, true, true, 0, vec![10, 11]),
498                (2, 2..3, true, false, 1, vec![20]),
499                (2, 0..1, false, true, 1, vec![21]),
500                (3, 1..2, true, true, 2, vec![30]),
501            ]
502        );
503    }
504
505    #[test]
506    fn skips_empty_batches_without_emitting_empty_chunks() {
507        let source = VecBatchSource::new(vec![
508            batch(Vec::new(), Vec::new()),
509            batch(vec![7], vec![70]),
510            batch(Vec::new(), Vec::new()),
511        ]);
512        let mut scan = PartitionScanner::new(source, vec![0]).unwrap();
513
514        let first = scan.next_chunk().unwrap().expect("non-empty partition");
515        assert_eq!(key_i64(&first), 7);
516        assert_eq!(chunk_values(&first), vec![70]);
517        assert!(first.is_partition_start);
518        assert!(first.is_partition_end);
519        assert!(scan.next_chunk().unwrap().is_none());
520
521        let mut empty =
522            PartitionScanner::new(VecBatchSource::new(vec![batch(vec![], vec![])]), vec![0])
523                .unwrap();
524        assert!(empty.next_chunk().unwrap().is_none());
525    }
526
527    #[test]
528    fn streams_skewed_partition_sizes_without_losing_boundaries() {
529        let mut batches = vec![batch(vec![1], vec![10])];
530        for offset in (0..10_000).step_by(1024) {
531            let len = (10_000 - offset).min(1024);
532            batches.push(batch(
533                vec![2; len],
534                (offset..offset + len).map(|v| v as i64).collect(),
535            ));
536        }
537        batches.push(batch(vec![3], vec![30]));
538
539        let mut scan = PartitionScanner::new(VecBatchSource::new(batches), vec![0]).unwrap();
540        let mut seen = Vec::new();
541        while let Some(chunk) = scan.next_chunk().unwrap() {
542            seen.push((
543                key_i64(&chunk),
544                chunk.row_count(),
545                chunk.is_partition_start,
546                chunk.is_partition_end,
547                chunk.partition_ordinal,
548            ));
549        }
550
551        assert_eq!(seen.first(), Some(&(1, 1, true, true, 0)));
552        assert_eq!(seen.last(), Some(&(3, 1, true, true, 2)));
553        let skewed_rows: usize = seen
554            .iter()
555            .filter(|(key, _, _, _, _)| *key == 2)
556            .map(|(_, rows, _, _, _)| *rows)
557            .sum();
558        assert_eq!(skewed_rows, 10_000);
559        assert_eq!(
560            seen.iter()
561                .filter(|(key, _, is_start, _, _)| *key == 2 && *is_start)
562                .count(),
563            1
564        );
565        assert_eq!(
566            seen.iter()
567                .filter(|(key, _, _, is_end, _)| *key == 2 && *is_end)
568                .count(),
569            1
570        );
571    }
572
573    #[test]
574    fn handles_thousands_of_partitions_in_input_order() {
575        let partition_count = 5_000usize;
576        let batches = (0..partition_count)
577            .collect::<Vec<_>>()
578            .chunks(257)
579            .map(|chunk| {
580                batch(
581                    chunk.iter().map(|v| *v as i64).collect(),
582                    chunk.iter().map(|v| (*v as i64) * 10).collect(),
583                )
584            })
585            .collect::<Vec<_>>();
586
587        let mut scan = PartitionScanner::new(VecBatchSource::new(batches), vec![0]).unwrap();
588        let mut count = 0usize;
589        while let Some(chunk) = scan.next_chunk().unwrap() {
590            assert_eq!(key_i64(&chunk), count as i64);
591            assert_eq!(chunk.partition_ordinal, count);
592            assert!(chunk.is_partition_start);
593            assert!(chunk.is_partition_end);
594            count += 1;
595        }
596
597        assert_eq!(count, partition_count);
598    }
599
600    #[test]
601    fn rejects_non_contiguous_partition_key_reappearance() {
602        let source = VecBatchSource::new(vec![batch(vec![1, 1, 2, 1], vec![10, 11, 20, 12])]);
603        let mut scan = PartitionScanner::new(source, vec![0]).unwrap();
604
605        assert!(scan.next_chunk().unwrap().is_some());
606        assert!(scan.next_chunk().unwrap().is_some());
607
608        match scan.next_chunk().unwrap_err() {
609            Error::InvalidFormat(message) => {
610                assert!(message.contains("partition key"));
611                assert!(message.contains("sorted"));
612            }
613            other => panic!("expected InvalidFormat, got {other:?}"),
614        }
615    }
616
617    #[test]
618    fn reads_partitions_from_segment_reader_row_groups() {
619        let config = SegmentConfigV2 {
620            row_group_size: 2,
621            ..SegmentConfigV2::default()
622        };
623        let mut writer = SegmentWriterV2::new(config);
624        writer
625            .write_batch(batch(vec![1, 1, 2, 2, 3], vec![10, 11, 20, 21, 30]))
626            .unwrap();
627        let segment = writer.finish().unwrap();
628        let reader =
629            SegmentReaderV2::open(Box::new(InMemorySegmentSource::new(segment.data))).unwrap();
630
631        let source = SegmentBatchSource::new(&reader, vec![0, 1]);
632        let mut scan = PartitionScanner::new(source, vec![0]).unwrap();
633        let mut partitions = Vec::new();
634        scan.for_each_partition(|partition| {
635            let row_count: usize = partition
636                .chunks()
637                .iter()
638                .map(PartitionChunk::row_count)
639                .sum();
640            partitions.push((
641                partition.key().clone(),
642                partition.partition_ordinal(),
643                row_count,
644            ));
645            Ok(())
646        })
647        .unwrap();
648
649        assert_eq!(partitions.len(), 3);
650        assert_eq!(partitions[0].0.values(), &[PartitionValue::Int64(1)]);
651        assert_eq!(partitions[0].1, 0);
652        assert_eq!(partitions[0].2, 2);
653        assert_eq!(partitions[1].0.values(), &[PartitionValue::Int64(2)]);
654        assert_eq!(partitions[1].1, 1);
655        assert_eq!(partitions[1].2, 2);
656        assert_eq!(partitions[2].0.values(), &[PartitionValue::Int64(3)]);
657        assert_eq!(partitions[2].1, 2);
658        assert_eq!(partitions[2].2, 1);
659    }
660}
alopex_core/dataframe/partition_scan.rs

alopex_core/dataframe/
partition_scan.rs
