Igosuki commented on a change in pull request #910: URL: https://github.com/apache/arrow-datafusion/pull/910#discussion_r696622294
########## File path: datafusion/src/physical_plan/avro.rs ########## @@ -0,0 +1,395 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Execution plan for reading line-delimited Avro files +use async_trait::async_trait; +use futures::Stream; + +use super::{common, source::Source, ExecutionPlan, Partitioning, RecordBatchStream}; +use crate::avro_to_arrow::infer_avro_schema_from_reader; +use crate::error::{DataFusionError, Result}; +use arrow::{ + datatypes::{Schema, SchemaRef}, + error::Result as ArrowResult, + record_batch::RecordBatch, +}; +use std::fs::File; +use std::{any::Any, io::Seek}; +use std::{ + io::Read, + pin::Pin, + sync::{Arc, Mutex}, + task::{Context, Poll}, +}; + +/// Line-delimited Avro read options +#[derive(Clone)] +pub struct AvroReadOptions<'a> { + /// The data source schema. + pub schema: Option<SchemaRef>, + + /// File extension; only files with this extension are selected for data input. + /// Defaults to ".avro". + pub file_extension: &'a str, +} + +impl<'a> Default for AvroReadOptions<'a> { + fn default() -> Self { + Self { + schema: None, + file_extension: ".avro", + } + } +} + +trait SeekRead: Read + Seek {} + +impl<T: Seek + Read> SeekRead for T {} +/// Execution plan for scanning Avro data source +#[derive(Debug)] +pub struct AvroExec { + source: Source<Box<dyn SeekRead + Send + Sync>>, + schema: SchemaRef, + projection: Option<Vec<usize>>, + projected_schema: SchemaRef, + file_extension: String, + batch_size: usize, + limit: Option<usize>, +} + +impl AvroExec { + /// Create a new execution plan for reading from a path + pub fn try_from_path( + path: &str, + options: AvroReadOptions, + projection: Option<Vec<usize>>, + batch_size: usize, + limit: Option<usize>, + ) -> Result<Self> { + let file_extension = options.file_extension.to_string(); + + let filenames = common::build_file_list(path, &file_extension)?; + + if filenames.is_empty() { + return Err(DataFusionError::Execution(format!( + "No files found at {path} with file extension {file_extension}", + path = path, + file_extension = file_extension.as_str() + ))); + } + + let schema = match options.schema { + Some(s) => s, + None => Arc::new(AvroExec::try_infer_schema(filenames.as_slice())?), + }; + + let projected_schema = match &projection { + None => schema.clone(), + Some(p) => Arc::new(Schema::new( + p.iter().map(|i| schema.field(*i).clone()).collect(), + )), + }; + + Ok(Self { + source: Source::PartitionedFiles { + path: path.to_string(), + filenames, + }, + schema, + projected_schema, + file_extension, + projection, + batch_size, + limit, + }) + } + /// Create a new execution plan for reading from a reader + pub fn try_new_from_reader( + reader: impl Read + Seek + Send + Sync + 'static, + options: AvroReadOptions, + projection: Option<Vec<usize>>, + batch_size: usize, + limit: Option<usize>, + ) -> Result<Self> { + let schema = match options.schema { + Some(s) => s, + None => { + return Err(DataFusionError::Execution( + "The schema must be provided in options when reading from a reader" + .to_string(), + )); + } + }; + + let projected_schema = match &projection { + None => schema.clone(), + Some(p) => Arc::new(Schema::new( + p.iter().map(|i| schema.field(*i).clone()).collect(), + )), + }; + + Ok(Self { + source: Source::Reader(Mutex::new(Some(Box::new(reader)))), + schema, + file_extension: String::new(), + projection, + projected_schema, + batch_size, + limit, + }) + } + + /// Path to directory containing partitioned CSV files with the same schema + pub fn path(&self) -> &str { + self.source.path() + } + + /// The individual files under path + pub fn filenames(&self) -> &[String] { + self.source.filenames() + } + + /// File extension + pub fn file_extension(&self) -> &str { + &self.file_extension + } + + /// Get the schema of the avro file + pub fn file_schema(&self) -> SchemaRef { + self.schema.clone() + } + + /// Optional projection for which columns to load + pub fn projection(&self) -> Option<&Vec<usize>> { + self.projection.as_ref() + } + + /// Batch size + pub fn batch_size(&self) -> usize { + self.batch_size + } + + /// Limit + pub fn limit(&self) -> Option<usize> { + self.limit + } + + /// Infer schema for given Avro dataset + pub fn try_infer_schema(filenames: &[String]) -> Result<Schema> { + let mut schemas = Vec::new(); + for filename in filenames { + let mut file = File::open(filename)?; + let schema = infer_avro_schema_from_reader(&mut file)?; + schemas.push(schema); + } + + Ok(Schema::try_merge(schemas)?) + } +} + +#[async_trait] +impl ExecutionPlan for AvroExec { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.projected_schema.clone() + } + + fn output_partitioning(&self) -> Partitioning { + Partitioning::UnknownPartitioning(match &self.source { + Source::PartitionedFiles { filenames, .. } => filenames.len(), + Source::Reader(_) => 1, + }) + } + + fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> { + Vec::new() + } + + fn with_new_children( Review comment: Added a test for that in `sql.rs` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
