reifydb_routine/function/json/
array.rs1use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
5use reifydb_type::{
6 util::bitvec::BitVec,
7 value::{Value, r#type::Type},
8};
9
10use crate::routine::{Function, FunctionKind, Routine, RoutineInfo, context::FunctionContext, error::RoutineError};
11
12pub struct JsonArray {
13 info: RoutineInfo,
14}
15
16impl Default for JsonArray {
17 fn default() -> Self {
18 Self::new()
19 }
20}
21
22impl JsonArray {
23 pub fn new() -> Self {
24 Self {
25 info: RoutineInfo::new("json::array"),
26 }
27 }
28}
29
30impl<'a> Routine<FunctionContext<'a>> for JsonArray {
31 fn info(&self) -> &RoutineInfo {
32 &self.info
33 }
34
35 fn return_type(&self, _input_types: &[Type]) -> Type {
36 Type::Any
37 }
38
39 fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
40 if args.is_empty() {
41 return Ok(Columns::new(vec![ColumnWithName::new(
42 ctx.fragment.clone(),
43 ColumnBuffer::any(vec![Box::new(Value::List(vec![]))]),
44 )]));
45 }
46
47 let mut unwrapped: Vec<_> = Vec::with_capacity(args.len());
49 let mut combined_bv: Option<BitVec> = None;
50
51 for col in args.iter() {
52 let (data, bitvec) = col.data().unwrap_option();
53 if let Some(bv) = bitvec {
54 combined_bv = Some(match combined_bv {
55 Some(existing) => existing.and(bv),
56 None => bv.clone(),
57 });
58 }
59 unwrapped.push(data);
60 }
61
62 let row_count = unwrapped[0].len();
63 let mut results: Vec<Box<Value>> = Vec::with_capacity(row_count);
64
65 for row in 0..row_count {
66 let mut items = Vec::with_capacity(unwrapped.len());
67 for col_data in unwrapped.iter() {
68 items.push(col_data.get_value(row));
69 }
70 results.push(Box::new(Value::List(items)));
71 }
72
73 let result_data = ColumnBuffer::any(results);
74 let final_data = match combined_bv {
75 Some(bv) => ColumnBuffer::Option {
76 inner: Box::new(result_data),
77 bitvec: bv,
78 },
79 None => result_data,
80 };
81
82 Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), final_data)]))
83 }
84}
85
86impl Function for JsonArray {
87 fn kinds(&self) -> &[FunctionKind] {
88 &[FunctionKind::Scalar]
89 }
90}