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::{Column, columns::Columns, data::ColumnData};
5use reifydb_type::{
6	util::bitvec::BitVec,
7	value::{Value, r#type::Type},
8};
9
10use crate::function::{Function, FunctionCapability, FunctionContext, FunctionInfo, error::FunctionError};
11
12pub struct JsonArray {
13	info: FunctionInfo,
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: FunctionInfo::new("json::array"),
26		}
27	}
28}
29
30impl Function for JsonArray {
31	fn info(&self) -> &FunctionInfo {
32		&self.info
33	}
34
35	fn capabilities(&self) -> &[FunctionCapability] {
36		&[FunctionCapability::Scalar]
37	}
38
39	fn return_type(&self, _input_types: &[Type]) -> Type {
40		Type::Any
41	}
42
43	fn execute(&self, ctx: &FunctionContext, args: &Columns) -> Result<Columns, FunctionError> {
44		if args.is_empty() {
45			return Ok(Columns::new(vec![Column::new(
46				ctx.fragment.clone(),
47				ColumnData::any(vec![Box::new(Value::List(vec![]))]),
48			)]));
49		}
50
51		// Check for any option columns and unwrap them
52		let mut unwrapped: Vec<_> = Vec::with_capacity(args.len());
53		let mut combined_bv: Option<BitVec> = None;
54
55		for col in args.iter() {
56			let (data, bitvec) = col.data().unwrap_option();
57			if let Some(bv) = bitvec {
58				combined_bv = Some(match combined_bv {
59					Some(existing) => existing.and(bv),
60					None => bv.clone(),
61				});
62			}
63			unwrapped.push(data);
64		}
65
66		let row_count = unwrapped[0].len();
67		let mut results: Vec<Box<Value>> = Vec::with_capacity(row_count);
68
69		for row in 0..row_count {
70			let mut items = Vec::with_capacity(unwrapped.len());
71			for col_data in unwrapped.iter() {
72				items.push(col_data.get_value(row));
73			}
74			results.push(Box::new(Value::List(items)));
75		}
76
77		let result_data = ColumnData::any(results);
78		let final_data = match combined_bv {
79			Some(bv) => ColumnData::Option {
80				inner: Box::new(result_data),
81				bitvec: bv,
82			},
83			None => result_data,
84		};
85
86		Ok(Columns::new(vec![Column::new(ctx.fragment.clone(), final_data)]))
87	}
88}