support jsonb_extract_path and jsonb_extract_path_text

Signed-off-by: Runji Wang <[email protected]>
risingwavelabs · Oct 30, 2023 · 15cf22d · 15cf22d
1 parent e392db0
commit 15cf22d
Show file tree

Hide file tree

Showing 7 changed files with 58 additions and 15 deletions.
diff --git 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
@@ -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<Item = DatumRef<'_>> {
+        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
@@ -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<JsonbRef<'_>> {
 /// NULL
 /// ```
 #[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, ...) -> jsonb")]
+pub fn jsonb_access_multi<'a>(v: JsonbRef<'a>, path: impl Row) -> Option<JsonbRef<'a>> {
     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
@@ -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
@@ -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
@@ -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
@@ -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'];