refactor: remove Vectors from RecordBatch completely

Signed-off-by: luofucong <[email protected]>

Author: MichaelScofield

src/client/src/database.rs

@@ -435,10 +435,10 @@ impl Database {
                             .context(ExternalSnafu)?;
                         match flight_message {
                             FlightMessage::RecordBatch(arrow_batch) => {
-                                yield RecordBatch::try_from_df_record_batch(
+                                yield Ok(RecordBatch::from_df_record_batch(
                                     schema_cloned.clone(),
                                     arrow_batch,
-                                )
+                                ))
                             }
                             FlightMessage::Metrics(_) => {}
                             FlightMessage::AffectedRows(_) | FlightMessage::Schema(_) => {

src/client/src/region.rs

@@ -182,10 +182,8 @@ impl RegionRequester {
 
                 match flight_message {
                     FlightMessage::RecordBatch(record_batch) => {
-                        let result_to_yield = RecordBatch::try_from_df_record_batch(
-                            schema_cloned.clone(),
-                            record_batch,
-                        );
+                        let result_to_yield =
+                            RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch);
 
                         // get the next message from the stream. normally it should be a metrics message.
                         if let Some(next_flight_message_result) = flight_message_stream.next().await
@@ -219,7 +217,7 @@ impl RegionRequester {
                             stream_ended = true;
                         }
 
-                        yield result_to_yield;
+                        yield Ok(result_to_yield);
                     }
                     FlightMessage::Metrics(s) => {
                         // just a branch in case of some metrics message comes after other things.

src/common/query/src/error.rs

@@ -52,9 +52,6 @@ pub enum Error {
         data_type: ArrowDatatype,
     },
 
-    #[snafu(display("Failed to downcast vector: {}", err_msg))]
-    DowncastVector { err_msg: String },
-
     #[snafu(display("Invalid input type: {}", err_msg))]
     InvalidInputType {
         #[snafu(implicit)]
@@ -209,8 +206,7 @@ pub type Result<T> = std::result::Result<T, Error>;
 impl ErrorExt for Error {
     fn status_code(&self) -> StatusCode {
         match self {
-            Error::DowncastVector { .. }
-            | Error::InvalidInputState { .. }
+            Error::InvalidInputState { .. }
             | Error::ToScalarValue { .. }
             | Error::GetScalarVector { .. }
             | Error::ArrowCompute { .. }

src/common/recordbatch/src/adapter.rs

@@ -314,10 +314,10 @@ impl Stream for RecordBatchStreamAdapter {
                         metric_collector.record_batch_metrics,
                     );
                 }
-                Poll::Ready(Some(RecordBatch::try_from_df_record_batch(
+                Poll::Ready(Some(Ok(RecordBatch::from_df_record_batch(
                     self.schema(),
                     df_record_batch,
-                )))
+                ))))
             }
             Poll::Ready(None) => {
                 if let Metrics::Unresolved(df_plan) | Metrics::PartialResolved(df_plan, _) =

src/common/recordbatch/src/error.rs

@@ -133,18 +133,6 @@ pub enum Error {
         source: datatypes::error::Error,
     },
 
-    #[snafu(display(
-        "Failed to downcast vector of type '{:?}' to type '{:?}'",
-        from_type,
-        to_type
-    ))]
-    DowncastVector {
-        from_type: ConcreteDataType,
-        to_type: ConcreteDataType,
-        #[snafu(implicit)]
-        location: Location,
-    },
-
     #[snafu(display("Error occurs when performing arrow computation"))]
     ArrowCompute {
         #[snafu(source)]
@@ -217,8 +205,6 @@ impl ErrorExt for Error {
             | Error::PhysicalExpr { .. }
             | Error::RecordBatchSliceIndexOverflow { .. } => StatusCode::Internal,
 
-            Error::DowncastVector { .. } => StatusCode::Unexpected,
-
             Error::PollStream { .. } => StatusCode::EngineExecuteQuery,
 
             Error::ArrowCompute { .. } => StatusCode::IllegalState,

src/common/recordbatch/src/lib.rs

@@ -30,19 +30,20 @@ use adapter::RecordBatchMetrics;
 use arc_swap::ArcSwapOption;
 use common_base::readable_size::ReadableSize;
 pub use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
+use datatypes::arrow::array::{ArrayRef, AsArray, StringBuilder};
 use datatypes::arrow::compute::SortOptions;
 pub use datatypes::arrow::record_batch::RecordBatch as DfRecordBatch;
 use datatypes::arrow::util::pretty;
 use datatypes::prelude::{ConcreteDataType, VectorRef};
-use datatypes::scalars::{ScalarVector, ScalarVectorBuilder};
 use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
 use datatypes::types::{JsonFormat, jsonb_to_string};
-use datatypes::vectors::{BinaryVector, StringVectorBuilder};
 use error::Result;
 use futures::task::{Context, Poll};
 use futures::{Stream, TryStreamExt};
 pub use recordbatch::RecordBatch;
-use snafu::{OptionExt, ResultExt, ensure};
+use snafu::{ResultExt, ensure};
+
+use crate::error::NewDfRecordBatchSnafu;
 
 pub trait RecordBatchStream: Stream<Item = Result<RecordBatch>> {
     fn name(&self) -> &str {
@@ -92,32 +93,26 @@ pub fn map_json_type_to_string(
     mapped_schema: &SchemaRef,
 ) -> Result<RecordBatch> {
     let mut vectors = Vec::with_capacity(original_schema.column_schemas().len());
-    for (vector, schema) in batch.columns.iter().zip(original_schema.column_schemas()) {
+    for (vector, schema) in batch.columns().iter().zip(original_schema.column_schemas()) {
         if let ConcreteDataType::Json(j) = &schema.data_type {
             if matches!(&j.format, JsonFormat::Jsonb) {
-                let mut string_vector_builder = StringVectorBuilder::with_capacity(vector.len());
-                let binary_vector = vector
-                    .as_any()
-                    .downcast_ref::<BinaryVector>()
-                    .with_context(|| error::DowncastVectorSnafu {
-                        from_type: schema.data_type.clone(),
-                        to_type: ConcreteDataType::binary_datatype(),
-                    })?;
-                for value in binary_vector.iter_data() {
+                let mut string_vector_builder = StringBuilder::new();
+                let binary_vector = vector.as_binary::<i32>();
+                for value in binary_vector.iter() {
                     let Some(value) = value else {
-                        string_vector_builder.push(None);
+                        string_vector_builder.append_null();
                         continue;
                     };
                     let string_value =
                         jsonb_to_string(value).with_context(|_| error::CastVectorSnafu {
                             from_type: schema.data_type.clone(),
                             to_type: ConcreteDataType::string_datatype(),
                         })?;
-                    string_vector_builder.push(Some(string_value.as_str()));
+                    string_vector_builder.append_value(string_value);
                 }
 
                 let string_vector = string_vector_builder.finish();
-                vectors.push(Arc::new(string_vector) as VectorRef);
+                vectors.push(Arc::new(string_vector) as ArrayRef);
             } else {
                 vectors.push(vector.clone());
             }
@@ -126,7 +121,15 @@ pub fn map_json_type_to_string(
         }
     }
 
-    RecordBatch::new(mapped_schema.clone(), vectors)
+    let record_batch = datatypes::arrow::record_batch::RecordBatch::try_new(
+        mapped_schema.arrow_schema().clone(),
+        vectors,
+    )
+    .context(NewDfRecordBatchSnafu)?;
+    Ok(RecordBatch::from_df_record_batch(
+        mapped_schema.clone(),
+        record_batch,
+    ))
 }
 
 /// Maps the json type to string in the schema.
@@ -758,7 +761,7 @@ impl Stream for MemoryTrackedStream {
                 let additional = batch
                     .columns()
                     .iter()
-                    .map(|c| c.memory_size())
+                    .map(|c| c.get_array_memory_size())
                     .sum::<usize>();
 
                 if let Err(e) = self.permit.track(additional, self.total_tracked) {