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/// Reads XML records selected by an absolute path into Arrow batches.
29pub struct XmlReader<'a> {
30    batches: BoxStream<'a, Result<RecordBatch>>,
31    // Keep the schema available after EOF.
32    schema: SchemaRef,
33}
34
35impl<'a> XmlReader<'a> {
36    /// Creates a reader without reading input. Path prefixes use schema bindings.
37    pub fn new<R: AsyncBufRead + Unpin + Send + 'a>(
38        input: R,
39        schema: XmlSchema,
40        record_path: &str,
41        options: ReaderOptions,
42    ) -> Result<Self> {
43        Ok(Self::from_state(ReaderState::new(
44            input,
45            schema,
46            record_path,
47            options,
48        )?))
49    }
50
51    /// Infers a fixed schema from a bounded sample, then replays sampled records.
52    /// Bindings resolve record-path prefixes; other namespaces get generated prefixes.
53    pub async fn infer<R: AsyncBufRead + Unpin + Send + 'a>(
54        input: R,
55        record_path: &str,
56        bindings: &[NamespaceBinding<'_>],
57        options: ReaderOptions,
58        sampling: InferenceOptions,
59    ) -> Result<Self> {
60        Ok(Self::from_state(
61            ReaderState::infer(input, record_path, bindings, options, sampling).await?,
62        ))
63    }
64
65    fn from_state<R: AsyncBufRead + Unpin + Send + 'a>(state: ReaderState<R>) -> Self {
66        let schema = state.schema();
67        let batches = stream::try_unfold(state, |mut state| async move {
68            Ok(state.next_batch().await?.map(|batch| (batch, state)))
69        })
70        .boxed();
71        Self { batches, schema }
72    }
73
74    /// Projected output schema.
75    pub fn schema(&self) -> SchemaRef {
76        self.schema.clone()
77    }
78
79    /// Returns a batch or EOF. The first error discards pending rows and ends input.
80    /// Cancelled calls can resume; calls after an error return `None`.
81    pub async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
82        self.batches.try_next().await
83    }
84
85    /// Consumes the reader as a batch stream.
86    pub fn into_stream(self) -> BoxStream<'a, Result<RecordBatch>> {
87        self.batches
88    }
89}