From 750a99632701d4e893582ea16c3a7b12dab6ffef Mon Sep 17 00:00:00 2001 From: Xiangjin Date: Fri, 17 Feb 2023 15:05:37 +0800 Subject: [PATCH 1/3] feat(expr): jsonb access --- dashboard/proto/gen/expr.ts | 16 +- e2e_test/batch/types/jsonb.slt.part | 23 +++ proto/expr.proto | 9 +- src/common/src/array/jsonb_array.rs | 40 ++++ src/expr/src/expr/build_expr_from_prost.rs | 4 +- src/expr/src/expr/expr_binary_nonnull.rs | 39 +++- src/expr/src/expr/expr_jsonb_access.rs | 222 +++++++++++++++++++++ src/expr/src/expr/mod.rs | 1 + src/expr/src/sig/func.rs | 5 + src/frontend/src/binder/expr/function.rs | 5 + 10 files changed, 359 insertions(+), 5 deletions(-) create mode 100644 src/expr/src/expr/expr_jsonb_access.rs diff --git a/dashboard/proto/gen/expr.ts b/dashboard/proto/gen/expr.ts index d5e49cb52885f..bf277974f9609 100644 --- a/dashboard/proto/gen/expr.ts +++ b/dashboard/proto/gen/expr.ts @@ -130,8 +130,12 @@ export const ExprNode_Type = { ARRAY_APPEND: "ARRAY_APPEND", ARRAY_PREPEND: "ARRAY_PREPEND", FORMAT_TYPE: "FORMAT_TYPE", + /** JSONB_ACCESS_INNER - jsonb -> int, jsonb -> text, jsonb #> text[] that returns jsonb */ + JSONB_ACCESS_INNER: "JSONB_ACCESS_INNER", + /** JSONB_ACCESS_STR - jsonb ->> int, jsonb ->> text, jsonb #>> text[] that returns text */ + JSONB_ACCESS_STR: "JSONB_ACCESS_STR", /** - * VNODE - Non-pure functions below (> 600) + * VNODE - Non-pure functions below (> 1000) * ------------------------ * Internal functions */ @@ -402,6 +406,12 @@ export function exprNode_TypeFromJSON(object: any): ExprNode_Type { case 534: case "FORMAT_TYPE": return ExprNode_Type.FORMAT_TYPE; + case 600: + case "JSONB_ACCESS_INNER": + return ExprNode_Type.JSONB_ACCESS_INNER; + case 601: + case "JSONB_ACCESS_STR": + return ExprNode_Type.JSONB_ACCESS_STR; case 1101: case "VNODE": return ExprNode_Type.VNODE; @@ -590,6 +600,10 @@ export function exprNode_TypeToJSON(object: ExprNode_Type): string { return "ARRAY_PREPEND"; case ExprNode_Type.FORMAT_TYPE: return "FORMAT_TYPE"; + case ExprNode_Type.JSONB_ACCESS_INNER: + return "JSONB_ACCESS_INNER"; + case ExprNode_Type.JSONB_ACCESS_STR: + return "JSONB_ACCESS_STR"; case ExprNode_Type.VNODE: return "VNODE"; case ExprNode_Type.NOW: diff --git a/e2e_test/batch/types/jsonb.slt.part b/e2e_test/batch/types/jsonb.slt.part index f0d890bc9d877..e7c0db15bc7e1 100644 --- a/e2e_test/batch/types/jsonb.slt.part +++ b/e2e_test/batch/types/jsonb.slt.part @@ -83,5 +83,28 @@ select null::jsonb::bool; ---- NULL +query T +select jsonb_array_element(jsonb_object_field('{"k2":[2,true,4]}', 'k2'), -2)::bool; +---- +t + +query TT +with t(v1) as (values (null::jsonb), ('null'), ('true'), ('1'), ('"a"'), ('[]'), ('{}')), + j(v1) as (select ('{"k":' || v1::varchar || '}')::jsonb from t) +select jsonb_object_field_text(v1, 'k'), jsonb_object_field(v1, 'k')::varchar from j order by 2; +---- +a "a" +1 1 +[] [] +NULL null +true true +{} {} +NULL NULL + +query T +select jsonb_array_element_text('true'::jsonb, 2); +---- +NULL + statement ok drop table t; diff --git a/proto/expr.proto b/proto/expr.proto index 6909ccbf53e0d..701d14f1b66e0 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -120,7 +120,14 @@ message ExprNode { ARRAY_PREPEND = 533; FORMAT_TYPE = 534; - // Non-pure functions below (> 600) + // Jsonb functions + + // jsonb -> int, jsonb -> text, jsonb #> text[] that returns jsonb + JSONB_ACCESS_INNER = 600; + // jsonb ->> int, jsonb ->> text, jsonb #>> text[] that returns text + JSONB_ACCESS_STR = 601; + + // Non-pure functions below (> 1000) // ------------------------ // Internal functions VNODE = 1101; diff --git a/src/common/src/array/jsonb_array.rs b/src/common/src/array/jsonb_array.rs index 09ba62b11e554..628971d571207 100644 --- a/src/common/src/array/jsonb_array.rs +++ b/src/common/src/array/jsonb_array.rs @@ -173,6 +173,10 @@ impl JsonbRef<'_> { output.freeze().into() } + pub fn is_jsonb_null(&self) -> bool { + matches!(self.0, Value::Null) + } + pub fn type_name(&self) -> &'static str { match self.0 { Value::Null => "null", @@ -184,6 +188,16 @@ impl JsonbRef<'_> { } } + pub fn array_len(&self) -> Result { + match self.0 { + Value::Array(v) => Ok(v.len()), + _ => Err(format!( + "cannot get array length of a jsonb {}", + self.type_name() + )), + } + } + pub fn as_bool(&self) -> Result { match self.0 { Value::Bool(v) => Ok(*v), @@ -207,6 +221,32 @@ impl JsonbRef<'_> { )), } } + + /// This is part of the `->>` or `#>>` syntax to access a child as string. + /// + /// * It is not `as_str`, because there is no runtime error when the jsonb type is not string. + /// * It is not same as [`Display`] or [`ToText`] (cast to string) in the following 2 cases: + /// * Jsonb null is displayed as 4-letter `null` but treated as sql null here. + /// * This function writes nothing and the caller is responsible for checking + /// [`is_jsonb_null`] to differentiate it from an empty string. + /// * Jsonb string is displayed with quotes but treated as its inner value here. + pub fn force_str(&self, writer: &mut W) -> std::fmt::Result { + match self.0 { + Value::String(v) => writer.write_str(v), + Value::Null => Ok(()), + Value::Bool(_) | Value::Number(_) | Value::Array(_) | Value::Object(_) => { + write!(writer, "{}", self.0) + } + } + } + + pub fn access_object_field(&self, field: &str) -> Option { + self.0.get(field).map(Self) + } + + pub fn access_array_element(&self, idx: usize) -> Option { + self.0.get(idx).map(Self) + } } #[derive(Debug)] diff --git a/src/expr/src/expr/build_expr_from_prost.rs b/src/expr/src/expr/build_expr_from_prost.rs index b7f6ef27d337d..50578e84033dd 100644 --- a/src/expr/src/expr/build_expr_from_prost.rs +++ b/src/expr/src/expr/build_expr_from_prost.rs @@ -68,7 +68,9 @@ pub fn build_from_prost(prost: &ExprNode) -> Result { Equal | NotEqual | LessThan | LessThanOrEqual | GreaterThan | GreaterThanOrEqual | Add | Subtract | Multiply | Divide | Modulus | Extract | RoundDigit | Pow | TumbleStart | Position | BitwiseShiftLeft | BitwiseShiftRight | BitwiseAnd | BitwiseOr | BitwiseXor - | ConcatOp | AtTimeZone | CastWithTimeZone => build_binary_expr_prost(prost), + | ConcatOp | AtTimeZone | CastWithTimeZone | JsonbAccessInner | JsonbAccessStr => { + build_binary_expr_prost(prost) + } And | Or | IsDistinctFrom | IsNotDistinctFrom | ArrayAccess | FormatType => { build_nullable_binary_expr_prost(prost) } diff --git a/src/expr/src/expr/expr_binary_nonnull.rs b/src/expr/src/expr/expr_binary_nonnull.rs index 910b563af67f8..7d1e7488153c8 100644 --- a/src/expr/src/expr/expr_binary_nonnull.rs +++ b/src/expr/src/expr/expr_binary_nonnull.rs @@ -13,14 +13,17 @@ // limitations under the License. use risingwave_common::array::{ - Array, BoolArray, DecimalArray, F64Array, I32Array, I64Array, IntervalArray, ListArray, - NaiveDateArray, NaiveDateTimeArray, StructArray, Utf8Array, + Array, BoolArray, DecimalArray, F64Array, I32Array, I64Array, IntervalArray, JsonbArrayBuilder, + ListArray, NaiveDateArray, NaiveDateTimeArray, StructArray, Utf8Array, Utf8ArrayBuilder, }; use risingwave_common::types::*; use risingwave_pb::expr::expr_node::Type; use super::Expression; use crate::expr::expr_binary_bytes::new_concat_op; +use crate::expr::expr_jsonb_access::{ + jsonb_array_element, jsonb_object_field, JsonbAccessExpression, +}; use crate::expr::template::{BinaryBytesExpression, BinaryExpression}; use crate::expr::{template_fast, BoxedExpression}; use crate::vector_op::arithmetic_op::*; @@ -680,6 +683,38 @@ pub fn new_binary_expr( )), Type::TumbleStart => new_tumble_start(l, r, ret)?, Type::ConcatOp => new_concat_op(l, r, ret), + Type::JsonbAccessInner => match r.return_type() { + DataType::Varchar => { + JsonbAccessExpression::::new_expr( + l, + r, + jsonb_object_field, + ) + .boxed() + } + DataType::Int32 => JsonbAccessExpression::::new_expr( + l, + r, + jsonb_array_element, + ) + .boxed(), + t => return Err(ExprError::UnsupportedFunction(format!("jsonb -> {t}"))), + }, + Type::JsonbAccessStr => match r.return_type() { + DataType::Varchar => JsonbAccessExpression::::new_expr( + l, + r, + jsonb_object_field, + ) + .boxed(), + DataType::Int32 => JsonbAccessExpression::::new_expr( + l, + r, + jsonb_array_element, + ) + .boxed(), + t => return Err(ExprError::UnsupportedFunction(format!("jsonb ->> {t}"))), + }, tp => { return Err(ExprError::UnsupportedFunction(format!( diff --git a/src/expr/src/expr/expr_jsonb_access.rs b/src/expr/src/expr/expr_jsonb_access.rs new file mode 100644 index 0000000000000..e4a725d75f3a5 --- /dev/null +++ b/src/expr/src/expr/expr_jsonb_access.rs @@ -0,0 +1,222 @@ +// Copyright 2023 RisingWave Labs +// +// 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 either::Either; +use risingwave_common::array::{ + Array, ArrayBuilder, ArrayImpl, ArrayRef, DataChunk, JsonbArray, JsonbArrayBuilder, JsonbRef, + Utf8ArrayBuilder, +}; +use risingwave_common::row::OwnedRow; +use risingwave_common::types::{DataType, Datum, Scalar, ScalarRef}; +use risingwave_common::util::iter_util::ZipEqFast; + +use super::{BoxedExpression, Expression}; + +/// This is forked from [`BinaryExpression`] for the following reasons: +/// * Optimize for the case when rhs path is const. (not implemented yet) +/// * It can return null when neither input is null. +/// * We could `append(RefItem)` directly rather than getting a `OwnedItem` first. +pub struct JsonbAccessExpression { + input: BoxedExpression, + path: Either, + func: F, + _phantom: std::marker::PhantomData, +} + +impl std::fmt::Debug for JsonbAccessExpression { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("JsonbAccessExpression") + .field("input", &self.input) + .field("path", &self.path) + .finish() + } +} + +impl JsonbAccessExpression +where + F: Send + Sync + for<'a> Fn(JsonbRef<'a>, A::RefItem<'_>) -> Option>, +{ + #[expect(dead_code)] + pub fn new_const(input: BoxedExpression, path: A::OwnedItem, func: F) -> Self { + Self { + input, + path: Either::Right(path), + func, + _phantom: std::marker::PhantomData, + } + } + + pub fn new_expr(input: BoxedExpression, path: BoxedExpression, func: F) -> Self { + Self { + input, + path: Either::Left(path), + func, + _phantom: std::marker::PhantomData, + } + } + + pub fn eval_strict<'a>( + &self, + v: Option>, + p: Option>, + ) -> Option> { + match (v, p) { + (Some(v), Some(p)) => (self.func)(v, p), + _ => None, + } + } +} + +impl Expression for JsonbAccessExpression +where + A: Array, + for<'a> &'a A: From<&'a ArrayImpl>, + O: AccessOutput, + F: Send + Sync + for<'a> Fn(JsonbRef<'a>, A::RefItem<'_>) -> Option>, +{ + fn return_type(&self) -> DataType { + O::return_type() + } + + fn eval(&self, input: &DataChunk) -> crate::Result { + let Either::Left(path_expr) = &self.path else { + unreachable!("optimization for const path not implemented yet"); + }; + let path_array = path_expr.eval_checked(input)?; + let path_array: &A = path_array.as_ref().into(); + + let input_array = self.input.eval_checked(input)?; + let input_array: &JsonbArray = input_array.as_ref().into(); + + let mut builder = O::new(input.capacity()); + match input.visibility() { + // We could ignore visibility and always evaluate access path for all values, because it + // never returns runtime error. But using visibility could save us some clone cost, + // unless we adjust [`JsonbArray`] to make sure all clones are on [`Arc`]. + Some(visibility) => { + for ((v, p), visible) in input_array + .iter() + .zip_eq_fast(path_array.iter()) + .zip_eq_fast(visibility.iter()) + { + let r = visible.then(|| self.eval_strict(v, p)).flatten(); + builder.output_nullable(r)?; + } + } + None => { + for (v, p) in input_array.iter().zip_eq_fast(path_array.iter()) { + builder.output_nullable(self.eval_strict(v, p))?; + } + } + }; + Ok(std::sync::Arc::new(builder.finish().into())) + } + + fn eval_row(&self, input: &OwnedRow) -> crate::Result { + let Either::Left(path_expr) = &self.path else { + unreachable!("optimization for const path not implemented yet"); + }; + let p = path_expr.eval_row(input)?; + let p = p + .as_ref() + .map(|p| p.as_scalar_ref_impl().try_into().unwrap()); + + let v = self.input.eval_row(input)?; + let v = v + .as_ref() + .map(|v| v.as_scalar_ref_impl().try_into().unwrap()); + + let r = self.eval_strict(v, p); + Ok(r.and_then(O::output_datum)) + } +} + +pub fn jsonb_object_field<'a>(v: JsonbRef<'a>, p: &str) -> Option> { + v.access_object_field(p) +} + +pub fn jsonb_array_element(v: JsonbRef<'_>, p: i32) -> Option> { + let idx = if p < 0 { + let Ok(len) = v.array_len() else { + return None; + }; + if ((-p) as usize) > len { + return None; + } else { + len - ((-p) as usize) + } + } else { + p as usize + }; + v.access_array_element(idx) +} + +trait AccessOutput: ArrayBuilder { + fn return_type() -> DataType; + fn output(&mut self, v: JsonbRef<'_>) -> crate::Result<()>; + fn output_datum(v: JsonbRef<'_>) -> Datum; + fn output_nullable(&mut self, v: Option>) -> crate::Result<()> { + match v { + Some(v) => self.output(v)?, + None => self.append_null(), + }; + Ok(()) + } +} + +impl AccessOutput for JsonbArrayBuilder { + fn return_type() -> DataType { + DataType::Jsonb + } + + fn output(&mut self, v: JsonbRef<'_>) -> crate::Result<()> { + self.append(Some(v)); + Ok(()) + } + + fn output_datum(v: JsonbRef<'_>) -> Datum { + Some(v.to_owned_scalar().to_scalar_value()) + } +} + +impl AccessOutput for Utf8ArrayBuilder { + fn return_type() -> DataType { + DataType::Varchar + } + + fn output(&mut self, v: JsonbRef<'_>) -> crate::Result<()> { + match v.is_jsonb_null() { + true => self.append_null(), + false => { + let mut writer = self.writer().begin(); + v.force_str(&mut writer) + .map_err(|e| crate::ExprError::Internal(e.into()))?; + writer.finish(); + } + }; + Ok(()) + } + + fn output_datum(v: JsonbRef<'_>) -> Datum { + match v.is_jsonb_null() { + true => None, + false => { + let mut s = String::new(); + v.force_str(&mut s).unwrap(); + let s: Box = s.into(); + Some(s.to_scalar_value()) + } + } + } +} diff --git a/src/expr/src/expr/mod.rs b/src/expr/src/expr/mod.rs index 08fa3b06adb8f..8ff9e80a01033 100644 --- a/src/expr/src/expr/mod.rs +++ b/src/expr/src/expr/mod.rs @@ -44,6 +44,7 @@ mod expr_field; mod expr_in; mod expr_input_ref; mod expr_is_null; +mod expr_jsonb_access; mod expr_literal; mod expr_nested_construct; mod expr_quaternary_bytes; diff --git a/src/expr/src/sig/func.rs b/src/expr/src/sig/func.rs index 5af8eced4b0a7..89353dac39ee9 100644 --- a/src/expr/src/sig/func.rs +++ b/src/expr/src/sig/func.rs @@ -309,6 +309,11 @@ fn build_type_derive_map() -> FuncSigMap { T::Varchar, ); + map.insert(E::JsonbAccessInner, vec![T::Jsonb, T::Int32], T::Jsonb); + map.insert(E::JsonbAccessInner, vec![T::Jsonb, T::Varchar], T::Jsonb); + map.insert(E::JsonbAccessStr, vec![T::Jsonb, T::Int32], T::Varchar); + map.insert(E::JsonbAccessStr, vec![T::Jsonb, T::Varchar], T::Varchar); + map } diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs index 2a77882f30658..ed6f105f7ad28 100644 --- a/src/frontend/src/binder/expr/function.rs +++ b/src/frontend/src/binder/expr/function.rs @@ -375,6 +375,11 @@ impl Binder { ("array_join", raw_call(ExprType::ArrayToString)), ("array_prepend", raw_call(ExprType::ArrayPrepend)), ("array_to_string", raw_call(ExprType::ArrayToString)), + // jsonb + ("jsonb_object_field", raw_call(ExprType::JsonbAccessInner)), + ("jsonb_array_element", raw_call(ExprType::JsonbAccessInner)), + ("jsonb_object_field_text", raw_call(ExprType::JsonbAccessStr)), + ("jsonb_array_element_text", raw_call(ExprType::JsonbAccessStr)), // System information operations. ( "pg_typeof", From 8b0b489b108c96ce3a8883d871106e9982584ea7 Mon Sep 17 00:00:00 2001 From: Xiangjin Date: Fri, 17 Feb 2023 23:50:08 +0800 Subject: [PATCH 2/3] jsonb_typeof and jsonb_array_length --- dashboard/proto/gen/expr.ts | 12 ++++++++ e2e_test/batch/types/jsonb.slt.part | 31 +++++++++++++------ proto/expr.proto | 2 ++ src/expr/src/expr/build_expr_from_prost.rs | 4 ++- src/expr/src/expr/expr_unary.rs | 13 ++++++++ src/expr/src/sig/func.rs | 2 ++ src/expr/src/vector_op/jsonb_info.rs | 36 ++++++++++++++++++++++ src/expr/src/vector_op/mod.rs | 1 + src/frontend/src/binder/expr/function.rs | 2 ++ 9 files changed, 93 insertions(+), 10 deletions(-) create mode 100644 src/expr/src/vector_op/jsonb_info.rs diff --git a/dashboard/proto/gen/expr.ts b/dashboard/proto/gen/expr.ts index bf277974f9609..66bc96aea0d75 100644 --- a/dashboard/proto/gen/expr.ts +++ b/dashboard/proto/gen/expr.ts @@ -134,6 +134,8 @@ export const ExprNode_Type = { JSONB_ACCESS_INNER: "JSONB_ACCESS_INNER", /** JSONB_ACCESS_STR - jsonb ->> int, jsonb ->> text, jsonb #>> text[] that returns text */ JSONB_ACCESS_STR: "JSONB_ACCESS_STR", + JSONB_TYPEOF: "JSONB_TYPEOF", + JSONB_ARRAY_LENGTH: "JSONB_ARRAY_LENGTH", /** * VNODE - Non-pure functions below (> 1000) * ------------------------ @@ -412,6 +414,12 @@ export function exprNode_TypeFromJSON(object: any): ExprNode_Type { case 601: case "JSONB_ACCESS_STR": return ExprNode_Type.JSONB_ACCESS_STR; + case 602: + case "JSONB_TYPEOF": + return ExprNode_Type.JSONB_TYPEOF; + case 603: + case "JSONB_ARRAY_LENGTH": + return ExprNode_Type.JSONB_ARRAY_LENGTH; case 1101: case "VNODE": return ExprNode_Type.VNODE; @@ -604,6 +612,10 @@ export function exprNode_TypeToJSON(object: ExprNode_Type): string { return "JSONB_ACCESS_INNER"; case ExprNode_Type.JSONB_ACCESS_STR: return "JSONB_ACCESS_STR"; + case ExprNode_Type.JSONB_TYPEOF: + return "JSONB_TYPEOF"; + case ExprNode_Type.JSONB_ARRAY_LENGTH: + return "JSONB_ARRAY_LENGTH"; case ExprNode_Type.VNODE: return "VNODE"; case ExprNode_Type.NOW: diff --git a/e2e_test/batch/types/jsonb.slt.part b/e2e_test/batch/types/jsonb.slt.part index e7c0db15bc7e1..54d29abaae6b7 100644 --- a/e2e_test/batch/types/jsonb.slt.part +++ b/e2e_test/batch/types/jsonb.slt.part @@ -88,23 +88,36 @@ select jsonb_array_element(jsonb_object_field('{"k2":[2,true,4]}', 'k2'), -2)::b ---- t -query TT +# Note the difference between access text directly vs access jsonb then cast to text. +query TTT with t(v1) as (values (null::jsonb), ('null'), ('true'), ('1'), ('"a"'), ('[]'), ('{}')), j(v1) as (select ('{"k":' || v1::varchar || '}')::jsonb from t) -select jsonb_object_field_text(v1, 'k'), jsonb_object_field(v1, 'k')::varchar from j order by 2; +select + jsonb_object_field_text(v1, 'k'), + jsonb_object_field(v1, 'k')::varchar, + jsonb_typeof(jsonb_object_field(v1, 'k')) +from j order by 2; ---- -a "a" -1 1 -[] [] -NULL null -true true -{} {} -NULL NULL +a "a" string +1 1 number +[] [] array +NULL null null +true true boolean +{} {} object +NULL NULL NULL query T select jsonb_array_element_text('true'::jsonb, 2); ---- NULL +query I +select jsonb_array_length('[7, 2]'); +---- +2 + +statement error cannot get array length +select jsonb_array_length('null'); + statement ok drop table t; diff --git a/proto/expr.proto b/proto/expr.proto index 701d14f1b66e0..89b2049dffdb4 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -126,6 +126,8 @@ message ExprNode { JSONB_ACCESS_INNER = 600; // jsonb ->> int, jsonb ->> text, jsonb #>> text[] that returns text JSONB_ACCESS_STR = 601; + JSONB_TYPEOF = 602; + JSONB_ARRAY_LENGTH = 603; // Non-pure functions below (> 1000) // ------------------------ diff --git a/src/expr/src/expr/build_expr_from_prost.rs b/src/expr/src/expr/build_expr_from_prost.rs index 50578e84033dd..34064d2cdafe4 100644 --- a/src/expr/src/expr/build_expr_from_prost.rs +++ b/src/expr/src/expr/build_expr_from_prost.rs @@ -64,7 +64,9 @@ pub fn build_from_prost(prost: &ExprNode) -> Result { // Fixed number of arguments and based on `Unary/Binary/Ternary/...Expression` Cast | Upper | Lower | Md5 | Not | IsTrue | IsNotTrue | IsFalse | IsNotFalse | IsNull | IsNotNull | Neg | Ascii | Abs | Ceil | Floor | Round | Exp | BitwiseNot | CharLength - | BoolOut | OctetLength | BitLength | ToTimestamp => build_unary_expr_prost(prost), + | BoolOut | OctetLength | BitLength | ToTimestamp | JsonbTypeof | JsonbArrayLength => { + build_unary_expr_prost(prost) + } Equal | NotEqual | LessThan | LessThanOrEqual | GreaterThan | GreaterThanOrEqual | Add | Subtract | Multiply | Divide | Modulus | Extract | RoundDigit | Pow | TumbleStart | Position | BitwiseShiftLeft | BitwiseShiftRight | BitwiseAnd | BitwiseOr | BitwiseXor diff --git a/src/expr/src/expr/expr_unary.rs b/src/expr/src/expr/expr_unary.rs index 00f9151078a0c..b957fb74d9628 100644 --- a/src/expr/src/expr/expr_unary.rs +++ b/src/expr/src/expr/expr_unary.rs @@ -30,6 +30,7 @@ use crate::vector_op::cast::*; use crate::vector_op::cmp::{is_false, is_not_false, is_not_true, is_true}; use crate::vector_op::conjunction; use crate::vector_op::exp::exp_f64; +use crate::vector_op::jsonb_info::{jsonb_array_length, jsonb_typeof}; use crate::vector_op::length::{bit_length, length_default, octet_length}; use crate::vector_op::lower::lower; use crate::vector_op::ltrim::ltrim; @@ -307,6 +308,18 @@ pub fn new_unary_expr( f64_sec_to_timestamptz, )) } + (ProstType::JsonbTypeof, DataType::Varchar, DataType::Jsonb) => { + UnaryBytesExpression::::new(child_expr, return_type, jsonb_typeof) + .boxed() + } + (ProstType::JsonbArrayLength, DataType::Int32, DataType::Jsonb) => { + UnaryExpression::::new( + child_expr, + return_type, + jsonb_array_length, + ) + .boxed() + } (expr, ret, child) => { return Err(ExprError::UnsupportedFunction(format!( "{:?}({:?}) -> {:?}", diff --git a/src/expr/src/sig/func.rs b/src/expr/src/sig/func.rs index 89353dac39ee9..1b160861ca666 100644 --- a/src/expr/src/sig/func.rs +++ b/src/expr/src/sig/func.rs @@ -313,6 +313,8 @@ fn build_type_derive_map() -> FuncSigMap { map.insert(E::JsonbAccessInner, vec![T::Jsonb, T::Varchar], T::Jsonb); map.insert(E::JsonbAccessStr, vec![T::Jsonb, T::Int32], T::Varchar); map.insert(E::JsonbAccessStr, vec![T::Jsonb, T::Varchar], T::Varchar); + map.insert(E::JsonbTypeof, vec![T::Jsonb], T::Varchar); + map.insert(E::JsonbArrayLength, vec![T::Jsonb], T::Int32); map } diff --git a/src/expr/src/vector_op/jsonb_info.rs b/src/expr/src/vector_op/jsonb_info.rs new file mode 100644 index 0000000000000..064e54959b3bd --- /dev/null +++ b/src/expr/src/vector_op/jsonb_info.rs @@ -0,0 +1,36 @@ +// Copyright 2023 RisingWave Labs +// +// 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::fmt::Write; + +use risingwave_common::array::JsonbRef; + +use crate::{ExprError, Result}; + +#[inline(always)] +pub fn jsonb_typeof(v: JsonbRef<'_>, writer: &mut dyn Write) -> Result<()> { + writer + .write_str(v.type_name()) + .map_err(|e| ExprError::Internal(e.into())) +} + +#[inline(always)] +pub fn jsonb_array_length(v: JsonbRef<'_>) -> Result { + v.array_len() + .map(|n| n as i32) + .map_err(|e| ExprError::InvalidParam { + name: "", + reason: e, + }) +} diff --git a/src/expr/src/vector_op/mod.rs b/src/expr/src/vector_op/mod.rs index d3f07ad7fed9d..164e433ad9dde 100644 --- a/src/expr/src/vector_op/mod.rs +++ b/src/expr/src/vector_op/mod.rs @@ -25,6 +25,7 @@ pub mod date_trunc; pub mod exp; pub mod extract; pub mod format_type; +pub mod jsonb_info; pub mod length; pub mod like; pub mod lower; diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs index ed6f105f7ad28..8969e2ac4a201 100644 --- a/src/frontend/src/binder/expr/function.rs +++ b/src/frontend/src/binder/expr/function.rs @@ -380,6 +380,8 @@ impl Binder { ("jsonb_array_element", raw_call(ExprType::JsonbAccessInner)), ("jsonb_object_field_text", raw_call(ExprType::JsonbAccessStr)), ("jsonb_array_element_text", raw_call(ExprType::JsonbAccessStr)), + ("jsonb_typeof", raw_call(ExprType::JsonbTypeof)), + ("jsonb_array_length", raw_call(ExprType::JsonbArrayLength)), // System information operations. ( "pg_typeof", From 6bef8af793bdb69a38b4fb4d21675a29b8e11431 Mon Sep 17 00:00:00 2001 From: Xiangjin Date: Wed, 22 Feb 2023 17:10:54 +0800 Subject: [PATCH 3/3] rename output_datum -> to_datum --- src/expr/src/expr/expr_jsonb_access.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/expr/src/expr/expr_jsonb_access.rs b/src/expr/src/expr/expr_jsonb_access.rs index e4a725d75f3a5..7c26779a39109 100644 --- a/src/expr/src/expr/expr_jsonb_access.rs +++ b/src/expr/src/expr/expr_jsonb_access.rs @@ -138,7 +138,7 @@ where .map(|v| v.as_scalar_ref_impl().try_into().unwrap()); let r = self.eval_strict(v, p); - Ok(r.and_then(O::output_datum)) + Ok(r.and_then(O::to_datum)) } } @@ -165,7 +165,7 @@ pub fn jsonb_array_element(v: JsonbRef<'_>, p: i32) -> Option> { trait AccessOutput: ArrayBuilder { fn return_type() -> DataType; fn output(&mut self, v: JsonbRef<'_>) -> crate::Result<()>; - fn output_datum(v: JsonbRef<'_>) -> Datum; + fn to_datum(v: JsonbRef<'_>) -> Datum; fn output_nullable(&mut self, v: Option>) -> crate::Result<()> { match v { Some(v) => self.output(v)?, @@ -185,7 +185,7 @@ impl AccessOutput for JsonbArrayBuilder { Ok(()) } - fn output_datum(v: JsonbRef<'_>) -> Datum { + fn to_datum(v: JsonbRef<'_>) -> Datum { Some(v.to_owned_scalar().to_scalar_value()) } } @@ -208,7 +208,7 @@ impl AccessOutput for Utf8ArrayBuilder { Ok(()) } - fn output_datum(v: JsonbRef<'_>) -> Datum { + fn to_datum(v: JsonbRef<'_>) -> Datum { match v.is_jsonb_null() { true => None, false => {