laminar_sql/datafusion/
complex_type_lambda.rs

1//! Lambda higher-order functions for arrays and maps (F-SCHEMA-015 Tier 3).
2//!
3//! Provides vectorized lambda evaluation over Arrow arrays:
4//!
5//! | Function | Lambda | Strategy |
6//! |----------|--------|----------|
7//! | `array_transform(arr, lambda)` | `x -> expr` | Flatten, eval, re-group |
8//! | `array_filter(arr, lambda)` | `x -> bool` | Flatten, eval, filter+rebuild |
9//! | `array_reduce(arr, init, lambda)` | `(acc, x) -> expr` | Sequential fold |
10//! | `map_filter(map, lambda)` | `(k, v) -> bool` | Eval on k+v, filter entries |
11//! | `map_transform_values(map, lambda)` | `(k, v) -> expr` | Eval on k+v, replace vals |
12//!
13//! Lambda expressions are specified as string literal SQL expressions.
14//! They are evaluated using DataFusion's SQL engine against a temporary
15//! table containing the element values. Native lambda syntax is deferred.
16
17use std::hash::{Hash, Hasher};
18use std::sync::Arc;
19
20use arrow::datatypes::DataType;
21use arrow_array::{Array, ArrayRef, BooleanArray, ListArray, MapArray, StructArray};
22use arrow_schema::{Field, Fields, Schema};
23use datafusion_common::Result;
24use datafusion_expr::{
25    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
26};
27
28use super::json_udf::expand_args;
29
30/// Registers all lambda HOFs with the given session context.
31pub fn register_lambda_functions(ctx: &datafusion::prelude::SessionContext) {
32    use datafusion_expr::ScalarUDF;
33
34    ctx.register_udf(ScalarUDF::new_from_impl(ArrayTransform::new()));
35    ctx.register_udf(ScalarUDF::new_from_impl(ArrayFilter::new()));
36    ctx.register_udf(ScalarUDF::new_from_impl(ArrayReduce::new()));
37    ctx.register_udf(ScalarUDF::new_from_impl(MapFilter::new()));
38    ctx.register_udf(ScalarUDF::new_from_impl(MapTransformValues::new()));
39}
40
41thread_local! {
42    // This standalone expression context is outside per-DB reservation accounting;
43    // direct expression/Arrow allocations are not governed by DataFusion's memory pool.
44    /// Cached `SessionContext` for lambda evaluation.
45    ///
46    /// Creating a `SessionContext` is expensive because it registers all
47    /// built-in functions and sets up catalogs. We cache it per-thread
48    /// and reuse it across lambda invocations.
49    static LAMBDA_CTX: std::cell::RefCell<Option<datafusion::prelude::SessionContext>> =
50        const { std::cell::RefCell::new(None) };
51}
52
53/// Evaluate a SQL expression against a `RecordBatch`, returning the result column.
54///
55/// The expression can reference columns by name from the batch schema.
56fn eval_expr_on_batch(sql_expr: &str, batch: &arrow_array::RecordBatch) -> Result<ArrayRef> {
57    // Reuse a thread-local SessionContext to avoid the cost of
58    // registering built-in functions on every invocation.
59    let ctx = LAMBDA_CTX.with(|cell| {
60        let mut opt = cell.borrow_mut();
61        opt.get_or_insert_with(datafusion::prelude::SessionContext::new)
62            .clone()
63    });
64
65    let provider =
66        datafusion::datasource::MemTable::try_new(batch.schema(), vec![vec![batch.clone()]])?;
67    let rt = tokio::runtime::Handle::try_current().map_err(|e| {
68        datafusion_common::DataFusionError::Internal(format!(
69            "lambda eval requires tokio runtime: {e}"
70        ))
71    })?;
72    // Use block_in_place to allow blocking inside an already-running tokio runtime.
73    tokio::task::block_in_place(|| {
74        rt.block_on(async {
75            ctx.register_table("__lambda_data", Arc::new(provider))?;
76            let df = ctx
77                .sql(&format!("SELECT {sql_expr} FROM __lambda_data"))
78                .await?;
79            let batches = df.collect().await?;
80            // Deregister the ephemeral table so the batch data is freed.
81            let _ = ctx.deregister_table("__lambda_data");
82            if batches.is_empty() {
83                Err(datafusion_common::DataFusionError::Internal(
84                    "lambda expression returned no data".into(),
85                ))
86            } else {
87                // Concatenate all result batches and return the first column.
88                let result = arrow::compute::concat_batches(&batches[0].schema(), &batches)?;
89                Ok(result.column(0).clone())
90            }
91        })
92    })
93}
94
95fn scalar_string_value(cv: &ColumnarValue) -> Result<String> {
96    match cv {
97        ColumnarValue::Scalar(s) => {
98            let arr = s.to_array_of_size(1)?;
99            let str_arr = arr
100                .as_any()
101                .downcast_ref::<arrow_array::StringArray>()
102                .ok_or_else(|| {
103                    datafusion_common::DataFusionError::Internal("expected Utf8 argument".into())
104                })?;
105            Ok(str_arr.value(0).to_string())
106        }
107        ColumnarValue::Array(arr) => {
108            let str_arr = arr
109                .as_any()
110                .downcast_ref::<arrow_array::StringArray>()
111                .ok_or_else(|| {
112                    datafusion_common::DataFusionError::Internal("expected Utf8 argument".into())
113                })?;
114            Ok(str_arr.value(0).to_string())
115        }
116    }
117}
118
119// ══════════════════════════════════════════════════════════════════
120// array_transform(arr, lambda_str) -> List
121// ══════════════════════════════════════════════════════════════════
122
123/// `array_transform(arr, 'x + 1')` — apply a lambda to each element.
124#[derive(Debug)]
125pub struct ArrayTransform {
126    signature: Signature,
127}
128
129impl ArrayTransform {
130    /// Creates a new `array_transform` UDF.
131    #[must_use]
132    pub fn new() -> Self {
133        Self {
134            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
135        }
136    }
137}
138
139impl Default for ArrayTransform {
140    fn default() -> Self {
141        Self::new()
142    }
143}
144impl PartialEq for ArrayTransform {
145    fn eq(&self, _: &Self) -> bool {
146        true
147    }
148}
149impl Eq for ArrayTransform {}
150impl Hash for ArrayTransform {
151    fn hash<H: Hasher>(&self, s: &mut H) {
152        "array_transform".hash(s);
153    }
154}
155
156impl ScalarUDFImpl for ArrayTransform {
157    fn as_any(&self) -> &dyn std::any::Any {
158        self
159    }
160
161    fn name(&self) -> &'static str {
162        "array_transform"
163    }
164    fn signature(&self) -> &Signature {
165        &self.signature
166    }
167
168    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
169        match &arg_types[0] {
170            DataType::List(f) => Ok(DataType::List(Arc::clone(f))),
171            _ => Ok(DataType::List(Arc::new(Field::new(
172                "item",
173                DataType::Utf8,
174                true,
175            )))),
176        }
177    }
178
179    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
180        let expanded = expand_args(&args.args)?;
181        let list_arr = expanded[0]
182            .as_any()
183            .downcast_ref::<ListArray>()
184            .ok_or_else(|| {
185                datafusion_common::DataFusionError::Internal(
186                    "array_transform: first arg must be List".into(),
187                )
188            })?;
189
190        let lambda_str = scalar_string_value(&args.args[1])?;
191        let flat_values = list_arr.values();
192
193        let schema = Arc::new(Schema::new(vec![Field::new(
194            "x",
195            flat_values.data_type().clone(),
196            true,
197        )]));
198        let batch = arrow_array::RecordBatch::try_new(schema, vec![Arc::clone(flat_values)])?;
199
200        let result_arr = eval_expr_on_batch(&lambda_str, &batch)?;
201
202        let new_field = Arc::new(Field::new("item", result_arr.data_type().clone(), true));
203        let new_list = ListArray::try_new(
204            new_field,
205            list_arr.offsets().clone(),
206            result_arr,
207            list_arr.nulls().cloned(),
208        )?;
209        Ok(ColumnarValue::Array(Arc::new(new_list)))
210    }
211}
212
213// ══════════════════════════════════════════════════════════════════
214// array_filter(arr, lambda_str) -> List
215// ══════════════════════════════════════════════════════════════════
216
217/// `array_filter(arr, 'x > 0')` — filter elements by a boolean lambda.
218#[derive(Debug)]
219pub struct ArrayFilter {
220    signature: Signature,
221}
222
223impl ArrayFilter {
224    /// Creates a new `array_filter` UDF.
225    #[must_use]
226    pub fn new() -> Self {
227        Self {
228            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
229        }
230    }
231}
232
233impl Default for ArrayFilter {
234    fn default() -> Self {
235        Self::new()
236    }
237}
238impl PartialEq for ArrayFilter {
239    fn eq(&self, _: &Self) -> bool {
240        true
241    }
242}
243impl Eq for ArrayFilter {}
244impl Hash for ArrayFilter {
245    fn hash<H: Hasher>(&self, s: &mut H) {
246        "array_filter".hash(s);
247    }
248}
249
250impl ScalarUDFImpl for ArrayFilter {
251    fn as_any(&self) -> &dyn std::any::Any {
252        self
253    }
254
255    fn name(&self) -> &'static str {
256        "array_filter"
257    }
258    fn signature(&self) -> &Signature {
259        &self.signature
260    }
261
262    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
263        match &arg_types[0] {
264            DataType::List(f) => Ok(DataType::List(Arc::clone(f))),
265            _ => Ok(DataType::List(Arc::new(Field::new(
266                "item",
267                DataType::Utf8,
268                true,
269            )))),
270        }
271    }
272
273    #[allow(
274        clippy::cast_sign_loss,
275        clippy::cast_possible_wrap,
276        clippy::cast_possible_truncation
277    )]
278    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
279        let expanded = expand_args(&args.args)?;
280        let list_arr = expanded[0]
281            .as_any()
282            .downcast_ref::<ListArray>()
283            .ok_or_else(|| {
284                datafusion_common::DataFusionError::Internal(
285                    "array_filter: first arg must be List".into(),
286                )
287            })?;
288
289        let lambda_str = scalar_string_value(&args.args[1])?;
290        let flat_values = list_arr.values();
291        let elem_type = flat_values.data_type().clone();
292
293        let schema = Arc::new(Schema::new(vec![Field::new("x", elem_type.clone(), true)]));
294        let batch = arrow_array::RecordBatch::try_new(schema, vec![Arc::clone(flat_values)])?;
295
296        let mask_arr = eval_expr_on_batch(&lambda_str, &batch)?;
297        let mask = mask_arr
298            .as_any()
299            .downcast_ref::<BooleanArray>()
300            .ok_or_else(|| {
301                datafusion_common::DataFusionError::Internal(
302                    "array_filter: lambda must return Boolean".into(),
303                )
304            })?;
305
306        let mut offsets = vec![0i32];
307        let mut filtered_indices: Vec<usize> = Vec::new();
308
309        for row in 0..list_arr.len() {
310            let start = list_arr.value_offsets()[row] as usize;
311            let end = list_arr.value_offsets()[row + 1] as usize;
312
313            for i in start..end {
314                if !mask.is_null(i) && mask.value(i) {
315                    filtered_indices.push(i);
316                }
317            }
318            offsets.push(filtered_indices.len() as i32);
319        }
320
321        let indices = arrow_array::UInt32Array::from(
322            filtered_indices
323                .iter()
324                .map(|&i| i as u32)
325                .collect::<Vec<_>>(),
326        );
327        let filtered_values = arrow::compute::take(flat_values.as_ref(), &indices, None)?;
328
329        let new_field = Arc::new(Field::new("item", elem_type, true));
330        let new_offsets =
331            arrow::buffer::OffsetBuffer::new(arrow::buffer::ScalarBuffer::from(offsets));
332        let new_list = ListArray::try_new(
333            new_field,
334            new_offsets,
335            filtered_values,
336            list_arr.nulls().cloned(),
337        )?;
338        Ok(ColumnarValue::Array(Arc::new(new_list)))
339    }
340}
341
342// ══════════════════════════════════════════════════════════════════
343// array_reduce(arr, init, lambda_str) -> scalar
344// ══════════════════════════════════════════════════════════════════
345
346/// `array_reduce(arr, init, '(acc + x)')` — fold/reduce array elements.
347#[derive(Debug)]
348pub struct ArrayReduce {
349    signature: Signature,
350}
351
352impl ArrayReduce {
353    /// Creates a new `array_reduce` UDF.
354    #[must_use]
355    pub fn new() -> Self {
356        Self {
357            signature: Signature::new(TypeSignature::Any(3), Volatility::Immutable),
358        }
359    }
360}
361
362impl Default for ArrayReduce {
363    fn default() -> Self {
364        Self::new()
365    }
366}
367impl PartialEq for ArrayReduce {
368    fn eq(&self, _: &Self) -> bool {
369        true
370    }
371}
372impl Eq for ArrayReduce {}
373impl Hash for ArrayReduce {
374    fn hash<H: Hasher>(&self, s: &mut H) {
375        "array_reduce".hash(s);
376    }
377}
378
379impl ScalarUDFImpl for ArrayReduce {
380    fn as_any(&self) -> &dyn std::any::Any {
381        self
382    }
383
384    fn name(&self) -> &'static str {
385        "array_reduce"
386    }
387    fn signature(&self) -> &Signature {
388        &self.signature
389    }
390
391    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
392        Ok(arg_types.get(1).cloned().unwrap_or(DataType::Int64))
393    }
394
395    #[allow(clippy::cast_sign_loss)]
396    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
397        let expanded = expand_args(&args.args)?;
398        let list_arr = expanded[0]
399            .as_any()
400            .downcast_ref::<ListArray>()
401            .ok_or_else(|| {
402                datafusion_common::DataFusionError::Internal(
403                    "array_reduce: first arg must be List".into(),
404                )
405            })?;
406
407        let init_arr = &expanded[1];
408        let lambda_str = scalar_string_value(&args.args[2])?;
409
410        let elem_type = list_arr.values().data_type().clone();
411        let acc_type = init_arr.data_type().clone();
412
413        let schema = Arc::new(Schema::new(vec![
414            Field::new("acc", acc_type, true),
415            Field::new("x", elem_type, true),
416        ]));
417
418        let mut result_builder: Vec<ArrayRef> = Vec::new();
419
420        for row in 0..list_arr.len() {
421            let start = list_arr.value_offsets()[row] as usize;
422            let end = list_arr.value_offsets()[row + 1] as usize;
423
424            let mut acc: ArrayRef = init_arr.slice(row, 1);
425
426            for i in start..end {
427                let x = list_arr.values().slice(i, 1);
428                let batch = arrow_array::RecordBatch::try_new(
429                    Arc::clone(&schema),
430                    vec![Arc::clone(&acc), x],
431                )?;
432                let result_col = eval_expr_on_batch(&lambda_str, &batch)?;
433                acc = result_col;
434            }
435
436            result_builder.push(acc);
437        }
438
439        if result_builder.is_empty() {
440            return Ok(ColumnarValue::Array(Arc::clone(init_arr)));
441        }
442
443        let refs: Vec<&dyn Array> = result_builder
444            .iter()
445            .map(std::convert::AsRef::as_ref)
446            .collect();
447        let result = arrow::compute::concat(&refs)?;
448        Ok(ColumnarValue::Array(result))
449    }
450}
451
452// ══════════════════════════════════════════════════════════════════
453// map_filter(map, lambda_str) -> Map
454// ══════════════════════════════════════════════════════════════════
455
456/// `map_filter(map, '(k <> ''temp'')')` — filter map entries by key+value.
457#[derive(Debug)]
458pub struct MapFilter {
459    signature: Signature,
460}
461
462impl MapFilter {
463    /// Creates a new `map_filter` UDF.
464    #[must_use]
465    pub fn new() -> Self {
466        Self {
467            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
468        }
469    }
470}
471
472impl Default for MapFilter {
473    fn default() -> Self {
474        Self::new()
475    }
476}
477impl PartialEq for MapFilter {
478    fn eq(&self, _: &Self) -> bool {
479        true
480    }
481}
482impl Eq for MapFilter {}
483impl Hash for MapFilter {
484    fn hash<H: Hasher>(&self, s: &mut H) {
485        "map_filter".hash(s);
486    }
487}
488
489impl ScalarUDFImpl for MapFilter {
490    fn as_any(&self) -> &dyn std::any::Any {
491        self
492    }
493
494    fn name(&self) -> &'static str {
495        "map_filter"
496    }
497    fn signature(&self) -> &Signature {
498        &self.signature
499    }
500
501    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
502        Ok(arg_types[0].clone())
503    }
504
505    #[allow(
506        clippy::cast_sign_loss,
507        clippy::cast_possible_wrap,
508        clippy::cast_possible_truncation
509    )]
510    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
511        let expanded = expand_args(&args.args)?;
512        let map_arr = expanded[0]
513            .as_any()
514            .downcast_ref::<MapArray>()
515            .ok_or_else(|| {
516                datafusion_common::DataFusionError::Internal(
517                    "map_filter: first arg must be Map".into(),
518                )
519            })?;
520
521        let lambda_str = scalar_string_value(&args.args[1])?;
522
523        let entries = map_arr.entries();
524        let key_col = entries.column(0);
525        let val_col = entries.column(1);
526
527        let key_type = key_col.data_type().clone();
528        let val_type = val_col.data_type().clone();
529
530        let schema = Arc::new(Schema::new(vec![
531            Field::new("k", key_type.clone(), true),
532            Field::new("v", val_type.clone(), true),
533        ]));
534        let batch = arrow_array::RecordBatch::try_new(
535            schema,
536            vec![Arc::clone(key_col), Arc::clone(val_col)],
537        )?;
538
539        let mask_arr = eval_expr_on_batch(&lambda_str, &batch)?;
540        let mask_bool = mask_arr
541            .as_any()
542            .downcast_ref::<BooleanArray>()
543            .ok_or_else(|| {
544                datafusion_common::DataFusionError::Internal(
545                    "map_filter: lambda must return Boolean".into(),
546                )
547            })?;
548
549        let mut offsets = vec![0i32];
550        let mut keep_indices: Vec<usize> = Vec::new();
551
552        for row in 0..map_arr.len() {
553            let start = map_arr.value_offsets()[row] as usize;
554            let end = map_arr.value_offsets()[row + 1] as usize;
555
556            for i in start..end {
557                if !mask_bool.is_null(i) && mask_bool.value(i) {
558                    keep_indices.push(i);
559                }
560            }
561            offsets.push(keep_indices.len() as i32);
562        }
563
564        let indices = arrow_array::UInt32Array::from(
565            keep_indices.iter().map(|&i| i as u32).collect::<Vec<_>>(),
566        );
567        let new_keys = arrow::compute::take(key_col.as_ref(), &indices, None)?;
568        let new_vals = arrow::compute::take(val_col.as_ref(), &indices, None)?;
569
570        let struct_fields = Fields::from(vec![
571            Field::new("key", key_type, false),
572            Field::new("value", val_type, true),
573        ]);
574        let new_entries = StructArray::try_new(struct_fields, vec![new_keys, new_vals], None)?;
575
576        let entries_field = Field::new("entries", new_entries.data_type().clone(), false);
577        let new_offsets =
578            arrow::buffer::OffsetBuffer::new(arrow::buffer::ScalarBuffer::from(offsets));
579        let new_map = MapArray::try_new(
580            Arc::new(entries_field),
581            new_offsets,
582            new_entries,
583            map_arr.nulls().cloned(),
584            false,
585        )?;
586        Ok(ColumnarValue::Array(Arc::new(new_map)))
587    }
588}
589
590// ══════════════════════════════════════════════════════════════════
591// map_transform_values(map, lambda_str) -> Map
592// ══════════════════════════════════════════════════════════════════
593
594/// `map_transform_values(map, 'v * 2')` — transform map values.
595#[derive(Debug)]
596pub struct MapTransformValues {
597    signature: Signature,
598}
599
600impl MapTransformValues {
601    /// Creates a new `map_transform_values` UDF.
602    #[must_use]
603    pub fn new() -> Self {
604        Self {
605            signature: Signature::new(TypeSignature::Any(2), Volatility::Immutable),
606        }
607    }
608}
609
610impl Default for MapTransformValues {
611    fn default() -> Self {
612        Self::new()
613    }
614}
615impl PartialEq for MapTransformValues {
616    fn eq(&self, _: &Self) -> bool {
617        true
618    }
619}
620impl Eq for MapTransformValues {}
621impl Hash for MapTransformValues {
622    fn hash<H: Hasher>(&self, s: &mut H) {
623        "map_transform_values".hash(s);
624    }
625}
626
627impl ScalarUDFImpl for MapTransformValues {
628    fn as_any(&self) -> &dyn std::any::Any {
629        self
630    }
631
632    fn name(&self) -> &'static str {
633        "map_transform_values"
634    }
635    fn signature(&self) -> &Signature {
636        &self.signature
637    }
638
639    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
640        Ok(arg_types[0].clone())
641    }
642
643    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
644        let expanded = expand_args(&args.args)?;
645        let map_arr = expanded[0]
646            .as_any()
647            .downcast_ref::<MapArray>()
648            .ok_or_else(|| {
649                datafusion_common::DataFusionError::Internal(
650                    "map_transform_values: first arg must be Map".into(),
651                )
652            })?;
653
654        let lambda_str = scalar_string_value(&args.args[1])?;
655
656        let entries = map_arr.entries();
657        let key_col = entries.column(0);
658        let val_col = entries.column(1);
659
660        let key_type = key_col.data_type().clone();
661
662        let schema = Arc::new(Schema::new(vec![
663            Field::new("k", key_type.clone(), true),
664            Field::new("v", val_col.data_type().clone(), true),
665        ]));
666        let batch = arrow_array::RecordBatch::try_new(
667            schema,
668            vec![Arc::clone(key_col), Arc::clone(val_col)],
669        )?;
670
671        let new_vals = eval_expr_on_batch(&lambda_str, &batch)?;
672
673        let struct_fields = Fields::from(vec![
674            Field::new("key", key_type, false),
675            Field::new("value", new_vals.data_type().clone(), true),
676        ]);
677        let new_entries =
678            StructArray::try_new(struct_fields, vec![Arc::clone(key_col), new_vals], None)?;
679
680        let entries_field = Field::new("entries", new_entries.data_type().clone(), false);
681        let new_map = MapArray::try_new(
682            Arc::new(entries_field),
683            map_arr.offsets().clone(),
684            new_entries,
685            map_arr.nulls().cloned(),
686            false,
687        )?;
688        Ok(ColumnarValue::Array(Arc::new(new_map)))
689    }
690}
691
692#[cfg(test)]
693mod tests {
694    use super::*;
695    use crate::datafusion::create_session_context;
696    use arrow_array::*;
697    use datafusion_common::config::ConfigOptions;
698
699    // ── array_transform ─────────────────────────────────────────
700
701    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
702    async fn test_array_transform_add_one() {
703        let values = Int64Array::from(vec![1, 2, 3, 4, 5, 6]);
704        let offsets =
705            arrow::buffer::OffsetBuffer::new(arrow::buffer::ScalarBuffer::from(vec![0i32, 3, 6]));
706        let list = ListArray::try_new(
707            Arc::new(Field::new("item", DataType::Int64, true)),
708            offsets,
709            Arc::new(values),
710            None,
711        )
712        .unwrap();
713
714        let udf = ArrayTransform::new();
715        let result = udf
716            .invoke_with_args(ScalarFunctionArgs {
717                args: vec![
718                    ColumnarValue::Array(Arc::new(list)),
719                    ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(
720                        "x + 1".into(),
721                    ))),
722                ],
723                number_rows: 0,
724                arg_fields: vec![],
725                return_field: Arc::new(Field::new(
726                    "output",
727                    DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
728                    true,
729                )),
730                config_options: Arc::new(ConfigOptions::default()),
731            })
732            .unwrap();
733
734        if let ColumnarValue::Array(arr) = result {
735            let la = arr.as_any().downcast_ref::<ListArray>().unwrap();
736            assert_eq!(la.len(), 2);
737            let row0 = la.value(0);
738            let r0 = row0.as_any().downcast_ref::<Int64Array>().unwrap();
739            assert_eq!(r0.value(0), 2);
740            assert_eq!(r0.value(1), 3);
741            assert_eq!(r0.value(2), 4);
742        } else {
743            panic!("expected Array");
744        }
745    }
746
747    // ── array_filter ────────────────────────────────────────────
748
749    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
750    async fn test_array_filter_positive() {
751        let values = Int64Array::from(vec![-1, 2, -3, 4]);
752        let offsets =
753            arrow::buffer::OffsetBuffer::new(arrow::buffer::ScalarBuffer::from(vec![0i32, 4]));
754        let list = ListArray::try_new(
755            Arc::new(Field::new("item", DataType::Int64, true)),
756            offsets,
757            Arc::new(values),
758            None,
759        )
760        .unwrap();
761
762        let udf = ArrayFilter::new();
763        let result = udf
764            .invoke_with_args(ScalarFunctionArgs {
765                args: vec![
766                    ColumnarValue::Array(Arc::new(list)),
767                    ColumnarValue::Scalar(datafusion_common::ScalarValue::Utf8(Some(
768                        "x > 0".into(),
769                    ))),
770                ],
771                number_rows: 0,
772                arg_fields: vec![],
773                return_field: Arc::new(Field::new(
774                    "output",
775                    DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
776                    true,
777                )),
778                config_options: Arc::new(ConfigOptions::default()),
779            })
780            .unwrap();
781
782        if let ColumnarValue::Array(arr) = result {
783            let la = arr.as_any().downcast_ref::<ListArray>().unwrap();
784            let row0 = la.value(0);
785            let r0 = row0.as_any().downcast_ref::<Int64Array>().unwrap();
786            assert_eq!(r0.len(), 2);
787            assert_eq!(r0.value(0), 2);
788            assert_eq!(r0.value(1), 4);
789        }
790    }
791
792    // ── Registration ────────────────────────────────────────────
793
794    #[test]
795    fn test_register_lambda_functions() {
796        use datafusion::execution::FunctionRegistry;
797
798        let ctx = create_session_context();
799        register_lambda_functions(&ctx);
800        assert!(ctx.udf("array_transform").is_ok());
801        assert!(ctx.udf("array_filter").is_ok());
802        assert!(ctx.udf("array_reduce").is_ok());
803        assert!(ctx.udf("map_filter").is_ok());
804        assert!(ctx.udf("map_transform_values").is_ok());
805    }
806}
laminar_sql/datafusion/complex_type_lambda.rs

laminar_sql/datafusion/
complex_type_lambda.rs
