Skip to main content

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(
58        stream::try_unfold((state, Vec::new()), |(mut state, mut buffer)| async move {
59            Ok(state
60                .next_record(&mut buffer)
61                .await?
62                .map(|record| (record, (state, buffer))))
63        })
64        .boxed(),
65    )
66}
67
68/// Reads XML records selected by an absolute path into Arrow batches.
69pub struct XmlReader<'a> {
70    batches: BoxStream<'a, Result<RecordBatch>>,
71    // Keep the schema available after EOF.
72    schema: SchemaRef,
73}
74
75impl<'a> XmlReader<'a> {
76    /// Creates a reader without reading input. Path prefixes use schema bindings.
77    pub fn new<R: AsyncBufRead + Unpin + Send + 'a>(
78        input: R,
79        schema: XmlSchema,
80        record_path: &str,
81        options: ReaderOptions,
82    ) -> Result<Self> {
83        Ok(Self::from_state(ReaderState::new(
84            input,
85            schema,
86            record_path,
87            options,
88        )?))
89    }
90
91    /// Infers a fixed schema from a bounded sample, then replays sampled records.
92    /// Bindings resolve record-path prefixes; other namespaces get generated prefixes.
93    pub async fn infer<R: AsyncBufRead + Unpin + Send + 'a>(
94        input: R,
95        record_path: &str,
96        bindings: &[NamespaceBinding<'_>],
97        options: ReaderOptions,
98        sampling: InferenceOptions,
99    ) -> Result<Self> {
100        Ok(Self::from_state(
101            ReaderState::infer(input, record_path, bindings, options, sampling).await?,
102        ))
103    }
104
105    fn from_state<R: AsyncBufRead + Unpin + Send + 'a>(state: ReaderState<R>) -> Self {
106        let schema = state.schema();
107        let batches = stream::try_unfold(state, |mut state| async move {
108            Ok(state.next_batch().await?.map(|batch| (batch, state)))
109        })
110        .boxed();
111        Self { batches, schema }
112    }
113
114    /// Projected output schema.
115    pub fn schema(&self) -> SchemaRef {
116        self.schema.clone()
117    }
118
119    /// Returns a batch or EOF. The first error discards pending rows and ends input.
120    /// Cancelled calls can resume; calls after an error return `None`.
121    pub async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
122        self.batches.try_next().await
123    }
124
125    /// Consumes the reader as a batch stream.
126    pub fn into_stream(self) -> BoxStream<'a, Result<RecordBatch>> {
127        self.batches
128    }
129}