athena_driver/postgresql/
json_row_expr.rs1const MAX_JSONB_BUILD_OBJECT_PAIRS: usize = 50;
5
6pub fn merge_jsonb_build_object_expr(pairs: &[String]) -> String {
11 if pairs.is_empty() {
12 return "jsonb_build_object()".to_string();
13 }
14 let chunks: Vec<String> = pairs
15 .chunks(MAX_JSONB_BUILD_OBJECT_PAIRS)
16 .map(|chunk| format!("jsonb_build_object({})", chunk.join(", ")))
17 .collect();
18 match chunks.len() {
19 1 => chunks[0].clone(),
20 _ => format!("({})", chunks.join(" || ")),
21 }
22}
23
24#[cfg(test)]
25mod tests {
26 use super::merge_jsonb_build_object_expr;
27
28 fn sample_pairs(n: usize) -> Vec<String> {
29 (0..n).map(|i| format!("'c{i}', t.c{i}")).collect()
30 }
31
32 #[test]
33 fn single_chunk_no_merge_operator() {
34 let pairs = sample_pairs(50);
35 let expr = merge_jsonb_build_object_expr(&pairs);
36 assert!(expr.starts_with("jsonb_build_object("));
37 assert!(!expr.contains("||"));
38 assert_eq!(expr.matches("jsonb_build_object").count(), 1);
39 }
40
41 #[test]
42 fn fifty_one_pairs_merge_two_objects() {
43 let pairs = sample_pairs(51);
44 let expr = merge_jsonb_build_object_expr(&pairs);
45 assert!(expr.starts_with('('));
46 assert!(expr.contains("||"));
47 assert_eq!(expr.matches("jsonb_build_object").count(), 2);
48 }
49
50 #[test]
51 fn one_pair_single_call() {
52 let pairs = sample_pairs(1);
53 let expr = merge_jsonb_build_object_expr(&pairs);
54 assert_eq!(expr, "jsonb_build_object('c0', t.c0)");
55 }
56}