feat(expr): support jsonb_extract_path(_text) function (#13143)

Signed-off-by: Runji Wang <[email protected]>
risingwavelabs · Oct 31, 2023 · 8a90ce5 · 8a90ce5
1 parent c97e08c
commit 8a90ce5
Show file tree

Hide file tree

Showing 6 changed files with 67 additions and 25 deletions.
diff --git a/proto/expr.proto b/proto/expr.proto
@@ -218,9 +218,9 @@ message ExprNode {
     // jsonb ->> int, jsonb ->> text that returns text
     JSONB_ACCESS_STR = 601;
     // jsonb #> text[] -> jsonb
-    JSONB_ACCESS_MULTI = 613;
+    JSONB_EXTRACT_PATH = 613;
     // jsonb #>> text[] -> text
-    JSONB_ACCESS_MULTI_STR = 614;
+    JSONB_EXTRACT_PATH_TEXT = 614;
     JSONB_TYPEOF = 602;
     JSONB_ARRAY_LENGTH = 603;
     IS_JSON = 604;

diff --git a/src/expr/impl/src/scalar/jsonb_access.rs b/src/expr/impl/src/scalar/jsonb_access.rs
@@ -71,7 +71,8 @@ pub fn jsonb_array_element(v: JsonbRef<'_>, p: i32) -> Option<JsonbRef<'_>> {
 
 /// Extracts JSON sub-object at the specified path, where path elements can be either field keys or array indexes.
 ///
-/// `jsonb #> text[] → jsonb`
+/// - `jsonb #> text[] → jsonb`
+/// - `jsonb_extract_path ( from_json jsonb, VARIADIC path_elems text[] ) → jsonb`
 ///
 /// # Examples
 ///
@@ -85,9 +86,14 @@ pub fn jsonb_array_element(v: JsonbRef<'_>, p: i32) -> Option<JsonbRef<'_>> {
 /// select '{"a": {"b": ["foo","bar"]}}'::jsonb #> '{a,b,null}'::text[];
 /// ----
 /// NULL
+///
+/// query T
+/// select jsonb_extract_path('{"a": {"b": ["foo","bar"]}}', 'a', 'b', '1');
+/// ----
+/// "bar"
 /// ```
-#[function("jsonb_access_multi(jsonb, varchar[]) -> jsonb")]
-pub fn jsonb_access_multi<'a>(v: JsonbRef<'a>, path: ListRef<'_>) -> Option<JsonbRef<'a>> {
+#[function("jsonb_extract_path(jsonb, varchar[]) -> jsonb")]
+pub fn jsonb_extract_path<'a>(v: JsonbRef<'a>, path: ListRef<'_>) -> Option<JsonbRef<'a>> {
     let mut jsonb = v;
     for key in path.iter() {
         // return null if any element is null
@@ -161,7 +167,8 @@ pub fn jsonb_array_element_str(v: JsonbRef<'_>, p: i32, writer: &mut impl Write)
 
 /// Extracts JSON sub-object at the specified path as text.
 ///
-/// `jsonb #>> text[] → text`
+/// - `jsonb #>> text[] → text`
+/// - `jsonb_extract_path_text ( from_json jsonb, VARIADIC path_elems text[] ) → text`
 ///
 /// # Examples
 ///
@@ -180,14 +187,19 @@ pub fn jsonb_array_element_str(v: JsonbRef<'_>, p: i32, writer: &mut impl Write)
 /// select '{"a": {"b": ["foo","bar"]}}'::jsonb #>> '{a,b,null}'::text[];
 /// ----
 /// NULL
+///
+/// query T
+/// select jsonb_extract_path_text('{"a": {"b": ["foo","bar"]}}', 'a', 'b', '1');
+/// ----
+/// bar
 /// ```
-#[function("jsonb_access_multi_str(jsonb, varchar[]) -> varchar")]
-pub fn jsonb_access_multi_str(
+#[function("jsonb_extract_path_text(jsonb, varchar[]) -> varchar")]
+pub fn jsonb_extract_path_text(
     v: JsonbRef<'_>,
     path: ListRef<'_>,
     writer: &mut impl Write,
 ) -> Option<()> {
-    let jsonb = jsonb_access_multi(v, path)?;
+    let jsonb = jsonb_extract_path(v, path)?;
     if jsonb.is_jsonb_null() {
         return None;
     }

diff --git a/src/frontend/src/binder/expr/binary_op.rs b/src/frontend/src/binder/expr/binary_op.rs
@@ -92,8 +92,8 @@ impl Binder {
             BinaryOperator::Arrow => ExprType::JsonbAccess,
             BinaryOperator::LongArrow => ExprType::JsonbAccessStr,
             BinaryOperator::HashMinus => ExprType::JsonbDeletePath,
-            BinaryOperator::HashArrow => ExprType::JsonbAccessMulti,
-            BinaryOperator::HashLongArrow => ExprType::JsonbAccessMultiStr,
+            BinaryOperator::HashArrow => ExprType::JsonbExtractPath,
+            BinaryOperator::HashLongArrow => ExprType::JsonbExtractPathText,
             BinaryOperator::Prefix => ExprType::StartsWith,
             BinaryOperator::Contains => ExprType::JsonbContains,
             BinaryOperator::Contained => ExprType::JsonbContained,

diff --git a/src/frontend/src/binder/expr/function.rs b/src/frontend/src/binder/expr/function.rs
@@ -882,6 +882,36 @@ 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(|_binder, mut inputs| {
+                    // rewrite: jsonb_extract_path(jsonb, s1, s2...)
+                    // to:      jsonb_extract_path(jsonb, array[s1, s2...])
+                    if inputs.len() < 2 {
+                        return Err(ErrorCode::ExprError("unexpected arguments number".into()).into());
+                    }
+                    inputs[0].cast_implicit_mut(DataType::Jsonb)?;
+                    let mut variadic_inputs = inputs.split_off(1);
+                    for input in &mut variadic_inputs {
+                        input.cast_implicit_mut(DataType::Varchar)?;
+                    }
+                    let array = FunctionCall::new_unchecked(ExprType::Array, variadic_inputs, DataType::List(Box::new(DataType::Varchar)));
+                    inputs.push(array.into());
+                    Ok(FunctionCall::new_unchecked(ExprType::JsonbExtractPath, inputs, DataType::Jsonb).into())
+                })),
+                ("jsonb_extract_path_text", raw(|_binder, mut inputs| {
+                    // rewrite: jsonb_extract_path_text(jsonb, s1, s2...)
+                    // to:      jsonb_extract_path_text(jsonb, array[s1, s2...])
+                    if inputs.len() < 2 {
+                        return Err(ErrorCode::ExprError("unexpected arguments number".into()).into());
+                    }
+                    inputs[0].cast_implicit_mut(DataType::Jsonb)?;
+                    let mut variadic_inputs = inputs.split_off(1);
+                    for input in &mut variadic_inputs {
+                        input.cast_implicit_mut(DataType::Varchar)?;
+                    }
+                    let array = FunctionCall::new_unchecked(ExprType::Array, variadic_inputs, DataType::List(Box::new(DataType::Varchar)));
+                    inputs.push(array.into());
+                    Ok(FunctionCall::new_unchecked(ExprType::JsonbExtractPathText, inputs, DataType::Varchar).into())
+                })),
                 ("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
@@ -176,8 +176,8 @@ impl ExprVisitor for ImpureAnalyzer {
             | expr_node::Type::JsonbCat
             | expr_node::Type::JsonbAccess
             | 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/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'];