ahri_tre_tabular/
table.rs

1use std::sync::Arc;
2
3use arrow_array::RecordBatch;
4use arrow_schema::{Schema, SchemaRef};
5
6use crate::{Result, TabularError};
7
8#[derive(Debug, Clone)]
9pub struct ArrowTable {
10    schema: SchemaRef,
11    batches: Vec<RecordBatch>,
12}
13
14impl ArrowTable {
15    pub fn try_new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Result<Self> {
16        validate_batch_schemas(&schema, &batches)?;
17        Ok(Self { schema, batches })
18    }
19
20    pub fn from_batches(batches: Vec<RecordBatch>) -> Result<Self> {
21        let schema = batches
22            .first()
23            .map(RecordBatch::schema)
24            .ok_or(TabularError::EmptyBatches)?;
25        Self::try_new(schema, batches)
26    }
27
28    pub fn empty(schema: Schema) -> Self {
29        Self {
30            schema: Arc::new(schema),
31            batches: Vec::new(),
32        }
33    }
34
35    pub fn schema(&self) -> &SchemaRef {
36        &self.schema
37    }
38
39    pub fn batches(&self) -> &[RecordBatch] {
40        &self.batches
41    }
42
43    pub fn into_batches(self) -> Vec<RecordBatch> {
44        self.batches
45    }
46
47    pub fn with_schema(&self, schema: Schema) -> Result<Self> {
48        let schema = Arc::new(schema);
49        let batches = self
50            .batches
51            .iter()
52            .map(|batch| RecordBatch::try_new(schema.clone(), batch.columns().to_vec()))
53            .collect::<std::result::Result<Vec<_>, _>>()?;
54        Self::try_new(schema, batches)
55    }
56
57    pub fn row_count(&self) -> usize {
58        self.batches.iter().map(RecordBatch::num_rows).sum()
59    }
60
61    pub fn column_count(&self) -> usize {
62        self.schema.fields().len()
63    }
64
65    pub fn is_empty(&self) -> bool {
66        self.row_count() == 0
67    }
68}
69
70fn validate_batch_schemas(schema: &SchemaRef, batches: &[RecordBatch]) -> Result<()> {
71    for (index, batch) in batches.iter().enumerate() {
72        if batch.schema().as_ref() != schema.as_ref() {
73            return Err(TabularError::SchemaMismatch { batch_index: index });
74        }
75    }
76    Ok(())
77}
ahri_tre_tabular/table.rs

ahri_tre_tabular/
table.rs
