feat: replace arrow2 with official implementation 🎉

Cargo.toml

@@ -20,7 +20,6 @@ members = [
     "src/common/time",
     "src/datanode",
     "src/datatypes",
-    "src/datatypes2",
     "src/frontend",
     "src/log-store",
     "src/meta-client",

benchmarks/Cargo.toml

@@ -5,10 +5,10 @@ edition = "2021"
 license = "Apache-2.0"
 
 [dependencies]
-arrow = "10"
+arrow = "26.0.0"
 clap = { version = "4.0", features = ["derive"] }
 client = { path = "../src/client" }
 indicatif = "0.17.1"
 itertools = "0.10.5"
-parquet = { version = "*" }
+parquet = "26.0.0"
 tokio = { version = "1.21", features = ["full"] }

benchmarks/src/bin/nyc-taxi.rs

@@ -20,7 +20,6 @@
 
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
-use std::sync::Arc;
 use std::time::Instant;
 
 use arrow::array::{ArrayRef, PrimitiveArray, StringArray, TimestampNanosecondArray};
@@ -32,9 +31,7 @@ use client::api::v1::column::Values;
 use client::api::v1::{Column, ColumnDataType, ColumnDef, CreateExpr, InsertExpr};
 use client::{Client, Database, Select};
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
-use parquet::arrow::{ArrowReader, ParquetFileArrowReader};
-use parquet::file::reader::FileReader;
-use parquet::file::serialized_reader::SerializedFileReader;
+use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 use tokio::task::JoinSet;
 
 const DATABASE_NAME: &str = "greptime";
@@ -86,10 +83,14 @@ async fn write_data(
     pb_style: ProgressStyle,
 ) -> u128 {
     let file = std::fs::File::open(&path).unwrap();
-    let file_reader = Arc::new(SerializedFileReader::new(file).unwrap());
-    let row_num = file_reader.metadata().file_metadata().num_rows();
-    let record_batch_reader = ParquetFileArrowReader::new(file_reader)
-        .get_record_reader(batch_size)
+    let record_batch_reader_builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+    let row_num = record_batch_reader_builder
+        .metadata()
+        .file_metadata()
+        .num_rows();
+    let record_batch_reader = record_batch_reader_builder
+        .with_batch_size(batch_size)
+        .build()
         .unwrap();
     let progress_bar = mpb.add(ProgressBar::new(row_num as _));
     progress_bar.set_style(pb_style);
@@ -210,9 +211,10 @@ fn build_values(column: &ArrayRef) -> Values {
         | DataType::FixedSizeList(_, _)
         | DataType::LargeList(_)
         | DataType::Struct(_)
-        | DataType::Union(_, _)
+        | DataType::Union(_, _, _)
         | DataType::Dictionary(_, _)
-        | DataType::Decimal(_, _)
+        | DataType::Decimal128(_, _)
+        | DataType::Decimal256(_, _)
         | DataType::Map(_, _) => todo!(),
     }
 }

src/api/greptime/v1/column.proto

@@ -32,7 +32,10 @@ message Column {
 
     repeated int32 date_values = 14;
     repeated int64 datetime_values = 15;
-    repeated int64 ts_millis_values = 16;
+    repeated int64 ts_second_values = 16;
+    repeated int64 ts_millisecond_values = 17;
+    repeated int64 ts_microsecond_values = 18;
+    repeated int64 ts_nanosecond_values = 19;
   }
   // The array of non-null values in this column.
   //
@@ -75,5 +78,8 @@ enum ColumnDataType {
   STRING = 12;
   DATE = 13;
   DATETIME = 14;
-  TIMESTAMP = 15;
+  TIMESTAMP_SECOND = 15;
+  TIMESTAMP_MILLISECOND = 16;
+  TIMESTAMP_MICROSECOND = 17;
+  TIMESTAMP_NANOSECOND = 18;
 }

src/api/src/helper.rs

@@ -15,6 +15,7 @@
 use common_base::BitVec;
 use common_time::timestamp::TimeUnit;
 use datatypes::prelude::ConcreteDataType;
+use datatypes::types::TimestampType;
 use datatypes::value::Value;
 use datatypes::vectors::VectorRef;
 use snafu::prelude::*;
@@ -56,7 +57,16 @@ impl From<ColumnDataTypeWrapper> for ConcreteDataType {
             ColumnDataType::String => ConcreteDataType::string_datatype(),
             ColumnDataType::Date => ConcreteDataType::date_datatype(),
             ColumnDataType::Datetime => ConcreteDataType::datetime_datatype(),
-            ColumnDataType::Timestamp => ConcreteDataType::timestamp_millis_datatype(),
+            ColumnDataType::TimestampSecond => ConcreteDataType::timestamp_second_datatype(),
+            ColumnDataType::TimestampMillisecond => {
+                ConcreteDataType::timestamp_millisecond_datatype()
+            }
+            ColumnDataType::TimestampMicrosecond => {
+                ConcreteDataType::timestamp_microsecond_datatype()
+            }
+            ColumnDataType::TimestampNanosecond => {
+                ConcreteDataType::timestamp_nanosecond_datatype()
+            }
         }
     }
 }
