feat(expr): access jsonb object field and array element (#8023)

Adds the following expressions:
* `jsonb_object_field(jsonb, varchar) -> jsonb`
* `jsonb_array_element(jsonb, int) -> jsonb`
* `jsonb_object_field_text(jsonb, varchar) -> varchar`
* `jsonb_array_element_text(jsonb, int) -> varchar`
* `jsonb_typeof(jsonb) -> varchar`
* `jsonb_array_length(jsonb) -> int`

The first two are actually operator `->` in PostgreSQL, and the two in the middle are operator `->>` in PostgreSQL. But our parser does not support parsing this syntax yet.

The optimization of constant rhs will be added in a followup.

Approved-By: BugenZhao

Co-Authored-By: Xiangjin <xiangjin@singularity-data.com>
Co-Authored-By: Xiangjin <xiangjin@risingwave-labs.com>

e2e_test/batch/types/jsonb.slt.part

@@ -83,5 +83,41 @@ select null::jsonb::bool;
 ----
 NULL
 
+query T
+select jsonb_array_element(jsonb_object_field('{"k2":[2,true,4]}', 'k2'), -2)::bool;
+----
+t
+
+# 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,
+    jsonb_typeof(jsonb_object_field(v1, 'k'))
+from j order by 2;
+----
+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;

proto/expr.proto

@@ -120,7 +120,16 @@ 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;
+    JSONB_TYPEOF = 602;
+    JSONB_ARRAY_LENGTH = 603;
+
+    // Non-pure functions below (> 1000)
     // ------------------------
     // Internal functions
     VNODE = 1101;

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<usize, String> {
+        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<bool, String> {
         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<W: std::fmt::Write>(&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> {
+        self.0.get(field).map(Self)
+    }
+
+    pub fn access_array_element(&self, idx: usize) -> Option<Self> {
+        self.0.get(idx).map(Self)
+    }
 }
 
 #[derive(Debug)]

src/expr/src/expr/build_expr_from_prost.rs

@@ -64,11 +64,15 @@ pub fn build_from_prost(prost: &ExprNode) -> Result<BoxedExpression> {
         // 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
-        | 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)
         }

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::<Utf8Array, JsonbArrayBuilder, _>::new_expr(
+                    l,
+                    r,
+                    jsonb_object_field,
+                )
+                .boxed()
+            }
+            DataType::Int32 => JsonbAccessExpression::<I32Array, JsonbArrayBuilder, _>::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::<Utf8Array, Utf8ArrayBuilder, _>::new_expr(
+                l,
+                r,
+                jsonb_object_field,
+            )
+            .boxed(),
+            DataType::Int32 => JsonbAccessExpression::<I32Array, Utf8ArrayBuilder, _>::new_expr(
+                l,
+                r,
+                jsonb_array_element,
+            )
+            .boxed(),
+            t => return Err(ExprError::UnsupportedFunction(format!("jsonb ->> {t}"))),
+        },
 
         tp => {
             return Err(ExprError::UnsupportedFunction(format!(