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#[derive(Debug)]
30pub struct XmlRecord {
31 pub xml: String,
33 pub byte_offset: u64,
35}
36
37pub 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 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
68pub struct XmlReader<'a> {
70 batches: BoxStream<'a, Result<RecordBatch>>,
71 schema: SchemaRef,
73}
74
75impl<'a> XmlReader<'a> {
76 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 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 pub fn schema(&self) -> SchemaRef {
116 self.schema.clone()
117 }
118
119 pub async fn next_batch(&mut self) -> Result<Option<RecordBatch>> {
122 self.batches.try_next().await
123 }
124
125 pub fn into_stream(self) -> BoxStream<'a, Result<RecordBatch>> {
127 self.batches
128 }
129}