micromegas_datafusion_extensions/jsonb/
each.rs

1use crate::jsonb::extract::object_or_array_entries;
2use async_trait::async_trait;
3use datafusion::arrow::array::{Array, ArrayRef, BinaryArray, DictionaryArray, GenericBinaryArray};
4use datafusion::arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef};
5use datafusion::arrow::record_batch::RecordBatch;
6use datafusion::catalog::Session;
7use datafusion::catalog::TableFunctionArgs;
8use datafusion::catalog::TableFunctionImpl;
9use datafusion::catalog::TableProvider;
10use datafusion::datasource::TableType;
11use datafusion::datasource::memory::{DataSourceExec, MemorySourceConfig};
12use datafusion::error::DataFusionError;
13use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder};
14use datafusion::physical_plan::ExecutionPlan;
15use datafusion::prelude::Expr;
16use datafusion::scalar::ScalarValue;
17use std::sync::Arc;
18
19/// A DataFusion `TableFunctionImpl` that expands a JSONB object or array into rows of (key, value).
20///
21/// For objects, `key` is the field name. For arrays, `key` is the element index (as a string).
22///
23/// Usage:
24/// ```sql
25/// SELECT key, jsonb_as_string(value)
26/// FROM jsonb_each(
27///   (SELECT properties FROM processes WHERE process_id = '...')
28/// )
29/// ```
30#[derive(Debug)]
31pub struct JsonbEachTableFunction {}
32
33impl JsonbEachTableFunction {
34    pub fn new() -> Self {
35        Self {}
36    }
37}
38
39impl Default for JsonbEachTableFunction {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45/// The source of JSONB data — either a literal value or a subquery to evaluate.
46#[derive(Debug, Clone)]
47enum JsonbSource {
48    Literal(ScalarValue),
49    Subquery(Arc<LogicalPlan>),
50}
51
52impl TableFunctionImpl for JsonbEachTableFunction {
53    fn call_with_args(
54        &self,
55        args: TableFunctionArgs,
56    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
57        let args = args.exprs();
58        if args.len() != 1 {
59            return Err(DataFusionError::Plan(
60                "jsonb_each requires exactly one argument (a JSONB object)".into(),
61            ));
62        }
63
64        let source = match &args[0] {
65            Expr::Literal(scalar, _metadata) => JsonbSource::Literal(scalar.clone()),
66            Expr::ScalarSubquery(subquery) => JsonbSource::Subquery(subquery.subquery.clone()),
67            other => {
68                let plan = LogicalPlanBuilder::empty(true)
69                    .project(vec![other.clone()])?
70                    .build()?;
71                JsonbSource::Subquery(Arc::new(plan))
72            }
73        };
74
75        Ok(Arc::new(JsonbEachTableProvider { source }))
76    }
77}
78
79fn output_schema() -> SchemaRef {
80    Arc::new(Schema::new(vec![
81        Field::new("key", DataType::Utf8, false),
82        Field::new("value", DataType::Binary, false),
83    ]))
84}
85
86/// Extract key-value entries from JSONB bytes.
87///
88/// For objects, keys are field names. For arrays, keys are element indices (as strings).
89/// Errors with the UDTF's usual message if the input is neither an object nor an array.
90fn extract_entries_from_jsonb(
91    jsonb_bytes: &[u8],
92) -> Result<Vec<(String, Vec<u8>)>, DataFusionError> {
93    object_or_array_entries(jsonb_bytes)?.ok_or_else(|| {
94        DataFusionError::Execution("jsonb_each: input is not a JSONB object or array".into())
95    })
96}
97
98fn entries_to_batch(entries: &[(String, Vec<u8>)]) -> Result<RecordBatch, DataFusionError> {
99    if entries.is_empty() {
100        return Ok(RecordBatch::new_empty(output_schema()));
101    }
102
103    let keys: Vec<&str> = entries.iter().map(|(k, _)| k.as_str()).collect();
104    let values: Vec<&[u8]> = entries.iter().map(|(_, v)| v.as_slice()).collect();
105
106    let key_array: ArrayRef = Arc::new(datafusion::arrow::array::StringArray::from(keys));
107    let value_array: ArrayRef = Arc::new(BinaryArray::from(values));
108
109    RecordBatch::try_new(output_schema(), vec![key_array, value_array])
110        .map_err(|e| DataFusionError::External(e.into()))
111}
112
113fn scalar_to_entries(scalar: &ScalarValue) -> Result<Vec<(String, Vec<u8>)>, DataFusionError> {
114    match scalar {
115        ScalarValue::Binary(Some(bytes)) => extract_entries_from_jsonb(bytes),
116        ScalarValue::Binary(None) => Ok(vec![]),
117        ScalarValue::Dictionary(_, inner) => scalar_to_entries(inner.as_ref()),
118        _ => Err(DataFusionError::Plan(format!(
119            "jsonb_each argument must be Binary (JSONB), got: {:?}",
120            scalar.data_type()
121        ))),
122    }
123}
124
125/// Extract JSONB bytes from all rows of a column, handling both plain Binary
126/// and Dictionary<Int32, Binary> encodings.
127fn extract_all_jsonb_bytes_from_column(column: &ArrayRef) -> Result<Vec<Vec<u8>>, DataFusionError> {
128    match column.data_type() {
129        DataType::Binary => {
130            let binary_array = column
131                .as_any()
132                .downcast_ref::<GenericBinaryArray<i32>>()
133                .ok_or_else(|| {
134                    DataFusionError::Internal("failed to cast column to BinaryArray".into())
135                })?;
136            Ok((0..binary_array.len())
137                .filter(|&i| !binary_array.is_null(i))
138                .map(|i| binary_array.value(i).to_vec())
139                .collect())
140        }
141        DataType::Dictionary(_, value_type) if matches!(value_type.as_ref(), DataType::Binary) => {
142            let dict_array = column
143                .as_any()
144                .downcast_ref::<DictionaryArray<Int32Type>>()
145                .ok_or_else(|| {
146                    DataFusionError::Internal(
147                        "failed to cast column to DictionaryArray<Int32, Binary>".into(),
148                    )
149                })?;
150            let binary_values = dict_array
151                .values()
152                .as_any()
153                .downcast_ref::<GenericBinaryArray<i32>>()
154                .ok_or_else(|| {
155                    DataFusionError::Internal("dictionary values are not a binary array".into())
156                })?;
157            Ok((0..dict_array.len())
158                .filter(|&i| !dict_array.is_null(i))
159                .map(|i| {
160                    let key_index = dict_array.keys().value(i) as usize;
161                    binary_values.value(key_index).to_vec()
162                })
163                .collect())
164        }
165        other => Err(DataFusionError::Execution(format!(
166            "jsonb_each subquery must return a Binary or Dictionary<Int32, Binary> column, got: {other:?}"
167        ))),
168    }
169}
170
171/// Table provider for expanding JSONB objects into key-value rows.
172#[derive(Debug)]
173pub struct JsonbEachTableProvider {
174    source: JsonbSource,
175}
176
177impl JsonbEachTableProvider {
178    /// Creates a new provider from a JSONB scalar value (for testing).
179    pub fn from_scalar(scalar: ScalarValue) -> Result<Self, DataFusionError> {
180        if !matches!(&scalar, ScalarValue::Binary(Some(_))) {
181            return Err(DataFusionError::Plan(format!(
182                "jsonb_each argument must be Binary (JSONB), got: {:?}",
183                scalar.data_type()
184            )));
185        }
186        Ok(Self {
187            source: JsonbSource::Literal(scalar),
188        })
189    }
190}
191
192#[async_trait]
193impl TableProvider for JsonbEachTableProvider {
194    fn schema(&self) -> SchemaRef {
195        output_schema()
196    }
197
198    fn table_type(&self) -> TableType {
199        TableType::Temporary
200    }
201
202    async fn scan(
203        &self,
204        state: &dyn Session,
205        projection: Option<&Vec<usize>>,
206        _filters: &[Expr],
207        limit: Option<usize>,
208    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
209        let entries = match &self.source {
210            JsonbSource::Literal(scalar) => scalar_to_entries(scalar)?,
211            JsonbSource::Subquery(plan) => {
212                let physical_plan = state.create_physical_plan(plan).await?;
213                let task_ctx = state.task_ctx();
214                let batches = datafusion::physical_plan::collect(physical_plan, task_ctx).await?;
215
216                if batches.is_empty() || batches.iter().all(|b| b.num_rows() == 0) {
217                    return Err(DataFusionError::Execution(
218                        "jsonb_each subquery returned no rows".into(),
219                    ));
220                }
221
222                let mut all_entries = Vec::new();
223                for batch in &batches {
224                    if batch.num_columns() != 1 {
225                        return Err(DataFusionError::Execution(format!(
226                            "jsonb_each subquery must return exactly one column, got {}",
227                            batch.num_columns()
228                        )));
229                    }
230                    for jsonb_bytes in extract_all_jsonb_bytes_from_column(batch.column(0))? {
231                        all_entries.extend(extract_entries_from_jsonb(&jsonb_bytes)?);
232                    }
233                }
234                all_entries
235            }
236        };
237
238        let mut record_batch = entries_to_batch(&entries)?;
239
240        // Apply limit if specified
241        if let Some(n) = limit
242            && n < record_batch.num_rows()
243        {
244            record_batch = record_batch.slice(0, n);
245        }
246
247        let source = MemorySourceConfig::try_new(
248            &[vec![record_batch]],
249            self.schema(),
250            projection.map(|v| v.to_owned()),
251        )?;
252        Ok(DataSourceExec::from_data_source(source))
253    }
254}
micromegas_datafusion_extensions/jsonb/each.rs

micromegas_datafusion_extensions/jsonb/
each.rs
