Skip to content

Commit

Permalink
expr: fix panic in json_array_append (#16929)
Browse files Browse the repository at this point in the history 
close #16930

mysql> drop table if exists t;
Query OK, 0 rows affected, 1 warning (0.05 sec)

mysql> create table t(a json, index idx((cast(json_array_append(a, '$.a', 1) as char(300)))));
Query OK, 0 rows affected (0.12 sec)

mysql> insert into t (a) values ('{\"a\": 1}');
Query OK, 1 row affected (0.07 sec)

mysql> insert into t (a) values ('{\"b\": 1}');
Query OK, 1 row affected (0.05 sec)

mysql> insert into t (a) values ('{\"a\": {\"b\": 1}}');
Query OK, 1 row affected (0.07 sec)

mysql> insert into t (a) values ('{\"a\": {\"b\": [1,2,3]}}');
Query OK, 1 row affected (0.07 sec)

mysql> insert into t (a) values ('[1,2,3]');
Query OK, 1 row affected (0.06 sec)

mysql> insert into t (a) values ('[\"a\",\"b\"]');
Query OK, 1 row affected (0.03 sec)

mysql> insert into t (a) values ('\"abc\"');
Query OK, 1 row affected (0.07 sec)

mysql> insert into t (a) values ('123');
Query OK, 1 row affected (0.07 sec)

mysql> insert into t (a) values ('true');
Query OK, 1 row affected (0.08 sec)

mysql> insert into t (a) values ('false');
Query OK, 1 row affected (0.04 sec)

mysql> insert into t (a) values (NULL);
Query OK, 1 row affected (0.08 sec)

mysql>
mysql> select * from t ignore index(idx) where cast(json_array_append(a, '$.a', 1) as char(300)) = '{"a": [1, 1]}';
+----------+
| a        |
+----------+
| {"a": 1} |
+----------+
1 row in set (0.06 sec)

Signed-off-by: dbsid <chenhuansheng@pingcap.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
  • Loading branch information
@dbsid @ti-chi-bot
dbsid and ti-chi-bot[bot] committed May 10, 2024
1 parent e157965 commit 69f13af
Showing 1 changed file with 166 additions and 24 deletions.
190 changes: 166 additions & 24 deletions components/tidb_query_expr/src/impl_json.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,40 +70,45 @@ fn json_modify(args: &[ScalarValueRef], mt: ModifyType) -> Result<Option<Json>>
#[inline]
fn json_array_append(args: &[ScalarValueRef]) -> Result<Option<Json>> {
assert!(args.len() >= 2);
// Returns None if Base is None
if args[0].to_owned().is_none() {
return Ok(None);
}
// base Json argument
let base: Option<JsonRef> = args[0].as_json();
let base = base.map_or(Json::none(), |json| Ok(json.to_owned()))?;

let buf_size = args.len() / 2;

let mut path_expr_list = Vec::with_capacity(buf_size);
let mut values = Vec::with_capacity(buf_size);
let mut base = base.map_or(Json::none(), |json| Ok(json.to_owned()))?;

for chunk in args[1..].chunks(2) {
let path: Option<BytesRef> = chunk[0].as_bytes();
let value: Option<JsonRef> = chunk[1].as_json();

path_expr_list.push(try_opt!(parse_json_path(path)));

let value = value
.as_ref()
.map_or(Json::none(), |json| Ok(json.to_owned()))?;
// extract the element from the path, then merge the value into the element
// 1. extrace the element from the path
let tmp_path_expr_list = vec![path_expr_list.last().unwrap().to_owned()];
let element = base.as_ref().extract(&tmp_path_expr_list)?;
// change element to JsonRef
let element_ref = element.as_ref().map(|e| e.as_ref());
let tmp_path_expr_list = vec![try_opt!(parse_json_path(path))];
let element: Option<Json> = base.as_ref().extract(&tmp_path_expr_list)?;
// 2. merge the value into the element
let tmp_values: Vec<JsonRef> = vec![element_ref.unwrap(), value.as_ref()];

values.push(Json::merge(tmp_values)?);
if let Some(elem) = element {
// if both elem and value are json object, wrap elem into a vector
if elem.get_type() == JsonType::Object && value.get_type() == JsonType::Object {
let array_json: Json = Json::from_array(vec![elem.clone()])?;
let tmp_values = vec![array_json.as_ref(), value.as_ref()];
let tmp_value = Json::merge(tmp_values)?;
base =
base.as_ref()
.modify(&tmp_path_expr_list, vec![tmp_value], ModifyType::Set)?;
} else {
let tmp_values = vec![elem.as_ref(), value.as_ref()];
let tmp_value = Json::merge(tmp_values)?;
base =
base.as_ref()
.modify(&tmp_path_expr_list, vec![tmp_value], ModifyType::Set)?;
}
}
}
Ok(Some(base.as_ref().modify(
&path_expr_list,
values,
ModifyType::Set,
)?))
Ok(Some(base))
}

/// validate the arguments are `(Option<JsonRef>, &[(Option<Bytes>,
Expand Down Expand Up @@ -1557,22 +1562,114 @@ mod tests {
#[test]
fn test_json_array_append() {
let cases: Vec<(Vec<ScalarValue>, _)> = vec![
// use exact testcase from TiDB repo
(
vec![
Some(Json::from_str(r#"{"a": 1, "b": [2, 3], "c": 4}"#).unwrap()).into(),
Some(b"$.d".to_vec()).into(),
Some(Json::from_str(r#""z""#).unwrap()).into(),
],
Some(r#"{"a": 1, "b": [2, 3], "c": 4}"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1, "b": [2, 3], "c": 4}"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
Some(Json::from_str(r#""w""#).unwrap()).into(),
],
Some(r#"[{"a": 1, "b": [2, 3], "c": 4}, "w"]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1, "b": [2, 3], "c": 4}"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
None::<Json>.into(),
None::<Bytes>.into(),
],
Some(r#"[{"a": 1, "b": [2, 3], "c": 4}, null]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1}"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
Some(Json::from_str(r#"{"b": 2}"#).unwrap()).into(),
],
Some(r#"[{"a": 1}, {"b": 2}]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1}"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
Some(Json::from_str(r#"{"b": 2}"#).unwrap()).into(),
],
Some(r#"[{"a": 1}, {"b": 2}]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1}"#).unwrap()).into(),
Some(b"$.a".to_vec()).into(),
Some(Json::from_str(r#"{"b": 2}"#).unwrap()).into(),
],
Some(r#"{"a": [1, {"b": 2}]}"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1}"#).unwrap()).into(),
Some(b"$.a".to_vec()).into(),
Some(Json::from_str(r#"{"b": 2}"#).unwrap()).into(),
Some(b"$.a[1]".to_vec()).into(),
Some(Json::from_str(r#"{"b": 2}"#).unwrap()).into(),
],
Some(r#"{"a": [1, [{"b": 2}, {"b": 2}]]}"#.parse().unwrap()),
),
(
vec![
None::<Json>.into(),
Some(b"$".to_vec()).into(),
None::<Json>.into(),
],
None::<Json>,
),
(
vec![
Some(Json::from_i64(9).unwrap()).into(),
None::<Json>.into(),
Some(b"$".to_vec()).into(),
Some(Json::from_u64(3).unwrap()).into(),
Some(Json::from_str(r#""a""#).unwrap()).into(),
],
Some(r#"[9,3]"#.parse().unwrap()),
None::<Json>,
),
(
vec![
Some(Json::from_str(r#"null"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
None::<Json>.into(),
],
Some(r#"[null, null]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"[]"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
None::<Json>.into(),
],
Some(r#"[null]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{}"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
None::<Json>.into(),
],
Some(r#"[{}, null]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1, "b": [2, 3], "c": 4}"#).unwrap()).into(),
None::<Bytes>.into(),
None::<Json>.into(),
],
None::<Json>,
),
// Following tests come from MySQL doc.
(
vec![
Some(Json::from_str(r#"["a", ["b", "c"], "d"]"#).unwrap()).into(),
Expand All @@ -1597,6 +1694,51 @@ mod tests {
],
Some(r#"["a", [["b", 3], "c"], "d"]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1, "b": [2, 3], "c": 4}"#).unwrap()).into(),
Some(b"$.b".to_vec()).into(),
Some(Json::from_str(r#""x""#).unwrap()).into(),
],
Some(r#"{"a": 1, "b": [2, 3, "x"], "c": 4}"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"{"a": 1, "b": [2, 3], "c": 4}"#).unwrap()).into(),
Some(b"$.c".to_vec()).into(),
Some(Json::from_str(r#""y""#).unwrap()).into(),
],
Some(r#"{"a": 1, "b": [2, 3], "c": [4, "y"]}"#.parse().unwrap()),
),
// Following tests come from MySQL test.
(
vec![
Some(Json::from_str(r#"[1,2,3, {"a":[4,5,6]}]"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
Some(Json::from_u64(7).unwrap()).into(),
],
Some(r#"[1, 2, 3, {"a": [4, 5, 6]}, 7]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"[1,2,3, {"a":[4,5,6]}]"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
Some(Json::from_u64(7).unwrap()).into(),
Some(b"$[3].a".to_vec()).into(),
Some(Json::from_f64(3.15).unwrap()).into(),
],
Some(r#"[1, 2, 3, {"a": [4, 5, 6, 3.15]}, 7]"#.parse().unwrap()),
),
(
vec![
Some(Json::from_str(r#"[1,2,3, {"a":[4,5,6]}]"#).unwrap()).into(),
Some(b"$".to_vec()).into(),
Some(Json::from_u64(7).unwrap()).into(),
Some(b"$[3].b".to_vec()).into(),
Some(Json::from_u64(8).unwrap()).into(),
],
Some(r#"[1, 2, 3, {"a": [4, 5, 6]}, 7]"#.parse().unwrap()),
),
];
for (args, expect_output) in cases {
let output: Option<Json> = RpnFnScalarEvaluator::new()
Expand Down

0 comments on commit 69f13af

Please sign in to comment.