From 15cf22d497fe8b99c216d8e8901832f0e5ce3a75 Mon Sep 17 00:00:00 2001 From: Runji Wang Date: Mon, 30 Oct 2023 15:12:24 +0800 Subject: [PATCH] support `jsonb_extract_path` and `jsonb_extract_path_text` Signed-off-by: Runji Wang --- proto/expr.proto | 2 ++ src/common/src/array/list_array.rs | 18 +++++++++++++++ src/expr/impl/src/scalar/jsonb_access.rs | 9 +++++--- src/frontend/src/binder/expr/function.rs | 2 ++ src/frontend/src/expr/pure.rs | 2 ++ src/frontend/src/expr/type_inference/func.rs | 16 +++++++++++++ src/tests/regress/data/sql/jsonb.sql | 24 ++++++++++---------- 7 files changed, 58 insertions(+), 15 deletions(-) diff --git a/proto/expr.proto b/proto/expr.proto index fecefc12b4ee7..06bfecc66d4d0 100644 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -236,6 +236,8 @@ message ExprNode { JSONB_EXISTS_ANY = 611; // jsonb ?& text[] JSONB_EXISTS_ALL = 612; + JSONB_EXTRACT_PATH = 616; + JSONB_EXTRACT_PATH_TEXT = 617; // Non-pure functions below (> 1000) // ------------------------ diff --git a/src/common/src/array/list_array.rs b/src/common/src/array/list_array.rs index 7eaaffff98534..d6a8aba766da4 100644 --- a/src/common/src/array/list_array.rs +++ b/src/common/src/array/list_array.rs @@ -538,6 +538,24 @@ impl<'a> ListRef<'a> { } } +impl Row for ListRef<'_> { + fn len(&self) -> usize { + self.len() + } + + fn datum_at(&self, index: usize) -> DatumRef<'_> { + self.get(index).unwrap() + } + + unsafe fn datum_at_unchecked(&self, index: usize) -> DatumRef<'_> { + self.get(index).unwrap() + } + + fn iter(&self) -> impl Iterator> { + self.clone().iter() + } +} + impl PartialEq for ListRef<'_> { fn eq(&self, other: &Self) -> bool { iter_elems_ref!(*self, lhs, { diff --git a/src/expr/impl/src/scalar/jsonb_access.rs b/src/expr/impl/src/scalar/jsonb_access.rs index 08e36bedf83cc..12d517ec31f4d 100644 --- a/src/expr/impl/src/scalar/jsonb_access.rs +++ b/src/expr/impl/src/scalar/jsonb_access.rs @@ -14,7 +14,8 @@ use std::fmt::Write; -use risingwave_common::types::{JsonbRef, ListRef}; +use risingwave_common::row::Row; +use risingwave_common::types::JsonbRef; use risingwave_expr::function; /// Extracts JSON object field with the given key. @@ -87,7 +88,8 @@ pub fn jsonb_array_element(v: JsonbRef<'_>, p: i32) -> Option> { /// NULL /// ``` #[function("jsonb_access_multi(jsonb, varchar[]) -> jsonb")] -pub fn jsonb_access_multi<'a>(v: JsonbRef<'a>, path: ListRef<'_>) -> Option> { +#[function("jsonb_extract_path(jsonb, ...) -> jsonb")] +pub fn jsonb_access_multi<'a>(v: JsonbRef<'a>, path: impl Row) -> Option> { let mut jsonb = v; for key in path.iter() { // return null if any element is null @@ -182,9 +184,10 @@ pub fn jsonb_array_element_str(v: JsonbRef<'_>, p: i32, writer: &mut impl Write) /// NULL /// ``` #[function("jsonb_access_multi_str(jsonb, varchar[]) -> varchar")] +#[function("jsonb_extract_path_text(jsonb, ...) -> varchar")] pub fn jsonb_access_multi_str( v: JsonbRef<'_>, - path: ListRef<'_>, + path: impl Row, writer: &mut impl Write, ) -> Option<()> { let jsonb = jsonb_access_multi(v, path)?; diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs index 6dfab2c3bc283..39da7540795a5 100644 --- a/src/frontend/src/binder/expr/function.rs +++ b/src/frontend/src/binder/expr/function.rs @@ -875,6 +875,8 @@ impl Binder { ("jsonb_array_element", raw_call(ExprType::JsonbAccess)), ("jsonb_object_field_text", raw_call(ExprType::JsonbAccessStr)), ("jsonb_array_element_text", raw_call(ExprType::JsonbAccessStr)), + ("jsonb_extract_path", raw_call(ExprType::JsonbExtractPath)), + ("jsonb_extract_path_text", raw_call(ExprType::JsonbExtractPathText)), ("jsonb_typeof", raw_call(ExprType::JsonbTypeof)), ("jsonb_array_length", raw_call(ExprType::JsonbArrayLength)), ("jsonb_object", raw_call(ExprType::JsonbObject)), diff --git a/src/frontend/src/expr/pure.rs b/src/frontend/src/expr/pure.rs index 7a63c7f95ae99..ffba850eb7874 100644 --- a/src/frontend/src/expr/pure.rs +++ b/src/frontend/src/expr/pure.rs @@ -177,6 +177,8 @@ impl ExprVisitor for ImpureAnalyzer { | expr_node::Type::JsonbAccessStr | expr_node::Type::JsonbAccessMulti | expr_node::Type::JsonbAccessMultiStr + | expr_node::Type::JsonbExtractPath + | expr_node::Type::JsonbExtractPathText | expr_node::Type::JsonbTypeof | expr_node::Type::JsonbArrayLength | expr_node::Type::JsonbObject diff --git a/src/frontend/src/expr/type_inference/func.rs b/src/frontend/src/expr/type_inference/func.rs index 84e315dacae45..b3fe1e8dd4766 100644 --- a/src/frontend/src/expr/type_inference/func.rs +++ b/src/frontend/src/expr/type_inference/func.rs @@ -541,6 +541,22 @@ fn infer_type_for_special( ensure_arity!("greatest/least", 1 <= | inputs |); Ok(Some(align_types(inputs.iter_mut())?)) } + ExprType::JsonbExtractPath => { + ensure_arity!("jsonb_extract_path", 2 <= | inputs |); + inputs[0].cast_implicit_mut(DataType::Jsonb)?; + for input in inputs.iter_mut().skip(1) { + input.cast_explicit_mut(DataType::Varchar)?; + } + Ok(Some(DataType::Jsonb)) + } + ExprType::JsonbExtractPathText => { + ensure_arity!("jsonb_extract_path_text", 2 <= | inputs |); + inputs[0].cast_implicit_mut(DataType::Jsonb)?; + for input in inputs.iter_mut().skip(1) { + input.cast_explicit_mut(DataType::Varchar)?; + } + Ok(Some(DataType::Varchar)) + } _ => Ok(None), } } diff --git a/src/tests/regress/data/sql/jsonb.sql b/src/tests/regress/data/sql/jsonb.sql index 69bd9a928a0f5..bf44f328c8ebb 100644 --- a/src/tests/regress/data/sql/jsonb.sql +++ b/src/tests/regress/data/sql/jsonb.sql @@ -449,20 +449,20 @@ SELECT jsonb_typeof('"1.0"') AS string; -- extract_path, extract_path_as_text ---@ SELECT jsonb_extract_path('{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}','f4','f6'); ---@ SELECT jsonb_extract_path('{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}','f2'); ---@ SELECT jsonb_extract_path('{"f2":["f3",1],"f4":{"f5":99,"f6":"stringy"}}','f2',0::text); ---@ SELECT jsonb_extract_path('{"f2":["f3",1],"f4":{"f5":99,"f6":"stringy"}}','f2',1::text); ---@ SELECT jsonb_extract_path_text('{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}','f4','f6'); ---@ SELECT jsonb_extract_path_text('{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}','f2'); ---@ SELECT jsonb_extract_path_text('{"f2":["f3",1],"f4":{"f5":99,"f6":"stringy"}}','f2',0::text); ---@ SELECT jsonb_extract_path_text('{"f2":["f3",1],"f4":{"f5":99,"f6":"stringy"}}','f2',1::text); +SELECT jsonb_extract_path('{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}','f4','f6'); +SELECT jsonb_extract_path('{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}','f2'); +SELECT jsonb_extract_path('{"f2":["f3",1],"f4":{"f5":99,"f6":"stringy"}}','f2',0::text); +SELECT jsonb_extract_path('{"f2":["f3",1],"f4":{"f5":99,"f6":"stringy"}}','f2',1::text); +SELECT jsonb_extract_path_text('{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}','f4','f6'); +SELECT jsonb_extract_path_text('{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}','f2'); +SELECT jsonb_extract_path_text('{"f2":["f3",1],"f4":{"f5":99,"f6":"stringy"}}','f2',0::text); +SELECT jsonb_extract_path_text('{"f2":["f3",1],"f4":{"f5":99,"f6":"stringy"}}','f2',1::text); -- extract_path nulls ---@ SELECT jsonb_extract_path('{"f2":{"f3":1},"f4":{"f5":null,"f6":"stringy"}}','f4','f5') IS NULL AS expect_false; ---@ SELECT jsonb_extract_path_text('{"f2":{"f3":1},"f4":{"f5":null,"f6":"stringy"}}','f4','f5') IS NULL AS expect_true; ---@ SELECT jsonb_extract_path('{"f2":{"f3":1},"f4":[0,1,2,null]}','f4','3') IS NULL AS expect_false; ---@ SELECT jsonb_extract_path_text('{"f2":{"f3":1},"f4":[0,1,2,null]}','f4','3') IS NULL AS expect_true; +SELECT jsonb_extract_path('{"f2":{"f3":1},"f4":{"f5":null,"f6":"stringy"}}','f4','f5') IS NULL AS expect_false; +SELECT jsonb_extract_path_text('{"f2":{"f3":1},"f4":{"f5":null,"f6":"stringy"}}','f4','f5') IS NULL AS expect_true; +SELECT jsonb_extract_path('{"f2":{"f3":1},"f4":[0,1,2,null]}','f4','3') IS NULL AS expect_false; +SELECT jsonb_extract_path_text('{"f2":{"f3":1},"f4":[0,1,2,null]}','f4','3') IS NULL AS expect_true; -- extract_path operators SELECT '{"f2":{"f3":1},"f4":{"f5":99,"f6":"stringy"}}'::jsonb#>array['f4','f6'];