@@ -81,7 +91,12 @@ impl TryFrom<ConcreteDataType> for ColumnDataTypeWrapper {
             ConcreteDataType::String(_) => ColumnDataType::String,
             ConcreteDataType::Date(_) => ColumnDataType::Date,
             ConcreteDataType::DateTime(_) => ColumnDataType::Datetime,
-            ConcreteDataType::Timestamp(_) => ColumnDataType::Timestamp,
+            ConcreteDataType::Timestamp(unit) => match unit {
+                TimestampType::Second(_) => ColumnDataType::TimestampSecond,
+                TimestampType::Millisecond(_) => ColumnDataType::TimestampMillisecond,
+                TimestampType::Microsecond(_) => ColumnDataType::TimestampMicrosecond,
+                TimestampType::Nanosecond(_) => ColumnDataType::TimestampNanosecond,
+            },
             ConcreteDataType::Null(_) | ConcreteDataType::List(_) => {
                 return error::IntoColumnDataTypeSnafu { from: datatype }.fail()
             }
@@ -153,8 +168,20 @@ impl Values {
                 datetime_values: Vec::with_capacity(capacity),
                 ..Default::default()
             },
-            ColumnDataType::Timestamp => Values {
-                ts_millis_values: Vec::with_capacity(capacity),
+            ColumnDataType::TimestampSecond => Values {
+                ts_second_values: Vec::with_capacity(capacity),
+                ..Default::default()
+            },
+            ColumnDataType::TimestampMillisecond => Values {
+                ts_millisecond_values: Vec::with_capacity(capacity),
+                ..Default::default()
+            },
+            ColumnDataType::TimestampMicrosecond => Values {
+                ts_microsecond_values: Vec::with_capacity(capacity),
+                ..Default::default()
+            },
+            ColumnDataType::TimestampNanosecond => Values {
+                ts_nanosecond_values: Vec::with_capacity(capacity),
                 ..Default::default()
             },
         }
@@ -187,9 +214,12 @@ impl Column {
             Value::Binary(val) => values.binary_values.push(val.to_vec()),
             Value::Date(val) => values.date_values.push(val.val()),
             Value::DateTime(val) => values.datetime_values.push(val.val()),
-            Value::Timestamp(val) => values
-                .ts_millis_values
-                .push(val.convert_to(TimeUnit::Millisecond)),
+            Value::Timestamp(val) => match val.unit() {
+                TimeUnit::Second => values.ts_second_values.push(val.value()),
+                TimeUnit::Millisecond => values.ts_millisecond_values.push(val.value()),
+                TimeUnit::Microsecond => values.ts_microsecond_values.push(val.value()),
+                TimeUnit::Nanosecond => values.ts_nanosecond_values.push(val.value()),
+            },
             Value::List(_) => unreachable!(),
         });
         self.null_mask = null_mask.into_vec();
