rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-21 20:55:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: remove constant vector and replicate operation (#8999)

Browse Source 
* refactor: remove constant vector

Signed-off-by: evenyag <realevenyag@gmail.com>

* refactor: remove vector replicate operation

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix: preserve scalar vector types with optional type hints

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix: remove obsolete mutable vector helper and import

Signed-off-by: evenyag <realevenyag@gmail.com>

* fix: preserve typed nulls in struct scalar conversion

Signed-off-by: evenyag <realevenyag@gmail.com>

---------

Signed-off-by: evenyag <realevenyag@gmail.com>
This commit is contained in:
Yingwen
2026-09-14 09:25:33 +00:00
committed by GitHub
parent a23fe1c2dc
commit 23a1a5ecaa
33 changed files with 534 additions and 1185 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/catalog/src/system_schema/information_schema/columns.rs
+4 -13
View File
@@ -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<RecordBatch> {
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<VectorRef> = vec![
Arc::new(self.catalog_names.finish()),
src/catalog/src/system_schema/information_schema/key_column_usage.rs
+2 -5
View File
@@ -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<RecordBatch> {
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<VectorRef> = vec![
Arc::new(self.constraint_catalog.finish()),
Arc::new(self.constraint_schema.finish()),
src/catalog/src/system_schema/information_schema/partitions.rs
+7 -18
View File
@@ -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<RecordBatch> {
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<VectorRef> = vec![
Arc::new(self.catalog_names.finish()),
src/catalog/src/system_schema/information_schema/table_constraints.rs
+3 -9
View File
@@ -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<RecordBatch> {
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<VectorRef> = vec![
constraint_catalogs,
src/common/function/src/aggrs/vector/avg.rs
+5 -9
View File
@@ -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<ArrayRef> = 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!(
src/common/function/src/aggrs/vector/product.rs
+5 -10
View File
@@ -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<ArrayRef> = 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();
src/common/function/src/aggrs/vector/sum.rs
+5 -9
View File
@@ -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<ArrayRef> = 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!(
src/common/function/src/scalars/expression.rs
-2
View File
@@ -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;
src/common/function/src/scalars/expression/binary.rs
-90
View File
@@ -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: Scalar, R: Scalar, O: Scalar, F>(
l: &VectorRef,
r: &VectorRef,
f: F,
ctx: &mut EvalContext,
) -> Result<<O as Scalar>::VectorType>
where
F: Fn(Option<L::RefType<'_>>, Option<R::RefType<'_>>, &mut EvalContext) -> Option<O>,
{
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: &<L as Scalar>::VectorType = unsafe { Helper::static_cast(l) };
let right: &ConstantVector = unsafe { Helper::static_cast(r) };
let right: &<R as Scalar>::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));
<O as Scalar>::VectorType::from_owned_iterator(it)
}
(false, false) => {
let left: &<L as Scalar>::VectorType = unsafe { Helper::static_cast(l) };
let right: &<R as Scalar>::VectorType = unsafe { Helper::static_cast(r) };
let it = left
.iter_data()
.zip(right.iter_data())
.map(|(a, b)| f(a, b, ctx));
<O as Scalar>::VectorType::from_owned_iterator(it)
}
(true, false) => {
let left: &ConstantVector = unsafe { Helper::static_cast(l) };
let left: &<L as Scalar>::VectorType = unsafe { Helper::static_cast(left.inner()) };
let a = left.get_data(0);
let right: &<R as Scalar>::VectorType = unsafe { Helper::static_cast(r) };
let it = right.iter_data().map(|b| f(a.clone(), b, ctx));
<O as Scalar>::VectorType::from_owned_iterator(it)
}
(true, true) => {
let left: &ConstantVector = unsafe { Helper::static_cast(l) };
let left: &<L as Scalar>::VectorType = unsafe { Helper::static_cast(left.inner()) };
let a = left.get_data(0);
let right: &ConstantVector = unsafe { Helper::static_cast(r) };
let right: &<R as Scalar>::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());
<O as Scalar>::VectorType::from_owned_iterator(it)
}
};
if let Some(error) = ctx.error.take() {
return Err(error);
}
Ok(result)
}
src/common/function/src/scalars/test.rs
+31 -1
View File
@@ -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: Scalar, R: Scalar, O: Scalar, F>(
l: &VectorRef,
r: &VectorRef,
f: F,
ctx: &mut EvalContext,
) -> Result<<O as Scalar>::VectorType>
where
F: Fn(Option<L::RefType<'_>>, Option<R::RefType<'_>>, &mut EvalContext) -> Option<O>,
{
debug_assert!(
l.len() == r.len(),
"Size of vectors must match to apply binary expression"
);
let left: &<L as Scalar>::VectorType = unsafe { Helper::static_cast(l) };
let right: &<R as Scalar>::VectorType = unsafe { Helper::static_cast(r) };
let result = <O as Scalar>::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 {
src/common/macro/src/admin_fn.rs
+1 -1
View File
@@ -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)
}
})
src/datatypes/src/schema/constraint.rs
+91 -10
View File
@@ -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());
src/datatypes/src/types/cast.rs
+1 -1
View File
@@ -26,7 +26,7 @@ pub fn cast(src_value: Value, dest_type: &ConcreteDataType) -> Result<Value> {
}
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)
src/datatypes/src/value.rs
+79 -3
View File
@@ -1085,11 +1085,21 @@ impl StructValue {
}
fn try_to_scalar_value(&self, output_type: &StructType) -> Result<ScalarValue> {
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::<Result<Vec<Arc<dyn Array>>>>()?;
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!(
src/datatypes/src/vectors.rs
-7
View File
@@ -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;
src/datatypes/src/vectors/binary.rs
-1
View File
@@ -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());
src/datatypes/src/vectors/boolean.rs
-5
View File
@@ -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<Vec<bool>> 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());
src/datatypes/src/vectors/constant.rs
-252
View File
@@ -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<VectorRef> {
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<VectorRef> {
Ok(Arc::new(ConstantVector::new(
self.inner().cast(to_type)?,
self.length,
)))
}
pub(crate) fn take_vector(&self, indices: &UInt32Vector) -> Result<VectorRef> {
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::<UInt32Array>().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<dyn Array> {
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<Vec<serde_json::Value>> {
std::iter::repeat_n(self.get(0), self.len())
.map(serde_json::Value::try_from)
.collect::<serde_json::Result<_>>()
.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]");
}
}
src/datatypes/src/vectors/decimal.rs
-37
View File
@@ -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;
src/datatypes/src/vectors/dictionary.rs
+1 -47
View File
@@ -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<K: ArrowDictionaryKeyType> VectorOp for DictionaryVector<K> {
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<VectorRef> {
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();
src/datatypes/src/vectors/eq.rs
-68
View File
@@ -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::<ConstantVector>()
.unwrap()
.inner(),
&**rhs
.as_any()
.downcast_ref::<ConstantVector>()
.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(
src/datatypes/src/vectors/helper.rs
+163 -143
View File
@@ -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<VectorRef> {
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<VectorRef> {
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::<Result<Vec<Value>>>()?;
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();
src/datatypes/src/vectors/null.rs
-11
View File
@@ -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());
src/datatypes/src/vectors/operations.rs
-48
View File
@@ -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<VectorRef> {
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<VectorRef> {
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<VectorRef> {
filter::filter_non_constant!(self, Decimal128Vector, filter)
}
@@ -128,10 +106,6 @@ impl VectorOp for Decimal128Vector {
}
impl<T: LogicalPrimitiveType> VectorOp for PrimitiveVector<T> {
fn replicate(&self, offsets: &[usize]) -> VectorRef {
std::sync::Arc::new(replicate::replicate_primitive(self, offsets))
}
fn filter(&self, filter: &BooleanVector) -> Result<VectorRef> {
filter::filter_non_constant!(self, PrimitiveVector<T>, filter)
}
@@ -146,10 +120,6 @@ impl<T: LogicalPrimitiveType> VectorOp for PrimitiveVector<T> {
}
impl VectorOp for NullVector {
fn replicate(&self, offsets: &[usize]) -> VectorRef {
replicate::replicate_null(self, offsets)
}
fn filter(&self, filter: &BooleanVector) -> Result<VectorRef> {
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<VectorRef> {
self.filter_vector(filter)
}
fn cast(&self, to_type: &ConcreteDataType) -> Result<VectorRef> {
self.cast_vector(to_type)
}
fn take(&self, indices: &UInt32Vector) -> Result<VectorRef> {
self.take_vector(indices)
}
}
src/datatypes/src/vectors/operations/filter.rs
-21
View File
@@ -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"]);
src/datatypes/src/vectors/operations/replicate.rs
-232
View File
@@ -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: ScalarVector>(c: &C, offsets: &[usize]) -> VectorRef {
assert_eq!(offsets.len(), c.len());
if offsets.is_empty() {
return c.slice(0, 0);
}
let mut builder = <<C as ScalarVector>::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::[<Timestamp $unit Vector>];
use $crate::timestamp::[<Timestamp $unit>];
let v = [<Timestamp $unit Vector>]::from_iterator((0..5).map([<Timestamp $unit>]::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);
}
}
src/datatypes/src/vectors/operations/take.rs
+1 -25
View File
@@ -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<T>(
@@ -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() {
src/datatypes/src/vectors/primitive.rs
-35
View File
@@ -366,40 +366,6 @@ where
}
}
pub(crate) fn replicate_primitive<T: LogicalPrimitiveType>(
vector: &PrimitiveVector<T>,
offsets: &[usize],
) -> PrimitiveVector<T> {
assert_eq!(offsets.len(), vector.len());
if offsets.is_empty() {
return vector.get_slice(0, 0);
}
let mut builder = PrimitiveVectorBuilder::<T>::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());
src/datatypes/src/vectors/string.rs
-1
View File
@@ -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());
src/datatypes/src/vectors/struct_vector.rs
+2 -64
View File
@@ -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<VectorRef> {
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::<Int8Type>::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
src/flow/src/expr/scalar.rs
+1
View File
@@ -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",
src/mito2/src/read/flat_projection.rs
+71
View File
@@ -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::<UInt32Type>::try_new(
UInt32Array::from(vec![0, 0, 0]),
values.to_vector().to_arrow_array(),
)
.unwrap(),
);
let arrays: Vec<Arc<dyn Array>> = mapper
.batch_schema()
.iter()
.map(|(id, _)| match id {
0 => dictionary.clone() as Arc<dyn Array>,
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::<Vec<_>>();
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();
src/mito2/src/read/projection.rs
+61 -7
View File
@@ -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<VectorRef> {
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()