feat(expr): allow jsonb_strip_nulls to remove null array elements

e2e_test/batch/types/jsonb_delete.slt.part

@@ -0,0 +1,83 @@
+query T
+select jsonb_strip_nulls(null);
+----
+NULL
+
+query T
+select jsonb_strip_nulls('1');
+----
+1
+
+query T
+select jsonb_strip_nulls('"a string"');
+----
+"a string"
+
+query T
+select jsonb_strip_nulls('null');
+----
+null
+
+query T
+select jsonb_strip_nulls('[1,2,null,3,4]');
+----
+[1, 2, null, 3, 4]
+
+query T
+select jsonb_strip_nulls('{"a":1,"b":null,"c":[2,null,3],"d":{"e":4,"f":null}}');
+----
+{"a": 1, "c": [2, null, 3], "d": {"e": 4}}
+
+query T
+select jsonb_strip_nulls('[1,{"a":1,"b":null,"c":2},3]');
+----
+[1, {"a": 1, "c": 2}, 3]
+
+# an empty object is not null and should not be stripped
+query T
+select jsonb_strip_nulls('{"a": {"b": null, "c": null}, "d": {} }');
+----
+{"a": {}, "d": {}}
+
+# jsonb_strip_nulls with strip_in_arrays=true
+
+query T
+select jsonb_strip_nulls(null, true);
+----
+NULL
+
+query T
+select jsonb_strip_nulls('1', true);
+----
+1
+
+query T
+select jsonb_strip_nulls('"a string"', true);
+----
+"a string"
+
+query T
+select jsonb_strip_nulls('null', true);
+----
+null
+
+query T
+select jsonb_strip_nulls('[1,2,null,3,4]', true);
+----
+[1, 2, 3, 4]
+
+query T
+select jsonb_strip_nulls('{"a":1,"b":null,"c":[2,null,3],"d":{"e":4,"f":null}}', true);
+----
+{"a": 1, "c": [2, 3], "d": {"e": 4}}
+
+query T
+select jsonb_strip_nulls('[1,{"a":1,"b":null,"c":2},3]', true);
+----
+[1, {"a": 1, "c": 2}, 3]
+
+# an empty object is not null and should not be stripped
+query T
+select jsonb_strip_nulls('{"a": {"b": null, "c": null}, "d": {} }', true);
+----
+{"a": {}, "d": {}}

src/expr/impl/src/scalar/jsonb_delete.rs

@@ -383,12 +383,20 @@ fn normalize_array_index(len: usize, index: i32) -> Option<usize> {
 fn jsonb_strip_nulls(v: JsonbRef<'_>) -> JsonbVal {
     let jsonb: ValueRef<'_> = v.into();
     let mut builder = jsonbb::Builder::<Vec<u8>>::with_capacity(jsonb.capacity());
-    jsonbb_strip_nulls(jsonb, &mut builder);
+    jsonbb_strip_nulls(jsonb, &mut builder, false);
+    JsonbVal::from(builder.finish())
+}
+
+#[function("jsonb_strip_nulls(jsonb, boolean) -> jsonb")]
+fn jsonb_strip_nulls_with_arrays(v: JsonbRef<'_>, strip_in_arrays: bool) -> JsonbVal {
+    let jsonb: ValueRef<'_> = v.into();
+    let mut builder = jsonbb::Builder::<Vec<u8>>::with_capacity(jsonb.capacity());
+    jsonbb_strip_nulls(jsonb, &mut builder, strip_in_arrays);
     JsonbVal::from(builder.finish())
 }
 
 /// Recursively removes all object fields that have null values from the given JSON value.
-fn jsonbb_strip_nulls(jsonb: ValueRef<'_>, builder: &mut jsonbb::Builder) {
+fn jsonbb_strip_nulls(jsonb: ValueRef<'_>, builder: &mut jsonbb::Builder, strip_in_arrays: bool) {
     match jsonb {
         ValueRef::Object(obj) => {
             builder.begin_object();
@@ -397,14 +405,17 @@ fn jsonbb_strip_nulls(jsonb: ValueRef<'_>, builder: &mut jsonbb::Builder) {
                     continue;
                 }
                 builder.add_string(k);
-                jsonbb_strip_nulls(v, builder);
+                jsonbb_strip_nulls(v, builder, strip_in_arrays);
             }
             builder.end_object();
         }
         ValueRef::Array(array) => {
             builder.begin_array();
             for v in array.iter() {
-                jsonbb_strip_nulls(v, builder);
+                match (strip_in_arrays, v) {
+                    (true, ValueRef::Null) => (),
+                    _ => jsonbb_strip_nulls(v, builder, strip_in_arrays),
+                }
             }
             builder.end_array();
         }

src/frontend/src/expr/type_inference/func.rs

@@ -789,6 +789,23 @@ fn infer_type_for_special(
             }
             Ok(Some(DataType::Varchar))
         }
+        ExprType::JsonbStripNulls => {
+            ensure_arity!("jsonb_strip_nulls", 1 <= | inputs | <= 2);
+            inputs[0].cast_implicit_mut(&DataType::Jsonb)?;
+
+            if let Some(strip_in_arrays) = inputs.get_mut(1) {
+                strip_in_arrays
+                    .cast_implicit_mut(&DataType::Boolean)
+                    .map_err(|_| {
+                        ErrorCode::BindError(format!(
+                            "Cannot cast {} to boolean for parameter strip_in_arrays",
+                            strip_in_arrays.return_type()
+                        ))
+                    })?;
+            }
+
+            Ok(Some(DataType::Jsonb))
+        }
         _ => Ok(None),
     }
 }

