Skip to content

Commit

Permalink
Flatten array
Browse files Browse the repository at this point in the history
  • Loading branch information
@trueleo
trueleo committed Mar 30, 2023
1 parent c5b7312 commit 73ed48b
Showing 1 changed file with 168 additions and 36 deletions.
204 changes: 168 additions & 36 deletions server/src/utils/json/flatten.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
*
*/

use itertools::Itertools;
use serde_json::map::Map;
use serde_json::value::Value;

Expand Down Expand Up @@ -57,17 +58,12 @@ pub fn flatten_object(
match value {
Value::Object(obj) => flatten_object(map, Some(&new_key), obj, separator)?,
Value::Array(arr) => {
let mut new_arr = Vec::with_capacity(arr.len());
for elem in arr {
if let Value::Object(object) = elem {
let mut map = Map::new();
flatten_object_disallow_array(&mut map, None, object, separator)?;
new_arr.push(Value::Object(map))
} else {
new_arr.push(elem)
}
// if value is object then decompose this list into lists
if arr.iter().any(|value| value.is_object()) {
flatten_array_objects(map, &new_key, arr, separator)?;
} else {
map.insert(new_key, Value::Array(arr));
}
map.insert(new_key, Value::Array(new_arr));
}
x => {
map.insert(new_key, x);
Expand All @@ -77,34 +73,71 @@ pub fn flatten_object(
Ok(())
}

fn flatten_object_disallow_array(
pub fn flatten_array_objects(
map: &mut Map<String, Value>,
parent_key: Option<&str>,
object: Map<String, Value>,
parent_key: &str,
arr: Vec<Value>,
separator: &str,
) -> Result<(), ()> {
for (key, value) in object.into_iter() {
let new_key = parent_key.map_or_else(
|| key.clone(),
|parent_key| format!("{}{}{}", parent_key, separator, key),
);
match value {
Value::Object(obj) => {
flatten_object_disallow_array(map, Some(&new_key), obj, separator)?;
let mut columns: Vec<(String, Vec<Value>)> = Vec::new();
let mut len = 0;
for value in arr {
if let Value::Object(object) = value {
let mut flattened_object = Map::new();
flatten_object(&mut flattened_object, None, object, separator)?;
let mut col_index = 0;
for (key, value) in flattened_object.into_iter().sorted_by(|a, b| a.0.cmp(&b.0)) {
loop {
if let Some((column_name, column)) = columns.get_mut(col_index) {
match (*column_name).cmp(&key) {
std::cmp::Ordering::Less => {
column.push(Value::Null);
col_index += 1;
continue;
}
std::cmp::Ordering::Equal => column.push(value),
std::cmp::Ordering::Greater => {
let mut list = vec![Value::Null; len];
list.push(value);
columns.insert(col_index, (key, list));
}
}
} else {
let mut list = vec![Value::Null; len];
list.push(value);
columns.push((key, list));
}
col_index += 1;
break;
}
}
Value::Array(_) => return Err(()),
x => {
map.insert(new_key, x);
for (_, column) in &mut columns[col_index..] {
column.push(Value::Null)
}
} else if value.is_null() {
for (_, column) in &mut columns {
column.push(Value::Null)
}
} else {
return Err(());
}
len += 1;
}

for (key, arr) in columns {
let new_key = format!("{}{}{}", parent_key, separator, key);
map.insert(new_key, Value::Array(arr));
}

Ok(())
}

#[cfg(test)]
mod tests {
use crate::utils::json::flatten::flatten_array_objects;

use super::flatten;
use serde_json::json;
use serde_json::{json, Map, Value};

#[test]
fn flatten_single_key_string() {
Expand Down Expand Up @@ -152,29 +185,128 @@ mod tests {
}

#[test]
fn nested_array() {
let obj = json!({"key": ["value1", "value2"]});
assert_eq!(obj.clone(), flatten(obj, ".").unwrap());
fn nested_obj_array() {
let obj = json!({"key": [{"a": "value0"}, {"a": "value1"}]});
assert_eq!(
json!({"key.a": ["value0", "value1"]}),
flatten(obj, ".").unwrap()
);
}

#[test]
fn nested_obj_array() {
let obj = json!({"key": ["value1", {"key": "value2"}]});
assert_eq!(obj.clone(), flatten(obj, ".").unwrap());
fn nested_obj_array_nulls() {
let obj = json!({"key": [{"a": "value0"}, {"a": "value1", "b": "value1"}]});
assert_eq!(
json!({"key.a": ["value0", "value1"], "key.b": [null, "value1"]}),
flatten(obj, ".").unwrap()
);
}

#[test]
fn flatten_mixed_object() {
let obj = json!({"a": 42, "arr": ["1", {"key": "2"}, {"key": {"nested": "3"}}]});
fn nested_obj_array_nulls_reversed() {
let obj = json!({"key": [{"a": "value0", "b": "value0"}, {"a": "value1"}]});
assert_eq!(
json!({"key.a": ["value0", "value1"], "key.b": ["value0", null]}),
flatten(obj, ".").unwrap()
);
}

#[test]
fn nested_obj_array_nested_obj() {
let obj = json!({"key": [{"a": {"p": 0}, "b": "value0"}, {"b": "value1"}]});
assert_eq!(
json!({"a": 42, "arr": ["1", {"key": "2"}, {"key.nested": "3"}]}),
json!({"key.a.p": [0, null], "key.b": ["value0", "value1"]}),
flatten(obj, ".").unwrap()
);
}

#[test]
fn err_flatten_mixed_object_with_array_within_object() {
let obj = json!({"a": 42, "arr": ["1", {"key": [1,2,4]}]});
fn nested_obj_array_nested_obj_array() {
let obj = json!({"key": [{"a": [{"p": "value0", "q": "value0"}, {"p": "value1", "q": null}], "b": "value0"}, {"b": "value1"}]});
assert_eq!(
json!({"key.a.p": [["value0", "value1"], null], "key.a.q": [["value0", null], null], "key.b": ["value0", "value1"]}),
flatten(obj, ".").unwrap()
);
}

#[test]
fn flatten_mixed_object() {
let obj = json!({"a": 42, "arr": ["1", {"key": "2"}, {"key": {"nested": "3"}}]});
assert!(flatten(obj, ".").is_err());
}

#[test]
fn flatten_array_nulls_at_start() {
let Value::Array(arr) = json!([
null,
{"p": 2, "q": 2},
{"q": 3},
]) else { unreachable!() };

let mut map = Map::new();
flatten_array_objects(&mut map, "key", arr, ".").unwrap();

assert_eq!(map.len(), 2);
assert_eq!(map.get("key.p").unwrap(), &json!([null, 2, null]));
assert_eq!(map.get("key.q").unwrap(), &json!([null, 2, 3]));
}

#[test]
fn flatten_array_objects_nulls_at_end() {
let Value::Array(arr) = json!([{"a": 1, "b": 1}, {"a": 2}, null]) else { unreachable!() };

let mut map = Map::new();
flatten_array_objects(&mut map, "key", arr, ".").unwrap();

assert_eq!(map.len(), 2);
assert_eq!(map.get("key.a").unwrap(), &json!([1, 2, null]));
assert_eq!(map.get("key.b").unwrap(), &json!([1, null, null]));
}

#[test]
fn flatten_array_objects_nulls_in_middle() {
let Value::Array(arr) = json!([{"a": 1, "b": 1}, null, {"a": 3, "c": 3}]) else { unreachable!() };

let mut map = Map::new();
flatten_array_objects(&mut map, "key", arr, ".").unwrap();

assert_eq!(map.len(), 3);
assert_eq!(map.get("key.a").unwrap(), &json!([1, null, 3]));
assert_eq!(map.get("key.b").unwrap(), &json!([1, null, null]));
assert_eq!(map.get("key.c").unwrap(), &json!([null, null, 3]));
}

#[test]
fn flatten_array_test() {
let Value::Array(arr) = json!([
{"p": 1, "q": 1},
{"r": 2, "q": 2},
{"p": 3, "r": 3}
]) else { unreachable!() };

let mut map = Map::new();
flatten_array_objects(&mut map, "key", arr, ".").unwrap();

assert_eq!(map.len(), 3);
assert_eq!(map.get("key.p").unwrap(), &json!([1, null, 3]));
assert_eq!(map.get("key.q").unwrap(), &json!([1, 2, null]));
assert_eq!(map.get("key.r").unwrap(), &json!([null, 2, 3]));
}

#[test]
fn flatten_array_nested_test() {
let Value::Array(arr) = json!([
{"p": 1, "q": [{"x": 1}, {"x": 2}]},
{"r": 2, "q": [{"x": 1}]},
{"p": 3, "r": 3}
]) else { unreachable!() };

let mut map = Map::new();
flatten_array_objects(&mut map, "key", arr, ".").unwrap();

assert_eq!(map.len(), 3);
assert_eq!(map.get("key.p").unwrap(), &json!([1, null, 3]));
assert_eq!(map.get("key.q.x").unwrap(), &json!([[1, 2], [1], null]));
assert_eq!(map.get("key.r").unwrap(), &json!([null, 2, 3]));
}
}

0 comments on commit 73ed48b

Please sign in to comment.