laminar_sql/datafusion/json_extensions/
object_ops.rs1#[allow(clippy::disallowed_types)] use std::collections::HashSet;
14use std::hash::{Hash, Hasher};
15use std::sync::Arc;
16
17use arrow::datatypes::DataType;
18use arrow_array::{
19 builder::LargeBinaryBuilder, Array, LargeBinaryArray, ListArray, MapArray, StringArray,
20};
21use datafusion_common::Result;
22use datafusion_expr::{
23 ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
24};
25
26use super::super::json_types;
27use super::super::json_udf::expand_args;
28
29const MAX_DEPTH: usize = 64;
31
32#[derive(Debug)]
41pub struct JsonbMerge {
42 signature: Signature,
43}
44
45impl JsonbMerge {
46 #[must_use]
48 pub fn new() -> Self {
49 Self {
50 signature: Signature::new(
51 TypeSignature::Exact(vec![DataType::LargeBinary, DataType::LargeBinary]),
52 Volatility::Immutable,
53 ),
54 }
55 }
56}
57
58impl Default for JsonbMerge {
59 fn default() -> Self {
60 Self::new()
61 }
62}
63
64impl PartialEq for JsonbMerge {
65 fn eq(&self, _other: &Self) -> bool {
66 true
67 }
68}
69
70impl Eq for JsonbMerge {}
71
72impl Hash for JsonbMerge {
73 fn hash<H: Hasher>(&self, state: &mut H) {
74 "jsonb_merge".hash(state);
75 }
76}
77
78impl ScalarUDFImpl for JsonbMerge {
79 fn as_any(&self) -> &dyn std::any::Any {
80 self
81 }
82
83 fn name(&self) -> &'static str {
84 "jsonb_merge"
85 }
86
87 fn signature(&self) -> &Signature {
88 &self.signature
89 }
90
91 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
92 Ok(DataType::LargeBinary)
93 }
94
95 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
96 let expanded = expand_args(&args.args)?;
97 let left_arr = expanded[0]
98 .as_any()
99 .downcast_ref::<LargeBinaryArray>()
100 .ok_or_else(|| {
101 datafusion_common::DataFusionError::Internal(
102 "jsonb_merge: first arg must be LargeBinary".into(),
103 )
104 })?;
105 let right_arr = expanded[1]
106 .as_any()
107 .downcast_ref::<LargeBinaryArray>()
108 .ok_or_else(|| {
109 datafusion_common::DataFusionError::Internal(
110 "jsonb_merge: second arg must be LargeBinary".into(),
111 )
112 })?;
113
114 let mut builder = LargeBinaryBuilder::with_capacity(left_arr.len(), 256);
115 for i in 0..left_arr.len() {
116 if left_arr.is_null(i) || right_arr.is_null(i) {
117 builder.append_null();
118 } else {
119 let left_val = json_types::jsonb_to_value(left_arr.value(i));
120 let right_val = json_types::jsonb_to_value(right_arr.value(i));
121 match (left_val, right_val) {
122 (
123 Some(serde_json::Value::Object(mut l)),
124 Some(serde_json::Value::Object(r)),
125 ) => {
126 for (k, v) in r {
127 l.insert(k, v);
128 }
129 builder
130 .append_value(json_types::encode_jsonb(&serde_json::Value::Object(l)));
131 }
132 (_, Some(r)) => {
133 builder.append_value(json_types::encode_jsonb(&r));
134 }
135 _ => builder.append_null(),
136 }
137 }
138 }
139 Ok(ColumnarValue::Array(Arc::new(builder.finish())))
140 }
141}
142
143#[derive(Debug)]
152pub struct JsonbDeepMerge {
153 signature: Signature,
154}
155
156impl JsonbDeepMerge {
157 #[must_use]
159 pub fn new() -> Self {
160 Self {
161 signature: Signature::new(
162 TypeSignature::Exact(vec![DataType::LargeBinary, DataType::LargeBinary]),
163 Volatility::Immutable,
164 ),
165 }
166 }
167}
168
169impl Default for JsonbDeepMerge {
170 fn default() -> Self {
171 Self::new()
172 }
173}
174
175impl PartialEq for JsonbDeepMerge {
176 fn eq(&self, _other: &Self) -> bool {
177 true
178 }
179}
180
181impl Eq for JsonbDeepMerge {}
182
183impl Hash for JsonbDeepMerge {
184 fn hash<H: Hasher>(&self, state: &mut H) {
185 "jsonb_deep_merge".hash(state);
186 }
187}
188
189fn deep_merge(
190 left: serde_json::Value,
191 right: serde_json::Value,
192 depth: usize,
193) -> std::result::Result<serde_json::Value, String> {
194 if depth > MAX_DEPTH {
195 return Err("jsonb_deep_merge: max depth exceeded".into());
196 }
197 match (left, right) {
198 (serde_json::Value::Object(mut l), serde_json::Value::Object(r)) => {
199 for (k, rv) in r {
200 let merged = if let Some(lv) = l.remove(&k) {
201 deep_merge(lv, rv, depth + 1)?
202 } else {
203 rv
204 };
205 l.insert(k, merged);
206 }
207 Ok(serde_json::Value::Object(l))
208 }
209 (_, r) => Ok(r),
210 }
211}
212
213impl ScalarUDFImpl for JsonbDeepMerge {
214 fn as_any(&self) -> &dyn std::any::Any {
215 self
216 }
217
218 fn name(&self) -> &'static str {
219 "jsonb_deep_merge"
220 }
221
222 fn signature(&self) -> &Signature {
223 &self.signature
224 }
225
226 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
227 Ok(DataType::LargeBinary)
228 }
229
230 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
231 let expanded = expand_args(&args.args)?;
232 let left_arr = expanded[0]
233 .as_any()
234 .downcast_ref::<LargeBinaryArray>()
235 .ok_or_else(|| {
236 datafusion_common::DataFusionError::Internal(
237 "jsonb_deep_merge: first arg must be LargeBinary".into(),
238 )
239 })?;
240 let right_arr = expanded[1]
241 .as_any()
242 .downcast_ref::<LargeBinaryArray>()
243 .ok_or_else(|| {
244 datafusion_common::DataFusionError::Internal(
245 "jsonb_deep_merge: second arg must be LargeBinary".into(),
246 )
247 })?;
248
249 let mut builder = LargeBinaryBuilder::with_capacity(left_arr.len(), 256);
250 for i in 0..left_arr.len() {
251 if left_arr.is_null(i) || right_arr.is_null(i) {
252 builder.append_null();
253 } else {
254 let left_val = json_types::jsonb_to_value(left_arr.value(i));
255 let right_val = json_types::jsonb_to_value(right_arr.value(i));
256 match (left_val, right_val) {
257 (Some(l), Some(r)) => {
258 let merged = deep_merge(l, r, 0)
259 .map_err(datafusion_common::DataFusionError::Execution)?;
260 builder.append_value(json_types::encode_jsonb(&merged));
261 }
262 _ => builder.append_null(),
263 }
264 }
265 }
266 Ok(ColumnarValue::Array(Arc::new(builder.finish())))
267 }
268}
269
270#[derive(Debug)]
279pub struct JsonbStripNulls {
280 signature: Signature,
281}
282
283impl JsonbStripNulls {
284 #[must_use]
286 pub fn new() -> Self {
287 Self {
288 signature: Signature::new(
289 TypeSignature::Exact(vec![DataType::LargeBinary]),
290 Volatility::Immutable,
291 ),
292 }
293 }
294}
295
296impl Default for JsonbStripNulls {
297 fn default() -> Self {
298 Self::new()
299 }
300}
301
302impl PartialEq for JsonbStripNulls {
303 fn eq(&self, _other: &Self) -> bool {
304 true
305 }
306}
307
308impl Eq for JsonbStripNulls {}
309
310impl Hash for JsonbStripNulls {
311 fn hash<H: Hasher>(&self, state: &mut H) {
312 "jsonb_strip_nulls".hash(state);
313 }
314}
315
316fn strip_nulls(val: serde_json::Value) -> serde_json::Value {
317 match val {
318 serde_json::Value::Object(obj) => {
319 let filtered: serde_json::Map<String, serde_json::Value> = obj
320 .into_iter()
321 .filter(|(_, v)| !v.is_null())
322 .map(|(k, v)| (k, strip_nulls(v)))
323 .collect();
324 serde_json::Value::Object(filtered)
325 }
326 serde_json::Value::Array(arr) => {
327 serde_json::Value::Array(arr.into_iter().map(strip_nulls).collect())
328 }
329 other => other,
330 }
331}
332
333impl ScalarUDFImpl for JsonbStripNulls {
334 fn as_any(&self) -> &dyn std::any::Any {
335 self
336 }
337
338 fn name(&self) -> &'static str {
339 "jsonb_strip_nulls"
340 }
341
342 fn signature(&self) -> &Signature {
343 &self.signature
344 }
345
346 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
347 Ok(DataType::LargeBinary)
348 }
349
350 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
351 let expanded = expand_args(&args.args)?;
352 let jsonb_arr = expanded[0]
353 .as_any()
354 .downcast_ref::<LargeBinaryArray>()
355 .ok_or_else(|| {
356 datafusion_common::DataFusionError::Internal(
357 "jsonb_strip_nulls: arg must be LargeBinary".into(),
358 )
359 })?;
360
361 let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
362 for i in 0..jsonb_arr.len() {
363 if jsonb_arr.is_null(i) {
364 builder.append_null();
365 } else {
366 match json_types::jsonb_to_value(jsonb_arr.value(i)) {
367 Some(val) => {
368 builder.append_value(json_types::encode_jsonb(&strip_nulls(val)));
369 }
370 None => builder.append_null(),
371 }
372 }
373 }
374 Ok(ColumnarValue::Array(Arc::new(builder.finish())))
375 }
376}
377
378#[derive(Debug)]
387pub struct JsonbRenameKeys {
388 signature: Signature,
389}
390
391impl JsonbRenameKeys {
392 #[must_use]
394 pub fn new() -> Self {
395 Self {
396 signature: Signature::new(
397 TypeSignature::Exact(vec![
398 DataType::LargeBinary,
399 DataType::Map(
400 Arc::new(arrow_schema::Field::new(
401 "entries",
402 DataType::Struct(
403 vec![
404 arrow_schema::Field::new("key", DataType::Utf8, false),
405 arrow_schema::Field::new("value", DataType::Utf8, true),
406 ]
407 .into(),
408 ),
409 false,
410 )),
411 false,
412 ),
413 ]),
414 Volatility::Immutable,
415 ),
416 }
417 }
418}
419
420impl Default for JsonbRenameKeys {
421 fn default() -> Self {
422 Self::new()
423 }
424}
425
426impl PartialEq for JsonbRenameKeys {
427 fn eq(&self, _other: &Self) -> bool {
428 true
429 }
430}
431
432impl Eq for JsonbRenameKeys {}
433
434impl Hash for JsonbRenameKeys {
435 fn hash<H: Hasher>(&self, state: &mut H) {
436 "jsonb_rename_keys".hash(state);
437 }
438}
439
440impl ScalarUDFImpl for JsonbRenameKeys {
441 fn as_any(&self) -> &dyn std::any::Any {
442 self
443 }
444
445 fn name(&self) -> &'static str {
446 "jsonb_rename_keys"
447 }
448
449 fn signature(&self) -> &Signature {
450 &self.signature
451 }
452
453 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
454 Ok(DataType::LargeBinary)
455 }
456
457 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
458 let expanded = expand_args(&args.args)?;
459 let jsonb_arr = expanded[0]
460 .as_any()
461 .downcast_ref::<LargeBinaryArray>()
462 .ok_or_else(|| {
463 datafusion_common::DataFusionError::Internal(
464 "jsonb_rename_keys: first arg must be LargeBinary".into(),
465 )
466 })?;
467 let map_arr = expanded[1]
468 .as_any()
469 .downcast_ref::<MapArray>()
470 .ok_or_else(|| {
471 datafusion_common::DataFusionError::Internal(
472 "jsonb_rename_keys: second arg must be Map<Utf8,Utf8>".into(),
473 )
474 })?;
475
476 let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
477 for i in 0..jsonb_arr.len() {
478 if jsonb_arr.is_null(i) || map_arr.is_null(i) {
479 builder.append_null();
480 } else {
481 let rename_map = extract_string_map(map_arr, i);
482 match json_types::jsonb_to_value(jsonb_arr.value(i)) {
483 Some(serde_json::Value::Object(obj)) => {
484 let renamed: serde_json::Map<String, serde_json::Value> = obj
485 .into_iter()
486 .map(|(k, v)| {
487 let new_key = rename_map.get(k.as_str()).cloned().unwrap_or(k);
488 (new_key, v)
489 })
490 .collect();
491 builder.append_value(json_types::encode_jsonb(&serde_json::Value::Object(
492 renamed,
493 )));
494 }
495 Some(other) => {
496 builder.append_value(json_types::encode_jsonb(&other));
497 }
498 None => builder.append_null(),
499 }
500 }
501 }
502 Ok(ColumnarValue::Array(Arc::new(builder.finish())))
503 }
504}
505
506#[allow(clippy::disallowed_types)] fn extract_string_map(map_arr: &MapArray, row: usize) -> std::collections::HashMap<String, String> {
509 let mut result = std::collections::HashMap::new();
510 let entries = map_arr.value(row);
511 let struct_arr = entries
512 .as_any()
513 .downcast_ref::<arrow_array::StructArray>()
514 .unwrap();
515 let keys = struct_arr
516 .column(0)
517 .as_any()
518 .downcast_ref::<StringArray>()
519 .unwrap();
520 let vals = struct_arr
521 .column(1)
522 .as_any()
523 .downcast_ref::<StringArray>()
524 .unwrap();
525 for j in 0..keys.len() {
526 if !keys.is_null(j) && !vals.is_null(j) {
527 result.insert(keys.value(j).to_owned(), vals.value(j).to_owned());
528 }
529 }
530 result
531}
532
533#[derive(Debug)]
541pub struct JsonbPick {
542 signature: Signature,
543}
544
545impl JsonbPick {
546 #[must_use]
548 pub fn new() -> Self {
549 Self {
550 signature: Signature::new(
551 TypeSignature::Exact(vec![
552 DataType::LargeBinary,
553 DataType::List(Arc::new(arrow_schema::Field::new(
554 "item",
555 DataType::Utf8,
556 true,
557 ))),
558 ]),
559 Volatility::Immutable,
560 ),
561 }
562 }
563}
564
565impl Default for JsonbPick {
566 fn default() -> Self {
567 Self::new()
568 }
569}
570
571impl PartialEq for JsonbPick {
572 fn eq(&self, _other: &Self) -> bool {
573 true
574 }
575}
576
577impl Eq for JsonbPick {}
578
579impl Hash for JsonbPick {
580 fn hash<H: Hasher>(&self, state: &mut H) {
581 "jsonb_pick".hash(state);
582 }
583}
584
585impl ScalarUDFImpl for JsonbPick {
586 fn as_any(&self) -> &dyn std::any::Any {
587 self
588 }
589
590 fn name(&self) -> &'static str {
591 "jsonb_pick"
592 }
593
594 fn signature(&self) -> &Signature {
595 &self.signature
596 }
597
598 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
599 Ok(DataType::LargeBinary)
600 }
601
602 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
603 let expanded = expand_args(&args.args)?;
604 let jsonb_arr = expanded[0]
605 .as_any()
606 .downcast_ref::<LargeBinaryArray>()
607 .ok_or_else(|| {
608 datafusion_common::DataFusionError::Internal(
609 "jsonb_pick: first arg must be LargeBinary".into(),
610 )
611 })?;
612 let keys_arr = expanded[1]
613 .as_any()
614 .downcast_ref::<ListArray>()
615 .ok_or_else(|| {
616 datafusion_common::DataFusionError::Internal(
617 "jsonb_pick: second arg must be List<Utf8>".into(),
618 )
619 })?;
620
621 let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
622 for i in 0..jsonb_arr.len() {
623 if jsonb_arr.is_null(i) || keys_arr.is_null(i) {
624 builder.append_null();
625 } else {
626 let key_set = extract_string_set(keys_arr, i);
627 match json_types::jsonb_to_value(jsonb_arr.value(i)) {
628 Some(serde_json::Value::Object(obj)) => {
629 let picked: serde_json::Map<String, serde_json::Value> = obj
630 .into_iter()
631 .filter(|(k, _)| key_set.contains(k.as_str()))
632 .collect();
633 builder.append_value(json_types::encode_jsonb(&serde_json::Value::Object(
634 picked,
635 )));
636 }
637 Some(other) => {
638 builder.append_value(json_types::encode_jsonb(&other));
639 }
640 None => builder.append_null(),
641 }
642 }
643 }
644 Ok(ColumnarValue::Array(Arc::new(builder.finish())))
645 }
646}
647
648fn extract_string_set(list_arr: &ListArray, row: usize) -> HashSet<String> {
650 let values = list_arr.value(row);
651 let str_arr = values.as_any().downcast_ref::<StringArray>();
652 let mut result = HashSet::new();
653 if let Some(arr) = str_arr {
654 for j in 0..arr.len() {
655 if !arr.is_null(j) {
656 result.insert(arr.value(j).to_owned());
657 }
658 }
659 }
660 result
661}
662
663#[derive(Debug)]
671pub struct JsonbExcept {
672 signature: Signature,
673}
674
675impl JsonbExcept {
676 #[must_use]
678 pub fn new() -> Self {
679 Self {
680 signature: Signature::new(
681 TypeSignature::Exact(vec![
682 DataType::LargeBinary,
683 DataType::List(Arc::new(arrow_schema::Field::new(
684 "item",
685 DataType::Utf8,
686 true,
687 ))),
688 ]),
689 Volatility::Immutable,
690 ),
691 }
692 }
693}
694
695impl Default for JsonbExcept {
696 fn default() -> Self {
697 Self::new()
698 }
699}
700
701impl PartialEq for JsonbExcept {
702 fn eq(&self, _other: &Self) -> bool {
703 true
704 }
705}
706
707impl Eq for JsonbExcept {}
708
709impl Hash for JsonbExcept {
710 fn hash<H: Hasher>(&self, state: &mut H) {
711 "jsonb_except".hash(state);
712 }
713}
714
715impl ScalarUDFImpl for JsonbExcept {
716 fn as_any(&self) -> &dyn std::any::Any {
717 self
718 }
719
720 fn name(&self) -> &'static str {
721 "jsonb_except"
722 }
723
724 fn signature(&self) -> &Signature {
725 &self.signature
726 }
727
728 fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
729 Ok(DataType::LargeBinary)
730 }
731
732 fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
733 let expanded = expand_args(&args.args)?;
734 let jsonb_arr = expanded[0]
735 .as_any()
736 .downcast_ref::<LargeBinaryArray>()
737 .ok_or_else(|| {
738 datafusion_common::DataFusionError::Internal(
739 "jsonb_except: first arg must be LargeBinary".into(),
740 )
741 })?;
742 let keys_arr = expanded[1]
743 .as_any()
744 .downcast_ref::<ListArray>()
745 .ok_or_else(|| {
746 datafusion_common::DataFusionError::Internal(
747 "jsonb_except: second arg must be List<Utf8>".into(),
748 )
749 })?;
750
751 let mut builder = LargeBinaryBuilder::with_capacity(jsonb_arr.len(), 256);
752 for i in 0..jsonb_arr.len() {
753 if jsonb_arr.is_null(i) || keys_arr.is_null(i) {
754 builder.append_null();
755 } else {
756 let exclude_set = extract_string_set(keys_arr, i);
757 match json_types::jsonb_to_value(jsonb_arr.value(i)) {
758 Some(serde_json::Value::Object(obj)) => {
759 let filtered: serde_json::Map<String, serde_json::Value> = obj
760 .into_iter()
761 .filter(|(k, _)| !exclude_set.contains(k.as_str()))
762 .collect();
763 builder.append_value(json_types::encode_jsonb(&serde_json::Value::Object(
764 filtered,
765 )));
766 }
767 Some(other) => {
768 builder.append_value(json_types::encode_jsonb(&other));
769 }
770 None => builder.append_null(),
771 }
772 }
773 }
774 Ok(ColumnarValue::Array(Arc::new(builder.finish())))
775 }
776}