Skip to content

Commit 6dfea15

Browse files
BohuTANGBohuTANG
authored
Merge pull request #7572 from andylokandy/json2
feat(expr): support variant type
2 parents 6b620fb + 93851b3 commit 6dfea15

File tree

23 files changed

+1008
-563
lines changed

23 files changed

+1008
-563
lines changed

‎Cargo.lock

Lines changed: 32 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/query/expression/Cargo.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,13 +22,16 @@ chrono-tz = "0.6.1"
2222
comfy-table = "6"
2323
educe = "0.4"
2424
enum-as-inner = "0.4"
25+
hex = "0.4.3"
2526
itertools = "0.10"
2627
match-template = "0.0.1"
2728
num-traits = "0.2.15"
2829
# TODO(andylokandy): Use the version from crates.io once
2930
# https://github.com/reem/rust-ordered-float/pull/110 is released.
3031
ordered-float = { git = "https://github.com/andylokandy/rust-ordered-float.git", branch = "as", features = ["serde"] }
3132
serde = "1.0"
33+
# TODO: Switch to jsonb. bson is used for placeholder.
34+
bson = "2.4"
3235

3336
[dev-dependencies]
3437
common-ast = { path = "../ast" }

‎src/query/expression/src/display.rs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use std::fmt::Formatter;
1818
use std::time::Duration;
1919
use std::time::UNIX_EPOCH;
2020

21+
use bson::Document;
2122
use chrono::DateTime;
2223
use chrono::Utc;
2324
use comfy_table::Cell;
@@ -39,6 +40,7 @@ use crate::types::number::NumberDataType;
3940
use crate::types::number::NumberDomain;
4041
use crate::types::number::NumberScalar;
4142
use crate::types::number::SimpleDomain;
43+
use crate::types::string::StringColumn;
4244
use crate::types::string::StringDomain;
4345
use crate::types::timestamp::Timestamp;
4446
use crate::types::timestamp::TimestampDomain;
@@ -103,6 +105,7 @@ impl<'a> Debug for ScalarRef<'a> {
103105
fields.iter().map(ScalarRef::to_string).join(", ")
104106
)
105107
}
108+
ScalarRef::Variant(s) => write!(f, "0x{}", &hex::encode(s)),
106109
}
107110
}
108111
}
@@ -123,6 +126,7 @@ impl Debug for Column {
123126
.field("fields", fields)
124127
.field("len", len)
125128
.finish(),
129+
Column::Variant(col) => write!(f, "{col:?}"),
126130
}
127131
}
128132
}
@@ -144,6 +148,11 @@ impl<'a> Display for ScalarRef<'a> {
144148
fields.iter().map(ScalarRef::to_string).join(", ")
145149
)
146150
}
151+
ScalarRef::Variant(s) => {
152+
let doc = Document::from_reader(*s).map_err(|_| std::fmt::Error)?;
153+
let bson = doc.get("v").ok_or(std::fmt::Error)?;
154+
write!(f, "{bson}")
155+
}
147156
}
148157
}
149158
}
@@ -205,6 +214,15 @@ impl Debug for NumberColumn {
205214
}
206215
}
207216

217+
impl Debug for StringColumn {
218+
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
219+
f.debug_struct("StringColumn")
220+
.field("data", &format_args!("0x{}", &hex::encode(&*self.data)))
221+
.field("offsets", &self.offsets)
222+
.finish()
223+
}
224+
}
225+
208226
impl Display for RawExpr {
209227
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
210228
match self {
@@ -293,6 +311,7 @@ impl Display for DataType {
293311
write!(f, ")")
294312
}
295313
}
314+
DataType::Variant => write!(f, "Variant"),
296315
DataType::Generic(index) => write!(f, "T{index}"),
297316
}
298317
}

