|
| 1 | +// Copyright 2024 RisingWave Labs |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +use risingwave_common::types::{JsonbRef, StructRef, StructValue}; |
| 16 | +use risingwave_expr::expr::Context; |
| 17 | +use risingwave_expr::{function, ExprError, Result}; |
| 18 | + |
| 19 | +/// Expands the top-level JSON object to a row having the composite type of the base argument. |
| 20 | +/// The JSON object is scanned for fields whose names match column names of the output row type, |
| 21 | +/// and their values are inserted into those columns of the output. (Fields that do not correspond |
| 22 | +/// to any output column name are ignored.) In typical use, the value of base is just NULL, which |
| 23 | +/// means that any output columns that do not match any object field will be filled with nulls. |
| 24 | +/// However, if base isn't NULL then the values it contains will be used for unmatched columns. |
| 25 | +/// |
| 26 | +/// # Examples |
| 27 | +/// |
| 28 | +/// ```slt |
| 29 | +/// query ITT |
| 30 | +/// select (jsonb_populate_record( |
| 31 | +/// null::struct<a int, b text[], c struct<d int, e text>>, |
| 32 | +/// '{"a": 1, "b": ["2", "a b"], "c": {"d": 4, "e": "a b c"}, "x": "foo"}' |
| 33 | +/// )).*; |
| 34 | +/// ---- |
| 35 | +/// 1 {2,"a b"} (4,"a b c") |
| 36 | +/// |
| 37 | +/// query ITT |
| 38 | +/// select (jsonb_populate_record( |
| 39 | +/// row(1, null, row(4, '5'))::struct<a int, b text[], c struct<d int, e text>>, |
| 40 | +/// '{"b": ["2", "a b"], "c": {"e": "a b c"}, "x": "foo"}' |
| 41 | +/// )).*; |
| 42 | +/// ---- |
| 43 | +/// 1 {2,"a b"} (4,"a b c") |
| 44 | +/// ``` |
| 45 | +#[function("jsonb_populate_record(struct, jsonb) -> struct")] |
| 46 | +fn jsonb_populate_record( |
| 47 | + base: Option<StructRef<'_>>, |
| 48 | + jsonb: JsonbRef<'_>, |
| 49 | + ctx: &Context, |
| 50 | +) -> Result<StructValue> { |
| 51 | + let output_type = ctx.return_type.as_struct(); |
| 52 | + jsonb.populate_struct(output_type, base).map_err(parse_err) |
| 53 | +} |
| 54 | + |
| 55 | +/// Expands the top-level JSON array of objects to a set of rows having the composite type of the |
| 56 | +/// base argument. Each element of the JSON array is processed as described above for |
| 57 | +/// `jsonb_populate_record`. |
| 58 | +/// |
| 59 | +/// # Examples |
| 60 | +/// |
| 61 | +/// ```slt |
| 62 | +/// query II |
| 63 | +/// select * from jsonb_populate_recordset( |
| 64 | +/// null::struct<a int, b int>, |
| 65 | +/// '[{"a":1,"b":2}, {"a":3,"b":4}]'::jsonb |
| 66 | +/// ); |
| 67 | +/// ---- |
| 68 | +/// 1 2 |
| 69 | +/// 3 4 |
| 70 | +/// |
| 71 | +/// query II |
| 72 | +/// select * from jsonb_populate_recordset( |
| 73 | +/// row(0, 0)::struct<a int, b int>, |
| 74 | +/// '[{}, {"a":1}, {"b":2}, {"a":1,"b":2}]'::jsonb |
| 75 | +/// ); |
| 76 | +/// ---- |
| 77 | +/// 0 0 |
| 78 | +/// 1 0 |
| 79 | +/// 0 2 |
| 80 | +/// 1 2 |
| 81 | +/// ``` |
| 82 | +#[function("jsonb_populate_recordset(struct, jsonb) -> setof struct")] |
| 83 | +fn jsonb_populate_recordset<'a>( |
| 84 | + base: Option<StructRef<'a>>, |
| 85 | + jsonb: JsonbRef<'a>, |
| 86 | + ctx: &'a Context, |
| 87 | +) -> Result<impl Iterator<Item = Result<StructValue>> + 'a> { |
| 88 | + let output_type = ctx.return_type.as_struct(); |
| 89 | + Ok(jsonb |
| 90 | + .array_elements() |
| 91 | + .map_err(parse_err)? |
| 92 | + .map(move |elem| elem.populate_struct(output_type, base).map_err(parse_err))) |
| 93 | +} |
| 94 | + |
| 95 | +/// Expands the top-level JSON object to a row having the composite type defined by an AS clause. |
| 96 | +/// The output record is filled from fields of the JSON object, in the same way as described above |
| 97 | +/// for `jsonb_populate_record`. Since there is no input record value, unmatched columns are always |
| 98 | +/// filled with nulls. |
| 99 | +/// |
| 100 | +/// # Examples |
| 101 | +/// |
| 102 | +/// // FIXME(runji): this query is blocked by parser and frontend support. |
| 103 | +/// ```slt,ignore |
| 104 | +/// query T |
| 105 | +/// select * from jsonb_to_record('{"a":1,"b":[1,2,3],"c":[1,2,3],"e":"bar","r": {"a": 123, "b": "a b c"}}') |
| 106 | +/// as x(a int, b text, c int[], d text, r struct<a int, b text>); |
| 107 | +/// ---- |
| 108 | +/// 1 [1,2,3] {1,2,3} NULL (123,"a b c") |
| 109 | +/// ``` |
| 110 | +#[function("jsonb_to_record(jsonb) -> struct", type_infer = "panic")] |
| 111 | +fn jsonb_to_record(jsonb: JsonbRef<'_>, ctx: &Context) -> Result<StructValue> { |
| 112 | + let output_type = ctx.return_type.as_struct(); |
| 113 | + jsonb.to_struct(output_type).map_err(parse_err) |
| 114 | +} |
| 115 | + |
| 116 | +/// Expands the top-level JSON array of objects to a set of rows having the composite type defined |
| 117 | +/// by an AS clause. Each element of the JSON array is processed as described above for |
| 118 | +/// `jsonb_populate_record`. |
| 119 | +/// |
| 120 | +/// # Examples |
| 121 | +/// |
| 122 | +/// // FIXME(runji): this query is blocked by parser and frontend support. |
| 123 | +/// ```slt,ignore |
| 124 | +/// query IT |
| 125 | +/// select * from jsonb_to_recordset('[{"a":1,"b":"foo"}, {"a":"2","c":"bar"}]') as x(a int, b text); |
| 126 | +/// ---- |
| 127 | +/// 1 foo |
| 128 | +/// 2 NULL |
| 129 | +/// ``` |
| 130 | +#[function("jsonb_to_recordset(jsonb) -> setof struct", type_infer = "panic")] |
| 131 | +fn jsonb_to_recordset<'a>( |
| 132 | + jsonb: JsonbRef<'a>, |
| 133 | + ctx: &'a Context, |
| 134 | +) -> Result<impl Iterator<Item = Result<StructValue>> + 'a> { |
| 135 | + let output_type = ctx.return_type.as_struct(); |
| 136 | + Ok(jsonb |
| 137 | + .array_elements() |
| 138 | + .map_err(parse_err)? |
| 139 | + .map(|elem| elem.to_struct(output_type).map_err(parse_err))) |
| 140 | +} |
| 141 | + |
| 142 | +/// Construct a parse error from String. |
| 143 | +fn parse_err(s: String) -> ExprError { |
| 144 | + ExprError::Parse(s.into()) |
| 145 | +} |
0 commit comments