Skip to content

Commit

Permalink
feat(storage): Map data type support bloom filter
Browse files Browse the repository at this point in the history
  • Loading branch information
b41sh committed Mar 9, 2023
1 parent a5e5079 commit 8c324cc
Show file tree
Hide file tree
Showing 8 changed files with 207 additions and 22 deletions.
7 changes: 7 additions & 0 deletions src/query/expression/src/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -399,6 +399,13 @@ impl FunctionID {
FunctionID::Factory { id, .. } => *id,
}
}

pub fn name(&self) -> &String {
match self {
FunctionID::Builtin { name, .. } => name,
FunctionID::Factory { name, .. } => name,
}
}
}

pub fn wrap_nullable<F>(f: F) -> impl Fn(&[ValueRef<AnyType>], &mut EvalContext) -> Value<AnyType>
Expand Down
22 changes: 21 additions & 1 deletion src/query/expression/src/kernels/concat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ use itertools::Itertools;

use crate::types::array::ArrayColumnBuilder;
use crate::types::decimal::DecimalColumn;
use crate::types::map::KvColumnBuilder;
use crate::types::nullable::NullableColumn;
use crate::types::number::NumberColumn;
use crate::types::string::StringColumnBuilder;
Expand All @@ -28,6 +29,7 @@ use crate::types::BooleanType;
use crate::types::DateType;
use crate::types::EmptyArrayType;
use crate::types::EmptyMapType;
use crate::types::MapType;
use crate::types::NullType;
use crate::types::NullableType;
use crate::types::NumberType;
Expand Down Expand Up @@ -138,7 +140,7 @@ impl Column {
let builder = Vec::with_capacity(capacity);
Self::concat_value_types::<DateType>(builder, columns)
}
Column::Array(col) | Column::Map(col) => {
Column::Array(col) => {
let mut offsets = Vec::with_capacity(capacity + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
Expand All @@ -148,6 +150,24 @@ impl Column {
let builder = ArrayColumnBuilder { builder, offsets };
Self::concat_value_types::<ArrayType<AnyType>>(builder, columns)
}
Column::Map(col) => {
let mut offsets = Vec::with_capacity(capacity + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
TypeDeserializerImpl::with_capacity(&col.values.data_type(), capacity)
.finish_to_column(),
);
let (key_builder, val_builder) = match builder {
ColumnBuilder::Tuple { fields, .. } => (fields[0].clone(), fields[1].clone()),
_ => unreachable!(),
};
let builder = KvColumnBuilder {
keys: key_builder,
values: val_builder,
};
let builder = ArrayColumnBuilder { builder, offsets };
Self::concat_value_types::<MapType<AnyType, AnyType>>(builder, columns)
}
Column::Nullable(_) => {
let mut bitmaps = Vec::with_capacity(columns.len());
let mut inners = Vec::with_capacity(columns.len());
Expand Down
24 changes: 23 additions & 1 deletion src/query/expression/src/kernels/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,18 @@ use common_arrow::arrow::bitmap::MutableBitmap;
use common_arrow::arrow::buffer::Buffer;
use common_exception::Result;

use crate::types::array::ArrayColumn;
use crate::types::array::ArrayColumnBuilder;
use crate::types::decimal::DecimalColumn;
use crate::types::map::KvColumnBuilder;
use crate::types::nullable::NullableColumn;
use crate::types::number::NumberColumn;
use crate::types::string::StringColumn;
use crate::types::string::StringColumnBuilder;
use crate::types::AnyType;
use crate::types::ArrayType;
use crate::types::BooleanType;
use crate::types::MapType;
use crate::types::ValueType;
use crate::types::VariantType;
use crate::with_decimal_type;
Expand Down Expand Up @@ -133,7 +136,7 @@ impl Column {
let d = Self::filter_primitive_types(column, filter);
Column::Date(d)
}
Column::Array(column) | Column::Map(column) => {
Column::Array(column) => {
let mut offsets = Vec::with_capacity(length + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
Expand All @@ -143,6 +146,25 @@ impl Column {
let builder = ArrayColumnBuilder { builder, offsets };
Self::filter_scalar_types::<ArrayType<AnyType>>(column, builder, filter)
}
Column::Map(column) => {
let mut offsets = Vec::with_capacity(length + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
TypeDeserializerImpl::with_capacity(&column.values.data_type(), length)
.finish_to_column(),
);
let (key_builder, val_builder) = match builder {
ColumnBuilder::Tuple { fields, .. } => (fields[0].clone(), fields[1].clone()),
_ => unreachable!(),
};
let builder = KvColumnBuilder {
keys: key_builder,
values: val_builder,
};
let builder = ArrayColumnBuilder { builder, offsets };
let column = ArrayColumn::try_downcast(column).unwrap();
Self::filter_scalar_types::<MapType<AnyType, AnyType>>(&column, builder, filter)
}
Column::Nullable(c) => {
let column = Self::filter(&c.column, filter);
let validity = Self::filter_scalar_types::<BooleanType>(
Expand Down
29 changes: 28 additions & 1 deletion src/query/expression/src/kernels/scatter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,10 @@ use common_arrow::arrow::bitmap::MutableBitmap;
use common_exception::Result;
use itertools::Itertools;

use crate::types::array::ArrayColumn;
use crate::types::array::ArrayColumnBuilder;
use crate::types::decimal::DecimalColumn;
use crate::types::map::KvColumnBuilder;
use crate::types::nullable::NullableColumn;
use crate::types::number::NumberColumn;
use crate::types::string::StringColumnBuilder;
Expand All @@ -26,6 +28,7 @@ use crate::types::ArrayType;
use crate::types::BooleanType;
use crate::types::DataType;
use crate::types::DateType;
use crate::types::MapType;
use crate::types::NumberType;
use crate::types::StringType;
use crate::types::TimestampType;
Expand Down Expand Up @@ -191,7 +194,7 @@ impl Column {
indices,
scatter_size,
),
Column::Array(column) | Column::Map(column) => {
Column::Array(column) => {
let mut offsets = Vec::with_capacity(length + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
Expand All @@ -206,6 +209,30 @@ impl Column {
scatter_size,
)
}
Column::Map(column) => {
let mut offsets = Vec::with_capacity(length + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
TypeDeserializerImpl::with_capacity(&column.values.data_type(), length)
.finish_to_column(),
);
let (key_builder, val_builder) = match builder {
ColumnBuilder::Tuple { fields, .. } => (fields[0].clone(), fields[1].clone()),
_ => unreachable!(),
};
let builder = KvColumnBuilder {
keys: key_builder,
values: val_builder,
};
let builder = ArrayColumnBuilder { builder, offsets };
let column = ArrayColumn::try_downcast(column).unwrap();
Self::scatter_scalars::<MapType<AnyType, AnyType>, _>(
&column,
builder,
indices,
scatter_size,
)
}
Column::Nullable(c) => {
let columns = c.column.scatter(data_type, indices, scatter_size);
let validitys = Self::scatter_scalars::<BooleanType, _>(
Expand Down
24 changes: 23 additions & 1 deletion src/query/expression/src/kernels/take.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,17 @@
use common_exception::Result;
use itertools::Itertools;

use crate::types::array::ArrayColumn;
use crate::types::array::ArrayColumnBuilder;
use crate::types::decimal::DecimalColumn;
use crate::types::map::KvColumnBuilder;
use crate::types::nullable::NullableColumn;
use crate::types::number::NumberColumn;
use crate::types::AnyType;
use crate::types::ArgType;
use crate::types::ArrayType;
use crate::types::BooleanType;
use crate::types::MapType;
use crate::types::NumberType;
use crate::types::StringType;
use crate::types::ValueType;
Expand Down Expand Up @@ -102,7 +105,7 @@ impl Column {
.unwrap();
Column::Date(d)
}
Column::Array(column) | Column::Map(column) => {
Column::Array(column) => {
let mut offsets = Vec::with_capacity(length + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
Expand All @@ -112,6 +115,25 @@ impl Column {
let builder = ArrayColumnBuilder { builder, offsets };
Self::take_value_types::<ArrayType<AnyType>, _>(column, builder, indices)
}
Column::Map(column) => {
let mut offsets = Vec::with_capacity(length + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
TypeDeserializerImpl::with_capacity(&column.values.data_type(), self.len())
.finish_to_column(),
);
let (key_builder, val_builder) = match builder {
ColumnBuilder::Tuple { fields, .. } => (fields[0].clone(), fields[1].clone()),
_ => unreachable!(),
};
let builder = KvColumnBuilder {
keys: key_builder,
values: val_builder,
};
let builder = ArrayColumnBuilder { builder, offsets };
let column = ArrayColumn::try_downcast(column).unwrap();
Self::take_value_types::<MapType<AnyType, AnyType>, _>(&column, builder, indices)
}
Column::Nullable(c) => {
let column = c.column.take(indices);
let validity = Self::take_arg_types::<BooleanType, _>(&c.validity, indices);
Expand Down
22 changes: 21 additions & 1 deletion src/query/expression/src/kernels/take_chunks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use itertools::Itertools;

use crate::types::array::ArrayColumnBuilder;
use crate::types::decimal::DecimalColumn;
use crate::types::map::KvColumnBuilder;
use crate::types::nullable::NullableColumn;
use crate::types::number::NumberColumn;
use crate::types::AnyType;
Expand All @@ -25,6 +26,7 @@ use crate::types::ArrayType;
use crate::types::BooleanType;
use crate::types::DataType;
use crate::types::DateType;
use crate::types::MapType;
use crate::types::NumberType;
use crate::types::StringType;
use crate::types::TimestampType;
Expand Down Expand Up @@ -237,7 +239,7 @@ impl Column {
let builder = DateType::create_builder(result_size, &[]);
Self::take_block_value_types::<DateType>(columns, builder, indices)
}
Column::Array(column) | Column::Map(column) => {
Column::Array(column) => {
let mut offsets = Vec::with_capacity(result_size + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
Expand All @@ -247,6 +249,24 @@ impl Column {
let builder = ArrayColumnBuilder { builder, offsets };
Self::take_block_value_types::<ArrayType<AnyType>>(columns, builder, indices)
}
Column::Map(column) => {
let mut offsets = Vec::with_capacity(result_size + 1);
offsets.push(0);
let builder = ColumnBuilder::from_column(
TypeDeserializerImpl::with_capacity(&column.values.data_type(), result_size)
.finish_to_column(),
);
let (key_builder, val_builder) = match builder {
ColumnBuilder::Tuple { fields, .. } => (fields[0].clone(), fields[1].clone()),
_ => unreachable!(),
};
let builder = KvColumnBuilder {
keys: key_builder,
values: val_builder,
};
let builder = ArrayColumnBuilder { builder, offsets };
Self::take_block_value_types::<MapType<AnyType, AnyType>>(columns, builder, indices)
}
Column::Nullable(_) => {
let inner_ty = datatype.as_nullable().unwrap();
let inner_columns = columns
Expand Down
Loading

0 comments on commit 8c324cc

Please sign in to comment.