micromegas_datafusion_extensions/jsonb/
array_elements.rs

1use crate::jsonb::extract::array_elements;
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 array into rows with a single `value` column.
20///
21/// Usage:
22/// ```sql
23/// SELECT jsonb_as_string(elem.value)
24/// FROM jsonb_array_elements(jsonb_parse('[1, 2, 3]')) as elem
25/// ```
26#[derive(Debug)]
27pub struct JsonbArrayElementsTableFunction {}
28
29impl JsonbArrayElementsTableFunction {
30    pub fn new() -> Self {
31        Self {}
32    }
33}
34
35impl Default for JsonbArrayElementsTableFunction {
36    fn default() -> Self {
37        Self::new()
38    }
39}
40
41/// The source of JSONB data — either a literal value or a subquery/expression to evaluate.
42#[derive(Debug, Clone)]
43enum JsonbSource {
44    Literal(ScalarValue),
45    Subquery(Arc<LogicalPlan>),
46}
47
48impl TableFunctionImpl for JsonbArrayElementsTableFunction {
49    fn call_with_args(
50        &self,
51        args: TableFunctionArgs,
52    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
53        let args = args.exprs();
54        if args.len() != 1 {
55            return Err(DataFusionError::Plan(
56                "jsonb_array_elements requires exactly one argument (a JSONB array)".into(),
57            ));
58        }
59
60        let source = match &args[0] {
61            Expr::Literal(scalar, _metadata) => JsonbSource::Literal(scalar.clone()),
62            Expr::ScalarSubquery(subquery) => JsonbSource::Subquery(subquery.subquery.clone()),
63            other => {
64                let plan = LogicalPlanBuilder::empty(true)
65                    .project(vec![other.clone()])?
66                    .build()?;
67                JsonbSource::Subquery(Arc::new(plan))
68            }
69        };
70
71        Ok(Arc::new(JsonbArrayElementsTableProvider { source }))
72    }
73}
74
75fn output_schema() -> SchemaRef {
76    Arc::new(Schema::new(vec![Field::new(
77        "value",
78        DataType::Binary,
79        false,
80    )]))
81}
82
83/// Extract element values from a JSONB array. Errors with the UDTF's usual message if the input
84/// is not an array.
85fn extract_elements_from_jsonb(jsonb_bytes: &[u8]) -> Result<Vec<Vec<u8>>, DataFusionError> {
86    array_elements(jsonb_bytes)?.ok_or_else(|| {
87        DataFusionError::Execution("jsonb_array_elements: input is not a JSONB array".into())
88    })
89}
90
91fn empty_exec(projection: Option<&Vec<usize>>) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
92    let batch = RecordBatch::new_empty(output_schema());
93    let source = MemorySourceConfig::try_new(
94        &[vec![batch]],
95        output_schema(),
96        projection.map(|v| v.to_owned()),
97    )?;
98    Ok(DataSourceExec::from_data_source(source))
99}
100
101fn elements_to_batch(elements: &[Vec<u8>]) -> Result<RecordBatch, DataFusionError> {
102    if elements.is_empty() {
103        return Ok(RecordBatch::new_empty(output_schema()));
104    }
105
106    let values: Vec<&[u8]> = elements.iter().map(|v| v.as_slice()).collect();
107    let value_array: ArrayRef = Arc::new(BinaryArray::from(values));
108
109    RecordBatch::try_new(output_schema(), vec![value_array])
110        .map_err(|e| DataFusionError::External(e.into()))
111}
112
113fn scalar_to_elements(scalar: &ScalarValue) -> Result<Vec<Vec<u8>>, DataFusionError> {
114    match scalar {
115        ScalarValue::Binary(Some(bytes)) => extract_elements_from_jsonb(bytes),
116        ScalarValue::Binary(None) => Ok(vec![]),
117        ScalarValue::Dictionary(_, inner) => scalar_to_elements(inner.as_ref()),
118        _ => Err(DataFusionError::Plan(format!(
119            "jsonb_array_elements 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_array_elements subquery must return a Binary or Dictionary<Int32, Binary> column, got: {other:?}"
167        ))),
168    }
169}
170
171/// Table provider for expanding JSONB arrays into value rows.
172#[derive(Debug)]
173pub struct JsonbArrayElementsTableProvider {
174    source: JsonbSource,
175}
176
177impl JsonbArrayElementsTableProvider {
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_array_elements 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 JsonbArrayElementsTableProvider {
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 elements = match &self.source {
210            JsonbSource::Literal(scalar) => scalar_to_elements(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                let mut all_elements = Vec::new();
217                if batches.is_empty() || batches.iter().all(|b| b.num_rows() == 0) {
218                    return empty_exec(projection);
219                }
220                for batch in &batches {
221                    if batch.num_columns() != 1 {
222                        return Err(DataFusionError::Execution(format!(
223                            "jsonb_array_elements subquery must return exactly one column, got {}",
224                            batch.num_columns()
225                        )));
226                    }
227                    for jsonb_bytes in extract_all_jsonb_bytes_from_column(batch.column(0))? {
228                        all_elements.extend(extract_elements_from_jsonb(&jsonb_bytes)?);
229                    }
230                }
231                all_elements
232            }
233        };
234
235        let mut record_batch = elements_to_batch(&elements)?;
236
237        // Apply limit if specified
238        if let Some(n) = limit
239            && n < record_batch.num_rows()
240        {
241            record_batch = record_batch.slice(0, n);
242        }
243
244        let source = MemorySourceConfig::try_new(
245            &[vec![record_batch]],
246            self.schema(),
247            projection.map(|v| v.to_owned()),
248        )?;
249        Ok(DataSourceExec::from_data_source(source))
250    }
251}
micromegas_datafusion_extensions/jsonb/array_elements.rs

micromegas_datafusion_extensions/jsonb/
array_elements.rs