‎src/query/expression/src/kernels/concat.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ use crate::types::NumberType;
3131
use crate::types::StringType;
3232
use crate::types::TimestampType;
3333
use crate::types::ValueType;
34+
use crate::types::VariantType;
3435
use crate::with_number_mapped_type;
3536
use crate::Chunk;
3637
use crate::Column;
@@ -78,7 +79,7 @@ impl Column {
7879
match &columns[0] {
7980
Column::Null { .. } => Self::concat_arg_types::<NullType>(columns),
8081
Column::EmptyArray { .. } => Self::concat_arg_types::<EmptyArrayType>(columns),
81-
Column::Number(col) => with_number_mapped_type!(NUM_TYPE, match col {
82+
Column::Number(col) => with_number_mapped_type!(|NUM_TYPE| match col {
8283
NumberColumn::NUM_TYPE(_) => {
8384
Self::concat_arg_types::<NumberType<NUM_TYPE>>(columns)
8485
}
@@ -128,6 +129,7 @@ impl Column {
128129
len: capacity,
129130
}
130131
}
132+
Column::Variant(_) => Self::concat_arg_types::<VariantType>(columns),
131133
}
132134
}
133135

‎src/query/expression/src/kernels/filter.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,11 @@ impl Column {
158158
let fields = fields.iter().map(|c| c.filter(filter)).collect();
159159
Column::Tuple { fields, len }
160160
}
161+
Column::Variant(column) => Self::filter_scalar_types::<StringType>(
162+
column,
163+
StringColumnBuilder::with_capacity(length, 0),
164+
filter,
165+
),
161166
}
162167
}
163168

‎src/query/expression/src/kernels/scatter.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ use crate::types::NumberType;
2828
use crate::types::StringType;
2929
use crate::types::TimestampType;
3030
use crate::types::ValueType;
31+
use crate::types::VariantType;
3132
use crate::with_number_mapped_type;
3233
use crate::Chunk;
3334
use crate::Column;
@@ -91,7 +92,7 @@ impl Column {
9192
Column::Null { .. } => {
9293
Self::scatter_repeat_scalars::<I>(&Scalar::Null, indices, scatter_size)
9394
}
94-
Column::Number(column) => with_number_mapped_type!(NUM_TYPE, match column {
95+
Column::Number(column) => with_number_mapped_type!(|NUM_TYPE| match column {
9596
NumberColumn::NUM_TYPE(values) => Self::scatter_scalars::<NumberType<NUM_TYPE>, _>(
9697
values,
9798
Vec::with_capacity(length),
@@ -108,14 +109,12 @@ impl Column {
108109
indices,
109110
scatter_size,
110111
),
111-
112112
Column::String(column) => Self::scatter_scalars::<StringType, _>(
113113
column,
114114
StringColumnBuilder::with_capacity(length, 0),
115115
indices,
116116
scatter_size,
117117
),
118-
119118
Column::Timestamp(column) => Self::scatter_scalars::<TimestampType, _>(
120119
column,
121120
TimestampColumnBuilder::with_capacity(length),
@@ -170,6 +169,12 @@ impl Column {
170169
})
171170
.collect()
172171
}
172+
Column::Variant(column) => Self::scatter_scalars::<VariantType, _>(
173+
column,
174+
StringColumnBuilder::with_capacity(length, 0),
175+
indices,
176+
scatter_size,
177+
),
173178
}
174179
}
175180

‎src/query/expression/src/kernels/take.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ use crate::types::BooleanType;
2626
use crate::types::NumberType;
2727
use crate::types::StringType;
2828
use crate::types::ValueType;
29+
use crate::types::VariantType;
2930
use crate::with_number_mapped_type;
3031
use crate::Chunk;
3132
use crate::Column;
@@ -53,7 +54,7 @@ impl Column {
5354
let length = indices.len();
5455
match self {
5556
Column::Null { .. } | Column::EmptyArray { .. } => self.slice(0..length),
56-
Column::Number(column) => with_number_mapped_type!(NUM_TYPE, match column {
57+
Column::Number(column) => with_number_mapped_type!(|NUM_TYPE| match column {
5758
NumberColumn::NUM_TYPE(values) =>
5859
Self::take_arg_types::<NumberType<NUM_TYPE>, _>(values, indices),
5960
}),
@@ -90,6 +91,7 @@ impl Column {
9091
len: indices.len(),
9192
}
9293
}
94+
Column::Variant(column) => Self::take_arg_types::<VariantType, _>(column, indices),
9395
}
9496
}
9597

‎src/query/expression/src/types.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ pub mod nullable;
2424
pub mod number;
2525
pub mod string;
2626
pub mod timestamp;
27+
pub mod variant;
2728

2829
use std::fmt::Debug;
2930
use std::ops::Range;
@@ -45,6 +46,7 @@ pub use self::number::NumberDataType;
4546
pub use self::number::NumberType;
4647
pub use self::string::StringType;
4748
pub use self::timestamp::TimestampType;
49+
pub use self::variant::VariantType;
4850
use crate::property::Domain;
4951
use crate::values::Column;
5052
use crate::values::Scalar;
@@ -66,6 +68,7 @@ pub enum DataType {
6668
Array(Box<DataType>),
6769
Map(Box<DataType>),
6870
Tuple(Vec<DataType>),
71+
Variant,
6972
Generic(usize),
7073
}
7174

‎src/query/expression/src/types/number.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,8 @@ impl<Num: Number> ValueType for NumberType<Num> {
127127
builder.push(Num::default());
128128
}
129129

130-
fn append_builder(builder: &mut Self::ColumnBuilder, other_builder: &Self::ColumnBuilder) {
131-
builder.extend_from_slice(other_builder);
130+
fn append_builder(builder: &mut Self::ColumnBuilder, other: &Self::ColumnBuilder) {
131+
builder.extend_from_slice(other);
132132
}
133133

134134
fn build_column(builder: Self::ColumnBuilder) -> Self::Column {
@@ -829,7 +829,7 @@ macro_rules! with_number_type {
829829

830830
#[macro_export]
831831
macro_rules! with_number_mapped_type {
832-
($t:tt, $($tail:tt)*) => {
832+
(| $t:tt | $($tail:tt)*) => {
833833
match_template::match_template! {
834834
$t = [
835835
UInt8 => u8, UInt16 => u16, UInt32 => u32, UInt64 => u64,

‎src/query/expression/src/types/string.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,7 @@ impl ArgType for StringType {
134134
}
135135
}
136136

137-
#[derive(Debug, Clone, PartialEq)]
137+
#[derive(Clone, PartialEq)]
138138
pub struct StringColumn {
139139
pub data: Buffer<u8>,
140140
pub offsets: Buffer<u64>,

0 commit comments

Comments
 (0)