xml_arrow/reader/
mod.rs

1use crate::{NamespaceBinding, XmlSchema};
2use arrow_array::RecordBatch;
3use arrow_schema::SchemaRef;
4use futures_util::{
5    io::AsyncBufRead,
6    stream::{self, BoxStream},
7    StreamExt, TryStreamExt,
8};
9
10mod error;
11mod inference;
12mod options;
13mod state;
14mod validation;
15
16pub use error::{Limit, ReadError, Result as ReaderResult};
17pub use options::{InferenceOptions, ReaderOptions};
18
19use error::Result;
20use state::ReaderState;
21
22#[derive(PartialEq, Eq)]
23struct Name {
24    uri: Option<String>,
25    local: String,
26}
27
28/// One selected XML record, ready for independent Arrow conversion.
29#[derive(Debug)]
30pub struct XmlRecord {
31    /// Complete record element with inherited namespace declarations attached.
32    pub xml: String,
33    /// Start of the record in the original document, in bytes.
34    pub byte_offset: u64,
35}
36
37/// Stream complete XML records without converting field values to Arrow.
38///
39/// Uses the same path matching, namespaces, document validation and input limits
40/// as [`XmlReader`]. Construction does not read input. The first error ends the
41/// stream. Batch size and projection are validated but do not change record XML;
42/// consumers can use [`crate::XmlBatchBuilder::new_projected`] for conversion.
43pub fn read_records<'a, R: AsyncBufRead + Unpin + Send + 'a>(
44    input: R,
45    schema: XmlSchema,
46    record_path: &str,
47    options: ReaderOptions,
48) -> Result<BoxStream<'a, Result<XmlRecord>>> {
49    options.validate(&schema, record_path)?;
50    // Keep namespace bindings without allocating unused Arrow column builders.
51    let schema = schema.project(&[]).map_err(ReadError::configuration)?;
52    let options = ReaderOptions {
53        projection: None,
54        ..options
55    };
56    let state = ReaderState::new(input, schema, record_path, options)?;
57    Ok(record_stream(state))
58}
59
60/// Infer reusable XML mappings from a bounded sample, then replay every raw record.
61///
62/// The returned schema includes namespace bindings and all sampled fields, even
63/// when `options.projection` selects fewer columns. Sampled records and any byte-
64/// limit lookahead precede the unread input, preserving original byte offsets.
65/// Limits, sampling and document validation match [`XmlReader::infer`]; conversion
66/// is left to independent [`crate::XmlBatchBuilder`] instances.
67pub async fn infer_records<'a, R: AsyncBufRead + Unpin + Send + 'a>(
68    input: R,
69    record_path: &str,
70    bindings: &[NamespaceBinding<'_>],
71    options: ReaderOptions,
72    sampling: InferenceOptions,
73) -> Result<(XmlSchema, BoxStream<'a, Result<XmlRecord>>)> {
74    let (state, schema) =
75        ReaderState::infer(input, record_path, bindings, options, sampling).await?;
76    Ok((schema, record_stream(state)))
77}
78
79fn record_stream<'a, R: AsyncBufRead + Unpin + Send + 'a>(
80    state: ReaderState<R>,
81) -> BoxStream<'a, Result<XmlRecord>> {
82    stream::try_unfold((state, Vec::new()), |(mut state, mut buffer)| async move {
83        Ok(state
84            .next_replay_record(&mut buffer)
85            .await?
86            .map(|record| (record, (state, buffer))))
87    })
88    .boxed()
89}
90
91/// Reads XML records selected by an absolute path into Arrow batches.
92pub struct XmlReader<'a> {
93    batches: BoxStream<'a, Result<RecordBatch>>,
94    // Keep the schema available after EOF.
95    schema: SchemaRef,
96}
97
98impl<'a> XmlReader<'a> {
99    /// Creates a reader without reading input. Path prefixes use schema bindings.
100    pub fn new<R: AsyncBufRead + Unpin + Send + 'a>(
101        input: R,
102        schema: XmlSchema,
103        record_path: &str,
104        options: ReaderOptions,
105    ) -> Result<Self> {
106        Ok(Self::from_state(ReaderState::new(
107            input,
108            schema,
109            record_path,
110            options,
111        )?))
112    }
113
114    /// Infers a fixed schema from a bounded sample, then replays sampled records.
115    /// Bindings resolve record-path prefixes; other namespaces get generated prefixes.
116    pub async fn infer<R: AsyncBufRead + Unpin + Send + 'a>(
117        input: R,
118        record_path: &str,
119        bindings: &[NamespaceBinding<'_>],
120        options: ReaderOptions,
121        sampling: InferenceOptions,
122    ) -> Result<Self> {
123        let (state, schema) =
124            ReaderState::infer(input, record_path, bindings, options, sampling).await?;
125        Ok(Self::from_state(state.with_schema(schema)?))
126    }
127
128    fn from_state<R: AsyncBufRead + Unpin + Send + 'a>(state: ReaderState<R>) -> Self {
129        let schema = state.schema();
130        let batches = stream::try_unfold(state, |mut state| async move {
131            Ok(state.next_batch().await?.map(|batch| (batch, state)))
132        })
133        .boxed();
134        Self { batches, schema }
135    }
136
137    /// Projected output schema.
138    pub fn schema(&self) -> SchemaRef {
139        self.schema.clone()
140    }
141
142    /// Returns a batch or EOF. The first error discards pending rows and ends input.
143    /// Cancelled calls can resume; calls after an error return `None`.
144    pub async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
145        self.batches.try_next().await
146    }
147
148    /// Consumes the reader as a batch stream.
149    pub fn into_stream(self) -> BoxStream<'a, Result<RecordBatch>> {
150        self.batches
151    }
152}
xml_arrow/reader/mod.rs

xml_arrow/reader/
mod.rs
