From 23a1a5ecaab5be4e045a2b7e2198b5f267bca851 Mon Sep 17 00:00:00 2001 From: Yingwen Date: Mon, 14 Sep 2026 09:25:33 +0000 Subject: [PATCH] refactor: remove constant vector and replicate operation (#8999) * refactor: remove constant vector Signed-off-by: evenyag * refactor: remove vector replicate operation Signed-off-by: evenyag * fix: preserve scalar vector types with optional type hints Signed-off-by: evenyag * fix: remove obsolete mutable vector helper and import Signed-off-by: evenyag * fix: preserve typed nulls in struct scalar conversion Signed-off-by: evenyag --------- Signed-off-by: evenyag --- .../information_schema/columns.rs | 17 +- .../information_schema/key_column_usage.rs | 7 +- .../information_schema/partitions.rs | 25 +- .../information_schema/table_constraints.rs | 12 +- src/common/function/src/aggrs/vector/avg.rs | 14 +- .../function/src/aggrs/vector/product.rs | 15 +- src/common/function/src/aggrs/vector/sum.rs | 14 +- src/common/function/src/scalars/expression.rs | 2 - .../function/src/scalars/expression/binary.rs | 90 ------ src/common/function/src/scalars/test.rs | 32 +- src/common/macro/src/admin_fn.rs | 2 +- src/datatypes/src/schema/constraint.rs | 101 +++++- src/datatypes/src/types/cast.rs | 2 +- src/datatypes/src/value.rs | 82 ++++- src/datatypes/src/vectors.rs | 7 - src/datatypes/src/vectors/binary.rs | 1 - src/datatypes/src/vectors/boolean.rs | 5 - src/datatypes/src/vectors/constant.rs | 252 --------------- src/datatypes/src/vectors/decimal.rs | 37 --- src/datatypes/src/vectors/dictionary.rs | 48 +-- src/datatypes/src/vectors/eq.rs | 68 ---- src/datatypes/src/vectors/helper.rs | 306 ++++++++++-------- src/datatypes/src/vectors/null.rs | 11 - src/datatypes/src/vectors/operations.rs | 48 --- .../src/vectors/operations/filter.rs | 21 -- .../src/vectors/operations/replicate.rs | 232 ------------- src/datatypes/src/vectors/operations/take.rs | 26 +- src/datatypes/src/vectors/primitive.rs | 35 -- src/datatypes/src/vectors/string.rs | 1 - src/datatypes/src/vectors/struct_vector.rs | 66 +--- src/flow/src/expr/scalar.rs | 1 + src/mito2/src/read/flat_projection.rs | 71 ++++ src/mito2/src/read/projection.rs | 68 +++- 33 files changed, 534 insertions(+), 1185 deletions(-) delete mode 100644 src/common/function/src/scalars/expression/binary.rs delete mode 100644 src/datatypes/src/vectors/constant.rs delete mode 100644 src/datatypes/src/vectors/operations/replicate.rs diff --git a/src/catalog/src/system_schema/information_schema/columns.rs b/src/catalog/src/system_schema/information_schema/columns.rs index 4e68b093cd..e92efd3beb 100644 --- a/src/catalog/src/system_schema/information_schema/columns.rs +++ b/src/catalog/src/system_schema/information_schema/columns.rs @@ -31,7 +31,7 @@ use datatypes::scalars::ScalarVectorBuilder; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::value::Value; use datatypes::vectors::{ - ConstantVector, Int64Vector, Int64VectorBuilder, StringVector, StringVectorBuilder, VectorRef, + Int64Vector, Int64VectorBuilder, StringVector, StringVectorBuilder, VectorRef, }; use futures::TryStreamExt; use snafu::{OptionExt, ResultExt}; @@ -408,18 +408,9 @@ impl InformationSchemaColumnsBuilder { fn finish(&mut self) -> Result { let rows_num = self.collation_names.len(); - let privileges = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![DEFAULT_PRIVILEGES])), - rows_num, - )); - let empty_string = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![EMPTY_STR])), - rows_num, - )); - let srs_ids = Arc::new(ConstantVector::new( - Arc::new(Int64Vector::from(vec![None])), - rows_num, - )); + let privileges = Arc::new(StringVector::from(vec![DEFAULT_PRIVILEGES; rows_num])); + let empty_string = Arc::new(StringVector::from(vec![EMPTY_STR; rows_num])); + let srs_ids = Arc::new(Int64Vector::from(vec![None; rows_num])); let columns: Vec = vec![ Arc::new(self.catalog_names.finish()), diff --git a/src/catalog/src/system_schema/information_schema/key_column_usage.rs b/src/catalog/src/system_schema/information_schema/key_column_usage.rs index 99b0d7f6b7..4f8f116afa 100644 --- a/src/catalog/src/system_schema/information_schema/key_column_usage.rs +++ b/src/catalog/src/system_schema/information_schema/key_column_usage.rs @@ -26,7 +26,7 @@ use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream; use datatypes::prelude::{ConcreteDataType, MutableVector, ScalarVectorBuilder, VectorRef}; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::value::Value; -use datatypes::vectors::{ConstantVector, StringVector, StringVectorBuilder, UInt32VectorBuilder}; +use datatypes::vectors::{StringVector, StringVectorBuilder, UInt32VectorBuilder}; use futures_util::TryStreamExt; use snafu::{OptionExt, ResultExt}; use store_api::storage::{ScanRequest, TableId}; @@ -326,10 +326,7 @@ impl InformationSchemaKeyColumnUsageBuilder { fn finish(&mut self) -> Result { let rows_num = self.table_catalog.len(); - let null_string_vector = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![None as Option<&str>])), - rows_num, - )); + let null_string_vector = Arc::new(StringVector::from(vec![None as Option<&str>; rows_num])); let columns: Vec = vec![ Arc::new(self.constraint_catalog.finish()), Arc::new(self.constraint_schema.finish()), diff --git a/src/catalog/src/system_schema/information_schema/partitions.rs b/src/catalog/src/system_schema/information_schema/partitions.rs index b9396fe554..d2cce93150 100644 --- a/src/catalog/src/system_schema/information_schema/partitions.rs +++ b/src/catalog/src/system_schema/information_schema/partitions.rs @@ -29,8 +29,8 @@ use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::timestamp::TimestampSecond; use datatypes::value::Value; use datatypes::vectors::{ - ConstantVector, Int64Vector, Int64VectorBuilder, MutableVector, StringVector, - StringVectorBuilder, TimestampSecondVector, TimestampSecondVectorBuilder, UInt64VectorBuilder, + Int64Vector, Int64VectorBuilder, MutableVector, StringVector, StringVectorBuilder, + TimestampSecondVector, TimestampSecondVectorBuilder, UInt64VectorBuilder, }; use futures::{StreamExt, TryStreamExt}; use partition::manager::PartitionInfo; @@ -361,22 +361,11 @@ impl InformationSchemaPartitionsBuilder { fn finish(&mut self) -> Result { let rows_num = self.catalog_names.len(); - let null_string_vector = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![None as Option<&str>])), - rows_num, - )); - let null_i64_vector = Arc::new(ConstantVector::new( - Arc::new(Int64Vector::from(vec![None])), - rows_num, - )); - let null_timestamp_second_vector = Arc::new(ConstantVector::new( - Arc::new(TimestampSecondVector::from(vec![None])), - rows_num, - )); - let partition_methods = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec![Some("RANGE")])), - rows_num, - )); + let null_string_vector = Arc::new(StringVector::from(vec![None as Option<&str>; rows_num])); + let null_i64_vector = Arc::new(Int64Vector::from(vec![None; rows_num])); + let null_timestamp_second_vector = + Arc::new(TimestampSecondVector::from(vec![None; rows_num])); + let partition_methods = Arc::new(StringVector::from(vec![Some("RANGE"); rows_num])); let columns: Vec = vec![ Arc::new(self.catalog_names.finish()), diff --git a/src/catalog/src/system_schema/information_schema/table_constraints.rs b/src/catalog/src/system_schema/information_schema/table_constraints.rs index 8c975207f2..2208d303c2 100644 --- a/src/catalog/src/system_schema/information_schema/table_constraints.rs +++ b/src/catalog/src/system_schema/information_schema/table_constraints.rs @@ -27,7 +27,7 @@ use datatypes::prelude::{ConcreteDataType, MutableVector}; use datatypes::scalars::ScalarVectorBuilder; use datatypes::schema::{ColumnSchema, Schema, SchemaRef}; use datatypes::value::Value; -use datatypes::vectors::{ConstantVector, StringVector, StringVectorBuilder, VectorRef}; +use datatypes::vectors::{StringVector, StringVectorBuilder, VectorRef}; use futures::TryStreamExt; use snafu::{OptionExt, ResultExt}; use store_api::storage::{ScanRequest, TableId}; @@ -242,14 +242,8 @@ impl InformationSchemaTableConstraintsBuilder { fn finish(&mut self) -> Result { let rows_num = self.constraint_names.len(); - let constraint_catalogs = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec!["def"])), - rows_num, - )); - let enforceds = Arc::new(ConstantVector::new( - Arc::new(StringVector::from(vec!["YES"])), - rows_num, - )); + let constraint_catalogs = Arc::new(StringVector::from(vec!["def"; rows_num])); + let enforceds = Arc::new(StringVector::from(vec!["YES"; rows_num])); let columns: Vec = vec![ constraint_catalogs, diff --git a/src/common/function/src/aggrs/vector/avg.rs b/src/common/function/src/aggrs/vector/avg.rs index ddf1823d28..822a7c16ef 100644 --- a/src/common/function/src/aggrs/vector/avg.rs +++ b/src/common/function/src/aggrs/vector/avg.rs @@ -183,8 +183,6 @@ mod tests { use std::sync::Arc; use arrow::array::StringArray; - use datatypes::scalars::ScalarVector; - use datatypes::vectors::{ConstantVector, StringVector, Vector}; use super::*; @@ -252,14 +250,12 @@ mod tests { vec_avg.evaluate().unwrap() ); - // test update with constant vector + // test update with repeated values let mut vec_avg = VectorAvg::default(); - let v: Vec = vec![ - Arc::new(ConstantVector::new( - Arc::new(StringVector::from_vec(vec!["[1.0,2.0,3.0]".to_string()])), - 4, - )) - .to_arrow_array(), + let v = vec![ + ScalarValue::Utf8(Some("[1.0,2.0,3.0]".to_string())) + .to_array_of_size(4) + .unwrap(), ]; vec_avg.update_batch(&v).unwrap(); assert_eq!( diff --git a/src/common/function/src/aggrs/vector/product.rs b/src/common/function/src/aggrs/vector/product.rs index d8f201966a..11cfa4cb1a 100644 --- a/src/common/function/src/aggrs/vector/product.rs +++ b/src/common/function/src/aggrs/vector/product.rs @@ -159,9 +159,6 @@ impl Accumulator for VectorProduct { mod tests { use std::sync::Arc; - use datatypes::scalars::ScalarVector; - use datatypes::vectors::{ConstantVector, StringVector, Vector}; - use super::*; #[test] @@ -213,14 +210,12 @@ mod tests { vec_product.update_batch(&v).unwrap(); assert_eq!(ScalarValue::Binary(None), vec_product.evaluate().unwrap()); - // test update with constant vector + // test update with repeated values let mut vec_product = VectorProduct::default(); - let v: Vec = vec![ - Arc::new(ConstantVector::new( - Arc::new(StringVector::from_vec(vec!["[1.0,2.0,3.0]".to_string()])), - 4, - )) - .to_arrow_array(), + let v = vec![ + ScalarValue::Utf8(Some("[1.0,2.0,3.0]".to_string())) + .to_array_of_size(4) + .unwrap(), ]; vec_product.update_batch(&v).unwrap(); diff --git a/src/common/function/src/aggrs/vector/sum.rs b/src/common/function/src/aggrs/vector/sum.rs index 3f875a8412..0ee8f0df7b 100644 --- a/src/common/function/src/aggrs/vector/sum.rs +++ b/src/common/function/src/aggrs/vector/sum.rs @@ -171,8 +171,6 @@ mod tests { use std::sync::Arc; use arrow::array::StringArray; - use datatypes::scalars::ScalarVector; - use datatypes::vectors::{ConstantVector, StringVector, Vector}; use super::*; @@ -225,14 +223,12 @@ mod tests { vec_sum.update_batch(&v).unwrap(); assert_eq!(ScalarValue::Binary(None), vec_sum.evaluate().unwrap()); - // test update with constant vector + // test update with repeated values let mut vec_sum = VectorSum::default(); - let v: Vec = vec![ - Arc::new(ConstantVector::new( - Arc::new(StringVector::from_vec(vec!["[1.0,2.0,3.0]".to_string()])), - 4, - )) - .to_arrow_array(), + let v = vec![ + ScalarValue::Utf8(Some("[1.0,2.0,3.0]".to_string())) + .to_array_of_size(4) + .unwrap(), ]; vec_sum.update_batch(&v).unwrap(); assert_eq!( diff --git a/src/common/function/src/scalars/expression.rs b/src/common/function/src/scalars/expression.rs index 63ed40fc8f..777c91e7c0 100644 --- a/src/common/function/src/scalars/expression.rs +++ b/src/common/function/src/scalars/expression.rs @@ -12,13 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod binary; mod ctx; mod if_func; mod is_null; mod unary; -pub use binary::scalar_binary_op; pub use ctx::EvalContext; pub use unary::scalar_unary_op; diff --git a/src/common/function/src/scalars/expression/binary.rs b/src/common/function/src/scalars/expression/binary.rs deleted file mode 100644 index 4d9f6268a8..0000000000 --- a/src/common/function/src/scalars/expression/binary.rs +++ /dev/null @@ -1,90 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::iter; - -use common_query::error::Result; -use datatypes::prelude::*; -use datatypes::vectors::{ConstantVector, Helper}; - -use crate::scalars::expression::ctx::EvalContext; - -pub fn scalar_binary_op( - l: &VectorRef, - r: &VectorRef, - f: F, - ctx: &mut EvalContext, -) -> Result<::VectorType> -where - F: Fn(Option>, Option>, &mut EvalContext) -> Option, -{ - debug_assert!( - l.len() == r.len(), - "Size of vectors must match to apply binary expression" - ); - - let result = match (l.is_const(), r.is_const()) { - (false, true) => { - let left: &::VectorType = unsafe { Helper::static_cast(l) }; - let right: &ConstantVector = unsafe { Helper::static_cast(r) }; - let right: &::VectorType = unsafe { Helper::static_cast(right.inner()) }; - let b = right.get_data(0); - - let it = left.iter_data().map(|a| f(a, b.clone(), ctx)); - ::VectorType::from_owned_iterator(it) - } - - (false, false) => { - let left: &::VectorType = unsafe { Helper::static_cast(l) }; - let right: &::VectorType = unsafe { Helper::static_cast(r) }; - - let it = left - .iter_data() - .zip(right.iter_data()) - .map(|(a, b)| f(a, b, ctx)); - ::VectorType::from_owned_iterator(it) - } - - (true, false) => { - let left: &ConstantVector = unsafe { Helper::static_cast(l) }; - let left: &::VectorType = unsafe { Helper::static_cast(left.inner()) }; - let a = left.get_data(0); - - let right: &::VectorType = unsafe { Helper::static_cast(r) }; - let it = right.iter_data().map(|b| f(a.clone(), b, ctx)); - ::VectorType::from_owned_iterator(it) - } - - (true, true) => { - let left: &ConstantVector = unsafe { Helper::static_cast(l) }; - let left: &::VectorType = unsafe { Helper::static_cast(left.inner()) }; - let a = left.get_data(0); - - let right: &ConstantVector = unsafe { Helper::static_cast(r) }; - let right: &::VectorType = unsafe { Helper::static_cast(right.inner()) }; - let b = right.get_data(0); - - let it = iter::repeat(a) - .zip(iter::repeat(b)) - .map(|(a, b)| f(a, b, ctx)) - .take(left.len()); - ::VectorType::from_owned_iterator(it) - } - }; - - if let Some(error) = ctx.error.take() { - return Err(error); - } - Ok(result) -} diff --git a/src/common/function/src/scalars/test.rs b/src/common/function/src/scalars/test.rs index 8768d546a3..e7f59a6592 100644 --- a/src/common/function/src/scalars/test.rs +++ b/src/common/function/src/scalars/test.rs @@ -14,13 +14,43 @@ use std::fmt; +use common_query::error::Result; use datafusion::logical_expr::ColumnarValue; use datafusion_expr::{ScalarFunctionArgs, Signature, Volatility}; use datatypes::arrow::datatypes::DataType; +use datatypes::prelude::{Scalar, ScalarVector, VectorRef}; use datatypes::vectors::{Helper, Vector}; use crate::function::{Function, extract_args}; -use crate::scalars::expression::{EvalContext, scalar_binary_op}; +use crate::scalars::expression::EvalContext; + +fn scalar_binary_op( + l: &VectorRef, + r: &VectorRef, + f: F, + ctx: &mut EvalContext, +) -> Result<::VectorType> +where + F: Fn(Option>, Option>, &mut EvalContext) -> Option, +{ + debug_assert!( + l.len() == r.len(), + "Size of vectors must match to apply binary expression" + ); + + let left: &::VectorType = unsafe { Helper::static_cast(l) }; + let right: &::VectorType = unsafe { Helper::static_cast(r) }; + let result = ::VectorType::from_owned_iterator( + left.iter_data() + .zip(right.iter_data()) + .map(|(a, b)| f(a, b, ctx)), + ); + + if let Some(error) = ctx.error.take() { + return Err(error); + } + Ok(result) +} #[derive(Clone)] pub(crate) struct TestAndFunction { diff --git a/src/common/macro/src/admin_fn.rs b/src/common/macro/src/admin_fn.rs index c580389f3c..053a2ad3f4 100644 --- a/src/common/macro/src/admin_fn.rs +++ b/src/common/macro/src/admin_fn.rs @@ -304,7 +304,7 @@ fn build_struct( .and_then(|cv| match cv { common_query::prelude::ColumnarValue::Vector(v) => Ok(v), common_query::prelude::ColumnarValue::Scalar(s) => { - datatypes::vectors::Helper::try_from_scalar_value(s, args.number_rows) + datatypes::vectors::Helper::try_from_scalar_value(s, args.number_rows, None) .context(common_query::error::FromScalarValueSnafu) } }) diff --git a/src/datatypes/src/schema/constraint.rs b/src/datatypes/src/schema/constraint.rs index 60025be68e..d946a92bec 100644 --- a/src/datatypes/src/schema/constraint.rs +++ b/src/datatypes/src/schema/constraint.rs @@ -23,7 +23,7 @@ use crate::error::{self, Result}; use crate::types::cast; use crate::value::Value; use crate::vectors::operations::VectorOp; -use crate::vectors::{TimestampMillisecondVector, VectorRef}; +use crate::vectors::{Helper, TimestampMillisecondVector, VectorRef}; pub const CURRENT_TIMESTAMP: &str = "current_timestamp"; pub const CURRENT_TIMESTAMP_FN: &str = "current_timestamp()"; @@ -151,15 +151,20 @@ impl ColumnDefaultConstraint { ColumnDefaultConstraint::Value(v) => { ensure!(is_nullable || !v.is_null(), error::NullDefaultSnafu); - // TODO(yingwen): - // 1. For null value, we could use NullVector once it supports custom logical type. - // 2. For non null value, we could use ConstantVector, but it would cause all codes - // attempt to downcast the vector fail if they don't check whether the vector is const - // first. - let mut mutable_vector = data_type.create_mutable_vector(1); - mutable_vector.try_push_value_ref(&v.as_value_ref())?; - let base_vector = mutable_vector.to_vector(); - Ok(base_vector.replicate(&[num_rows])) + if let Ok(vector) = v.try_to_scalar_value(data_type).and_then(|scalar| { + Helper::try_from_scalar_value(scalar, num_rows, Some(data_type)) + }) { + return Ok(vector); + } + + // Some extension values, such as JSON nested in a struct, cannot safely + // round-trip through ScalarValue. Preserve their logical type with the + // type-specific vector builder instead. + let mut mutable_vector = data_type.create_mutable_vector(num_rows); + for _ in 0..num_rows { + mutable_vector.try_push_value_ref(&v.as_value_ref())?; + } + Ok(mutable_vector.to_vector()) } } } @@ -391,6 +396,82 @@ mod tests { assert_eq!(Value::Int32(10), v); } + #[test] + fn test_struct_default_null_fields_and_json() { + use crate::types::{StructField, StructType}; + use crate::value::StructValue; + + let inner_type = StructType::from([StructField::new( + "x", + ConcreteDataType::int32_datatype(), + true, + )]); + let inner = Value::Struct(StructValue::new(vec![Value::Null], inner_type)); + let json = crate::json::JsonSettings::default() + .encode(serde_json::json!({"answer": 42})) + .unwrap(); + let nested_type = StructType::from([StructField::new("nested", inner.data_type(), true)]); + let json_type = StructType::from([StructField::new("json", json.data_type(), true)]); + let values = [ + inner.clone(), + Value::Struct(StructValue::new(vec![inner], nested_type)), + Value::Struct(StructValue::new(vec![], StructType::default())), + Value::Struct(StructValue::new(vec![json], json_type)), + ]; + for value in values { + let data_type = value.data_type(); + // JSON children are read back as their underlying struct values. + let expected = serde_json::Value::try_from(value.clone()).unwrap(); + for num_rows in [1, 3] { + let vector = ColumnDefaultConstraint::Value(value.clone()) + .create_default_vector(&data_type, false, num_rows) + .unwrap(); + assert_eq!(data_type, vector.data_type()); + assert_eq!( + data_type.as_arrow_type(), + *vector.to_arrow_array().data_type() + ); + assert_eq!(num_rows, vector.len()); + assert_eq!(0, vector.null_count()); + for row in 0..num_rows { + assert_eq!( + expected, + serde_json::Value::try_from(vector.get(row)).unwrap() + ); + } + } + } + } + + #[test] + fn test_string_default_preserves_batch_schema() { + use arrow::datatypes::{Field, Schema}; + use arrow::record_batch::RecordBatch; + + for data_type in [ + ConcreteDataType::large_string_datatype(), + ConcreteDataType::utf8_view_datatype(), + ] { + let schema = Arc::new(Schema::new(vec![Field::new( + "tag", + data_type.as_arrow_type(), + true, + )])); + for value in [Value::from("greptime"), Value::Null] { + let vector = ColumnDefaultConstraint::Value(value.clone()) + .create_default_vector(&data_type, true, 3) + .unwrap(); + let batch = + RecordBatch::try_new(schema.clone(), vec![vector.to_arrow_array()]).unwrap(); + assert_eq!(3, batch.num_rows()); + assert_eq!(data_type, vector.data_type()); + for row in 0..batch.num_rows() { + assert_eq!(value, vector.get(row)); + } + } + } + } + #[test] fn test_create_default_vector_by_func() { let constraint = ColumnDefaultConstraint::Function(CURRENT_TIMESTAMP.to_string()); diff --git a/src/datatypes/src/types/cast.rs b/src/datatypes/src/types/cast.rs index d7a27a45e1..2f92713c26 100644 --- a/src/datatypes/src/types/cast.rs +++ b/src/datatypes/src/types/cast.rs @@ -26,7 +26,7 @@ pub fn cast(src_value: Value, dest_type: &ConcreteDataType) -> Result { } let src_type = src_value.data_type(); let scalar_value = src_value.try_to_scalar_value(&src_type)?; - let new_value = Helper::try_from_scalar_value(scalar_value, 1)? + let new_value = Helper::try_from_scalar_value(scalar_value, 1, None)? .cast(dest_type)? .get(0); Ok(new_value) diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 9ad9e1d46e..35a6acabd3 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -1085,11 +1085,21 @@ impl StructValue { } fn try_to_scalar_value(&self, output_type: &StructType) -> Result { + let output_fields = output_type.fields(); + ensure!( + self.items.len() == output_fields.len(), + InconsistentStructFieldsAndItemsSnafu { + field_len: output_fields.len(), + item_len: self.items.len() + } + ); let arrays = self .items .iter() - .map(|value| { - let scalar_value = value.try_to_scalar_value(&value.data_type())?; + .zip(output_fields.iter()) + .map(|(value, field)| { + // Null values need the declared field type to produce a typed null array. + let scalar_value = value.try_to_scalar_value(field.data_type())?; scalar_value .to_array() .context(ConvertScalarToArrowArraySnafu) @@ -1097,7 +1107,13 @@ impl StructValue { .collect::>>>()?; let fields = output_type.as_arrow_fields(); - let struct_array = StructArray::new(fields, arrays, None); + let struct_array = + StructArray::try_new_with_length(fields, arrays, None, 1).map_err(|error| { + error::ToScalarValueSnafu { + reason: error.to_string(), + } + .build() + })?; Ok(ScalarValue::Struct(Arc::new(struct_array))) } } @@ -3082,6 +3098,66 @@ pub(crate) mod tests { } } + #[test] + fn test_struct_scalar_null_fields() { + let struct_type = StructType::from([ + StructField::new("x", ConcreteDataType::int32_datatype(), true), + StructField::new("name", ConcreteDataType::string_datatype(), true), + ]); + let value = StructValue::new(vec![Value::Null, Value::from("hello")], struct_type.clone()); + let ScalarValue::Struct(array) = value.try_to_scalar_value(&struct_type).unwrap() else { + panic!("Expected struct scalar"); + }; + assert_eq!(1, array.len()); + assert_eq!(0, array.null_count()); + assert_eq!(&struct_type.as_arrow_fields(), array.fields()); + assert_eq!( + ScalarValue::Int32(None), + ScalarValue::try_from_array(array.column(0), 0).unwrap() + ); + + let nested_type = StructType::from([StructField::new( + "nested", + ConcreteDataType::struct_datatype(struct_type), + true, + )]); + for child in [Value::Struct(value), Value::Null] { + let nested = StructValue::new(vec![child.clone()], nested_type.clone()); + let ScalarValue::Struct(array) = nested.try_to_scalar_value(&nested_type).unwrap() + else { + panic!("Expected struct scalar"); + }; + let vector = crate::vectors::Helper::try_into_vector(array.column(0).clone()).unwrap(); + assert_eq!(child, vector.get(0)); + } + + let empty_type = StructType::default(); + let empty = StructValue::new(vec![], empty_type.clone()); + let ScalarValue::Struct(array) = empty.try_to_scalar_value(&empty_type).unwrap() else { + panic!("Expected struct scalar"); + }; + assert_eq!(1, array.len()); + assert_eq!(0, array.num_columns()); + assert_eq!(0, array.null_count()); + } + + #[test] + fn test_struct_scalar_invalid_fields() { + let struct_type = StructType::from([StructField::new( + "x", + ConcreteDataType::int32_datatype(), + false, + )]); + for child in [Value::Null, Value::from("wrong type")] { + let value = StructValue::new(vec![child], struct_type.clone()); + assert!(value.try_to_scalar_value(&struct_type).is_err()); + } + let value = StructValue::new(vec![Value::Int32(1)], struct_type.clone()); + assert!(value.try_to_scalar_value(&StructType::default()).is_err()); + let empty = StructValue::new(vec![], StructType::default()); + assert!(empty.try_to_scalar_value(&struct_type).is_err()); + } + #[test] fn test_timestamp_to_scalar_value() { assert_eq!( diff --git a/src/datatypes/src/vectors.rs b/src/datatypes/src/vectors.rs index 593bec833d..38db5940df 100644 --- a/src/datatypes/src/vectors.rs +++ b/src/datatypes/src/vectors.rs @@ -27,7 +27,6 @@ use crate::vectors::operations::VectorOp; mod binary; mod boolean; -mod constant; mod date; mod decimal; mod dictionary; @@ -48,7 +47,6 @@ mod validity; pub use binary::{BinaryVector, BinaryVectorBuilder}; pub use boolean::{BooleanVector, BooleanVectorBuilder}; -pub use constant::ConstantVector; pub use date::{DateVector, DateVectorBuilder}; pub use decimal::{Decimal128Vector, Decimal128VectorBuilder}; pub(crate) use dictionary::StringDictionaryVectorBuilder; @@ -126,11 +124,6 @@ pub trait Vector: Send + Sync + Serializable + Debug + VectorOp { /// This is `O(1)`. fn null_count(&self) -> usize; - /// Returns true when it's a ConstantColumn - fn is_const(&self) -> bool { - false - } - /// Returns whether row is null. fn is_null(&self, row: usize) -> bool; diff --git a/src/datatypes/src/vectors/binary.rs b/src/datatypes/src/vectors/binary.rs index 6d9954fb80..77344b4651 100644 --- a/src/datatypes/src/vectors/binary.rs +++ b/src/datatypes/src/vectors/binary.rs @@ -479,7 +479,6 @@ mod tests { assert_eq!(2, v.len()); assert_eq!("BinaryVector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_valid()); assert!(!v.only_null()); assert_eq!(128, v.memory_size()); diff --git a/src/datatypes/src/vectors/boolean.rs b/src/datatypes/src/vectors/boolean.rs index 4ea7a4591b..48625da516 100644 --- a/src/datatypes/src/vectors/boolean.rs +++ b/src/datatypes/src/vectors/boolean.rs @@ -40,10 +40,6 @@ impl BooleanVector { pub fn as_boolean_array(&self) -> &BooleanArray { &self.array } - - pub(crate) fn false_count(&self) -> usize { - self.array.false_count() - } } impl From> for BooleanVector { @@ -251,7 +247,6 @@ mod tests { let v = BooleanVector::from(bools.clone()); assert_eq!(9, v.len()); assert_eq!("BooleanVector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_valid()); assert!(!v.only_null()); assert_eq!(2, v.memory_size()); diff --git a/src/datatypes/src/vectors/constant.rs b/src/datatypes/src/vectors/constant.rs deleted file mode 100644 index 3a0adc7b62..0000000000 --- a/src/datatypes/src/vectors/constant.rs +++ /dev/null @@ -1,252 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use std::any::Any; -use std::fmt; -use std::sync::Arc; - -use arrow::array::{Array, ArrayRef, UInt32Array}; -use snafu::{ResultExt, ensure}; - -use crate::data_type::ConcreteDataType; -use crate::error::{self, Result, SerializeSnafu}; -use crate::serialize::Serializable; -use crate::value::{Value, ValueRef}; -use crate::vectors::{BooleanVector, Helper, UInt32Vector, Validity, Vector, VectorRef}; - -#[derive(Clone)] -pub struct ConstantVector { - length: usize, - vector: VectorRef, -} - -impl ConstantVector { - /// Create a new [ConstantVector]. - /// - /// # Panics - /// Panics if `vector.len() != 1`. - pub fn new(vector: VectorRef, length: usize) -> Self { - assert_eq!(1, vector.len()); - - // Avoid const recursion. - if vector.is_const() { - let vec: &ConstantVector = unsafe { Helper::static_cast(&vector) }; - return Self::new(vec.inner().clone(), length); - } - Self { vector, length } - } - - pub fn inner(&self) -> &VectorRef { - &self.vector - } - - /// Returns the constant value. - pub fn get_constant_ref(&self) -> ValueRef<'_> { - self.vector.get_ref(0) - } - - pub(crate) fn replicate_vector(&self, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), self.len()); - - if offsets.is_empty() { - return self.slice(0, 0); - } - - Arc::new(ConstantVector::new( - self.vector.clone(), - *offsets.last().unwrap(), - )) - } - - pub(crate) fn filter_vector(&self, filter: &BooleanVector) -> Result { - let length = self.len() - filter.false_count(); - if length == self.len() { - return Ok(Arc::new(self.clone())); - } - Ok(Arc::new(ConstantVector::new(self.inner().clone(), length))) - } - - pub(crate) fn cast_vector(&self, to_type: &ConcreteDataType) -> Result { - Ok(Arc::new(ConstantVector::new( - self.inner().cast(to_type)?, - self.length, - ))) - } - - pub(crate) fn take_vector(&self, indices: &UInt32Vector) -> Result { - if indices.is_empty() { - return Ok(self.slice(0, 0)); - } - ensure!( - indices.null_count() == 0, - error::UnsupportedOperationSnafu { - op: "taking a null index", - vector_type: self.vector_type_name(), - } - ); - - let len = self.len(); - let arr = indices.to_arrow_array(); - let indices_arr = arr.as_any().downcast_ref::().unwrap(); - if !arrow::compute::min_boolean( - &arrow::compute::kernels::cmp::lt(indices_arr, &UInt32Array::new_scalar(len as u32)) - .unwrap(), - ) - .unwrap() - { - panic!( - "Array index out of bounds, cannot take index out of the length of the array: {len}" - ); - } - - Ok(Arc::new(ConstantVector::new( - self.inner().clone(), - indices.len(), - ))) - } -} - -impl Vector for ConstantVector { - fn data_type(&self) -> ConcreteDataType { - self.vector.data_type() - } - - fn vector_type_name(&self) -> String { - "ConstantVector".to_string() - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn len(&self) -> usize { - self.length - } - - fn to_arrow_array(&self) -> ArrayRef { - let v = self.vector.replicate(&[self.length]); - v.to_arrow_array() - } - - fn to_boxed_arrow_array(&self) -> Box { - let v = self.vector.replicate(&[self.length]); - v.to_boxed_arrow_array() - } - - fn is_const(&self) -> bool { - true - } - - fn validity(&self) -> Validity { - if self.vector.is_null(0) { - Validity::all_null(self.length) - } else { - Validity::all_valid(self.length) - } - } - - fn memory_size(&self) -> usize { - self.vector.memory_size() - } - - fn is_null(&self, _row: usize) -> bool { - self.vector.is_null(0) - } - - fn only_null(&self) -> bool { - self.vector.is_null(0) - } - - fn slice(&self, _offset: usize, length: usize) -> VectorRef { - Arc::new(Self { - vector: self.vector.clone(), - length, - }) - } - - fn get(&self, _index: usize) -> Value { - self.vector.get(0) - } - - fn get_ref(&self, _index: usize) -> ValueRef<'_> { - self.vector.get_ref(0) - } - - fn null_count(&self) -> usize { - if self.only_null() { self.len() } else { 0 } - } -} - -impl fmt::Debug for ConstantVector { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "ConstantVector([{:?}; {}])", self.get(0), self.len()) - } -} - -impl Serializable for ConstantVector { - fn serialize_to_json(&self) -> Result> { - std::iter::repeat_n(self.get(0), self.len()) - .map(serde_json::Value::try_from) - .collect::>() - .context(SerializeSnafu) - } -} - -#[cfg(test)] -mod tests { - use arrow::datatypes::DataType as ArrowDataType; - - use super::*; - use crate::vectors::Int32Vector; - - #[test] - fn test_constant_vector_misc() { - let a = Int32Vector::from_slice(vec![1]); - let c = ConstantVector::new(Arc::new(a), 10); - - assert_eq!("ConstantVector", c.vector_type_name()); - assert!(c.is_const()); - assert_eq!(10, c.len()); - assert!(c.validity().is_all_valid()); - assert!(!c.only_null()); - assert_eq!(4, c.memory_size()); - - for i in 0..10 { - assert!(!c.is_null(i)); - assert_eq!(Value::Int32(1), c.get(i)); - } - - let arrow_arr = c.to_arrow_array(); - assert_eq!(10, arrow_arr.len()); - assert_eq!(&ArrowDataType::Int32, arrow_arr.data_type()); - } - - #[test] - fn test_debug_null_array() { - let a = Int32Vector::from_slice(vec![1]); - let c = ConstantVector::new(Arc::new(a), 10); - - let s = format!("{c:?}"); - assert_eq!(s, "ConstantVector([Int32(1); 10])"); - } - - #[test] - fn test_serialize_json() { - let a = Int32Vector::from_slice(vec![1]); - let c = ConstantVector::new(Arc::new(a), 10); - - let s = serde_json::to_string(&c.serialize_to_json().unwrap()).unwrap(); - assert_eq!(s, "[1,1,1,1,1,1,1,1,1,1]"); - } -} diff --git a/src/datatypes/src/vectors/decimal.rs b/src/datatypes/src/vectors/decimal.rs index 46fa73297a..34d066f072 100644 --- a/src/datatypes/src/vectors/decimal.rs +++ b/src/datatypes/src/vectors/decimal.rs @@ -385,43 +385,6 @@ impl Decimal128VectorBuilder { vectors::impl_try_from_arrow_array_for_vector!(Decimal128Array, Decimal128Vector); -pub(crate) fn replicate_decimal128( - vector: &Decimal128Vector, - offsets: &[usize], -) -> Decimal128Vector { - assert_eq!(offsets.len(), vector.len()); - - if offsets.is_empty() { - return vector.get_slice(0, 0); - } - - // Safety: safe to unwrap because we the vector ensures precision and scale are valid. - let mut builder = Decimal128VectorBuilder::with_capacity(*offsets.last().unwrap()) - .with_precision_and_scale(vector.precision(), vector.scale()) - .unwrap(); - - let mut previous_offset = 0; - - for (offset, value) in offsets.iter().zip(vector.array.iter()) { - let repeat_times = *offset - previous_offset; - match value { - Some(data) => { - unsafe { - // Safety: std::iter::Repeat and std::iter::Take implement TrustedLen. - builder - .mutable_array - .append_trusted_len_iter(std::iter::repeat_n(data, repeat_times)); - } - } - None => { - builder.mutable_array.append_nulls(repeat_times); - } - } - previous_offset = *offset; - } - builder.finish() -} - #[cfg(test)] pub mod tests { use arrow_array::Decimal128Array; diff --git a/src/datatypes/src/vectors/dictionary.rs b/src/datatypes/src/vectors/dictionary.rs index d41daf4638..5f117b0ae8 100644 --- a/src/datatypes/src/vectors/dictionary.rs +++ b/src/datatypes/src/vectors/dictionary.rs @@ -17,8 +17,7 @@ use std::fmt; use std::sync::Arc; use arrow::array::{ - Array, ArrayBuilder, ArrayRef, DictionaryArray, PrimitiveArray, PrimitiveBuilder, - StringDictionaryBuilder, + Array, ArrayBuilder, ArrayRef, DictionaryArray, PrimitiveArray, StringDictionaryBuilder, }; use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, UInt32Type}; use serde_json::Value as JsonValue; @@ -320,36 +319,6 @@ impl<'a, K: ArrowDictionaryKeyType> Iterator for DictionaryIter<'a, K> { } impl VectorOp for DictionaryVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - let keys = self.array.keys(); - let mut replicated_keys = PrimitiveBuilder::new(); - - let mut previous_offset = 0; - let mut key_iter = keys.iter().chain(std::iter::repeat(None)); - for &offset in offsets { - let key = key_iter.next().unwrap(); - - // repeat this key (offset - previous_offset) times - let repeat_count = offset - previous_offset; - for _ in 0..repeat_count { - replicated_keys.append_option(key); - } - - previous_offset = offset; - } - - let new_keys = replicated_keys.finish(); - let new_array = DictionaryArray::try_new(new_keys, self.values().clone()) - .expect("Failed to create replicated dictionary array"); - - Arc::new(Self { - array: new_array, - key_type: self.key_type.clone(), - item_type: self.item_type.clone(), - item_vector: self.item_vector.clone(), - }) - } - fn filter(&self, filter: &vectors::BooleanVector) -> Result { let key_array: ArrayRef = Arc::new(self.array.keys().clone()); let key_vector = Helper::try_into_vector(&key_array)?; @@ -485,21 +454,6 @@ mod tests { assert_eq!(sliced.get(2), Value::Null); } - #[test] - fn test_replicate() { - let dict_vec = create_test_dictionary(); - - // Replicate with offsets [0, 2, 5] - should get values at these indices - let offsets = vec![0, 2, 5]; - let replicated = dict_vec.replicate(&offsets); - assert_eq!(replicated.len(), 5); - assert_eq!(replicated.get(0), Value::String("b".to_string().into())); - assert_eq!(replicated.get(1), Value::String("b".to_string().into())); - assert_eq!(replicated.get(2), Value::String("c".to_string().into())); - assert_eq!(replicated.get(3), Value::String("c".to_string().into())); - assert_eq!(replicated.get(4), Value::String("c".to_string().into())); - } - #[test] fn test_filter() { let dict_vec = create_test_dictionary(); diff --git a/src/datatypes/src/vectors/eq.rs b/src/datatypes/src/vectors/eq.rs index e8e942d3ff..937e4587c3 100644 --- a/src/datatypes/src/vectors/eq.rs +++ b/src/datatypes/src/vectors/eq.rs @@ -18,7 +18,6 @@ use common_time::interval::IntervalUnit; use crate::data_type::DataType; use crate::types::{DurationType, TimeType, TimestampType}; -use crate::vectors::constant::ConstantVector; use crate::vectors::struct_vector::StructVector; use crate::vectors::{ BinaryVector, BooleanVector, DateVector, Decimal128Vector, DurationMicrosecondVector, @@ -58,23 +57,6 @@ fn equal(lhs: &dyn Vector, rhs: &dyn Vector) -> bool { return false; } - if lhs.is_const() || rhs.is_const() { - // Length has been checked before, so we only need to compare inner - // vector here. - return equal( - &**lhs - .as_any() - .downcast_ref::() - .unwrap() - .inner(), - &**rhs - .as_any() - .downcast_ref::() - .unwrap() - .inner(), - ); - } - use crate::data_type::ConcreteDataType::*; let lhs_type = lhs.data_type(); @@ -191,10 +173,6 @@ mod tests { Some(b"world".to_vec()), ]))); assert_vector_ref_eq(Arc::new(BooleanVector::from(vec![true, false]))); - assert_vector_ref_eq(Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - ))); assert_vector_ref_eq(Arc::new(BooleanVector::from(vec![true, false]))); assert_vector_ref_eq(Arc::new(DateVector::from(vec![Some(100), Some(120)]))); assert_vector_ref_eq(Arc::new(TimestampSecondVector::from_values([100, 120]))); @@ -259,22 +237,6 @@ mod tests { ]))); } - // Regression: second arm must downcast `rhs` (was `lhs`), or same-length ConstantVectors - // with different inners compare equal. - #[test] - fn test_constant_vector_eq_compares_both_inners() { - assert_vector_ref_ne( - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - )), - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![false])), - 5, - )), - ); - } - #[test] fn test_vector_ne() { assert_vector_ref_ne( @@ -289,36 +251,6 @@ mod tests { Arc::new(Int32Vector::from_slice([1, 2, 3, 4])), Arc::new(BooleanVector::from(vec![true, true])), ); - assert_vector_ref_ne( - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - )), - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 4, - )), - ); - assert_vector_ref_ne( - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - )), - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![false])), - 4, - )), - ); - assert_vector_ref_ne( - Arc::new(ConstantVector::new( - Arc::new(BooleanVector::from(vec![true])), - 5, - )), - Arc::new(ConstantVector::new( - Arc::new(Int32Vector::from_slice(vec![1])), - 4, - )), - ); assert_vector_ref_ne(Arc::new(NullVector::new(5)), Arc::new(NullVector::new(8))); assert_vector_ref_ne( diff --git a/src/datatypes/src/vectors/helper.rs b/src/datatypes/src/vectors/helper.rs index 685b07419f..3d9e71853c 100644 --- a/src/datatypes/src/vectors/helper.rs +++ b/src/datatypes/src/vectors/helper.rs @@ -30,21 +30,21 @@ use datafusion_common::ScalarValue; use snafu::{OptionExt, ResultExt}; use crate::data_type::ConcreteDataType; -use crate::error::{self, ConvertArrowArrayToScalarsSnafu, Result}; +use crate::error::{self, ConvertScalarToArrowArraySnafu, Result}; use crate::prelude::DataType; -use crate::scalars::{Scalar, ScalarVectorBuilder}; +use crate::scalars::Scalar; use crate::types::StructType; -use crate::value::{ListValue, ListValueRef, Value}; +use crate::value::Value; use crate::vectors::struct_vector::StructVector; use crate::vectors::{ - BinaryVector, BooleanVector, ConstantVector, DateVector, Decimal128Vector, DictionaryVector, + BinaryVector, BooleanVector, DateVector, Decimal128Vector, DictionaryVector, DurationMicrosecondVector, DurationMillisecondVector, DurationNanosecondVector, DurationSecondVector, Float32Vector, Float64Vector, Int8Vector, Int16Vector, Int32Vector, Int64Vector, IntervalDayTimeVector, IntervalMonthDayNanoVector, IntervalYearMonthVector, - ListVector, ListVectorBuilder, MutableVector, NullVector, StringVector, TimeMicrosecondVector, - TimeMillisecondVector, TimeNanosecondVector, TimeSecondVector, TimestampMicrosecondVector, - TimestampMillisecondVector, TimestampNanosecondVector, TimestampSecondVector, UInt8Vector, - UInt16Vector, UInt32Vector, UInt64Vector, Vector, VectorRef, + ListVector, NullVector, StringVector, TimeMicrosecondVector, TimeMillisecondVector, + TimeNanosecondVector, TimeSecondVector, TimestampMicrosecondVector, TimestampMillisecondVector, + TimestampNanosecondVector, TimestampSecondVector, UInt8Vector, UInt16Vector, UInt32Vector, + UInt64Vector, Vector, VectorRef, }; /// Helper functions for `Vector`. @@ -102,135 +102,50 @@ impl Helper { }) } - /// Try to cast an arrow scalar value into vector - pub fn try_from_scalar_value(value: ScalarValue, length: usize) -> Result { - let vector = match value { - ScalarValue::Null => ConstantVector::new(Arc::new(NullVector::new(1)), length), - ScalarValue::Boolean(v) => { - ConstantVector::new(Arc::new(BooleanVector::from(vec![v])), length) + /// Materializes an Arrow scalar as a vector of the given length. + /// + /// With a type hint, casts the array to the requested representation and returns + /// an error if the vector cannot preserve that concrete type. Without a hint, + /// normalizes the scalar to a supported GreptimeDB representation. + pub fn try_from_scalar_value( + value: ScalarValue, + length: usize, + type_hint: Option<&ConcreteDataType>, + ) -> Result { + if let Some(data_type) = type_hint { + let mut array = value + .to_array_of_size(length) + .context(ConvertScalarToArrowArraySnafu)?; + let arrow_type = data_type.as_arrow_type(); + if array.data_type() != &arrow_type { + array = compute::cast(&array, &arrow_type).context(error::ArrowComputeSnafu)?; } - ScalarValue::Float16(v) => ConstantVector::new( - Arc::new(Float32Vector::from(vec![v.map(f32::from)])), - length, - ), - ScalarValue::Float32(v) => { - ConstantVector::new(Arc::new(Float32Vector::from(vec![v])), length) + let vector = Self::try_into_vector(array)?; + if &vector.data_type() != data_type { + return error::CastTypeSnafu { + msg: format!( + "Scalar materialization produced {:?}, expected {data_type:?}", + vector.data_type() + ), + } + .fail(); } - ScalarValue::Float64(v) => { - ConstantVector::new(Arc::new(Float64Vector::from(vec![v])), length) + return Ok(vector); + } + + let value = match value { + // GreptimeDB doesn't support Float16 vectors. + ScalarValue::Float16(v) => ScalarValue::Float32(v.map(f32::from)), + ScalarValue::LargeUtf8(v) => ScalarValue::Utf8(v), + ScalarValue::LargeBinary(v) | ScalarValue::FixedSizeBinary(_, v) => { + ScalarValue::Binary(v) } - ScalarValue::Int8(v) => { - ConstantVector::new(Arc::new(Int8Vector::from(vec![v])), length) - } - ScalarValue::Int16(v) => { - ConstantVector::new(Arc::new(Int16Vector::from(vec![v])), length) - } - ScalarValue::Int32(v) => { - ConstantVector::new(Arc::new(Int32Vector::from(vec![v])), length) - } - ScalarValue::Int64(v) => { - ConstantVector::new(Arc::new(Int64Vector::from(vec![v])), length) - } - ScalarValue::UInt8(v) => { - ConstantVector::new(Arc::new(UInt8Vector::from(vec![v])), length) - } - ScalarValue::UInt16(v) => { - ConstantVector::new(Arc::new(UInt16Vector::from(vec![v])), length) - } - ScalarValue::UInt32(v) => { - ConstantVector::new(Arc::new(UInt32Vector::from(vec![v])), length) - } - ScalarValue::UInt64(v) => { - ConstantVector::new(Arc::new(UInt64Vector::from(vec![v])), length) - } - ScalarValue::Utf8(v) | ScalarValue::LargeUtf8(v) => { - ConstantVector::new(Arc::new(StringVector::from(vec![v])), length) - } - ScalarValue::Binary(v) - | ScalarValue::LargeBinary(v) - | ScalarValue::FixedSizeBinary(_, v) => { - ConstantVector::new(Arc::new(BinaryVector::from(vec![v])), length) - } - ScalarValue::List(array) => { - let item_type = Arc::new(ConcreteDataType::try_from(&array.value_type())?); - let mut builder = ListVectorBuilder::with_type_capacity(item_type.clone(), 1); - let scalar_values = ScalarValue::convert_array_to_scalar_vec(array.as_ref()) - .context(ConvertArrowArrayToScalarsSnafu)?; - let values = scalar_values - .into_iter() - .flat_map(|v| v.unwrap_or_else(|| vec![ScalarValue::Null])) - .map(ScalarValue::try_into) - .collect::>>()?; - builder.push(Some(ListValueRef::Ref { - val: &ListValue::new(values, item_type), - })); - let list_vector = builder.to_vector(); - ConstantVector::new(list_vector, length) - } - ScalarValue::Date32(v) => { - ConstantVector::new(Arc::new(DateVector::from(vec![v])), length) - } - ScalarValue::TimestampSecond(v, _) => { - // Timezone is unimplemented now. - ConstantVector::new(Arc::new(TimestampSecondVector::from(vec![v])), length) - } - ScalarValue::TimestampMillisecond(v, _) => { - // Timezone is unimplemented now. - ConstantVector::new(Arc::new(TimestampMillisecondVector::from(vec![v])), length) - } - ScalarValue::TimestampMicrosecond(v, _) => { - // Timezone is unimplemented now. - ConstantVector::new(Arc::new(TimestampMicrosecondVector::from(vec![v])), length) - } - ScalarValue::TimestampNanosecond(v, _) => { - // Timezone is unimplemented now. - ConstantVector::new(Arc::new(TimestampNanosecondVector::from(vec![v])), length) - } - ScalarValue::Time32Second(v) => { - ConstantVector::new(Arc::new(TimeSecondVector::from(vec![v])), length) - } - ScalarValue::Time32Millisecond(v) => { - ConstantVector::new(Arc::new(TimeMillisecondVector::from(vec![v])), length) - } - ScalarValue::Time64Microsecond(v) => { - ConstantVector::new(Arc::new(TimeMicrosecondVector::from(vec![v])), length) - } - ScalarValue::Time64Nanosecond(v) => { - ConstantVector::new(Arc::new(TimeNanosecondVector::from(vec![v])), length) - } - ScalarValue::IntervalYearMonth(v) => { - ConstantVector::new(Arc::new(IntervalYearMonthVector::from(vec![v])), length) - } - ScalarValue::IntervalDayTime(v) => { - ConstantVector::new(Arc::new(IntervalDayTimeVector::from(vec![v])), length) - } - ScalarValue::IntervalMonthDayNano(v) => { - ConstantVector::new(Arc::new(IntervalMonthDayNanoVector::from(vec![v])), length) - } - ScalarValue::DurationSecond(v) => { - ConstantVector::new(Arc::new(DurationSecondVector::from(vec![v])), length) - } - ScalarValue::DurationMillisecond(v) => { - ConstantVector::new(Arc::new(DurationMillisecondVector::from(vec![v])), length) - } - ScalarValue::DurationMicrosecond(v) => { - ConstantVector::new(Arc::new(DurationMicrosecondVector::from(vec![v])), length) - } - ScalarValue::DurationNanosecond(v) => { - ConstantVector::new(Arc::new(DurationNanosecondVector::from(vec![v])), length) - } - ScalarValue::Decimal128(v, p, s) => { - let vector = Decimal128Vector::from(vec![v]).with_precision_and_scale(p, s)?; - ConstantVector::new(Arc::new(vector), length) - } - ScalarValue::Struct(v) => { - let struct_type = StructType::from(v.fields()); - ConstantVector::new( - Arc::new(StructVector::try_new(struct_type, (*v).clone())?), - length, - ) - } - ScalarValue::Decimal32(_, _, _) + // Timezones are not supported by GreptimeDB vectors. + ScalarValue::TimestampSecond(v, _) => ScalarValue::TimestampSecond(v, None), + ScalarValue::TimestampMillisecond(v, _) => ScalarValue::TimestampMillisecond(v, None), + ScalarValue::TimestampMicrosecond(v, _) => ScalarValue::TimestampMicrosecond(v, None), + ScalarValue::TimestampNanosecond(v, _) => ScalarValue::TimestampNanosecond(v, None), + value @ (ScalarValue::Decimal32(_, _, _) | ScalarValue::Decimal64(_, _, _) | ScalarValue::Decimal256(_, _, _) | ScalarValue::FixedSizeList(_) @@ -241,15 +156,19 @@ impl Helper { | ScalarValue::BinaryView(_) | ScalarValue::Map(_) | ScalarValue::Date64(_) - | ScalarValue::RunEndEncoded(_, _, _) => { + | ScalarValue::RunEndEncoded(_, _, _)) => { return error::ConversionSnafu { from: format!("Unsupported scalar value: {value}"), } .fail(); } + value => value, }; - Ok(Arc::new(vector)) + let array = value + .to_array_of_size(length) + .context(ConvertScalarToArrowArraySnafu)?; + Self::try_into_vector(array) } /// Try to cast an arrow array into vector @@ -457,12 +376,15 @@ mod tests { }; use arrow::buffer::Buffer; use arrow::datatypes::{Int32Type, IntervalMonthDayNano}; - use arrow_array::{BinaryArray, DictionaryArray, FixedSizeBinaryArray, LargeStringArray}; - use arrow_schema::DataType; + use arrow_array::{ + BinaryArray, DictionaryArray, FixedSizeBinaryArray, LargeStringArray, StructArray, + }; + use arrow_schema::{DataType, Field, Fields}; use common_decimal::Decimal128; use common_time::time::Time; use common_time::timestamp::TimeUnit; use common_time::{Date, Duration}; + use datafusion_common::scalar::ScalarStructBuilder; use super::*; use crate::value::Value; @@ -496,7 +418,7 @@ mod tests { #[test] fn test_try_from_scalar_date_value() { - let vector = Helper::try_from_scalar_value(ScalarValue::Date32(Some(42)), 3).unwrap(); + let vector = Helper::try_from_scalar_value(ScalarValue::Date32(Some(42)), 3, None).unwrap(); assert_eq!(ConcreteDataType::date_datatype(), vector.data_type()); assert_eq!(3, vector.len()); for i in 0..vector.len() { @@ -507,7 +429,7 @@ mod tests { #[test] fn test_try_from_scalar_duration_value() { let vector = - Helper::try_from_scalar_value(ScalarValue::DurationSecond(Some(42)), 3).unwrap(); + Helper::try_from_scalar_value(ScalarValue::DurationSecond(Some(42)), 3, None).unwrap(); assert_eq!( ConcreteDataType::duration_second_datatype(), vector.data_type() @@ -524,7 +446,8 @@ mod tests { #[test] fn test_try_from_scalar_decimal128_value() { let vector = - Helper::try_from_scalar_value(ScalarValue::Decimal128(Some(42), 3, 1), 3).unwrap(); + Helper::try_from_scalar_value(ScalarValue::Decimal128(Some(42), 3, 1), 3, None) + .unwrap(); assert_eq!( ConcreteDataType::decimal128_datatype(3, 1), vector.data_type() @@ -542,7 +465,7 @@ mod tests { &ArrowDataType::Int32, true, )); - let vector = Helper::try_from_scalar_value(value, 3).unwrap(); + let vector = Helper::try_from_scalar_value(value, 3, None).unwrap(); assert_eq!( ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype())), vector.data_type() @@ -555,6 +478,101 @@ mod tests { } } + #[test] + fn test_try_from_scalar_value_materializes_values() { + let vector = Helper::try_from_scalar_value(ScalarValue::Int32(Some(42)), 4, None).unwrap(); + assert_eq!(ConcreteDataType::int32_datatype(), vector.data_type()); + assert_eq!(4, vector.len()); + assert_eq!(0, vector.null_count()); + for i in 0..vector.len() { + assert_eq!(Value::Int32(42), vector.get(i)); + } + + let empty = Helper::try_from_scalar_value(ScalarValue::Int32(Some(42)), 0, None).unwrap(); + assert_eq!(ConcreteDataType::int32_datatype(), empty.data_type()); + assert!(empty.is_empty()); + + let nulls = Helper::try_from_scalar_value(ScalarValue::Int32(None), 3, None).unwrap(); + assert_eq!(3, nulls.len()); + assert_eq!(3, nulls.null_count()); + for i in 0..nulls.len() { + assert_eq!(Value::Null, nulls.get(i)); + } + } + + #[test] + fn test_try_from_scalar_struct_value() { + let fields = Fields::from(vec![ + Field::new("id", ArrowDataType::Int32, false), + Field::new("name", ArrowDataType::Utf8, true), + ]); + let value = ScalarValue::Struct(Arc::new(StructArray::new( + fields.clone(), + vec![ + ScalarValue::Int32(Some(7)).to_array().unwrap(), + ScalarValue::Utf8(Some("greptime".to_string())) + .to_array() + .unwrap(), + ], + None, + ))); + + let vector = Helper::try_from_scalar_value(value, 3, None).unwrap(); + assert_eq!( + ConcreteDataType::struct_datatype(StructType::from(&fields)), + vector.data_type() + ); + assert_eq!(3, vector.len()); + for i in 0..vector.len() { + let Value::Struct(value) = vector.get(i) else { + panic!("expected struct value"); + }; + assert_eq!( + &[Value::Int32(7), Value::String("greptime".into())], + value.items() + ); + } + + let null = ScalarStructBuilder::new_null(fields); + let vector = Helper::try_from_scalar_value(null, 2, None).unwrap(); + assert_eq!(2, vector.len()); + assert_eq!(2, vector.null_count()); + assert_eq!(Value::Null, vector.get(0)); + assert_eq!(Value::Null, vector.get(1)); + } + + #[test] + fn test_try_from_scalar_value_normalizes_arrow_types() { + let string = Helper::try_from_scalar_value( + ScalarValue::LargeUtf8(Some("greptime".to_string())), + 2, + None, + ) + .unwrap(); + assert_eq!(ConcreteDataType::string_datatype(), string.data_type()); + assert_eq!(&ArrowDataType::Utf8, string.to_arrow_array().data_type()); + + let binary = Helper::try_from_scalar_value( + ScalarValue::FixedSizeBinary(2, Some(vec![1, 2])), + 2, + None, + ) + .unwrap(); + assert_eq!(ConcreteDataType::binary_datatype(), binary.data_type()); + assert_eq!(&ArrowDataType::Binary, binary.to_arrow_array().data_type()); + + let timestamp = Helper::try_from_scalar_value( + ScalarValue::TimestampMillisecond(Some(42), Some("UTC".into())), + 2, + None, + ) + .unwrap(); + assert_eq!( + &ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + timestamp.to_arrow_array().data_type() + ); + } + #[test] fn test_like_utf8() { fn assert_vector(expected: Vec<&str>, actual: &VectorRef) { @@ -745,7 +763,8 @@ mod tests { #[test] fn test_try_from_scalar_time_value() { - let vector = Helper::try_from_scalar_value(ScalarValue::Time32Second(Some(42)), 3).unwrap(); + let vector = + Helper::try_from_scalar_value(ScalarValue::Time32Second(Some(42)), 3, None).unwrap(); assert_eq!(ConcreteDataType::time_second_datatype(), vector.data_type()); assert_eq!(3, vector.len()); for i in 0..vector.len() { @@ -758,6 +777,7 @@ mod tests { let vector = Helper::try_from_scalar_value( ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano::new(1, 1, 2000))), 3, + None, ) .unwrap(); diff --git a/src/datatypes/src/vectors/null.rs b/src/datatypes/src/vectors/null.rs index e049eb205c..2510824ce0 100644 --- a/src/datatypes/src/vectors/null.rs +++ b/src/datatypes/src/vectors/null.rs @@ -260,16 +260,6 @@ impl ScalarVectorBuilder for NullVectorBuilder { } } -pub(crate) fn replicate_null(vector: &NullVector, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), vector.len()); - - if offsets.is_empty() { - return vector.slice(0, 0); - } - - Arc::new(NullVector::new(*offsets.last().unwrap())) -} - #[cfg(test)] mod tests { use serde_json; @@ -290,7 +280,6 @@ mod tests { assert_eq!(vector2.null_count(), 16); assert_eq!("NullVector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_null()); assert!(v.only_null()); diff --git a/src/datatypes/src/vectors/operations.rs b/src/datatypes/src/vectors/operations.rs index 1dc9e21a0d..f75ff060f3 100644 --- a/src/datatypes/src/vectors/operations.rs +++ b/src/datatypes/src/vectors/operations.rs @@ -14,14 +14,12 @@ mod cast; mod filter; -mod replicate; mod take; use std::sync::Arc; use crate::error::{self, Result}; use crate::types::LogicalPrimitiveType; -use crate::vectors::constant::ConstantVector; use crate::vectors::{ BinaryVector, BooleanVector, ConcreteDataType, Decimal128Vector, ListVector, NullVector, PrimitiveVector, StringVector, UInt32Vector, Vector, VectorRef, @@ -29,14 +27,6 @@ use crate::vectors::{ /// Vector compute operations. pub trait VectorOp { - /// Copies each element according `offsets` parameter. - /// - `i-th` element should be copied `offsets[i] - offsets[i - 1]` times - /// - `0-th` element would be copied `offsets[0]` times - /// - /// # Panics - /// Panics if `offsets.len() != self.len()`. - fn replicate(&self, offsets: &[usize]) -> VectorRef; - /// Filters the vector, returns elements matching the `filter` (i.e. where the values are true). /// /// Note that the nulls of `filter` are interpreted as `false` will lead to these elements being masked out. @@ -57,10 +47,6 @@ pub trait VectorOp { macro_rules! impl_scalar_vector_op { ($($VectorType: ident),+) => {$( impl VectorOp for $VectorType { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - replicate::replicate_scalar(self, offsets) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, $VectorType, filter) } @@ -92,10 +78,6 @@ macro_rules! impl_scalar_vector_op { impl_scalar_vector_op!(BinaryVector, BooleanVector, StringVector); impl VectorOp for ListVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - replicate::replicate_list(self, offsets) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, ListVector, filter) } @@ -110,10 +92,6 @@ impl VectorOp for ListVector { } impl VectorOp for Decimal128Vector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - std::sync::Arc::new(replicate::replicate_decimal128(self, offsets)) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, Decimal128Vector, filter) } @@ -128,10 +106,6 @@ impl VectorOp for Decimal128Vector { } impl VectorOp for PrimitiveVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - std::sync::Arc::new(replicate::replicate_primitive(self, offsets)) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, PrimitiveVector, filter) } @@ -146,10 +120,6 @@ impl VectorOp for PrimitiveVector { } impl VectorOp for NullVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - replicate::replicate_null(self, offsets) - } - fn filter(&self, filter: &BooleanVector) -> Result { filter::filter_non_constant!(self, NullVector, filter) } @@ -166,21 +136,3 @@ impl VectorOp for NullVector { take::take_indices!(self, NullVector, indices) } } - -impl VectorOp for ConstantVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - self.replicate_vector(offsets) - } - - fn filter(&self, filter: &BooleanVector) -> Result { - self.filter_vector(filter) - } - - fn cast(&self, to_type: &ConcreteDataType) -> Result { - self.cast_vector(to_type) - } - - fn take(&self, indices: &UInt32Vector) -> Result { - self.take_vector(indices) - } -} diff --git a/src/datatypes/src/vectors/operations/filter.rs b/src/datatypes/src/vectors/operations/filter.rs index 7dde6c1f94..b27bbb5de4 100644 --- a/src/datatypes/src/vectors/operations/filter.rs +++ b/src/datatypes/src/vectors/operations/filter.rs @@ -39,7 +39,6 @@ mod tests { TimestampMicrosecond, TimestampMillisecond, TimestampNanosecond, TimestampSecond, }; use crate::types::WrapperType; - use crate::vectors::constant::ConstantVector; use crate::vectors::{ BooleanVector, Int32Vector, NullVector, StringVector, VectorOp, VectorRef, }; @@ -67,26 +66,6 @@ mod tests { check_filter_primitive(&[5, 7], &[5, 6, 7], &[true, false, true]); } - fn check_filter_constant(expect_length: usize, input_length: usize, filter: &[bool]) { - let v = ConstantVector::new(Arc::new(Int32Vector::from_slice([123])), input_length); - let filter = BooleanVector::from_slice(filter); - let out = v.filter(&filter).unwrap(); - - assert!(out.is_const()); - assert_eq!(expect_length, out.len()); - } - - #[test] - fn test_filter_constant() { - check_filter_constant(0, 0, &[]); - check_filter_constant(1, 1, &[true]); - check_filter_constant(0, 1, &[false]); - check_filter_constant(1, 2, &[false, true]); - check_filter_constant(2, 2, &[true, true]); - check_filter_constant(1, 4, &[false, false, false, true]); - check_filter_constant(2, 4, &[false, true, false, true]); - } - #[test] fn test_filter_scalar() { let v = StringVector::from_slice(&["0", "1", "2", "3"]); diff --git a/src/datatypes/src/vectors/operations/replicate.rs b/src/datatypes/src/vectors/operations/replicate.rs deleted file mode 100644 index dee0b3855b..0000000000 --- a/src/datatypes/src/vectors/operations/replicate.rs +++ /dev/null @@ -1,232 +0,0 @@ -// Copyright 2023 Greptime Team -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use crate::prelude::*; -pub(crate) use crate::vectors::decimal::replicate_decimal128; -pub(crate) use crate::vectors::null::replicate_null; -pub(crate) use crate::vectors::primitive::replicate_primitive; -use crate::vectors::{ListVector, ListVectorBuilder}; - -pub(crate) fn replicate_scalar(c: &C, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), c.len()); - - if offsets.is_empty() { - return c.slice(0, 0); - } - let mut builder = <::Builder>::with_capacity(c.len()); - - let mut previous_offset = 0; - for (i, offset) in offsets.iter().enumerate() { - let data = c.get_data(i); - for _ in previous_offset..*offset { - builder.push(data.clone()); - } - previous_offset = *offset; - } - builder.to_vector() -} - -pub(crate) fn replicate_list(c: &ListVector, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), c.len()); - - if offsets.is_empty() { - return c.slice(0, 0); - } - let mut builder = - ListVectorBuilder::with_type_capacity(c.item_type(), *offsets.last().unwrap()); - - let mut previous_offset = 0; - for (i, offset) in offsets.iter().enumerate() { - let data = c.get_data(i); - for _ in previous_offset..*offset { - builder.push(data.clone()); - } - previous_offset = *offset; - } - builder.to_vector() -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use common_time::timestamp::TimeUnit; - use common_time::{Date, Timestamp}; - use paste::paste; - - use super::*; - use crate::value::{ListValue, ListValueRef}; - use crate::vectors::constant::ConstantVector; - use crate::vectors::{ - Decimal128Vector, Int32Vector, ListVectorBuilder, NullVector, StringVector, VectorOp, - }; - - #[test] - fn test_replicate_primitive() { - let v = Int32Vector::from_iterator(0..5); - let offsets = [0, 1, 2, 3, 4]; - - let v = v.replicate(&offsets); - assert_eq!(4, v.len()); - - for i in 0..4 { - assert_eq!(Value::Int32(i as i32 + 1), v.get(i)); - } - } - - #[test] - fn test_replicate_nullable_primitive() { - let v = Int32Vector::from(vec![None, Some(1), None, Some(2)]); - let offsets = [2, 4, 6, 8]; - let v = v.replicate(&offsets); - assert_eq!(8, v.len()); - - let expect: VectorRef = Arc::new(Int32Vector::from(vec![ - None, - None, - Some(1), - Some(1), - None, - None, - Some(2), - Some(2), - ])); - assert_eq!(expect, v); - } - - #[test] - fn test_replicate_scalar() { - let v = StringVector::from_slice(&["0", "1", "2", "3"]); - let offsets = [1, 3, 5, 6]; - - let v = v.replicate(&offsets); - assert_eq!(6, v.len()); - - let expect: VectorRef = Arc::new(StringVector::from_slice(&["0", "1", "1", "2", "2", "3"])); - assert_eq!(expect, v); - } - - #[test] - fn test_replicate_list() { - let item_type = Arc::new(ConcreteDataType::int32_datatype()); - let first = ListValue::new(vec![Value::Int32(1), Value::Int32(2)], item_type.clone()); - let second = ListValue::new(vec![Value::Int32(3)], item_type.clone()); - let mut builder = ListVectorBuilder::with_type_capacity(item_type, 2); - builder.push(Some(ListValueRef::Ref { val: &first })); - builder.push(Some(ListValueRef::Ref { val: &second })); - let v = builder.finish(); - - let v = v.replicate(&[1, 3]); - let mut expect_builder = - ListVectorBuilder::with_type_capacity(Arc::new(ConcreteDataType::int32_datatype()), 3); - expect_builder.push(Some(ListValueRef::Ref { val: &first })); - expect_builder.push(Some(ListValueRef::Ref { val: &second })); - expect_builder.push(Some(ListValueRef::Ref { val: &second })); - let expect = expect_builder.to_vector(); - assert_eq!(expect, v); - } - - #[test] - fn test_replicate_constant() { - let v = Arc::new(StringVector::from_slice(&["hello"])); - let cv = ConstantVector::new(v.clone(), 2); - let offsets = [1, 4]; - - let cv = cv.replicate(&offsets); - assert_eq!(4, cv.len()); - - let expect: VectorRef = Arc::new(ConstantVector::new(v, 4)); - assert_eq!(expect, cv); - } - - #[test] - fn test_replicate_null() { - let v = NullVector::new(0); - let offsets = []; - let v = v.replicate(&offsets); - assert!(v.is_empty()); - - let v = NullVector::new(3); - let offsets = [1, 3, 5]; - - let v = v.replicate(&offsets); - assert_eq!(5, v.len()); - } - - macro_rules! impl_replicate_date_like_test { - ($VectorType: ident, $ValueType: ident, $method: ident) => {{ - use $crate::vectors::$VectorType; - - let v = $VectorType::from_iterator((0..5).map($ValueType::$method)); - let offsets = [0, 1, 2, 3, 4]; - - let v = v.replicate(&offsets); - assert_eq!(4, v.len()); - - for i in 0..4 { - assert_eq!( - Value::$ValueType($ValueType::$method((i as i32 + 1).into())), - v.get(i) - ); - } - }}; - } - - macro_rules! impl_replicate_timestamp_test { - ($unit: ident) => {{ - paste!{ - use $crate::vectors::[]; - use $crate::timestamp::[]; - let v = []::from_iterator((0..5).map([]::from)); - let offsets = [0, 1, 2, 3, 4]; - let v = v.replicate(&offsets); - assert_eq!(4, v.len()); - for i in 0..4 { - assert_eq!( - Value::Timestamp(Timestamp::new(i as i64 + 1, TimeUnit::$unit)), - v.get(i) - ); - } - } - }}; - } - - #[test] - fn test_replicate_date_like() { - impl_replicate_date_like_test!(DateVector, Date, new); - impl_replicate_timestamp_test!(Second); - impl_replicate_timestamp_test!(Millisecond); - impl_replicate_timestamp_test!(Microsecond); - impl_replicate_timestamp_test!(Nanosecond); - } - - #[test] - fn test_replicate_decimal() { - let data = vec![100]; - // create a decimal vector - let v = Decimal128Vector::from_values(data.clone()) - .with_precision_and_scale(10, 2) - .unwrap(); - let offsets = [5]; - let v = v.replicate(&offsets); - assert_eq!(5, v.len()); - - let expect: VectorRef = Arc::new( - Decimal128Vector::from_values(vec![100; 5]) - .with_precision_and_scale(10, 2) - .unwrap(), - ); - assert_eq!(expect, v); - } -} diff --git a/src/datatypes/src/vectors/operations/take.rs b/src/datatypes/src/vectors/operations/take.rs index 41947d3ef5..167f650c8e 100644 --- a/src/datatypes/src/vectors/operations/take.rs +++ b/src/datatypes/src/vectors/operations/take.rs @@ -43,8 +43,7 @@ mod tests { use crate::types::{LogicalPrimitiveType, WrapperType}; use crate::vectors::operations::VectorOp; use crate::vectors::{ - BooleanVector, ConstantVector, Int32Vector, NullVector, PrimitiveVector, StringVector, - UInt32Vector, + BooleanVector, Int32Vector, NullVector, PrimitiveVector, StringVector, UInt32Vector, }; fn check_take_primitive( @@ -119,29 +118,6 @@ mod tests { take_time_like_test!(TimestampNanosecondVector, TimestampNanosecond, from_native); } - fn check_take_constant(expect_length: usize, input_length: usize, indices: &[u32]) { - let v = ConstantVector::new(Arc::new(Int32Vector::from_slice([111])), input_length); - let indices = UInt32Vector::from_slice(indices); - let out = v.take(&indices).unwrap(); - - assert!(out.is_const()); - assert_eq!(expect_length, out.len()); - } - - #[test] - fn test_take_constant() { - check_take_constant(2, 5, &[3, 4]); - check_take_constant(3, 10, &[1, 2, 3]); - check_take_constant(4, 10, &[1, 5, 3, 6]); - check_take_constant(5, 10, &[1, 9, 8, 7, 3]); - } - - #[test] - #[should_panic] - fn test_take_constant_out_of_index() { - check_take_constant(2, 5, &[3, 5]); - } - #[test] #[should_panic] fn test_take_out_of_index() { diff --git a/src/datatypes/src/vectors/primitive.rs b/src/datatypes/src/vectors/primitive.rs index dc85f94c4f..92c8394e71 100644 --- a/src/datatypes/src/vectors/primitive.rs +++ b/src/datatypes/src/vectors/primitive.rs @@ -366,40 +366,6 @@ where } } -pub(crate) fn replicate_primitive( - vector: &PrimitiveVector, - offsets: &[usize], -) -> PrimitiveVector { - assert_eq!(offsets.len(), vector.len()); - - if offsets.is_empty() { - return vector.get_slice(0, 0); - } - - let mut builder = PrimitiveVectorBuilder::::with_capacity(*offsets.last().unwrap()); - - let mut previous_offset = 0; - - for (offset, value) in offsets.iter().zip(vector.array.iter()) { - let repeat_times = *offset - previous_offset; - match value { - Some(data) => { - unsafe { - // Safety: std::iter::Repeat and std::iter::Take implement TrustedLen. - builder - .mutable_array - .append_trusted_len_iter(std::iter::repeat_n(data, repeat_times)); - } - } - None => { - builder.mutable_array.append_nulls(repeat_times); - } - } - previous_offset = *offset; - } - builder.finish() -} - #[cfg(test)] mod tests { use std::vec; @@ -433,7 +399,6 @@ mod tests { fn check_vec(v: Int32Vector) { assert_eq!(4, v.len()); assert_eq!("Int32Vector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_valid()); assert!(!v.only_null()); diff --git a/src/datatypes/src/vectors/string.rs b/src/datatypes/src/vectors/string.rs index 53324280c3..0476ecc644 100644 --- a/src/datatypes/src/vectors/string.rs +++ b/src/datatypes/src/vectors/string.rs @@ -604,7 +604,6 @@ mod tests { let v = StringVector::from(strs.clone()); assert_eq!(3, v.len()); assert_eq!("StringVector", v.vector_type_name()); - assert!(!v.is_const()); assert!(v.validity().is_all_valid()); assert!(!v.only_null()); assert_eq!(1040, v.memory_size()); diff --git a/src/datatypes/src/vectors/struct_vector.rs b/src/datatypes/src/vectors/struct_vector.rs index e6284351b4..4798c8fddc 100644 --- a/src/datatypes/src/vectors/struct_vector.rs +++ b/src/datatypes/src/vectors/struct_vector.rs @@ -15,7 +15,7 @@ use std::any::Any; use std::sync::Arc; -use arrow::array::{MutableArrayData, NullBufferBuilder}; +use arrow::array::NullBufferBuilder; use arrow::compute::TakeOptions; use arrow::datatypes::DataType as ArrowDataType; use arrow_array::{Array, ArrayRef, StructArray}; @@ -143,31 +143,6 @@ impl Vector for StructVector { } impl VectorOp for StructVector { - fn replicate(&self, offsets: &[usize]) -> VectorRef { - assert_eq!(offsets.len(), self.len()); - assert!(offsets.is_sorted(), "offsets must be non-decreasing"); - - let Some(&output_len) = offsets.last() else { - return self.slice(0, 0); - }; - - let source = self.array.to_data(); - let mut output = MutableArrayData::new(vec![&source], false, output_len); - let mut previous_offset = 0; - - for (index, &offset) in offsets.iter().enumerate() { - for _ in previous_offset..offset { - output.extend(0, index, index + 1); - } - previous_offset = offset; - } - - Arc::new(StructVector { - array: StructArray::from(output.freeze()), - fields: self.fields.clone(), - }) - } - fn cast(&self, _to_type: &ConcreteDataType) -> Result { UnsupportedOperationSnafu { op: "cast", @@ -466,10 +441,6 @@ impl ScalarVectorBuilder for StructVectorBuilder { #[cfg(test)] mod tests { - use arrow::array::{DictionaryArray, Int8Array, StringArray}; - use arrow::buffer::NullBuffer; - use arrow::datatypes::Int8Type; - use super::*; use crate::json::JsonSettings; use crate::schema::{ColumnDefaultConstraint, ColumnSchema}; @@ -555,7 +526,7 @@ mod tests { } #[test] - fn test_replicate_preserves_json2_identity() { + fn test_default_vector_preserves_json2_identity() { let json = JsonSettings::default() .encode(serde_json::json!({"answer": 42})) .unwrap(); @@ -576,39 +547,6 @@ mod tests { assert_eq!(replicated.len(), 2); } - #[test] - fn test_replicate_preserves_dictionary_and_nulls() { - let fields = StructType::new(Arc::new(vec![StructField::new( - "label", - ConcreteDataType::dictionary_datatype( - ConcreteDataType::int8_datatype(), - ConcreteDataType::string_datatype(), - ), - true, - )])); - let dictionary = DictionaryArray::::new( - Int8Array::from(vec![Some(0), Some(1)]), - Arc::new(StringArray::from(vec!["a", "b"])), - ); - let array = StructArray::new( - fields.as_arrow_fields(), - vec![Arc::new(dictionary)], - Some(NullBuffer::from(vec![true, false])), - ); - let vector = StructVector::try_new(fields.clone(), array).unwrap(); - - let replicated = vector.replicate(&[2, 3]); - - assert_eq!( - replicated.data_type(), - ConcreteDataType::struct_datatype(fields) - ); - assert_eq!(replicated.len(), 3); - assert_eq!(replicated.null_count(), 1); - assert_eq!(replicated.get(0), replicated.get(1)); - assert!(replicated.is_null(2)); - } - #[test] fn test_deep_nested_struct_list() { // level 1: struct diff --git a/src/flow/src/expr/scalar.rs b/src/flow/src/expr/scalar.rs index 18649bff12..af16381e4c 100644 --- a/src/flow/src/expr/scalar.rs +++ b/src/flow/src/expr/scalar.rs @@ -216,6 +216,7 @@ impl ScalarExpr { msg: "Failed to convert literal to scalar value", })?, batch.row_count(), + None, ) .context(DataTypeSnafu { msg: "Failed to convert scalar value to vector ref when parsing literal", diff --git a/src/mito2/src/read/flat_projection.rs b/src/mito2/src/read/flat_projection.rs index 19dc962c56..2abf6aa48a 100644 --- a/src/mito2/src/read/flat_projection.rs +++ b/src/mito2/src/read/flat_projection.rs @@ -667,6 +667,77 @@ mod tests { ); } + #[test] + fn test_single_value_dictionary_preserves_string_type() { + use datatypes::arrow::array::{DictionaryArray, TimestampMillisecondArray, UInt32Array}; + use datatypes::arrow::datatypes::UInt32Type; + + for data_type in [ + ConcreteDataType::large_string_datatype(), + ConcreteDataType::utf8_view_datatype(), + ] { + let mut builder = RegionMetadataBuilder::new(RegionId::new(1024, 0)); + builder + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new("tag", data_type.clone(), true), + semantic_type: SemanticType::Tag, + column_id: 0, + }) + .push_column_metadata(ColumnMetadata { + column_schema: ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + semantic_type: SemanticType::Timestamp, + column_id: 1, + }) + .primary_key(vec![0]); + let metadata = Arc::new(builder.build().unwrap()); + let mapper = FlatProjectionMapper::new(&metadata, [0, 1]).unwrap(); + + for value in [Value::from("greptime"), Value::Null] { + let mut values = data_type.create_mutable_vector(1); + values.try_push_value_ref(&value.as_value_ref()).unwrap(); + let dictionary = Arc::new( + DictionaryArray::::try_new( + UInt32Array::from(vec![0, 0, 0]), + values.to_vector().to_arrow_array(), + ) + .unwrap(), + ); + let arrays: Vec> = mapper + .batch_schema() + .iter() + .map(|(id, _)| match id { + 0 => dictionary.clone() as Arc, + 1 => Arc::new(TimestampMillisecondArray::from(vec![1, 2, 3])), + _ => unreachable!(), + }) + .collect(); + let fields = mapper + .batch_schema() + .iter() + .zip(&arrays) + .map(|((id, _), array)| { + Field::new(id.to_string(), array.data_type().clone(), true) + }) + .collect::>(); + let batch = DfRecordBatch::try_new( + Arc::new(datatypes::arrow::datatypes::Schema::new(fields)), + arrays, + ) + .unwrap(); + let output = mapper.convert(&batch, &CacheStrategy::Disabled).unwrap(); + let vector = Helper::try_into_vector(output.column(0)).unwrap(); + assert_eq!(data_type, vector.data_type()); + for row in 0..3 { + assert_eq!(value, vector.get(row)); + } + } + } + } + #[test] fn test_json_type_hint_does_not_concretize_legacy_json() { let metadata = metadata_with_legacy_json(); diff --git a/src/mito2/src/read/projection.rs b/src/mito2/src/read/projection.rs index c6d4baee3a..fbf6a29b73 100644 --- a/src/mito2/src/read/projection.rs +++ b/src/mito2/src/read/projection.rs @@ -19,7 +19,7 @@ use std::cmp::Ordering; use common_recordbatch::error::DataTypesSnafu; use datatypes::prelude::{ConcreteDataType, DataType}; use datatypes::value::Value; -use datatypes::vectors::VectorRef; +use datatypes::vectors::{Helper, VectorRef}; use snafu::{OptionExt, ResultExt}; use store_api::metadata::RegionMetadataRef; use store_api::storage::ColumnId; @@ -82,12 +82,21 @@ pub(crate) fn new_repeated_vector( value: &Value, num_rows: usize, ) -> common_recordbatch::error::Result { - let mut mutable_vector = data_type.create_mutable_vector(1); - mutable_vector - .try_push_value_ref(&value.as_value_ref()) - .context(DataTypesSnafu)?; - let base_vector = mutable_vector.to_vector(); - Ok(base_vector.replicate(&[num_rows])) + if let Ok(vector) = value + .try_to_scalar_value(data_type) + .and_then(|scalar| Helper::try_from_scalar_value(scalar, num_rows, Some(data_type))) + { + return Ok(vector); + } + + // Preserve extension types that cannot safely round-trip through ScalarValue. + let mut mutable_vector = data_type.create_mutable_vector(num_rows); + for _ in 0..num_rows { + mutable_vector + .try_push_value_ref(&value.as_value_ref()) + .context(DataTypesSnafu)?; + } + Ok(mutable_vector.to_vector()) } #[cfg(test)] @@ -110,6 +119,51 @@ mod tests { use crate::read::flat_projection::FlatProjectionMapper; use crate::read::read_columns::ReadColumns; + #[test] + fn test_repeated_struct_null_fields_and_json() { + use datatypes::types::{StructField, StructType}; + use datatypes::value::StructValue; + + let inner_type = StructType::from([StructField::new( + "x", + ConcreteDataType::int32_datatype(), + true, + )]); + let inner = Value::Struct(StructValue::new(vec![Value::Null], inner_type)); + let json = datatypes::json::JsonSettings::default() + .encode(serde_json::json!({"answer": 42})) + .unwrap(); + let nested_type = StructType::from([StructField::new("nested", inner.data_type(), true)]); + let json_type = StructType::from([StructField::new("json", json.data_type(), true)]); + let values = [ + inner.clone(), + Value::Struct(StructValue::new(vec![inner], nested_type)), + Value::Struct(StructValue::new(vec![], StructType::default())), + Value::Struct(StructValue::new(vec![json], json_type)), + ]; + for value in values { + let data_type = value.data_type(); + // JSON children are read back as their underlying struct values. + let expected = serde_json::Value::try_from(value.clone()).unwrap(); + for num_rows in [0, 1, 3] { + let vector = new_repeated_vector(&data_type, &value, num_rows).unwrap(); + assert_eq!(data_type, vector.data_type()); + assert_eq!( + data_type.as_arrow_type(), + *vector.to_arrow_array().data_type() + ); + assert_eq!(num_rows, vector.len()); + assert_eq!(0, vector.null_count()); + for row in 0..num_rows { + assert_eq!( + expected, + serde_json::Value::try_from(vector.get(row)).unwrap() + ); + } + } + } + } + fn print_record_batch(record_batch: RecordBatch) -> String { pretty::pretty_format_batches(&[record_batch.into_df_record_batch()]) .unwrap()