@@ -200,7 +230,10 @@ impl Column {
 mod tests {
     use std::sync::Arc;
 
-    use datatypes::vectors::BooleanVector;
+    use datatypes::vectors::{
+        BooleanVector, TimestampMicrosecondVector, TimestampMillisecondVector,
+        TimestampNanosecondVector, TimestampSecondVector,
+    };
 
     use super::*;
 
@@ -258,8 +291,8 @@ mod tests {
         let values = values.datetime_values;
         assert_eq!(2, values.capacity());
 
-        let values = Values::with_capacity(ColumnDataType::Timestamp, 2);
-        let values = values.ts_millis_values;
+        let values = Values::with_capacity(ColumnDataType::TimestampMillisecond, 2);
+        let values = values.ts_millisecond_values;
         assert_eq!(2, values.capacity());
     }
 
@@ -326,8 +359,8 @@ mod tests {
             ColumnDataTypeWrapper(ColumnDataType::Datetime).into()
         );
         assert_eq!(
-            ConcreteDataType::timestamp_millis_datatype(),
-            ColumnDataTypeWrapper(ColumnDataType::Timestamp).into()
+            ConcreteDataType::timestamp_millisecond_datatype(),
+            ColumnDataTypeWrapper(ColumnDataType::TimestampMillisecond).into()
         );
     }
 
@@ -394,8 +427,8 @@ mod tests {
             ConcreteDataType::datetime_datatype().try_into().unwrap()
         );
         assert_eq!(
-            ColumnDataTypeWrapper(ColumnDataType::Timestamp),
-            ConcreteDataType::timestamp_millis_datatype()
+            ColumnDataTypeWrapper(ColumnDataType::TimestampMillisecond),
+            ConcreteDataType::timestamp_millisecond_datatype()
                 .try_into()
                 .unwrap()
         );
@@ -412,7 +445,48 @@ mod tests {
         assert!(result.is_err());
         assert_eq!(
             result.unwrap_err().to_string(),
-            "Failed to create column datatype from List(ListType { inner: Boolean(BooleanType) })"
+            "Failed to create column datatype from List(ListType { item_type: Boolean(BooleanType) })"
+        );
+    }
+
+    #[test]
+    fn test_column_put_timestamp_values() {
+        let mut column = Column {
+            column_name: "test".to_string(),
+            semantic_type: 0,
+            values: Some(Values {
+                ..Default::default()
+            }),
+            null_mask: vec![],
+            datatype: 0,
+        };
+
+        let vector = Arc::new(TimestampNanosecondVector::from_vec(vec![1, 2, 3]));
+        column.push_vals(3, vector);
+        assert_eq!(
+            vec![1, 2, 3],
+            column.values.as_ref().unwrap().ts_nanosecond_values
+        );
+
+        let vector = Arc::new(TimestampMillisecondVector::from_vec(vec![4, 5, 6]));
+        column.push_vals(3, vector);
+        assert_eq!(
+            vec![4, 5, 6],
+            column.values.as_ref().unwrap().ts_millisecond_values
+        );
+
+        let vector = Arc::new(TimestampMicrosecondVector::from_vec(vec![7, 8, 9]));
+        column.push_vals(3, vector);
+        assert_eq!(
+            vec![7, 8, 9],
+            column.values.as_ref().unwrap().ts_microsecond_values
+        );
+
+        let vector = Arc::new(TimestampSecondVector::from_vec(vec![10, 11, 12]));
+        column.push_vals(3, vector);
+        assert_eq!(
+            vec![10, 11, 12],
+            column.values.as_ref().unwrap().ts_second_values
         );
     }
 

src/catalog/Cargo.toml

@@ -19,9 +19,7 @@ common-recordbatch = { path = "../common/recordbatch" }
 common-runtime = { path = "../common/runtime" }
 common-telemetry = { path = "../common/telemetry" }
 common-time = { path = "../common/time" }
-datafusion = { git = "https://github.com/apache/arrow-datafusion.git", branch = "arrow2", features = [
-    "simd",
-] }
+datafusion = "14.0.0"
 datatypes = { path = "../datatypes" }
 futures = "0.3"
 futures-util = "0.3"

src/catalog/src/error.rs

@@ -17,7 +17,7 @@ use std::any::Any;
 use common_error::ext::{BoxedError, ErrorExt};
 use common_error::prelude::{Snafu, StatusCode};
 use datafusion::error::DataFusionError;
-use datatypes::arrow;
+use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::RawSchema;
 use snafu::{Backtrace, ErrorCompat};
 
@@ -51,14 +51,12 @@ pub enum Error {
     SystemCatalog { msg: String, backtrace: Backtrace },
 
     #[snafu(display(
-        "System catalog table type mismatch, expected: binary, found: {:?} source: {}",
+        "System catalog table type mismatch, expected: binary, found: {:?}",
         data_type,
-        source
     ))]
     SystemCatalogTypeMismatch {
-        data_type: arrow::datatypes::DataType,
-        #[snafu(backtrace)]
-        source: datatypes::error::Error,
+        data_type: ConcreteDataType,
+        backtrace: Backtrace,
     },
 
     #[snafu(display("Invalid system catalog entry type: {:?}", entry_type))]
