arrow_avro/reader/async_reader/
builder.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::codec::{AvroFieldBuilder, Tz};
19use crate::errors::AvroError;
20use crate::reader::async_reader::ReaderState;
21use crate::reader::header::{Header, HeaderDecoder, HeaderInfo};
22use crate::reader::record::RecordDecoder;
23use crate::reader::{AsyncAvroFileReader, AsyncFileReader, Decoder};
24use crate::schema::{AvroSchema, FingerprintAlgorithm};
25use indexmap::IndexMap;
26use std::ops::Range;
27
28const DEFAULT_HEADER_SIZE_HINT: u64 = 16 * 1024; // 16 KB
29
30/// Builder for an asynchronous Avro file reader.
31pub struct ReaderBuilder<R> {
32    reader: R,
33    file_size: u64,
34    batch_size: usize,
35    range: Option<Range<u64>>,
36    reader_schema: Option<AvroSchema>,
37    projection: Option<Vec<usize>>,
38    header_size_hint: Option<u64>,
39    utf8_view: bool,
40    strict_mode: bool,
41    tz: Tz,
42}
43
44impl<R> ReaderBuilder<R> {
45    pub(super) fn new(reader: R, file_size: u64, batch_size: usize) -> Self {
46        Self {
47            reader,
48            file_size,
49            batch_size,
50            range: None,
51            reader_schema: None,
52            projection: None,
53            header_size_hint: None,
54            utf8_view: false,
55            strict_mode: false,
56            tz: Default::default(),
57        }
58    }
59
60    /// Specify a byte range to read from the Avro file.
61    /// If this is provided, the reader will read all the blocks within the specified range,
62    /// if the range ends mid-block, it will attempt to fetch the remaining bytes to complete the block,
63    /// but no further blocks will be read.
64    /// If this is omitted, the full file will be read.
65    pub fn with_range(self, range: Range<u64>) -> Self {
66        Self {
67            range: Some(range),
68            ..self
69        }
70    }
71
72    /// Specify a reader schema to use when reading the Avro file.
73    /// This can be useful to project specific columns or handle schema evolution.
74    /// If this is not provided, the schema will be derived from the Arrow schema provided.
75    pub fn with_reader_schema(self, reader_schema: AvroSchema) -> Self {
76        Self {
77            reader_schema: Some(reader_schema),
78            ..self
79        }
80    }
81
82    /// Specify a projection of column indices to read from the Avro file.
83    /// This can help optimize reading by only fetching the necessary columns.
84    pub fn with_projection(self, projection: Vec<usize>) -> Self {
85        Self {
86            projection: Some(projection),
87            ..self
88        }
89    }
90
91    /// Provide a hint for the expected size of the Avro header in bytes.
92    /// This can help optimize the initial read operation when fetching the header.
93    pub fn with_header_size_hint(self, hint: u64) -> Self {
94        Self {
95            header_size_hint: Some(hint),
96            ..self
97        }
98    }
99
100    /// Enable usage of Utf8View types when reading string data.
101    pub fn with_utf8_view(self, utf8_view: bool) -> Self {
102        Self { utf8_view, ..self }
103    }
104
105    /// Enable strict mode for schema validation and data reading.
106    pub fn with_strict_mode(self, strict_mode: bool) -> Self {
107        Self {
108            strict_mode,
109            ..self
110        }
111    }
112
113    /// Sets the timezone representation for Avro timestamp fields.
114    ///
115    /// The default is `Tz::OffsetZero`, meaning the "+00:00" time zone ID.
116    pub fn with_tz(mut self, tz: Tz) -> Self {
117        self.tz = tz;
118        self
119    }
120}
121
122/// Reads the Avro file header (magic, metadata, sync marker) asynchronously from `reader`.
123///
124/// On success, returns the parsed [`HeaderInfo`] containing the header and its length in bytes.
125pub async fn read_header_info<R>(
126    reader: &mut R,
127    file_size: u64,
128    header_size_hint: Option<u64>,
129) -> Result<HeaderInfo, AvroError>
130where
131    R: AsyncFileReader,
132{
133    read_header(reader, file_size, header_size_hint)
134        .await
135        .map(|(header, header_len)| HeaderInfo::new(header, header_len))
136}
137
138async fn read_header<R>(
139    reader: &mut R,
140    file_size: u64,
141    header_size_hint: Option<u64>,
142) -> Result<(Header, u64), AvroError>
143where
144    R: AsyncFileReader,
145{
146    let mut decoder = HeaderDecoder::default();
147    let mut position = 0;
148    loop {
149        let range_to_fetch = position
150            ..(position + header_size_hint.unwrap_or(DEFAULT_HEADER_SIZE_HINT)).min(file_size);
151
152        // Maybe EOF after the header, no actual data
153        if range_to_fetch.is_empty() {
154            break;
155        }
156
157        let current_data = reader.get_bytes(range_to_fetch.clone()).await?;
158        if current_data.is_empty() {
159            return Err(AvroError::EOF(
160                "Unexpected EOF while fetching header data".into(),
161            ));
162        }
163
164        let read = current_data.len();
165        let decoded = decoder.decode(&current_data)?;
166        if decoded != read {
167            position += decoded as u64;
168            break;
169        }
170        position += read as u64;
171    }
172
173    decoder
174        .flush()
175        .map(|header| (header, position))
176        .ok_or_else(|| AvroError::EOF("Unexpected EOF while reading Avro header".into()))
177}
178
179impl<R: AsyncFileReader> ReaderBuilder<R> {
180    /// Build the asynchronous Avro reader with the provided parameters.
181    /// This reads the header first to initialize the reader state.
182    pub async fn try_build(mut self) -> Result<AsyncAvroFileReader<R>, AvroError> {
183        if self.file_size == 0 {
184            return Err(AvroError::InvalidArgument("File size cannot be 0".into()));
185        }
186
187        // Start by reading the header from the beginning of the avro file
188        // take the writer schema from the header
189        let header_info =
190            read_header_info(&mut self.reader, self.file_size, self.header_size_hint).await?;
191
192        self.build_with_header(header_info)
193    }
194
195    /// Build the asynchronous Avro reader with the provided header.
196    ///
197    /// This allows initializing the reader with pre-parsed header information.
198    /// Note that this method is not async because it does not need to perform any I/O operations.
199    ///
200    /// Note: Any `header_size_hint` set via [`Self::with_header_size_hint`] is not used
201    /// when building with a pre-parsed header, since no header fetching occurs.
202    pub fn build_with_header(
203        self,
204        header_info: HeaderInfo,
205    ) -> Result<AsyncAvroFileReader<R>, AvroError> {
206        let writer_schema = header_info.writer_schema()?;
207
208        // If projection exists, project the reader schema,
209        // if no reader schema is provided, parse it from the header(get the raw writer schema), and project that
210        // this projected schema will be the schema used for reading.
211        let projected_reader_schema = self
212            .projection
213            .as_deref()
214            .map(|projection| {
215                let base_schema = if let Some(reader_schema) = &self.reader_schema {
216                    reader_schema
217                } else {
218                    &writer_schema
219                };
220                base_schema.project(projection)
221            })
222            .transpose()?;
223
224        // Use either the projected reader schema or the original reader schema(if no projection)
225        // (both optional, at worst no reader schema is provided, in which case we read with the writer schema)
226        let effective_reader_schema = projected_reader_schema
227            .as_ref()
228            .or(self.reader_schema.as_ref())
229            .map(|s| s.schema())
230            .transpose()?;
231
232        let root = {
233            let writer_schema = writer_schema.schema()?;
234            let mut builder = AvroFieldBuilder::new(&writer_schema);
235            if let Some(reader_schema) = &effective_reader_schema {
236                builder = builder.with_reader_schema(reader_schema);
237            }
238            builder
239                .with_utf8view(self.utf8_view)
240                .with_strict_mode(self.strict_mode)
241                .with_tz(self.tz)
242                .build()
243        }?;
244
245        let record_decoder = RecordDecoder::try_new_with_options(root.data_type())?;
246        let decoder = Decoder::from_parts(
247            self.batch_size,
248            record_decoder,
249            None,
250            IndexMap::new(),
251            FingerprintAlgorithm::Rabin,
252        );
253        let header_len = header_info.header_len();
254        let range = match self.range {
255            Some(r) => {
256                // If this PartitionedFile's range starts at 0, we need to skip the header bytes.
257                // But then we need to seek back 16 bytes to include the sync marker for the first block,
258                // as the logic in this reader searches the data for the first sync marker(after which a block starts),
259                // then reads blocks from the count, size etc.
260                let start = r.start.max(header_len.checked_sub(16).ok_or(AvroError::ParseError("Avro header length overflow, header was not long enough to contain avro bytes".to_string()))?);
261                let end = r.end.max(start).min(self.file_size); // Ensure end is not less than start, worst case range is empty
262                start..end
263            }
264            None => 0..self.file_size,
265        };
266
267        // Determine if there is actually data to fetch, note that we subtract the header len from range.start,
268        // so we need to check if range.end == header_len to see if there's no data after the header
269        let reader_state = if range.start == range.end || header_len == range.end {
270            ReaderState::Finished
271        } else {
272            ReaderState::Idle {
273                reader: self.reader,
274            }
275        };
276
277        let codec = header_info.compression()?;
278        let sync_marker = header_info.sync();
279
280        Ok(AsyncAvroFileReader::new(
281            range,
282            self.file_size,
283            decoder,
284            codec,
285            sync_marker,
286            reader_state,
287        ))
288    }
289}
arrow_avro/reader/async_reader/builder.rs

arrow_avro/reader/async_reader/
builder.rs
