micromegas_datafusion_extensions/jsonb/
format_json.rs

1use crate::binary_column_accessor::create_binary_accessor;
2use datafusion::arrow::array::StringDictionaryBuilder;
3use datafusion::arrow::datatypes::{DataType, Int32Type};
4use datafusion::common::{Result, internal_err};
5use datafusion::logical_expr::{
6    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
7};
8use jsonb::RawJsonb;
9use std::sync::Arc;
10
11/// A scalar UDF that formats JSONB binary data as a JSON string.
12///
13/// Accepts both Binary and Dictionary<Int32, Binary> inputs, making it compatible
14/// with dictionary-encoded JSONB columns and the output of `properties_to_jsonb`.
15/// Returns Dictionary<Int32, Utf8> for memory efficiency.
16#[derive(Debug, PartialEq, Eq, Hash)]
17pub struct JsonbFormatJson {
18    signature: Signature,
19}
20
21impl JsonbFormatJson {
22    pub fn new() -> Self {
23        Self {
24            signature: Signature::any(1, Volatility::Immutable),
25        }
26    }
27}
28
29impl Default for JsonbFormatJson {
30    fn default() -> Self {
31        Self::new()
32    }
33}
34
35impl ScalarUDFImpl for JsonbFormatJson {
36    fn name(&self) -> &str {
37        "jsonb_format_json"
38    }
39
40    fn signature(&self) -> &Signature {
41        &self.signature
42    }
43
44    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
45        Ok(DataType::Dictionary(
46            Box::new(DataType::Int32),
47            Box::new(DataType::Utf8),
48        ))
49    }
50
51    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
52        let args = ColumnarValue::values_to_arrays(&args.args)?;
53        if args.len() != 1 {
54            return internal_err!("wrong number of arguments to jsonb_format_json");
55        }
56
57        // Use BinaryColumnAccessor to handle both Binary and Dictionary<Int32, Binary>
58        let binary_accessor = create_binary_accessor(&args[0])
59            .map_err(|e| datafusion::error::DataFusionError::Execution(
60                format!("Invalid input type for jsonb_format_json: {}. Expected Binary or Dictionary<Int32, Binary>", e)
61            ))?;
62
63        let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
64
65        for index in 0..binary_accessor.len() {
66            if binary_accessor.is_null(index) {
67                dict_builder.append_null();
68            } else {
69                let src_buffer = binary_accessor.value(index);
70                let jsonb = RawJsonb::new(src_buffer);
71                dict_builder.append_value(jsonb.to_string());
72            }
73        }
74
75        Ok(ColumnarValue::Array(Arc::new(dict_builder.finish())))
76    }
77}
78
79/// Creates a user-defined function to format a JSONB value as a JSON string.
80///
81/// This function accepts both `Binary` and `Dictionary<Int32, Binary>` inputs,
82/// allowing it to work seamlessly with dictionary-encoded JSONB columns.
83/// Returns `Dictionary<Int32, Utf8>` for memory efficiency.
84pub fn make_jsonb_format_json_udf() -> datafusion::logical_expr::ScalarUDF {
85    datafusion::logical_expr::ScalarUDF::new_from_impl(JsonbFormatJson::new())
86}
micromegas_datafusion_extensions/jsonb/format_json.rs

micromegas_datafusion_extensions/jsonb/
format_json.rs
