Skip to main content

uqa_execution/batch/
batches.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use super::{OwnedPhysicalRow, PhysicalRow, ResultRow, RowFragment, RowSchema, DEFAULT_BATCH_SIZE};
8
9/// A schema and bounded vector of physical rows flowing between operators.
10#[derive(Debug, Clone, PartialEq)]
11pub struct Batch {
12    pub schema: RowSchema,
13    pub rows: Vec<PhysicalRow>,
14}
15
16impl Batch {
17    /// Compatibility constructor for named rows entering the physical engine. The resulting batch is positional immediately; maps do not flow to the next operator.
18    pub fn new(schema: RowSchema, rows: Vec<ResultRow>) -> Self {
19        let rows = rows
20            .into_iter()
21            .map(|row| PhysicalRow::from_result_row(&schema, row))
22            .collect();
23        Self { schema, rows }
24    }
25
26    pub fn from_physical_rows(schema: RowSchema, rows: Vec<PhysicalRow>) -> Self {
27        debug_assert!(rows.iter().all(|row| {
28            row.fragments.iter().map(RowFragment::len).sum::<usize>() == schema.physical_width()
29        }));
30        Self { schema, rows }
31    }
32
33    pub fn empty(schema: RowSchema) -> Self {
34        Self {
35            schema,
36            rows: Vec::new(),
37        }
38    }
39
40    pub fn len(&self) -> usize {
41        self.rows.len()
42    }
43
44    pub fn is_empty(&self) -> bool {
45        self.rows.is_empty()
46    }
47
48    pub fn into_result_rows(self) -> Vec<ResultRow> {
49        let schema = self.schema;
50        if schema.index.cold.identity_layout {
51            return self
52                .rows
53                .into_iter()
54                .map(|row| schema.materialize_identity_result_row(row))
55                .collect();
56        }
57        schema.materialize_remapped_result_rows(self.rows)
58    }
59
60    /// Consume a batch without materializing named maps.
61    pub fn into_owned_rows(self) -> Vec<OwnedPhysicalRow> {
62        let schema = self.schema;
63        self.rows
64            .into_iter()
65            .map(|row| OwnedPhysicalRow::new(schema.clone(), row))
66            .collect()
67    }
68
69    /// Split named rows into batches of at most [`DEFAULT_BATCH_SIZE`].
70    pub fn chunked(schema: RowSchema, rows: Vec<ResultRow>) -> Vec<Batch> {
71        if rows.is_empty() {
72            return vec![Batch::empty(schema)];
73        }
74        let mut out = Vec::with_capacity(rows.len().div_ceil(DEFAULT_BATCH_SIZE));
75        let mut buf = Vec::with_capacity(DEFAULT_BATCH_SIZE);
76        for row in rows {
77            buf.push(row);
78            if buf.len() == DEFAULT_BATCH_SIZE {
79                out.push(Batch::new(schema.clone(), std::mem::take(&mut buf)));
80                buf.reserve(DEFAULT_BATCH_SIZE);
81            }
82        }
83        if !buf.is_empty() {
84            out.push(Batch::new(schema, buf));
85        }
86        out
87    }
88}