Skip to main content

reifydb_routine/function/json/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use 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		// Check for any option columns and unwrap them
48		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}