Skip to content

Turn jsonb stringify functions into separate functions #32055

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
antiguru merged 1 commit into from
Apr 1, 2025
Merged

Turn jsonb stringify functions into separate functions #32055

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/buf.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ breaking:
- compute-types/src/sinks.proto
# reason: Ignore because plans are currently not persisted.
- expr/src/relation.proto
# reason: Ignore because plans are currently not persisted.
- expr/src/scalar.proto
# reason: we very carefully evolve these protobuf definitions
- persist-client/src/internal/state.proto
# reason: does currently not require backward-compatibility
Expand Down
2 changes: 1 addition & 1 deletion src/expr-test-util/tests/testdata/scalar
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ build-scalar
case when (#0 > -2) then substr(#1, 1, 4) else "hello" end

build-scalar
(call_binary (jsonb_get_string true) #2 ("field1" string))
(call_binary (jsonb_get_string_stringify) #2 ("field1" string))
----
(#2 ->> "field1")

Expand Down
20 changes: 10 additions & 10 deletions src/expr/src/interpret.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -688,12 +688,12 @@ impl SpecialBinary {
}

match func {
BinaryFunc::JsonbGetString { stringify } => Some(SpecialBinary {
map_fn: if *stringify {
|l, r| jsonb_get_string(l, r, true)
} else {
|l, r| jsonb_get_string(l, r, false)
},
BinaryFunc::JsonbGetString => Some(SpecialBinary {
map_fn: |l, r| jsonb_get_string(l, r, false),
pushdownable: (true, false),
}),
BinaryFunc::JsonbGetStringStringify => Some(SpecialBinary {
map_fn: |l, r| jsonb_get_string(l, r, true),
pushdownable: (true, false),
}),
BinaryFunc::Eq => Some(SpecialBinary {
Expand Down Expand Up @@ -1177,8 +1177,8 @@ mod tests {
Lte,
Gte,
DateTruncTimestamp,
JsonbGetString { stringify: true },
JsonbGetString { stringify: false },
JsonbGetString,
JsonbGetStringStringify,
]
};

Expand All @@ -1200,7 +1200,7 @@ mod tests {
.scalar_type
.base_eq(&ScalarType::Timestamp { precision: None })
}
JsonbGetString { .. } => {
JsonbGetString | JsonbGetStringStringify => {
arg0.scalar_type.base_eq(&ScalarType::Jsonb)
&& arg1.scalar_type.base_eq(&ScalarType::String)
}
Expand Down Expand Up @@ -1567,7 +1567,7 @@ mod tests {
let expr = MirScalarExpr::CallUnary {
func: UnaryFunc::CastJsonbToNumeric(CastJsonbToNumeric(None)),
expr: Box::new(MirScalarExpr::CallBinary {
func: BinaryFunc::JsonbGetString { stringify: false },
func: BinaryFunc::JsonbGetString,
expr1: Box::new(MirScalarExpr::Column(0)),
expr2: Box::new(MirScalarExpr::Literal(
Ok(Row::pack_slice(&["ts".into()])),
Expand Down
11 changes: 8 additions & 3 deletions src/expr/src/scalar.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@

// See https://developers.google.com/protocol-buffers for what's going on here.

// buf breaking: ignore (Ignore because plans are currently not persisted.)

syntax = "proto3";

package mz_expr.scalar;
Expand Down Expand Up @@ -574,9 +576,12 @@ message ProtoBinaryFunc {
google.protobuf.Empty timezone_interval_timestamp_tz = 95;
google.protobuf.Empty timezone_interval_time = 96;
google.protobuf.Empty text_concat = 97;
bool jsonb_get_int64 = 98;
bool jsonb_get_string = 99;
bool jsonb_get_path = 100;
google.protobuf.Empty jsonb_get_int64 = 98;
google.protobuf.Empty jsonb_get_int64_stringify = 197;
google.protobuf.Empty jsonb_get_string = 99;
google.protobuf.Empty jsonb_get_string_stringify = 198;
google.protobuf.Empty jsonb_get_path = 100;
google.protobuf.Empty jsonb_get_path_stringify = 199;
Copy link
Contributor

@ggevay ggevay Apr 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CI is complaining about compatibility, but as discussed on Slack, you could put an ignore into this file.

google.protobuf.Empty jsonb_contains_string = 101;
google.protobuf.Empty jsonb_concat = 102;
google.protobuf.Empty jsonb_contains_jsonb = 103;
Expand Down
123 changes: 66 additions & 57 deletions src/expr/src/scalar/func.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2310,15 +2310,12 @@ pub enum BinaryFunc {
TimezoneIntervalTime,
TimezoneOffset,
TextConcat,
JsonbGetInt64 {
stringify: bool,
},
JsonbGetString {
stringify: bool,
},
JsonbGetPath {
stringify: bool,
},
JsonbGetInt64,
JsonbGetInt64Stringify,
JsonbGetString,
JsonbGetStringStringify,
JsonbGetPath,
JsonbGetPathStringify,
JsonbContainsString,
JsonbConcat,
JsonbContainsJsonb,
Expand Down Expand Up @@ -2577,15 +2574,12 @@ impl BinaryFunc {
BinaryFunc::TimezoneIntervalTime => timezone_interval_time(a, b),
BinaryFunc::TimezoneOffset => timezone_offset(a, b, temp_storage),
BinaryFunc::TextConcat => Ok(text_concat_binary(a, b, temp_storage)),
BinaryFunc::JsonbGetInt64 { stringify } => {
Ok(jsonb_get_int64(a, b, temp_storage, *stringify))
}
BinaryFunc::JsonbGetString { stringify } => {
Ok(jsonb_get_string(a, b, temp_storage, *stringify))
}
BinaryFunc::JsonbGetPath { stringify } => {
Ok(jsonb_get_path(a, b, temp_storage, *stringify))
}
BinaryFunc::JsonbGetInt64 => Ok(jsonb_get_int64(a, b, temp_storage, false)),
BinaryFunc::JsonbGetInt64Stringify => Ok(jsonb_get_int64(a, b, temp_storage, true)),
BinaryFunc::JsonbGetString => Ok(jsonb_get_string(a, b, temp_storage, false)),
BinaryFunc::JsonbGetStringStringify => Ok(jsonb_get_string(a, b, temp_storage, true)),
BinaryFunc::JsonbGetPath => Ok(jsonb_get_path(a, b, temp_storage, false)),
BinaryFunc::JsonbGetPathStringify => Ok(jsonb_get_path(a, b, temp_storage, true)),
BinaryFunc::JsonbContainsString => Ok(jsonb_contains_string(a, b)),
BinaryFunc::JsonbConcat => Ok(jsonb_concat(a, b, temp_storage)),
BinaryFunc::JsonbContainsJsonb => Ok(jsonb_contains_jsonb(a, b)),
Expand Down Expand Up @@ -2783,13 +2777,13 @@ impl BinaryFunc {

MzRenderTypmod | TextConcat => ScalarType::String.nullable(in_nullable),

JsonbGetInt64 { stringify: true }
| JsonbGetString { stringify: true }
| JsonbGetPath { stringify: true } => ScalarType::String.nullable(true),
JsonbGetInt64Stringify
| JsonbGetStringStringify
| JsonbGetPathStringify => ScalarType::String.nullable(true),

JsonbGetInt64 { stringify: false }
| JsonbGetString { stringify: false }
| JsonbGetPath { stringify: false }
JsonbGetInt64
| JsonbGetString
| JsonbGetPath
| JsonbConcat
| JsonbDeleteInt64
| JsonbDeleteString => ScalarType::Jsonb.nullable(true),
Expand Down Expand Up @@ -3067,9 +3061,12 @@ impl BinaryFunc {
| RegexpReplace { .. }
| StartsWith => false,

JsonbGetInt64 { .. }
| JsonbGetString { .. }
| JsonbGetPath { .. }
JsonbGetInt64
| JsonbGetInt64Stringify
| JsonbGetString
| JsonbGetStringStringify
| JsonbGetPath
| JsonbGetPathStringify
| JsonbConcat
| JsonbDeleteInt64
| JsonbDeleteString
Expand Down Expand Up @@ -3184,9 +3181,12 @@ impl BinaryFunc {
| Gte
| JsonbConcat
| JsonbContainsJsonb
| JsonbGetInt64 { .. }
| JsonbGetString { .. }
| JsonbGetPath { .. }
| JsonbGetInt64
| JsonbGetInt64Stringify
| JsonbGetString
| JsonbGetStringStringify
| JsonbGetPath
| JsonbGetPathStringify
| JsonbContainsString
| JsonbDeleteInt64
| JsonbDeleteString
Expand Down Expand Up @@ -3330,9 +3330,12 @@ impl BinaryFunc {
| BinaryFunc::BitShiftRightUInt16
| BinaryFunc::BitShiftRightUInt32
| BinaryFunc::BitShiftRightUInt64 => false,
BinaryFunc::JsonbGetInt64 { .. }
| BinaryFunc::JsonbGetString { .. }
| BinaryFunc::JsonbGetPath { .. }
BinaryFunc::JsonbGetInt64
| BinaryFunc::JsonbGetInt64Stringify
| BinaryFunc::JsonbGetString
| BinaryFunc::JsonbGetStringStringify
| BinaryFunc::JsonbGetPath
| BinaryFunc::JsonbGetPathStringify
| BinaryFunc::JsonbContainsString
| BinaryFunc::JsonbConcat
| BinaryFunc::JsonbContainsJsonb
Expand Down Expand Up @@ -3493,9 +3496,12 @@ impl BinaryFunc {
| BinaryFunc::TimezoneIntervalTimestampTz
| BinaryFunc::TimezoneIntervalTime
| BinaryFunc::TimezoneOffset => (false, false),
BinaryFunc::JsonbGetInt64 { .. }
| BinaryFunc::JsonbGetString { .. }
| BinaryFunc::JsonbGetPath { .. }
BinaryFunc::JsonbGetInt64
| BinaryFunc::JsonbGetInt64Stringify
| BinaryFunc::JsonbGetString
| BinaryFunc::JsonbGetStringStringify
| BinaryFunc::JsonbGetPath
| BinaryFunc::JsonbGetPathStringify
| BinaryFunc::JsonbContainsString
| BinaryFunc::JsonbConcat
| BinaryFunc::JsonbContainsJsonb
Expand Down Expand Up @@ -3697,12 +3703,12 @@ impl fmt::Display for BinaryFunc {
BinaryFunc::TimezoneIntervalTime => f.write_str("timezoneit"),
BinaryFunc::TimezoneOffset => f.write_str("timezone_offset"),
BinaryFunc::TextConcat => f.write_str("||"),
BinaryFunc::JsonbGetInt64 { stringify: false } => f.write_str("->"),
BinaryFunc::JsonbGetInt64 { stringify: true } => f.write_str("->>"),
BinaryFunc::JsonbGetString { stringify: false } => f.write_str("->"),
BinaryFunc::JsonbGetString { stringify: true } => f.write_str("->>"),
BinaryFunc::JsonbGetPath { stringify: false } => f.write_str("#>"),
BinaryFunc::JsonbGetPath { stringify: true } => f.write_str("#>>"),
BinaryFunc::JsonbGetInt64 => f.write_str("->"),
BinaryFunc::JsonbGetInt64Stringify => f.write_str("->>"),
BinaryFunc::JsonbGetString => f.write_str("->"),
BinaryFunc::JsonbGetStringStringify => f.write_str("->>"),
BinaryFunc::JsonbGetPath => f.write_str("#>"),
BinaryFunc::JsonbGetPathStringify => f.write_str("#>>"),
BinaryFunc::JsonbContainsString | BinaryFunc::MapContainsKey => f.write_str("?"),
BinaryFunc::JsonbConcat => f.write_str("||"),
BinaryFunc::JsonbContainsJsonb | BinaryFunc::MapContainsMap => f.write_str("@>"),
Expand Down Expand Up @@ -3925,15 +3931,12 @@ impl Arbitrary for BinaryFunc {
Just(BinaryFunc::TimezoneIntervalTime).boxed(),
Just(BinaryFunc::TimezoneOffset).boxed(),
Just(BinaryFunc::TextConcat).boxed(),
bool::arbitrary()
.prop_map(|stringify| BinaryFunc::JsonbGetInt64 { stringify })
.boxed(),
bool::arbitrary()
.prop_map(|stringify| BinaryFunc::JsonbGetString { stringify })
.boxed(),
bool::arbitrary()
.prop_map(|stringify| BinaryFunc::JsonbGetPath { stringify })
.boxed(),
Just(BinaryFunc::JsonbGetInt64).boxed(),
Just(BinaryFunc::JsonbGetInt64Stringify).boxed(),
Just(BinaryFunc::JsonbGetString).boxed(),
Just(BinaryFunc::JsonbGetStringStringify).boxed(),
Just(BinaryFunc::JsonbGetPath).boxed(),
Just(BinaryFunc::JsonbGetPathStringify).boxed(),
Just(BinaryFunc::JsonbContainsString).boxed(),
Just(BinaryFunc::JsonbConcat).boxed(),
Just(BinaryFunc::JsonbContainsJsonb).boxed(),
Expand Down Expand Up @@ -4125,9 +4128,12 @@ impl RustType<ProtoBinaryFunc> for BinaryFunc {
BinaryFunc::TimezoneIntervalTime => TimezoneIntervalTime(()),
BinaryFunc::TimezoneOffset => TimezoneOffset(()),
BinaryFunc::TextConcat => TextConcat(()),
BinaryFunc::JsonbGetInt64 { stringify } => JsonbGetInt64(*stringify),
BinaryFunc::JsonbGetString { stringify } => JsonbGetString(*stringify),
BinaryFunc::JsonbGetPath { stringify } => JsonbGetPath(*stringify),
BinaryFunc::JsonbGetInt64 => JsonbGetInt64(()),
BinaryFunc::JsonbGetInt64Stringify => JsonbGetInt64Stringify(()),
BinaryFunc::JsonbGetString => JsonbGetString(()),
BinaryFunc::JsonbGetStringStringify => JsonbGetStringStringify(()),
BinaryFunc::JsonbGetPath => JsonbGetPath(()),
BinaryFunc::JsonbGetPathStringify => JsonbGetPathStringify(()),
BinaryFunc::JsonbContainsString => JsonbContainsString(()),
BinaryFunc::JsonbConcat => JsonbConcat(()),
BinaryFunc::JsonbContainsJsonb => JsonbContainsJsonb(()),
Expand Down Expand Up @@ -4344,9 +4350,12 @@ impl RustType<ProtoBinaryFunc> for BinaryFunc {
TimezoneIntervalTime(()) => Ok(BinaryFunc::TimezoneIntervalTime),
TimezoneOffset(()) => Ok(BinaryFunc::TimezoneOffset),
TextConcat(()) => Ok(BinaryFunc::TextConcat),
JsonbGetInt64(stringify) => Ok(BinaryFunc::JsonbGetInt64 { stringify }),
JsonbGetString(stringify) => Ok(BinaryFunc::JsonbGetString { stringify }),
JsonbGetPath(stringify) => Ok(BinaryFunc::JsonbGetPath { stringify }),
JsonbGetInt64(()) => Ok(BinaryFunc::JsonbGetInt64),
JsonbGetInt64Stringify(()) => Ok(BinaryFunc::JsonbGetInt64Stringify),
JsonbGetString(()) => Ok(BinaryFunc::JsonbGetString),
JsonbGetStringStringify(()) => Ok(BinaryFunc::JsonbGetStringStringify),
JsonbGetPath(()) => Ok(BinaryFunc::JsonbGetPath),
JsonbGetPathStringify(()) => Ok(BinaryFunc::JsonbGetPathStringify),
JsonbContainsString(()) => Ok(BinaryFunc::JsonbContainsString),
JsonbConcat(()) => Ok(BinaryFunc::JsonbConcat),
JsonbContainsJsonb(()) => Ok(BinaryFunc::JsonbContainsJsonb),
Expand Down
8 changes: 4 additions & 4 deletions src/expr/tests/testdata/interpret
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ interpret
[
[("{\"created_ms\": 1000}" jsonb) ("{\"created_ms\": 2000}" jsonb) ("{\"created_ms\": 3000}" jsonb)]
]
(call_binary (jsonb_get_string false) #0 ("created_ms" string))
(call_binary (jsonb_get_string) #0 ("created_ms" string))
[(0 numeric) (2050 numeric) (4000 numeric) null]
----
may contain: [2050]
Expand All @@ -70,7 +70,7 @@ interpret
[
[("{\"code\": \"00135\"}" jsonb) ("{\"code\": \"22122\"}" jsonb) ("{\"code\": \"34153\"}" jsonb)]
]
(call_binary (jsonb_get_string true) #0 ("code" string))
(call_binary (jsonb_get_string_stringify) #0 ("code" string))
["00000" "20000" "2" "80000" null]
----
may contain: ["20000" "2"]
Expand All @@ -82,7 +82,7 @@ interpret
[
[("{\"created_ms\": 1000}" jsonb) ("{\"created_ms\": 2000}" jsonb) ("{\"created_ms\": 3000}" jsonb)]
]
(call_binary (jsonb_get_string true) #0 ("created_ms" string))
(call_binary (jsonb_get_string_stringify) #0 ("created_ms" string))
["00000" "20000" "2" "80000" null]
----
may contain: ["00000" "20000" "2" "80000"]
Expand All @@ -94,7 +94,7 @@ interpret
[
[null null null]
]
(call_binary (jsonb_get_string false) #0 ("created_ms" string))
(call_binary (jsonb_get_string) #0 ("created_ms" string))
["00000" "foo" null true]
----
may contain: [null]
Expand Down
2 changes: 1 addition & 1 deletion src/expr/tests/testdata/reduce
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ false

reduce
(call_unary is_null
(call_binary (jsonb_get_int64 false) #1 #0))
(call_binary (jsonb_get_int64) #1 #0))
[(jsonb false) (int64 false)]
----
((#1 -> #0)) IS NULL
Expand Down
12 changes: 6 additions & 6 deletions src/sql/src/func.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4765,19 +4765,19 @@ pub static OP_IMPLS: LazyLock<BTreeMap<&'static str, Func>> = LazyLock::new(|| {

// JSON, MAP, RANGE, LIST, ARRAY
"->" => Scalar {
params!(Jsonb, Int64) => JsonbGetInt64 { stringify: false } => Jsonb, 3212;
params!(Jsonb, String) => JsonbGetString { stringify: false } => Jsonb, 3211;
params!(Jsonb, Int64) => JsonbGetInt64 => Jsonb, 3212;
params!(Jsonb, String) => JsonbGetString => Jsonb, 3211;
params!(MapAny, String) => MapGetValue => Any, oid::OP_GET_VALUE_MAP_OID;
},
"->>" => Scalar {
params!(Jsonb, Int64) => JsonbGetInt64 { stringify: true } => String, 3481;
params!(Jsonb, String) => JsonbGetString { stringify: true } => String, 3477;
params!(Jsonb, Int64) => JsonbGetInt64Stringify => String, 3481;
params!(Jsonb, String) => JsonbGetStringStringify => String, 3477;
},
"#>" => Scalar {
params!(Jsonb, ScalarType::Array(Box::new(ScalarType::String))) => JsonbGetPath { stringify: false } => Jsonb, 3213;
params!(Jsonb, ScalarType::Array(Box::new(ScalarType::String))) => JsonbGetPath => Jsonb, 3213;
},
"#>>" => Scalar {
params!(Jsonb, ScalarType::Array(Box::new(ScalarType::String))) => JsonbGetPath { stringify: true } => String, 3206;
params!(Jsonb, ScalarType::Array(Box::new(ScalarType::String))) => JsonbGetPathStringify => String, 3206;
},
"@>" => Scalar {
params!(Jsonb, Jsonb) => JsonbContainsJsonb => Bool, 3246;
Expand Down
2 changes: 1 addition & 1 deletion src/sql/src/plan/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4374,7 +4374,7 @@ fn plan_subscript_jsonb(
},
exprs,
},
BinaryFunc::JsonbGetPath { stringify: false },
BinaryFunc::JsonbGetPath,
);
Ok(expr.into())
}
Expand Down
4 changes: 2 additions & 2 deletions test/sqllogictest/distinct_arrangements.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -984,8 +984,8 @@ Arrange␠export␠iterative 2
Arrange␠export␠iterative␠err 2
Arrange␠recursive␠err 3
ArrangeAccumulable␠[val:␠empty] 10
ArrangeBy[[CallBinary␠{␠func:␠JsonbGetString␠{␠stringify:␠true␠},␠expr1:␠Column(1),␠expr2:␠Literal(Ok(Row{[String("id")]}),␠ColumnType␠{␠scalar_type:␠String,␠nullable:␠false␠})␠}]] 2
ArrangeBy[[CallBinary␠{␠func:␠JsonbGetString␠{␠stringify:␠true␠},␠expr1:␠Column(2),␠expr2:␠Literal(Ok(Row{[String("id")]}),␠ColumnType␠{␠scalar_type:␠String,␠nullable:␠false␠})␠}]] 1
ArrangeBy[[CallBinary␠{␠func:␠JsonbGetStringStringify,␠expr1:␠Column(1),␠expr2:␠Literal(Ok(Row{[String("id")]}),␠ColumnType␠{␠scalar_type:␠String,␠nullable:␠false␠})␠}]] 2
ArrangeBy[[CallBinary␠{␠func:␠JsonbGetStringStringify,␠expr1:␠Column(2),␠expr2:␠Literal(Ok(Row{[String("id")]}),␠ColumnType␠{␠scalar_type:␠String,␠nullable:␠false␠})␠}]] 1
ArrangeBy[[CallVariadic␠{␠func:␠Coalesce,␠exprs:␠[Column(2),␠Column(3)]␠}]] 2
ArrangeBy[[Column(0),␠CallUnary␠{␠func:␠CastInt32ToNumeric(CastInt32ToNumeric(None)),␠expr:␠Column(1)␠}]] 1
ArrangeBy[[Column(0),␠CallUnary␠{␠func:␠CastUint64ToNumeric(CastUint64ToNumeric(None)),␠expr:␠Column(2)␠}]] 1
Expand Down