refactor: upgrade DataFusion, Arrow and Sqlparser (#1074)

* refactor: upgrade DataFusion, Arrow and Sqlparser

* fix: resolve PR comments

src/common/base/src/bytes.rs

@@ -20,6 +20,12 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer};
 #[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize)]
 pub struct Bytes(bytes::Bytes);
 
+impl From<Bytes> for bytes::Bytes {
+    fn from(value: Bytes) -> Self {
+        value.0
+    }
+}
+
 impl From<bytes::Bytes> for Bytes {
     fn from(bytes: bytes::Bytes) -> Bytes {
         Bytes(bytes)

src/common/grpc/Cargo.toml

@@ -16,7 +16,7 @@ common-runtime = { path = "../runtime" }
 dashmap = "5.4"
 datafusion.workspace = true
 datatypes = { path = "../../datatypes" }
-flatbuffers = "22"
+flatbuffers = "23.1"
 futures = "0.3"
 prost.workspace = true
 snafu = { version = "0.7", features = ["backtraces"] }

src/common/grpc/src/flight.rs

@@ -16,8 +16,9 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use api::v1::{AffectedRows, FlightMetadata};
-use arrow_flight::utils::{flight_data_from_arrow_batch, flight_data_to_arrow_batch};
+use arrow_flight::utils::flight_data_to_arrow_batch;
 use arrow_flight::{FlightData, IpcMessage, SchemaAsIpc};
+use common_base::bytes::Bytes;
 use common_recordbatch::{RecordBatch, RecordBatches};
 use datatypes::arrow;
 use datatypes::arrow::datatypes::Schema as ArrowSchema;
@@ -39,38 +40,58 @@ pub enum FlightMessage {
     AffectedRows(usize),
 }
 
-#[derive(Default)]
 pub struct FlightEncoder {
     write_options: writer::IpcWriteOptions,
+    data_gen: writer::IpcDataGenerator,
+    dictionary_tracker: writer::DictionaryTracker,
+}
+
+impl Default for FlightEncoder {
+    fn default() -> Self {
+        Self {
+            write_options: writer::IpcWriteOptions::default(),
+            data_gen: writer::IpcDataGenerator::default(),
+            dictionary_tracker: writer::DictionaryTracker::new(false),
+        }
+    }
 }
 
 impl FlightEncoder {
-    pub fn encode(&self, flight_message: FlightMessage) -> FlightData {
+    pub fn encode(&mut self, flight_message: FlightMessage) -> FlightData {
         match flight_message {
             FlightMessage::Schema(schema) => {
                 SchemaAsIpc::new(schema.arrow_schema(), &self.write_options).into()
             }
             FlightMessage::Recordbatch(recordbatch) => {
-                let (flight_dictionaries, flight_batch) = flight_data_from_arrow_batch(
-                    recordbatch.df_record_batch(),
-                    &self.write_options,
-                );
+                let (encoded_dictionaries, encoded_batch) = self
+                    .data_gen
+                    .encoded_batch(
+                        recordbatch.df_record_batch(),
+                        &mut self.dictionary_tracker,
+                        &self.write_options,
+                    )
+                    .expect("DictionaryTracker configured above to not fail on replacement");
 
                 // TODO(LFC): Handle dictionary as FlightData here, when we supported Arrow's Dictionary DataType.
                 // Currently we don't have a datatype corresponding to Arrow's Dictionary DataType,
                 // so there won't be any "dictionaries" here. Assert to be sure about it, and
                 // perform a "testing guard" in case we forgot to handle the possible "dictionaries"
                 // here in the future.
-                debug_assert_eq!(flight_dictionaries.len(), 0);
+                debug_assert_eq!(encoded_dictionaries.len(), 0);
 
-                flight_batch
+                encoded_batch.into()
             }
             FlightMessage::AffectedRows(rows) => {
                 let metadata = FlightMetadata {
                     affected_rows: Some(AffectedRows { value: rows as _ }),
                 }
                 .encode_to_vec();
-                FlightData::new(None, IpcMessage(build_none_flight_msg()), metadata, vec![])
+                FlightData::new(
+                    None,
+                    IpcMessage(build_none_flight_msg().into()),
+                    metadata,
+                    vec![],
+                )
             }
         }
     }
@@ -83,7 +104,8 @@ pub struct FlightDecoder {
 
 impl FlightDecoder {
     pub fn try_decode(&mut self, flight_data: FlightData) -> Result<FlightMessage> {
-        let message = root_as_message(flight_data.data_header.as_slice()).map_err(|e| {
+        let bytes = flight_data.data_header.slice(..);
+        let message = root_as_message(&bytes).map_err(|e| {
             InvalidFlightDataSnafu {
                 reason: e.to_string(),
             }
@@ -91,7 +113,7 @@ impl FlightDecoder {
         })?;
         match message.header_type() {
             MessageHeader::NONE => {
-                let metadata = FlightMetadata::decode(flight_data.app_metadata.as_slice())
+                let metadata = FlightMetadata::decode(flight_data.app_metadata)
                     .context(DecodeFlightDataSnafu)?;
                 if let Some(AffectedRows { value }) = metadata.affected_rows {
                     return Ok(FlightMessage::AffectedRows(value as _));
@@ -176,7 +198,7 @@ pub fn flight_messages_to_recordbatches(messages: Vec<FlightMessage>) -> Result<
     }
 }
 
-fn build_none_flight_msg() -> Vec<u8> {
+fn build_none_flight_msg() -> Bytes {
     let mut builder = FlatBufferBuilder::new();
 
     let mut message = arrow::ipc::MessageBuilder::new(&mut builder);
@@ -187,7 +209,7 @@ fn build_none_flight_msg() -> Vec<u8> {
     let data = message.finish();
     builder.finish(data, None);
 
-    builder.finished_data().to_vec()
+    builder.finished_data().into()
 }
 
 #[cfg(test)]

src/common/query/src/error.rs

@@ -239,7 +239,6 @@ impl From<BoxedError> for Error {
 
 #[cfg(test)]
 mod tests {
-    use datatypes::arrow::error::ArrowError;
     use snafu::GenerateImplicitData;
 
     use super::*;
@@ -286,7 +285,7 @@ mod tests {
     fn test_convert_df_recordbatch_stream_error() {
         let result: std::result::Result<i32, common_recordbatch::error::Error> =
             Err(common_recordbatch::error::Error::PollStream {
-                source: ArrowError::DivideByZero,
+                source: DataFusionError::Internal("blabla".to_string()),
                 backtrace: Backtrace::generate(),
             });
         let error = result

src/common/query/src/physical_plan.rs

@@ -315,7 +315,11 @@ mod test {
                 .unwrap()
                 .build()
                 .unwrap();
-        let physical_plan = ctx.create_physical_plan(&logical_plan).await.unwrap();
+        let physical_plan = ctx
+            .state()
+            .create_physical_plan(&logical_plan)
+            .await
+            .unwrap();
         let df_recordbatches = collect(physical_plan, Arc::new(TaskContext::from(&ctx)))
             .await
             .unwrap();

src/common/recordbatch/src/adapter.rs

@@ -18,9 +18,9 @@ use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use datafusion::arrow::datatypes::SchemaRef as DfSchemaRef;
+use datafusion::error::Result as DfResult;
 use datafusion::physical_plan::RecordBatchStream as DfRecordBatchStream;
 use datafusion_common::DataFusionError;
-use datatypes::arrow::error::{ArrowError, Result as ArrowResult};
 use datatypes::schema::{Schema, SchemaRef};
 use futures::ready;
 use snafu::ResultExt;
@@ -57,14 +57,14 @@ impl DfRecordBatchStream for DfRecordBatchStreamAdapter {
 }
 
 impl Stream for DfRecordBatchStreamAdapter {
-    type Item = ArrowResult<DfRecordBatch>;
+    type Item = DfResult<DfRecordBatch>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         match Pin::new(&mut self.stream).poll_next(cx) {
             Poll::Pending => Poll::Pending,
             Poll::Ready(Some(recordbatch)) => match recordbatch {
                 Ok(recordbatch) => Poll::Ready(Some(Ok(recordbatch.into_df_record_batch()))),
-                Err(e) => Poll::Ready(Some(Err(ArrowError::ExternalError(Box::new(e))))),
+                Err(e) => Poll::Ready(Some(Err(DataFusionError::External(Box::new(e))))),
             },
             Poll::Ready(None) => Poll::Ready(None),
         }

src/common/recordbatch/src/error.rs

@@ -55,7 +55,7 @@ pub enum Error {
 
     #[snafu(display("Failed to poll stream, source: {}", source))]
     PollStream {
-        source: datatypes::arrow::error::ArrowError,
+        source: datafusion::error::DataFusionError,
         backtrace: Backtrace,
     },
 

src/common/substrait/Cargo.toml

@@ -5,6 +5,8 @@ edition.workspace = true
 license.workspace = true
 
 [dependencies]
+async-recursion = "1.0"
+async-trait.workspace = true
 bytes = "1.1"
 catalog = { path = "../../catalog" }
 common-catalog = { path = "../catalog" }
@@ -15,6 +17,7 @@ datafusion-expr.workspace = true
 datatypes = { path = "../../datatypes" }
 futures = "0.3"
 prost.workspace = true
+session = { path = "../../session" }
 snafu.workspace = true
 table = { path = "../../table" }
 

src/common/substrait/src/df_expr.rs

@@ -635,8 +635,6 @@ mod utils {
             Operator::Modulo => "modulo",
             Operator::And => "and",
             Operator::Or => "or",
-            Operator::Like => "like",
-            Operator::NotLike => "not_like",
             Operator::IsDistinctFrom => "is_distinct_from",
             Operator::IsNotDistinctFrom => "is_not_distinct_from",
             Operator::RegexMatch => "regex_match",
@@ -649,8 +647,6 @@ mod utils {
             Operator::BitwiseShiftRight => "bitwise_shift_right",
             Operator::BitwiseShiftLeft => "bitwise_shift_left",
             Operator::StringConcat => "string_concat",
-            Operator::ILike => "i_like",
-            Operator::NotILike => "not_i_like",
         }
     }
 

src/common/substrait/src/df_logical.rs

@@ -14,16 +14,20 @@
 
 use std::sync::Arc;
 
+use async_recursion::async_recursion;
+use async_trait::async_trait;
 use bytes::{Buf, Bytes, BytesMut};
+use catalog::table_source::DfTableSourceProvider;
 use catalog::CatalogManagerRef;
-use common_error::prelude::BoxedError;
+use common_catalog::format_full_table_name;
 use common_telemetry::debug;
 use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
-use datafusion::common::{DFField, DFSchema};
+use datafusion::common::{DFField, DFSchema, OwnedTableReference};
 use datafusion::datasource::DefaultTableSource;
 use datafusion::physical_plan::project_schema;
-use datafusion_expr::{Filter, LogicalPlan, TableScan, TableSource};
+use datafusion_expr::{Filter, LogicalPlan, TableScan};
 use prost::Message;
+use session::context::QueryContext;
 use snafu::{ensure, OptionExt, ResultExt};
 use substrait_proto::proto::expression::mask_expression::{StructItem, StructSelect};
 use substrait_proto::proto::expression::MaskExpression;
@@ -37,8 +41,8 @@ use table::table::adapter::DfTableProviderAdapter;
 use crate::context::ConvertorContext;
 use crate::df_expr::{expression_from_df_expr, to_df_expr};
 use crate::error::{
-    self, DFInternalSnafu, DecodeRelSnafu, EmptyPlanSnafu, EncodeRelSnafu, Error, InternalSnafu,
-    InvalidParametersSnafu, MissingFieldSnafu, SchemaNotMatchSnafu, TableNotFoundSnafu,
+    self, DFInternalSnafu, DecodeRelSnafu, EmptyPlanSnafu, EncodeRelSnafu, Error,
+    InvalidParametersSnafu, MissingFieldSnafu, ResolveTableSnafu, SchemaNotMatchSnafu,
     UnknownPlanSnafu, UnsupportedExprSnafu, UnsupportedPlanSnafu,
 };
 use crate::schema::{from_schema, to_schema};
@@ -46,18 +50,19 @@ use crate::SubstraitPlan;
 
 pub struct DFLogicalSubstraitConvertor;
 
+#[async_trait]
 impl SubstraitPlan for DFLogicalSubstraitConvertor {
     type Error = Error;
 
     type Plan = LogicalPlan;
 
-    fn decode<B: Buf + Send>(
+    async fn decode<B: Buf + Send>(
         &self,
         message: B,
         catalog_manager: CatalogManagerRef,
     ) -> Result<Self::Plan, Self::Error> {
         let plan = Plan::decode(message).context(DecodeRelSnafu)?;
-        self.convert_plan(plan, catalog_manager)
+        self.convert_plan(plan, catalog_manager).await
     }
 
     fn encode(&self, plan: Self::Plan) -> Result<Bytes, Self::Error> {
@@ -71,7 +76,7 @@ impl SubstraitPlan for DFLogicalSubstraitConvertor {
 }
 
 impl DFLogicalSubstraitConvertor {
-    fn convert_plan(
+    async fn convert_plan(
         &self,
         mut plan: Plan,
         catalog_manager: CatalogManagerRef,
@@ -102,20 +107,25 @@ impl DFLogicalSubstraitConvertor {
             .fail()?
         };
 
-        self.rel_to_logical_plan(&mut ctx, Box::new(rel), catalog_manager)
+        // TODO(LFC): Create table provider from outside, respect "disallow_cross_schema_query" option in query engine state.
+        let mut table_provider =
+            DfTableSourceProvider::new(catalog_manager, false, &QueryContext::new());
+        self.rel_to_logical_plan(&mut ctx, Box::new(rel), &mut table_provider)
+            .await
     }
 
-    fn rel_to_logical_plan(
+    #[async_recursion]
+    async fn rel_to_logical_plan(
         &self,
         ctx: &mut ConvertorContext,
         rel: Box<Rel>,
-        catalog_manager: CatalogManagerRef,
+        table_provider: &mut DfTableSourceProvider,
     ) -> Result<LogicalPlan, Error> {
         let rel_type = rel.rel_type.context(EmptyPlanSnafu)?;
 
         // build logical plan
         let logical_plan = match rel_type {
-            RelType::Read(read_rel) => self.convert_read_rel(ctx, read_rel, catalog_manager)?,
+            RelType::Read(read_rel) => self.convert_read_rel(ctx, read_rel, table_provider).await?,
             RelType::Filter(filter) => {
                 let FilterRel {
                     common: _,
@@ -128,7 +138,7 @@ impl DFLogicalSubstraitConvertor {
                     field: "input",
                     plan: "Filter",
                 })?;
-                let input = Arc::new(self.rel_to_logical_plan(ctx, input, catalog_manager)?);
+                let input = Arc::new(self.rel_to_logical_plan(ctx, input, table_provider).await?);
 
                 let condition = condition.context(MissingFieldSnafu {
                     field: "condition",
@@ -191,11 +201,11 @@ impl DFLogicalSubstraitConvertor {
         Ok(logical_plan)
     }
 
-    fn convert_read_rel(
+    async fn convert_read_rel(
         &self,
         ctx: &mut ConvertorContext,
         read_rel: Box<ReadRel>,
-        catalog_manager: CatalogManagerRef,
+        table_provider: &mut DfTableSourceProvider,
     ) -> Result<LogicalPlan, Error> {
         // Extract the catalog, schema and table name from NamedTable. Assume the first three are those names.
         let read_type = read_rel.read_type.context(MissingFieldSnafu {
@@ -230,17 +240,17 @@ impl DFLogicalSubstraitConvertor {
             .projection
             .map(|mask_expr| self.convert_mask_expression(mask_expr));
 
-        // Get table handle from catalog manager
-        let table_ref = catalog_manager
-            .table(&catalog_name, &schema_name, &table_name)
-            .map_err(BoxedError::new)
-            .context(InternalSnafu)?
-            .context(TableNotFoundSnafu {
-                name: format!("{catalog_name}.{schema_name}.{table_name}"),
+        let table_ref = OwnedTableReference::Full {
+            catalog: catalog_name.clone(),
+            schema: schema_name.clone(),
+            table: table_name.clone(),
+        };
+        let adapter = table_provider
+            .resolve_table(table_ref)
+            .await
+            .with_context(|_| ResolveTableSnafu {
+                table_name: format_full_table_name(&catalog_name, &schema_name, &table_name),
             })?;
-        let adapter = Arc::new(DefaultTableSource::new(Arc::new(
-            DfTableProviderAdapter::new(table_ref),
-        )));
 
         // Get schema directly from the table, and compare it with the schema retrieved from substrait proto.
         let stored_schema = adapter.schema();
@@ -262,7 +272,7 @@ impl DFLogicalSubstraitConvertor {
         };
 
         // Calculate the projected schema
-        let qualified = &format!("{catalog_name}.{schema_name}.{table_name}");
+        let qualified = &format_full_table_name(&catalog_name, &schema_name, &table_name);
         let projected_schema = Arc::new(
             project_schema(&stored_schema, projection.as_ref())
                 .and_then(|x| {
@@ -281,7 +291,7 @@ impl DFLogicalSubstraitConvertor {
 
         // TODO(ruihang): Support limit(fetch)
         Ok(LogicalPlan::TableScan(TableScan {
-            table_name: format!("{catalog_name}.{schema_name}.{table_name}"),
+            table_name: qualified.to_string(),
             source: adapter,
             projection,
             projected_schema,
@@ -314,7 +324,7 @@ impl DFLogicalSubstraitConvertor {
             .fail()?,
             LogicalPlan::Filter(filter) => {
                 let input = Some(Box::new(
-                    self.logical_plan_to_rel(ctx, filter.input().clone())?,
+                    self.logical_plan_to_rel(ctx, filter.input.clone())?,
                 ));
 
                 let schema = plan
@@ -324,7 +334,7 @@ impl DFLogicalSubstraitConvertor {
                     .context(error::ConvertDfSchemaSnafu)?;
                 let condition = Some(Box::new(expression_from_df_expr(
                     ctx,
-                    filter.predicate(),
+                    &filter.predicate,
                     &schema,
                 )?));
 
@@ -396,7 +406,10 @@ impl DFLogicalSubstraitConvertor {
             | LogicalPlan::Explain(_)
             | LogicalPlan::Analyze(_)
             | LogicalPlan::Extension(_)
-            | LogicalPlan::Prepare(_) => InvalidParametersSnafu {
+            | LogicalPlan::Prepare(_)
+            | LogicalPlan::Dml(_)
+            | LogicalPlan::DescribeTable(_)
+            | LogicalPlan::Unnest(_) => InvalidParametersSnafu {
                 reason: format!(
                     "Trying to convert DDL/DML plan to substrait proto, plan: {plan:?}",
                 ),
@@ -524,6 +537,7 @@ mod test {
     use catalog::{CatalogList, CatalogProvider, RegisterTableRequest};
     use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
     use datafusion::common::{DFSchema, ToDFSchema};
+    use datafusion_expr::TableSource;
     use datatypes::schema::RawSchema;
     use table::requests::CreateTableRequest;
     use table::test_util::{EmptyTable, MockTableEngine};
@@ -572,7 +586,7 @@ mod test {
         let convertor = DFLogicalSubstraitConvertor;
 
         let proto = convertor.encode(plan.clone()).unwrap();
-        let tripped_plan = convertor.decode(proto, catalog).unwrap();
+        let tripped_plan = convertor.decode(proto, catalog).await.unwrap();
 
         assert_eq!(format!("{plan:?}"), format!("{tripped_plan:?}"));
     }

src/common/substrait/src/error.rs

@@ -105,6 +105,13 @@ pub enum Error {
         #[snafu(backtrace)]
         source: datatypes::error::Error,
     },
+
+    #[snafu(display("Unable to resolve table: {table_name}, error: {source}"))]
+    ResolveTable {
+        table_name: String,
+        #[snafu(backtrace)]
+        source: catalog::error::Error,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -127,6 +134,7 @@ impl ErrorExt for Error {
             | Error::SchemaNotMatch { .. } => StatusCode::InvalidArguments,
             Error::DFInternal { .. } | Error::Internal { .. } => StatusCode::Internal,
             Error::ConvertDfSchema { source } => source.status_code(),
+            Error::ResolveTable { source, .. } => source.status_code(),
         }
     }
 

src/common/substrait/src/lib.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #![feature(let_chains)]
+#![feature(trait_upcasting)]
 
 mod context;
 mod df_expr;
@@ -21,17 +22,19 @@ pub mod error;
 mod schema;
 mod types;
 
+use async_trait::async_trait;
 use bytes::{Buf, Bytes};
 use catalog::CatalogManagerRef;
 
 pub use crate::df_logical::DFLogicalSubstraitConvertor;
 
+#[async_trait]
 pub trait SubstraitPlan {
     type Error: std::error::Error;
 
     type Plan;
 
-    fn decode<B: Buf + Send>(
+    async fn decode<B: Buf + Send>(
         &self,
         message: B,
         catalog_manager: CatalogManagerRef,