acta/
batch.rs

1//! Native logical rows grouped by one Acta data block.
2
3use std::sync::Arc;
4
5use crate::array::Array;
6use crate::error::{Error, ErrorContext, Result};
7use crate::schema::{Column, LogicalType, Schema};
8
9/// A decoded block with one logical array per schema column.
10#[derive(Debug, Clone, PartialEq)]
11pub struct RecordBatch {
12    schema: Arc<Schema>,
13    columns: Vec<Array>,
14    row_count: usize,
15}
16
17impl RecordBatch {
18    /// Build a batch from arrays that already have one position per row.
19    ///
20    /// The row count is given rather than taken from the first array, so a
21    /// batch keeps its shape independently of how many columns it has, and
22    /// every array is checked against it.
23    pub fn try_new(schema: Arc<Schema>, columns: Vec<Array>, row_count: usize) -> Result<Self> {
24        if columns.len() != schema.column_count() {
25            return Err(invalid("the column count does not match the schema"));
26        }
27        for (column, array) in schema.columns().iter().zip(&columns) {
28            if array.len() != row_count {
29                return Err(invalid(format!(
30                    "column {} has {} values, expected {row_count}",
31                    column.name(),
32                    array.len()
33                )));
34            }
35            if !holds(column, array) {
36                return Err(invalid(format!(
37                    "column {} was decoded as {array:?}, which is not {:?}",
38                    column.name(),
39                    column.logical_type()
40                )));
41            }
42            if let Some(validity) = validity(array) {
43                if validity.len() != row_count {
44                    return Err(invalid(format!(
45                        "column {} has a validity bitmap with {} bits, expected {row_count}",
46                        column.name(),
47                        validity.len()
48                    )));
49                }
50                if !column.is_nullable() {
51                    return Err(invalid(format!(
52                        "non-nullable column {} has a validity bitmap",
53                        column.name()
54                    )));
55                }
56            }
57        }
58        Ok(Self {
59            schema,
60            columns,
61            row_count,
62        })
63    }
64
65    /// The schema shared by this batch.
66    pub fn schema(&self) -> &Schema {
67        &self.schema
68    }
69
70    /// The decoded arrays in schema order.
71    pub fn columns(&self) -> &[Array] {
72        &self.columns
73    }
74
75    /// The number of logical rows in the block.
76    pub fn row_count(&self) -> usize {
77        self.row_count
78    }
79
80    /// Find a decoded array by its schema column name.
81    pub fn column_by_name(&self, name: &str) -> Option<&Array> {
82        self.schema
83            .columns()
84            .iter()
85            .position(|column| column.name() == name)
86            .and_then(|index| self.columns.get(index))
87    }
88
89    /// Get a decoded array by schema position.
90    pub fn column(&self, index: usize) -> Option<&Array> {
91        self.columns.get(index)
92    }
93
94    /// Copy a row range for the writer's bounded block assembler.
95    ///
96    /// The writer slices an appended batch only where it crosses a block
97    /// boundary, so every row is copied at most once on its way into a block.
98    pub(crate) fn slice(&self, start: usize, end: usize) -> Self {
99        assert!(start <= end && end <= self.row_count);
100        Self {
101            schema: Arc::clone(&self.schema),
102            columns: self
103                .columns
104                .iter()
105                .map(|column| column.slice(start, end))
106                .collect(),
107            row_count: end - start,
108        }
109    }
110
111    /// Select rows in file order for a scan filter. This remains private to
112    /// the reader so callers do not receive a row-oriented mutation API.
113    pub(crate) fn take(&self, indices: &[usize]) -> Self {
114        assert!(indices.iter().all(|&index| index < self.row_count));
115        Self {
116            schema: Arc::clone(&self.schema),
117            columns: self
118                .columns
119                .iter()
120                .map(|column| column.take(indices))
121                .collect(),
122            row_count: indices.len(),
123        }
124    }
125}
126
127/// Whether an array is the representation its schema column calls for.
128fn holds(column: &Column, array: &Array) -> bool {
129    matches!(
130        (column.logical_type(), array),
131        (LogicalType::Bool, Array::Bool(_))
132            | (LogicalType::Int8, Array::Int8(_))
133            | (LogicalType::Int16, Array::Int16(_))
134            | (LogicalType::Int32, Array::Int32(_))
135            | (LogicalType::Int64, Array::Int64(_))
136            | (LogicalType::UInt8, Array::UInt8(_))
137            | (LogicalType::UInt16, Array::UInt16(_))
138            | (LogicalType::UInt32, Array::UInt32(_))
139            | (LogicalType::UInt64, Array::UInt64(_))
140            | (LogicalType::Float32, Array::Float32(_))
141            | (LogicalType::Float64, Array::Float64(_))
142            | (LogicalType::Decimal { .. }, Array::Decimal(_))
143            | (LogicalType::Timestamp { .. }, Array::Timestamp(_))
144            | (LogicalType::Utf8, Array::Utf8(_))
145            | (LogicalType::Categorical { .. }, Array::Categorical(_))
146            | (LogicalType::Binary, Array::Binary(_))
147            | (LogicalType::FixedBinary { .. }, Array::FixedBinary(_))
148            | (LogicalType::Date32, Array::Date32(_))
149    )
150}
151
152fn validity(array: &Array) -> Option<&[bool]> {
153    match array {
154        Array::Bool(array) => array.validity(),
155        Array::Int8(array) => array.validity(),
156        Array::Int16(array) => array.validity(),
157        Array::Int32(array) => array.validity(),
158        Array::Int64(array) => array.validity(),
159        Array::UInt8(array) => array.validity(),
160        Array::UInt16(array) => array.validity(),
161        Array::UInt32(array) => array.validity(),
162        Array::UInt64(array) => array.validity(),
163        Array::Float32(array) => array.validity(),
164        Array::Float64(array) => array.validity(),
165        Array::Decimal(array) => array.validity(),
166        Array::Timestamp(array) => array.validity(),
167        Array::Utf8(array) | Array::Categorical(array) => array.validity(),
168        Array::Binary(array) | Array::FixedBinary(array) => array.validity(),
169        Array::Date32(array) => array.validity(),
170    }
171}
172
173fn invalid(message: impl Into<String>) -> Error {
174    Error::invalid_argument(message).with_context(ErrorContext::Payload)
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use crate::array::{PrimitiveArray, ScalarValue, Utf8Array};
181
182    fn schema() -> Arc<Schema> {
183        Arc::new(Schema::new(
184            1,
185            vec![
186                Column::new(1, "value", LogicalType::Int64, true),
187                Column::new(2, "label", LogicalType::Utf8, false),
188            ],
189            None,
190        ))
191    }
192
193    fn batch() -> RecordBatch {
194        RecordBatch::try_new(
195            schema(),
196            vec![
197                Array::Int64(PrimitiveArray::new(
198                    vec![1, 2, 3, 4],
199                    Some(vec![true, false, true, true]),
200                )),
201                Array::Utf8(Utf8Array::new(
202                    vec!["a".into(), "b".into(), "c".into(), "d".into()],
203                    None,
204                )),
205            ],
206            4,
207        )
208        .expect("a well-formed batch")
209    }
210
211    #[test]
212    fn a_slice_reports_the_rows_of_its_range() {
213        assert_eq!(batch().slice(1, 3).row_count(), 2);
214    }
215
216    #[test]
217    fn a_slice_keeps_every_column() {
218        assert_eq!(batch().slice(1, 3).columns().len(), 2);
219    }
220
221    #[test]
222    fn a_slice_keeps_its_rows_aligned_across_columns() {
223        let sliced = batch().slice(1, 3);
224
225        assert_eq!(
226            (
227                sliced.column(0).unwrap().value_at(1),
228                sliced.column(1).unwrap().value_at(1)
229            ),
230            (Some(ScalarValue::Int64(3)), Some(ScalarValue::Utf8("c")))
231        );
232    }
233
234    #[test]
235    fn a_slice_keeps_the_nulls_of_its_range() {
236        let sliced = batch().slice(1, 3);
237
238        assert_eq!(sliced.column(0).unwrap().value_at(0), None);
239    }
240
241    #[test]
242    fn a_slice_shares_the_schema_it_came_from() {
243        assert_eq!(batch().slice(0, 1).schema(), schema().as_ref());
244    }
245
246    #[test]
247    fn slicing_a_whole_batch_reproduces_it() {
248        assert_eq!(batch().slice(0, 4), batch());
249    }
250}
acta/batch.rs

acta/
batch.rs
