diff --git a/Cargo.lock b/Cargo.lock index a041619ecef..5f657c3d5ee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1666,6 +1666,7 @@ dependencies = [ "bytes", "common-error", "common-macro", + "common-recordbatch", "common-runtime", "common-test-util", "datafusion", @@ -1943,6 +1944,7 @@ dependencies = [ "datatypes", "futures", "paste", + "pin-project", "serde", "serde_json", "snafu", diff --git a/src/common/datasource/Cargo.toml b/src/common/datasource/Cargo.toml index 3027052e4cc..193fc108da7 100644 --- a/src/common/datasource/Cargo.toml +++ b/src/common/datasource/Cargo.toml @@ -19,6 +19,7 @@ async-trait.workspace = true bytes.workspace = true common-error.workspace = true common-macro.workspace = true +common-recordbatch.workspace = true common-runtime.workspace = true datafusion.workspace = true datatypes.workspace = true diff --git a/src/common/datasource/src/file_format/orc.rs b/src/common/datasource/src/file_format/orc.rs index effbed98e90..f11798c3300 100644 --- a/src/common/datasource/src/file_format/orc.rs +++ b/src/common/datasource/src/file_format/orc.rs @@ -12,18 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::pin::Pin; use std::sync::Arc; -use std::task::{Context, Poll}; -use arrow::compute::cast; use arrow_schema::{ArrowError, Schema, SchemaRef}; use async_trait::async_trait; -use datafusion::arrow::record_batch::RecordBatch as DfRecordBatch; +use common_recordbatch::adapter::RecordBatchStreamTypeAdapter; use datafusion::datasource::physical_plan::{FileMeta, FileOpenFuture, FileOpener}; use datafusion::error::{DataFusionError, Result as DfResult}; -use datafusion::physical_plan::RecordBatchStream; -use futures::{Stream, StreamExt, TryStreamExt}; +use futures::{StreamExt, TryStreamExt}; use object_store::ObjectStore; use orc_rust::arrow_reader::{create_arrow_schema, Cursor}; use orc_rust::async_arrow_reader::ArrowStreamReader; @@ -61,73 +57,6 @@ pub async fn infer_orc_schema Ok(create_arrow_schema(&cursor)) } -pub struct OrcArrowStreamReaderAdapter { - output_schema: SchemaRef, - projection: Vec, - stream: ArrowStreamReader, -} - -impl OrcArrowStreamReaderAdapter { - pub fn new( - output_schema: SchemaRef, - stream: ArrowStreamReader, - projection: Option>, - ) -> Self { - let projection = if let Some(projection) = projection { - projection - } else { - (0..output_schema.fields().len()).collect() - }; - - Self { - output_schema, - projection, - stream, - } - } -} - -impl RecordBatchStream - for OrcArrowStreamReaderAdapter -{ - fn schema(&self) -> SchemaRef { - self.output_schema.clone() - } -} - -impl Stream for OrcArrowStreamReaderAdapter { - type Item = DfResult; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let batch = futures::ready!(Pin::new(&mut self.stream).poll_next(cx)) - .map(|r| r.map_err(|e| DataFusionError::External(Box::new(e)))); - - let projected_schema = self.output_schema.project(&self.projection)?; - let batch = batch.map(|b| { - b.and_then(|b| { - let mut columns = Vec::with_capacity(self.projection.len()); - for idx in self.projection.iter() { - let column = b.column(*idx); - let field = self.output_schema.field(*idx); - - if column.data_type() != field.data_type() { - let output = cast(&column, field.data_type())?; - columns.push(output) - } else { - columns.push(column.clone()) - } - } - - let record_batch = DfRecordBatch::try_new(projected_schema.into(), columns)?; - - Ok(record_batch) - }) - }); - - Poll::Ready(batch) - } -} - #[async_trait] impl FileFormat for OrcFormat { async fn infer_schema(&self, store: &ObjectStore, path: &str) -> Result { @@ -166,7 +95,15 @@ impl OrcOpener { impl FileOpener for OrcOpener { fn open(&self, meta: FileMeta) -> DfResult { let object_store = self.object_store.clone(); - let output_schema = self.output_schema.clone(); + let projected_schema = if let Some(projection) = &self.projection { + let projected_schema = self + .output_schema + .project(projection) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + Arc::new(projected_schema) + } else { + self.output_schema.clone() + }; let projection = self.projection.clone(); Ok(Box::pin(async move { let reader = object_store @@ -178,7 +115,8 @@ impl FileOpener for OrcOpener { .await .map_err(|e| DataFusionError::External(Box::new(e)))?; - let stream = OrcArrowStreamReaderAdapter::new(output_schema, stream_reader, projection); + let stream = + RecordBatchStreamTypeAdapter::new(projected_schema, stream_reader, projection); let adopted = stream.map_err(|e| ArrowError::ExternalError(Box::new(e))); Ok(adopted.boxed()) diff --git a/src/common/recordbatch/Cargo.toml b/src/common/recordbatch/Cargo.toml index 5425313bdb3..c90153b65ce 100644 --- a/src/common/recordbatch/Cargo.toml +++ b/src/common/recordbatch/Cargo.toml @@ -14,6 +14,7 @@ datafusion.workspace = true datatypes.workspace = true futures.workspace = true paste = "1.0" +pin-project.workspace = true serde.workspace = true serde_json.workspace = true snafu.workspace = true diff --git a/src/common/recordbatch/src/adapter.rs b/src/common/recordbatch/src/adapter.rs index 3cbb8eb4e2c..a2dfd4ebaa9 100644 --- a/src/common/recordbatch/src/adapter.rs +++ b/src/common/recordbatch/src/adapter.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::future::Future; +use std::marker::PhantomData; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -20,12 +21,13 @@ use std::task::{Context, Poll}; use datafusion::arrow::compute::cast; use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef; use datafusion::error::Result as DfResult; -use datafusion::parquet::arrow::async_reader::{AsyncFileReader, ParquetRecordBatchStream}; use datafusion::physical_plan::metrics::{BaselineMetrics, MetricValue}; use datafusion::physical_plan::{ExecutionPlan, RecordBatchStream as DfRecordBatchStream}; +use datafusion_common::arrow::error::ArrowError; use datafusion_common::DataFusionError; use datatypes::schema::{Schema, SchemaRef}; use futures::ready; +use pin_project::pin_project; use snafu::ResultExt; use crate::error::{self, Result}; @@ -37,56 +39,76 @@ use crate::{ type FutureStream = Pin> + Send>>; -/// ParquetRecordBatchStream -> DataFusion RecordBatchStream -pub struct ParquetRecordBatchStreamAdapter { - stream: ParquetRecordBatchStream, - output_schema: DfSchemaRef, +/// Casts the `RecordBatch`es of `stream` against the `output_schema`. +#[pin_project] +pub struct RecordBatchStreamTypeAdapter { + #[pin] + stream: T, + projected_schema: DfSchemaRef, projection: Vec, + phantom: PhantomData, } -impl ParquetRecordBatchStreamAdapter { - pub fn new( - output_schema: DfSchemaRef, - stream: ParquetRecordBatchStream, - projection: Option>, - ) -> Self { +impl RecordBatchStreamTypeAdapter +where + T: Stream>, + E: std::error::Error + Send + Sync + 'static, +{ + pub fn new(projected_schema: DfSchemaRef, stream: T, projection: Option>) -> Self { let projection = if let Some(projection) = projection { projection } else { - (0..output_schema.fields().len()).collect() + (0..projected_schema.fields().len()).collect() }; Self { stream, - output_schema, + projected_schema, projection, + phantom: Default::default(), } } } -impl DfRecordBatchStream - for ParquetRecordBatchStreamAdapter +impl DfRecordBatchStream for RecordBatchStreamTypeAdapter +where + T: Stream>, + E: std::error::Error + Send + Sync + 'static, { fn schema(&self) -> DfSchemaRef { - self.stream.schema().clone() + self.projected_schema.clone() } } -impl Stream for ParquetRecordBatchStreamAdapter { +impl Stream for RecordBatchStreamTypeAdapter +where + T: Stream>, + E: std::error::Error + Send + Sync + 'static, +{ type Item = DfResult; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let batch = futures::ready!(Pin::new(&mut self.stream).poll_next(cx)) + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.project(); + + let batch = futures::ready!(this.stream.poll_next(cx)) .map(|r| r.map_err(|e| DataFusionError::External(Box::new(e)))); - let projected_schema = self.output_schema.project(&self.projection)?; + let projected_schema = this.projected_schema.clone(); + let projection = this.projection.clone(); let batch = batch.map(|b| { b.and_then(|b| { - let mut columns = Vec::with_capacity(self.projection.len()); - for idx in self.projection.iter() { - let column = b.column(*idx); - let field = self.output_schema.field(*idx); + let projected_column = b.project(&projection)?; + if projected_column.schema().fields.len() != projected_schema.fields.len() { + return Err(DataFusionError::ArrowError(ArrowError::SchemaError(format!( + "Trying to cast a RecordBatch into an incompatible schema. RecordBatch: {}, Target: {}", + projected_column.schema(), + projected_schema, + )))); + } + let mut columns = Vec::with_capacity(projected_schema.fields.len()); + for (idx,field) in projected_schema.fields.iter().enumerate() { + let column = projected_column.column(idx); if column.data_type() != field.data_type() { let output = cast(&column, field.data_type())?; columns.push(output) @@ -94,9 +116,7 @@ impl Stream for ParquetRecordBatchS columns.push(column.clone()) } } - - let record_batch = DfRecordBatch::try_new(projected_schema.into(), columns)?; - + let record_batch = DfRecordBatch::try_new(projected_schema, columns)?; Ok(record_batch) }) }); diff --git a/src/operator/src/statement/copy_table_from.rs b/src/operator/src/statement/copy_table_from.rs index 61aeb5e6036..7e7b748cb7e 100644 --- a/src/operator/src/statement/copy_table_from.rs +++ b/src/operator/src/statement/copy_table_from.rs @@ -19,14 +19,12 @@ use std::sync::Arc; use common_base::readable_size::ReadableSize; use common_datasource::file_format::csv::{CsvConfigBuilder, CsvFormat, CsvOpener}; use common_datasource::file_format::json::{JsonFormat, JsonOpener}; -use common_datasource::file_format::orc::{ - infer_orc_schema, new_orc_stream_reader, OrcArrowStreamReaderAdapter, -}; +use common_datasource::file_format::orc::{infer_orc_schema, new_orc_stream_reader}; use common_datasource::file_format::{FileFormat, Format}; use common_datasource::lister::{Lister, Source}; use common_datasource::object_store::{build_backend, parse_url}; use common_datasource::util::find_dir_and_filename; -use common_recordbatch::adapter::ParquetRecordBatchStreamAdapter; +use common_recordbatch::adapter::RecordBatchStreamTypeAdapter; use common_recordbatch::DfSendableRecordBatchStream; use common_telemetry::{debug, tracing}; use datafusion::datasource::listing::PartitionedFile; @@ -206,44 +204,76 @@ impl StatementExecutor { async fn build_read_stream( &self, - schema: SchemaRef, + compat_schema: SchemaRef, object_store: &ObjectStore, file_metadata: &FileMetadata, projection: Vec, ) -> Result { match file_metadata { - FileMetadata::Csv { format, path, .. } => { + FileMetadata::Csv { + format, + path, + schema, + } => { + let projected_schema = Arc::new( + compat_schema + .project(&projection) + .context(error::ProjectSchemaSnafu)?, + ); let csv_conf = CsvConfigBuilder::default() .batch_size(DEFAULT_BATCH_SIZE) .file_schema(schema.clone()) - .file_projection(Some(projection)) + .file_projection(Some(projection.clone())) .build() .context(error::BuildCsvConfigSnafu)?; - self.build_file_stream( - CsvOpener::new(csv_conf, object_store.clone(), format.compression_type), - path, - schema, - ) - .await + let stream = self + .build_file_stream( + CsvOpener::new(csv_conf, object_store.clone(), format.compression_type), + path, + schema.clone(), + ) + .await?; + + Ok(Box::pin(RecordBatchStreamTypeAdapter::new( + projected_schema, + stream, + Some(projection), + ))) } - FileMetadata::Json { format, path, .. } => { - let projected_schema = Arc::new( + FileMetadata::Json { + format, + path, + schema, + } => { + let projected_file_schema = Arc::new( schema .project(&projection) .context(error::ProjectSchemaSnafu)?, ); - self.build_file_stream( - JsonOpener::new( - DEFAULT_BATCH_SIZE, - projected_schema, - object_store.clone(), - format.compression_type, - ), - path, - schema, - ) - .await + let projected_schema = Arc::new( + compat_schema + .project(&projection) + .context(error::ProjectSchemaSnafu)?, + ); + let stream = self + .build_file_stream( + JsonOpener::new( + DEFAULT_BATCH_SIZE, + projected_file_schema, + object_store.clone(), + format.compression_type, + ), + path, + schema.clone(), + ) + .await?; + + Ok(Box::pin(RecordBatchStreamTypeAdapter::new( + projected_schema, + stream, + Some(projection), + ))) } FileMetadata::Parquet { metadata, path, .. } => { let reader = object_store @@ -253,13 +283,18 @@ impl StatementExecutor { .context(error::ReadObjectSnafu { path })?; let builder = ParquetRecordBatchStreamBuilder::new_with_metadata(reader, metadata.clone()); - let upstream = builder + let stream = builder .build() .context(error::BuildParquetRecordBatchStreamSnafu)?; - Ok(Box::pin(ParquetRecordBatchStreamAdapter::new( - schema, - upstream, + let projected_schema = Arc::new( + compat_schema + .project(&projection) + .context(error::ProjectSchemaSnafu)?, + ); + Ok(Box::pin(RecordBatchStreamTypeAdapter::new( + projected_schema, + stream, Some(projection), ))) } @@ -272,9 +307,18 @@ impl StatementExecutor { let stream = new_orc_stream_reader(reader) .await .context(error::ReadOrcSnafu)?; - let stream = OrcArrowStreamReaderAdapter::new(schema, stream, Some(projection)); - Ok(Box::pin(stream)) + let projected_schema = Arc::new( + compat_schema + .project(&projection) + .context(error::ProjectSchemaSnafu)?, + ); + + Ok(Box::pin(RecordBatchStreamTypeAdapter::new( + projected_schema, + stream, + Some(projection), + ))) } } } @@ -420,7 +464,10 @@ fn ensure_schema_compatible(from: &SchemaRef, to: &SchemaRef) -> Result<()> { } } -/// Allows the file schema is a subset of table +/// Generates a maybe compatible schema of the file schema. +/// +/// If there is a field is found in table schema, +/// copy the field data type to maybe compatible schema(`compatible_fields`). fn generated_schema_projection_and_compatible_file_schema( file: &SchemaRef, table: &SchemaRef,