fix(Copy From): fix incorrect type casts (#3264)

* refactor: refactor RecordBatchStreamTypeAdapter

* fix(Copy From): fix incorrect type casts

* fix: unit tests

* chore: add comment

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",

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

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<R: AsyncRead + AsyncSeek + Unpin + Send + 'static>
     Ok(create_arrow_schema(&cursor))
 }
 
-pub struct OrcArrowStreamReaderAdapter<T: AsyncRead + AsyncSeek + Unpin + Send + 'static> {
-    output_schema: SchemaRef,
-    projection: Vec<usize>,
-    stream: ArrowStreamReader<T>,
-}
-
-impl<T: AsyncRead + AsyncSeek + Unpin + Send + 'static> OrcArrowStreamReaderAdapter<T> {
-    pub fn new(
-        output_schema: SchemaRef,
-        stream: ArrowStreamReader<T>,
-        projection: Option<Vec<usize>>,
-    ) -> Self {
-        let projection = if let Some(projection) = projection {
-            projection
-        } else {
-            (0..output_schema.fields().len()).collect()
-        };
-
-        Self {
-            output_schema,
-            projection,
-            stream,
-        }
-    }
-}
-
-impl<T: AsyncRead + AsyncSeek + Unpin + Send + 'static> RecordBatchStream
-    for OrcArrowStreamReaderAdapter<T>
-{
-    fn schema(&self) -> SchemaRef {
-        self.output_schema.clone()
-    }
-}
-
-impl<T: AsyncRead + AsyncSeek + Unpin + Send + 'static> Stream for OrcArrowStreamReaderAdapter<T> {
-    type Item = DfResult<DfRecordBatch>;
-
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        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<Schema> {
@@ -166,7 +95,15 @@ impl OrcOpener {
 impl FileOpener for OrcOpener {
     fn open(&self, meta: FileMeta) -> DfResult<FileOpenFuture> {
         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())

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

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<Box<dyn std::future::Future<Output = Result<SendableRecordBatchStream>> + Send>>;
 
-/// ParquetRecordBatchStream -> DataFusion RecordBatchStream
-pub struct ParquetRecordBatchStreamAdapter<T> {
-    stream: ParquetRecordBatchStream<T>,
-    output_schema: DfSchemaRef,
+/// Casts the `RecordBatch`es of `stream` against the `output_schema`.
+#[pin_project]
+pub struct RecordBatchStreamTypeAdapter<T, E> {
+    #[pin]
+    stream: T,
+    projected_schema: DfSchemaRef,
     projection: Vec<usize>,
+    phantom: PhantomData<E>,
 }
 
-impl<T: Unpin + AsyncFileReader + Send + 'static> ParquetRecordBatchStreamAdapter<T> {
-    pub fn new(
-        output_schema: DfSchemaRef,
-        stream: ParquetRecordBatchStream<T>,
-        projection: Option<Vec<usize>>,
-    ) -> Self {
+impl<T, E> RecordBatchStreamTypeAdapter<T, E>
+where
+    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    pub fn new(projected_schema: DfSchemaRef, stream: T, projection: Option<Vec<usize>>) -> 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<T: Unpin + AsyncFileReader + Send + 'static> DfRecordBatchStream
-    for ParquetRecordBatchStreamAdapter<T>
+impl<T, E> DfRecordBatchStream for RecordBatchStreamTypeAdapter<T, E>
+where
+    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
+    E: std::error::Error + Send + Sync + 'static,
 {
     fn schema(&self) -> DfSchemaRef {
-        self.stream.schema().clone()
+        self.projected_schema.clone()
     }
 }
 
-impl<T: Unpin + AsyncFileReader + Send + 'static> Stream for ParquetRecordBatchStreamAdapter<T> {
+impl<T, E> Stream for RecordBatchStreamTypeAdapter<T, E>
+where
+    T: Stream<Item = std::result::Result<DfRecordBatch, E>>,
+    E: std::error::Error + Send + Sync + 'static,
+{
     type Item = DfResult<DfRecordBatch>;
 
-    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let batch = futures::ready!(Pin::new(&mut self.stream).poll_next(cx))
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        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<T: Unpin + AsyncFileReader + Send + 'static> 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)
             })
         });

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<usize>,
     ) -> Result<DfSendableRecordBatchStream> {
         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,