datafusion_datasource_avro/
source.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
18//! Execution plan for reading line-delimited Avro files
19
20use std::sync::Arc;
21
22use arrow::datatypes::{Schema, SchemaRef};
23use arrow_avro::reader::{Reader, ReaderBuilder};
24use datafusion_common::error::Result;
25use datafusion_datasource::TableSchema;
26use datafusion_datasource::file::FileSource;
27use datafusion_datasource::file_scan_config::FileScanConfig;
28use datafusion_datasource::file_stream::FileOpener;
29use datafusion_datasource::projection::{ProjectionOpener, SplitProjection};
30use datafusion_physical_expr_adapter::BatchAdapterFactory;
31use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
32use datafusion_physical_plan::projection::ProjectionExprs;
33
34use object_store::ObjectStore;
35
36/// AvroSource holds the extra configuration that is necessary for opening avro files
37#[derive(Clone)]
38pub struct AvroSource {
39    table_schema: TableSchema,
40    batch_size: Option<usize>,
41    projection: SplitProjection,
42    metrics: ExecutionPlanMetricsSet,
43}
44
45impl AvroSource {
46    /// Initialize an AvroSource with the provided schema
47    pub fn new(table_schema: impl Into<TableSchema>) -> Self {
48        let table_schema = table_schema.into();
49        Self {
50            projection: SplitProjection::unprojected(&table_schema),
51            table_schema,
52            batch_size: None,
53            metrics: ExecutionPlanMetricsSet::new(),
54        }
55    }
56
57    fn open<R: std::io::BufRead>(
58        &self,
59        reader: R,
60        projection: Option<Vec<usize>>,
61    ) -> Result<Reader<R>> {
62        let mut builder = ReaderBuilder::new()
63            .with_batch_size(self.batch_size.expect("Batch size must set before open"));
64        if let Some(projection) = projection {
65            builder = builder.with_projection(projection);
66        }
67        builder.build(reader).map_err(Into::into)
68    }
69
70    fn projected_file_schema(&self) -> SchemaRef {
71        let file_schema = self.table_schema.file_schema();
72        if self.projection.file_indices.is_empty() {
73            return Arc::clone(file_schema);
74        }
75
76        Arc::new(Schema::new(
77            self.projection
78                .file_indices
79                .iter()
80                .map(|idx| file_schema.field(*idx).clone())
81                .collect::<Vec<_>>(),
82        ))
83    }
84
85    fn writer_projection_for_schema(
86        &self,
87        writer_schema: &Schema,
88        target_schema: &Schema,
89    ) -> Option<Vec<usize>> {
90        // `arrow-avro` accepts projection ordinals against the file's writer schema,
91        // while DataFusion plans projection against the logical table schema. Remap
92        // projected column names to writer ordinals so reader-level pushdown still
93        // preserves DataFusion's existing name-based projection semantics.
94        let projection = target_schema
95            .fields()
96            .iter()
97            .filter_map(|field| {
98                writer_schema
99                    .column_with_name(field.name())
100                    .map(|(idx, _)| idx)
101            })
102            .collect::<Vec<_>>();
103
104        let identity_projection = projection.len() == writer_schema.fields().len()
105            && projection
106                .iter()
107                .enumerate()
108                .all(|(idx, value)| idx == *value);
109
110        (!identity_projection).then_some(projection)
111    }
112}
113
114impl FileSource for AvroSource {
115    fn create_file_opener(
116        &self,
117        object_store: Arc<dyn ObjectStore>,
118        _base_config: &FileScanConfig,
119        _partition: usize,
120    ) -> Result<Arc<dyn FileOpener>> {
121        let mut opener = Arc::new(private::AvroOpener {
122            config: Arc::new(self.clone()),
123            object_store,
124        }) as Arc<dyn FileOpener>;
125        opener = ProjectionOpener::try_new(
126            self.projection.clone(),
127            Arc::clone(&opener),
128            self.table_schema.file_schema(),
129        )?;
130        Ok(opener)
131    }
132
133    fn table_schema(&self) -> &TableSchema {
134        &self.table_schema
135    }
136
137    fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
138        let mut conf = self.clone();
139        conf.batch_size = Some(batch_size);
140        Arc::new(conf)
141    }
142
143    fn try_pushdown_projection(
144        &self,
145        projection: &ProjectionExprs,
146    ) -> Result<Option<Arc<dyn FileSource>>> {
147        let mut source = self.clone();
148        let new_projection = self.projection.source.try_merge(projection)?;
149        let split_projection =
150            SplitProjection::new(self.table_schema.file_schema(), &new_projection);
151        source.projection = split_projection;
152        Ok(Some(Arc::new(source)))
153    }
154
155    fn projection(&self) -> Option<&ProjectionExprs> {
156        Some(&self.projection.source)
157    }
158
159    fn metrics(&self) -> &ExecutionPlanMetricsSet {
160        &self.metrics
161    }
162
163    fn file_type(&self) -> &str {
164        "avro"
165    }
166
167    fn supports_repartitioning(&self) -> bool {
168        // Avro OCF does not support safe byte-range splitting in this reader path.
169        false
170    }
171}
172
173mod private {
174    use super::*;
175    use std::io::BufReader;
176    use std::io::Seek;
177
178    use bytes::Buf;
179    use datafusion_datasource::{PartitionedFile, file_stream::FileOpenFuture};
180    use futures::StreamExt;
181    use object_store::{GetResultPayload, ObjectStore, ObjectStoreExt};
182
183    pub struct AvroOpener {
184        pub config: Arc<AvroSource>,
185        pub object_store: Arc<dyn ObjectStore>,
186    }
187
188    impl FileOpener for AvroOpener {
189        fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
190            let object_store = Arc::clone(&self.object_store);
191            let config = Arc::clone(&self.config);
192            let projected_file_schema = config.projected_file_schema();
193
194            Ok(Box::pin(async move {
195                let r = object_store
196                    .get(&partitioned_file.object_meta.location)
197                    .await?;
198                match r.payload {
199                    GetResultPayload::File(mut file, _) => {
200                        // Probe the writer schema first so logical projected columns can be
201                        // translated to the writer-schema ordinals expected by `arrow-avro`.
202                        let probe_reader =
203                            config.open(BufReader::new(file.try_clone()?), None)?;
204                        let writer_projection = config.writer_projection_for_schema(
205                            probe_reader.schema().as_ref(),
206                            projected_file_schema.as_ref(),
207                        );
208                        file.rewind()?;
209                        let reader =
210                            config.open(BufReader::new(file), writer_projection)?;
211                        let batch_adapter =
212                            BatchAdapterFactory::new(Arc::clone(&projected_file_schema))
213                                .make_adapter(&reader.schema())?;
214                        Ok(futures::stream::iter(reader)
215                            .map(move |r| {
216                                r.map_err(Into::into)
217                                    .and_then(|batch| batch_adapter.adapt_batch(&batch))
218                            })
219                            .boxed())
220                    }
221                    GetResultPayload::Stream(_) => {
222                        let bytes = r.bytes().await?;
223                        // As above, inspect the writer schema before constructing the real
224                        // reader so `with_projection` can use writer-schema ordinals.
225                        let probe_reader =
226                            config.open(BufReader::new(bytes.clone().reader()), None)?;
227                        let writer_projection = config.writer_projection_for_schema(
228                            probe_reader.schema().as_ref(),
229                            projected_file_schema.as_ref(),
230                        );
231                        let reader = config
232                            .open(BufReader::new(bytes.reader()), writer_projection)?;
233                        let batch_adapter =
234                            BatchAdapterFactory::new(Arc::clone(&projected_file_schema))
235                                .make_adapter(&reader.schema())?;
236                        Ok(futures::stream::iter(reader)
237                            .map(move |r| {
238                                r.map_err(Into::into)
239                                    .and_then(|batch| batch_adapter.adapt_batch(&batch))
240                            })
241                            .boxed())
242                    }
243                }
244            }))
245        }
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252    use arrow::array::*;
253    use arrow::array::{
254        BinaryArray, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array,
255        TimestampMicrosecondArray, TimestampMillisecondArray,
256    };
257    use arrow::datatypes::TimeUnit;
258    use arrow::datatypes::{DataType, Field};
259    use std::fs::File;
260    use std::io::BufReader;
261
262    fn build_reader(
263        name: &'_ str,
264        projection: Option<Vec<usize>>,
265    ) -> Reader<BufReader<File>> {
266        let testdata = datafusion_common::test_util::arrow_test_data();
267        let filename = format!("{testdata}/avro/{name}");
268        let mut builder = ReaderBuilder::new().with_batch_size(64);
269        if let Some(proj) = projection {
270            builder = builder.with_projection(proj);
271        }
272        builder
273            .build(BufReader::new(File::open(filename).unwrap()))
274            .unwrap()
275    }
276
277    fn get_col<'a, T: 'static>(
278        batch: &'a RecordBatch,
279        col: (usize, &Field),
280    ) -> Option<&'a T> {
281        batch.column(col.0).as_any().downcast_ref::<T>()
282    }
283
284    #[test]
285    fn test_avro_basic() {
286        let mut reader = build_reader("alltypes_dictionary.avro", None);
287        let batch = reader.next().unwrap().unwrap();
288
289        assert_eq!(11, batch.num_columns());
290        assert_eq!(2, batch.num_rows());
291
292        let schema = reader.schema();
293        let batch_schema = batch.schema();
294        assert_eq!(schema, batch_schema);
295
296        let id = schema.column_with_name("id").unwrap();
297        assert_eq!(0, id.0);
298        assert_eq!(&DataType::Int32, id.1.data_type());
299        let col = get_col::<Int32Array>(&batch, id).unwrap();
300        assert_eq!(0, col.value(0));
301        assert_eq!(1, col.value(1));
302        let bool_col = schema.column_with_name("bool_col").unwrap();
303        assert_eq!(1, bool_col.0);
304        assert_eq!(&DataType::Boolean, bool_col.1.data_type());
305        let col = get_col::<BooleanArray>(&batch, bool_col).unwrap();
306        assert!(col.value(0));
307        assert!(!col.value(1));
308        let tinyint_col = schema.column_with_name("tinyint_col").unwrap();
309        assert_eq!(2, tinyint_col.0);
310        assert_eq!(&DataType::Int32, tinyint_col.1.data_type());
311        let col = get_col::<Int32Array>(&batch, tinyint_col).unwrap();
312        assert_eq!(0, col.value(0));
313        assert_eq!(1, col.value(1));
314        let smallint_col = schema.column_with_name("smallint_col").unwrap();
315        assert_eq!(3, smallint_col.0);
316        assert_eq!(&DataType::Int32, smallint_col.1.data_type());
317        let col = get_col::<Int32Array>(&batch, smallint_col).unwrap();
318        assert_eq!(0, col.value(0));
319        assert_eq!(1, col.value(1));
320        let int_col = schema.column_with_name("int_col").unwrap();
321        assert_eq!(4, int_col.0);
322        let col = get_col::<Int32Array>(&batch, int_col).unwrap();
323        assert_eq!(0, col.value(0));
324        assert_eq!(1, col.value(1));
325        assert_eq!(&DataType::Int32, int_col.1.data_type());
326        let col = get_col::<Int32Array>(&batch, int_col).unwrap();
327        assert_eq!(0, col.value(0));
328        assert_eq!(1, col.value(1));
329        let bigint_col = schema.column_with_name("bigint_col").unwrap();
330        assert_eq!(5, bigint_col.0);
331        let col = get_col::<Int64Array>(&batch, bigint_col).unwrap();
332        assert_eq!(0, col.value(0));
333        assert_eq!(10, col.value(1));
334        assert_eq!(&DataType::Int64, bigint_col.1.data_type());
335        let float_col = schema.column_with_name("float_col").unwrap();
336        assert_eq!(6, float_col.0);
337        let col = get_col::<Float32Array>(&batch, float_col).unwrap();
338        assert_eq!(0.0, col.value(0));
339        assert_eq!(1.1, col.value(1));
340        assert_eq!(&DataType::Float32, float_col.1.data_type());
341        let col = get_col::<Float32Array>(&batch, float_col).unwrap();
342        assert_eq!(0.0, col.value(0));
343        assert_eq!(1.1, col.value(1));
344        let double_col = schema.column_with_name("double_col").unwrap();
345        assert_eq!(7, double_col.0);
346        assert_eq!(&DataType::Float64, double_col.1.data_type());
347        let col = get_col::<Float64Array>(&batch, double_col).unwrap();
348        assert_eq!(0.0, col.value(0));
349        assert_eq!(10.1, col.value(1));
350        let date_string_col = schema.column_with_name("date_string_col").unwrap();
351        assert_eq!(8, date_string_col.0);
352        assert_eq!(&DataType::Binary, date_string_col.1.data_type());
353        let col = get_col::<BinaryArray>(&batch, date_string_col).unwrap();
354        assert_eq!("01/01/09".as_bytes(), col.value(0));
355        assert_eq!("01/01/09".as_bytes(), col.value(1));
356        let string_col = schema.column_with_name("string_col").unwrap();
357        assert_eq!(9, string_col.0);
358        assert_eq!(&DataType::Binary, string_col.1.data_type());
359        let col = get_col::<BinaryArray>(&batch, string_col).unwrap();
360        assert_eq!("0".as_bytes(), col.value(0));
361        assert_eq!("1".as_bytes(), col.value(1));
362        let timestamp_col = schema.column_with_name("timestamp_col").unwrap();
363        assert_eq!(10, timestamp_col.0);
364        assert_eq!(
365            &DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
366            timestamp_col.1.data_type()
367        );
368        let col = get_col::<TimestampMicrosecondArray>(&batch, timestamp_col).unwrap();
369        assert_eq!(1230768000000000, col.value(0));
370        assert_eq!(1230768060000000, col.value(1));
371    }
372
373    #[test]
374    fn test_avro_with_projection() {
375        // Test projection to filter and reorder columns
376        let projection = vec![9, 7, 1]; // string_col, double_col, bool_col
377
378        let mut reader = build_reader("alltypes_dictionary.avro", Some(projection));
379        let batch = reader.next().unwrap().unwrap();
380
381        // Only 3 columns should be present (not all 11)
382        assert_eq!(3, batch.num_columns());
383        assert_eq!(2, batch.num_rows());
384
385        let schema = reader.schema();
386        let batch_schema = batch.schema();
387        assert_eq!(schema, batch_schema);
388
389        // Verify columns are in the order specified in projection
390        // First column should be string_col (was at index 9 in original)
391        assert_eq!("string_col", schema.field(0).name());
392        assert_eq!(&DataType::Binary, schema.field(0).data_type());
393        let col = batch
394            .column(0)
395            .as_any()
396            .downcast_ref::<BinaryArray>()
397            .unwrap();
398        assert_eq!("0".as_bytes(), col.value(0));
399        assert_eq!("1".as_bytes(), col.value(1));
400
401        // Second column should be double_col (was at index 7 in original)
402        assert_eq!("double_col", schema.field(1).name());
403        assert_eq!(&DataType::Float64, schema.field(1).data_type());
404        let col = batch
405            .column(1)
406            .as_any()
407            .downcast_ref::<Float64Array>()
408            .unwrap();
409        assert_eq!(0.0, col.value(0));
410        assert_eq!(10.1, col.value(1));
411
412        // Third column should be bool_col (was at index 1 in original)
413        assert_eq!("bool_col", schema.field(2).name());
414        assert_eq!(&DataType::Boolean, schema.field(2).data_type());
415        let col = batch
416            .column(2)
417            .as_any()
418            .downcast_ref::<BooleanArray>()
419            .unwrap();
420        assert!(col.value(0));
421        assert!(!col.value(1));
422    }
423
424    #[test]
425    fn test_avro_timestamp_logical_types() {
426        let mut reader = build_reader("timestamp_logical_types.avro", None);
427        let batch = reader.next().unwrap().unwrap();
428
429        assert_eq!(7, batch.num_columns());
430        assert_eq!(2, batch.num_rows());
431
432        let schema = reader.schema();
433        let ts_millis = schema.column_with_name("ts_millis").unwrap();
434        assert_eq!(1, ts_millis.0);
435        assert_eq!(
436            &DataType::Timestamp(TimeUnit::Millisecond, Some("+00:00".into())),
437            ts_millis.1.data_type()
438        );
439        let col = get_col::<TimestampMillisecondArray>(&batch, ts_millis).unwrap();
440        assert_eq!(0, col.value(0));
441        assert_eq!(1_000, col.value(1));
442
443        let ts_micros = schema.column_with_name("ts_micros").unwrap();
444        assert_eq!(2, ts_micros.0);
445        assert_eq!(
446            &DataType::Timestamp(TimeUnit::Microsecond, Some("+00:00".into())),
447            ts_micros.1.data_type()
448        );
449        let col = get_col::<TimestampMicrosecondArray>(&batch, ts_micros).unwrap();
450        assert_eq!(0, col.value(0));
451        assert_eq!(1_000_000, col.value(1));
452
453        let ts_nanos = schema.column_with_name("ts_nanos").unwrap();
454        assert_eq!(3, ts_nanos.0);
455        assert_eq!(
456            &DataType::Timestamp(TimeUnit::Nanosecond, Some("+00:00".into())),
457            ts_nanos.1.data_type()
458        );
459        let col = get_col::<TimestampNanosecondArray>(&batch, ts_nanos).unwrap();
460        assert_eq!(0, col.value(0));
461        assert_eq!(1_000_000_000, col.value(1));
462
463        let local_ts_millis = schema.column_with_name("local_ts_millis").unwrap();
464        assert_eq!(4, local_ts_millis.0);
465        assert_eq!(
466            &DataType::Timestamp(TimeUnit::Millisecond, None),
467            local_ts_millis.1.data_type()
468        );
469        let col = get_col::<TimestampMillisecondArray>(&batch, local_ts_millis).unwrap();
470        assert_eq!(0, col.value(0));
471        assert_eq!(1_000, col.value(1));
472
473        let local_ts_micros = schema.column_with_name("local_ts_micros").unwrap();
474        assert_eq!(5, local_ts_micros.0);
475        assert_eq!(
476            &DataType::Timestamp(TimeUnit::Microsecond, None),
477            local_ts_micros.1.data_type()
478        );
479        let col = get_col::<TimestampMicrosecondArray>(&batch, local_ts_micros).unwrap();
480        assert_eq!(0, col.value(0));
481        assert_eq!(1_000_000, col.value(1));
482
483        let local_ts_nanos = schema.column_with_name("local_ts_nanos").unwrap();
484        assert_eq!(6, local_ts_nanos.0);
485        assert_eq!(
486            &DataType::Timestamp(TimeUnit::Nanosecond, None),
487            local_ts_nanos.1.data_type()
488        );
489        let col = get_col::<TimestampNanosecondArray>(&batch, local_ts_nanos).unwrap();
490        assert_eq!(0, col.value(0));
491        assert_eq!(1_000_000_000, col.value(1));
492    }
493}
datafusion_datasource_avro/source.rs

datafusion_datasource_avro/
source.rs
