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 document = roxmltree::Document::parse(input)?;
57        let row = parse(&document, &self.schema)?;
58        for ((column, value), field) in self.columns.iter().zip(&row).zip(&self.schema.fields) {
59            validate_capacity(column.as_ref(), field, &[value])?;
60        }
61        for ((column, value), field) in self.columns.iter_mut().zip(row).zip(&self.schema.fields) {
62            append(column.as_mut(), field, value);
63        }
64        self.rows += 1;
65        Ok(())
66    }
67
68    /// Pending row count.
69    pub fn len(&self) -> usize {
70        self.rows
71    }
72
73    /// Whether the batch is empty.
74    pub fn is_empty(&self) -> bool {
75        self.rows == 0
76    }
77
78    /// Whether the row limit has been reached.
79    pub fn is_full(&self) -> bool {
80        self.rows == self.batch_size
81    }
82
83    /// Projected output schema.
84    pub fn schema(&self) -> SchemaRef {
85        self.schema.arrow_schema()
86    }
87
88    /// Flushes pending rows; returns `None` if empty.
89    pub fn flush(&mut self) -> Result<Option<RecordBatch>> {
90        if self.is_empty() {
91            return Ok(None);
92        }
93        let columns = self
94            .columns
95            .iter_mut()
96            .map(|column| column.finish())
97            .collect();
98        let options = RecordBatchOptions::new().with_row_count(Some(self.rows));
99        let batch = RecordBatch::try_new_with_options(self.schema(), columns, &options)?;
100        self.rows = 0;
101        Ok(Some(batch))
102    }
103}
xml_arrow/batch.rs

xml_arrow/
batch.rs