@@ -222,10 +220,11 @@ impl ErrorExt for Error {
             | Error::ValueDeserialize { .. }
             | Error::Io { .. } => StatusCode::StorageUnavailable,
 
-            Error::RegisterTable { .. } => StatusCode::Internal,
+            Error::RegisterTable { .. } | Error::SystemCatalogTypeMismatch { .. } => {
+                StatusCode::Internal
+            }
 
             Error::ReadSystemCatalog { source, .. } => source.status_code(),
-            Error::SystemCatalogTypeMismatch { source, .. } => source.status_code(),
             Error::InvalidCatalogValue { source, .. } => source.status_code(),
 
             Error::TableExists { .. } => StatusCode::TableAlreadyExists,
@@ -265,7 +264,6 @@ impl From<Error> for DataFusionError {
 #[cfg(test)]
 mod tests {
     use common_error::mock::MockError;
-    use datatypes::arrow::datatypes::DataType;
     use snafu::GenerateImplicitData;
 
     use super::*;
@@ -314,11 +312,8 @@ mod tests {
         assert_eq!(
             StatusCode::Internal,
             Error::SystemCatalogTypeMismatch {
-                data_type: DataType::Boolean,
-                source: datatypes::error::Error::UnsupportedArrowType {
-                    arrow_type: DataType::Boolean,
-                    backtrace: Backtrace::generate()
-                }
+                data_type: ConcreteDataType::binary_datatype(),
+                backtrace: Backtrace::generate(),
             }
             .status_code()
         );

src/catalog/src/helper.rs

@@ -138,7 +138,7 @@ impl TableGlobalKey {
 
 /// Table global info contains necessary info for a datanode to create table regions, including
 /// table id, table meta(schema...), region id allocation across datanodes.
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct TableGlobalValue {
     /// Id of datanode that created the global table info kv. only for debugging.
     pub node_id: u64,

src/catalog/src/local/manager.rs

@@ -145,27 +145,34 @@ impl LocalCatalogManager {
     /// Convert `RecordBatch` to a vector of `Entry`.
     fn record_batch_to_entry(rb: RecordBatch) -> Result<Vec<Entry>> {
         ensure!(
-            rb.df_recordbatch.columns().len() >= 6,
+            rb.num_columns() >= 6,
             SystemCatalogSnafu {
-                msg: format!("Length mismatch: {}", rb.df_recordbatch.columns().len())
+                msg: format!("Length mismatch: {}", rb.num_columns())
             }
         );
 
-        let entry_type = UInt8Vector::try_from_arrow_array(&rb.df_recordbatch.columns()[0])
-            .with_context(|_| SystemCatalogTypeMismatchSnafu {
-                data_type: rb.df_recordbatch.columns()[ENTRY_TYPE_INDEX]
-                    .data_type()
-                    .clone(),
+        let entry_type = rb
+            .column(ENTRY_TYPE_INDEX)
+            .as_any()
+            .downcast_ref::<UInt8Vector>()
+            .with_context(|| SystemCatalogTypeMismatchSnafu {
+                data_type: rb.column(ENTRY_TYPE_INDEX).data_type(),
             })?;
 
-        let key = BinaryVector::try_from_arrow_array(&rb.df_recordbatch.columns()[1])
-            .with_context(|_| SystemCatalogTypeMismatchSnafu {
-                data_type: rb.df_recordbatch.columns()[KEY_INDEX].data_type().clone(),
+        let key = rb
+            .column(KEY_INDEX)
+            .as_any()
+            .downcast_ref::<BinaryVector>()
+            .with_context(|| SystemCatalogTypeMismatchSnafu {
+                data_type: rb.column(KEY_INDEX).data_type(),
             })?;
 
-        let value = BinaryVector::try_from_arrow_array(&rb.df_recordbatch.columns()[3])
-            .with_context(|_| SystemCatalogTypeMismatchSnafu {
-                data_type: rb.df_recordbatch.columns()[VALUE_INDEX].data_type().clone(),
+        let value = rb
+            .column(VALUE_INDEX)
+            .as_any()
+            .downcast_ref::<BinaryVector>()
+            .with_context(|| SystemCatalogTypeMismatchSnafu {
+                data_type: rb.column(VALUE_INDEX).data_type(),
             })?;
 
         let mut res = Vec::with_capacity(rb.num_rows());