src/tests/regress/data/expected/jsonb.out

@@ -3989,6 +3989,56 @@ select jsonb_strip_nulls('{"a": {"b": null, "c": null}, "d": {} }');
  {"a": {}, "d": {}}
 (1 row)
 
+-- jsonb_strip_nulls (strip_in_arrays=true)
+select jsonb_strip_nulls(null, true);
+ jsonb_strip_nulls 
+-------------------
+
+(1 row)
+
+select jsonb_strip_nulls('1', true);
+ jsonb_strip_nulls 
+-------------------
+ 1
+(1 row)
+
+select jsonb_strip_nulls('"a string"', true);
+ jsonb_strip_nulls 
+-------------------
+ "a string"
+(1 row)
+
+select jsonb_strip_nulls('null', true);
+ jsonb_strip_nulls 
+-------------------
+ null
+(1 row)
+
+select jsonb_strip_nulls('[1,2,null,3,4]', true);
+ jsonb_strip_nulls 
+-------------------
+ [1, 2, 3, 4]
+(1 row)
+
+select jsonb_strip_nulls('{"a":1,"b":null,"c":[2,null,3],"d":{"e":4,"f":null}}', true);
+ jsonb_strip_nulls 
+--------------------------------------
+ {"a": 1, "c": [2, 3], "d": {"e": 4}}
+(1 row)
+
+select jsonb_strip_nulls('[1,{"a":1,"b":null,"c":2},3]', true);
+ jsonb_strip_nulls 
+--------------------------
+ [1, {"a": 1, "c": 2}, 3]
+(1 row)
+
+-- an empty object is not null and should not be stripped
+select jsonb_strip_nulls('{"a": {"b": null, "c": null}, "d": {} }', true);
+ jsonb_strip_nulls 
+--------------------
+ {"a": {}, "d": {}}
+(1 row)
+
 select jsonb_pretty('{"a": "test", "b": [1, 2, 3], "c": "test3", "d":{"dd": "test4", "dd2":{"ddd": "test5"}}}');
         jsonb_pretty        
 ----------------------------

src/tests/regress/data/sql/jsonb.sql

@@ -1046,6 +1046,24 @@ select jsonb_strip_nulls('[1,{"a":1,"b":null,"c":2},3]');
 -- an empty object is not null and should not be stripped
 select jsonb_strip_nulls('{"a": {"b": null, "c": null}, "d": {} }');
 
+-- jsonb_strip_nulls (strip_in_arrays=true)
+
+--@ select jsonb_strip_nulls(null, true);
+
+--@ select jsonb_strip_nulls('1', true);
+
+--@ select jsonb_strip_nulls('"a string"', true);
+
+--@ select jsonb_strip_nulls('null', true);
+
+--@ select jsonb_strip_nulls('[1,2,null,3,4]', true);
+
+--@ select jsonb_strip_nulls('{"a":1,"b":null,"c":[2,null,3],"d":{"e":4,"f":null}}', true);
+
+--@ select jsonb_strip_nulls('[1,{"a":1,"b":null,"c":2},3]', true);
+
+-- an empty object is not null and should not be stripped
+--@ select jsonb_strip_nulls('{"a": {"b": null, "c": null}, "d": {} }', true);
 
 select jsonb_pretty('{"a": "test", "b": [1, 2, 3], "c": "test3", "d":{"dd": "test4", "dd2":{"ddd": "test5"}}}');
 select jsonb_pretty('[{"f1":1,"f2":null},2,null,[[{"x":true},6,7],8],3]');