uqa_execution/batch/
batches.rs1use super::{OwnedPhysicalRow, PhysicalRow, ResultRow, RowFragment, RowSchema, DEFAULT_BATCH_SIZE};
8
9#[derive(Debug, Clone, PartialEq)]
11pub struct Batch {
12 pub schema: RowSchema,
13 pub rows: Vec<PhysicalRow>,
14}
15
16impl Batch {
17 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 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 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}