xml_arrow/
batch.rs

1use arrow_array::{
2    builder::{make_builder, ArrayBuilder},
3    RecordBatch, RecordBatchOptions,
4};
5use arrow_schema::SchemaRef;
6
7use crate::{
8    builders::{append, validate_capacity},
9    parse::parse,
10    Error, Result, XmlSchema,
11};
12
13/// Builds Arrow batches from complete XML record elements.
14pub struct XmlBatchBuilder {
15    schema: XmlSchema,
16    columns: Vec<Box<dyn ArrayBuilder>>,
17    batch_size: usize,
18    rows: usize,
19}
20
21impl XmlBatchBuilder {
22    /// Creates a builder with a positive row limit.
23    pub fn new(schema: XmlSchema, batch_size: usize) -> Result<Self> {
24        if batch_size == 0 {
25            return Err(Error::Configuration("batch_size must be positive".into()));
26        }
27        let columns = schema
28            .fields
29            .iter()
30            .map(|f| make_builder(f.field.data_type(), batch_size.min(1024)))
31            .collect();
32        Ok(Self {
33            schema,
34            columns,
35            batch_size,
36            rows: 0,
37        })
38    }
39
40    /// Selects columns in order; duplicate or invalid indices fail.
41    /// Only selected fields are converted; an empty projection counts rows.
42    /// XML syntax is always checked.
43    pub fn new_projected(
44        schema: XmlSchema,
45        batch_size: usize,
46        projection: &[usize],
47    ) -> Result<Self> {
48        Self::new(schema.project(projection)?, batch_size)
49    }
50
51    /// Appends a complete XML record. Errors leave pending rows unchanged.
52    pub fn append_xml(&mut self, input: &str) -> Result<()> {
53        if self.is_full() {
54            return Err(Error::BatchFull);
55        }
56        let row = parse(input, &self.schema)?;
57        for ((column, value), field) in self.columns.iter().zip(&row).zip(&self.schema.fields) {
58            validate_capacity(column.as_ref(), field, &[value])?;
59        }
60        for ((column, value), field) in self.columns.iter_mut().zip(row).zip(&self.schema.fields) {
61            append(column.as_mut(), field, value);
62        }
63        self.rows += 1;
64        Ok(())
65    }
66
67    /// Pending row count.
68    pub fn len(&self) -> usize {
69        self.rows
70    }
71
72    /// Whether the batch is empty.
73    pub fn is_empty(&self) -> bool {
74        self.rows == 0
75    }
76
77    /// Whether the row limit has been reached.
78    pub fn is_full(&self) -> bool {
79        self.rows == self.batch_size
80    }
81
82    /// Projected output schema.
83    pub fn schema(&self) -> SchemaRef {
84        self.schema.arrow_schema()
85    }
86
87    /// Flushes pending rows; returns `None` if empty.
88    pub fn flush(&mut self) -> Result<Option<RecordBatch>> {
89        if self.is_empty() {
90            return Ok(None);
91        }
92        let columns = self
93            .columns
94            .iter_mut()
95            .map(|column| column.finish())
96            .collect();
97        let options = RecordBatchOptions::new().with_row_count(Some(self.rows));
98        let batch = RecordBatch::try_new_with_options(self.schema(), columns, &options)?;
99        self.rows = 0;
100        Ok(Some(batch))
101    }
102}
xml_arrow/batch.rs

xml_arrow/
batch.rs
