lance_datafusion/udf/
json.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use arrow_array::builder::{
5    BooleanBuilder, Float64Builder, Int64Builder, LargeBinaryBuilder, StringBuilder,
6};
7use arrow_array::{Array, ArrayRef, LargeBinaryArray, StringArray};
8use arrow_schema::DataType;
9use datafusion::error::{DataFusionError, Result};
10use datafusion::logical_expr::{ScalarUDF, Volatility};
11use datafusion::physical_plan::ColumnarValue;
12use datafusion::prelude::create_udf;
13use lance_arrow::json::{json_path_exists, select_json_path, select_json_path_values};
14use std::sync::Arc;
15
16/// Represents the type of a JSONB value
17#[repr(u8)]
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum JsonbType {
20    Null = 0,
21    Boolean = 1,
22    Int64 = 2,
23    Float64 = 3,
24    String = 4,
25    Array = 5,
26    Object = 6,
27}
28
29impl JsonbType {
30    /// Convert from u8 value
31    pub fn from_u8(value: u8) -> Option<Self> {
32        match value {
33            0 => Some(Self::Null),
34            1 => Some(Self::Boolean),
35            2 => Some(Self::Int64),
36            3 => Some(Self::Float64),
37            4 => Some(Self::String),
38            5 => Some(Self::Array),
39            6 => Some(Self::Object),
40            _ => None,
41        }
42    }
43
44    /// Convert to u8 value for storage in Arrow arrays
45    pub fn as_u8(self) -> u8 {
46        self as u8
47    }
48}
49
50/// Common helper functions and types for JSON UDFs
51mod common {
52    use super::*;
53
54    /// Convert ColumnarValue arguments to ArrayRef vector
55    ///
56    /// Note: This implementation currently broadcasts scalars to arrays.
57    /// Future optimization: handle scalars directly without broadcasting
58    /// to improve performance for scalar inputs.
59    pub fn columnar_to_arrays(args: &[ColumnarValue]) -> Vec<ArrayRef> {
60        args.iter()
61            .map(|arg| match arg {
62                ColumnarValue::Array(arr) => arr.clone(),
63                ColumnarValue::Scalar(scalar) => scalar.to_array().unwrap(),
64            })
65            .collect()
66    }
67
68    /// Create DataFusionError for execution failures (simplified error wrapping)
69    pub fn execution_error(msg: impl Into<String>) -> DataFusionError {
70        DataFusionError::Execution(msg.into())
71    }
72
73    /// Validate argument count for UDF
74    pub fn validate_arg_count(
75        args: &[ArrayRef],
76        expected: usize,
77        function_name: &str,
78    ) -> Result<()> {
79        if args.len() != expected {
80            return Err(execution_error(format!(
81                "{} requires exactly {} arguments",
82                function_name, expected
83            )));
84        }
85        Ok(())
86    }
87
88    /// Extract and validate LargeBinaryArray from first argument
89    pub fn extract_jsonb_array(args: &[ArrayRef]) -> Result<&LargeBinaryArray> {
90        args[0]
91            .as_any()
92            .downcast_ref::<LargeBinaryArray>()
93            .ok_or_else(|| execution_error("First argument must be LargeBinary"))
94    }
95
96    /// Extract and validate StringArray from specified argument
97    pub fn extract_string_array(args: &[ArrayRef], arg_index: usize) -> Result<&StringArray> {
98        args[arg_index]
99            .as_any()
100            .downcast_ref::<StringArray>()
101            .ok_or_else(|| execution_error(format!("Argument {} must be String", arg_index + 1)))
102    }
103
104    /// Get string value at index, handling scalar broadcast case
105    /// When a scalar is converted to an array, it becomes a single-element array
106    /// This function handles accessing that value repeatedly for all rows
107    pub fn get_string_value_at(string_array: &StringArray, index: usize) -> Option<&str> {
108        // Handle scalar broadcast case: if array has only 1 element, always use index 0
109        let actual_index = if string_array.len() == 1 { 0 } else { index };
110
111        if string_array.is_null(actual_index) {
112            None
113        } else {
114            Some(string_array.value(actual_index))
115        }
116    }
117
118    /// Get a JSON field or array element by key.
119    pub fn get_json_value_by_key(
120        raw_jsonb: &lance_jsonb::RawJsonb,
121        key: &str,
122    ) -> Result<Option<lance_jsonb::OwnedJsonb>> {
123        if raw_jsonb.is_object().unwrap_or(false) {
124            raw_jsonb
125                .get_by_name(key, false)
126                .map_err(|e| execution_error(format!("Failed to get field '{}': {}", key, e)))
127        } else if raw_jsonb.is_array().unwrap_or(false) {
128            match key.parse::<usize>() {
129                Ok(index) => raw_jsonb.get_by_index(index).map_err(|e| {
130                    execution_error(format!("Failed to get array element [{}]: {}", index, e))
131                }),
132                Err(_) => Ok(None),
133            }
134        } else {
135            Ok(None)
136        }
137    }
138
139    /// Parse JSONPath with proper error handling (no false returns)
140    pub fn parse_json_path(path: &str) -> Result<lance_jsonb::jsonpath::JsonPath<'_>> {
141        lance_arrow::json::parse_json_path(path).map_err(|e| execution_error(e.to_string()))
142    }
143}
144
145/// Convert JSONB value to string using jsonb's built-in serde (strict mode)
146fn json_value_to_string(value: lance_jsonb::OwnedJsonb) -> Result<Option<String>> {
147    let raw_jsonb = value.as_raw();
148
149    // Check for null first
150    if raw_jsonb
151        .is_null()
152        .map_err(|e| common::execution_error(format!("Failed to check null: {}", e)))?
153    {
154        return Ok(None);
155    }
156
157    // Use jsonb's built-in to_str() method - strict conversion
158    raw_jsonb
159        .to_str()
160        .map(Some)
161        .map_err(|e| common::execution_error(format!("Failed to convert to string: {}", e)))
162}
163
164/// Convert JSONB value to integer using jsonb's built-in serde (strict mode)
165fn json_value_to_int(value: lance_jsonb::OwnedJsonb) -> Result<Option<i64>> {
166    let raw_jsonb = value.as_raw();
167
168    // Check for null first
169    if raw_jsonb
170        .is_null()
171        .map_err(|e| common::execution_error(format!("Failed to check null: {}", e)))?
172    {
173        return Ok(None);
174    }
175
176    // Use jsonb's built-in to_i64() method - strict conversion
177    raw_jsonb
178        .to_i64()
179        .map(Some)
180        .map_err(|e| common::execution_error(format!("Failed to convert to integer: {}", e)))
181}
182
183/// Convert JSONB value to float using jsonb's built-in serde (strict mode)
184fn json_value_to_float(value: lance_jsonb::OwnedJsonb) -> Result<Option<f64>> {
185    let raw_jsonb = value.as_raw();
186
187    // Check for null first
188    if raw_jsonb
189        .is_null()
190        .map_err(|e| common::execution_error(format!("Failed to check null: {}", e)))?
191    {
192        return Ok(None);
193    }
194
195    // Use jsonb's built-in to_f64() method - strict conversion
196    raw_jsonb
197        .to_f64()
198        .map(Some)
199        .map_err(|e| common::execution_error(format!("Failed to convert to float: {}", e)))
200}
201
202/// Convert JSONB value to boolean using jsonb's built-in serde (strict mode)
203fn json_value_to_bool(value: lance_jsonb::OwnedJsonb) -> Result<Option<bool>> {
204    let raw_jsonb = value.as_raw();
205
206    // Check for null first
207    if raw_jsonb
208        .is_null()
209        .map_err(|e| common::execution_error(format!("Failed to check null: {}", e)))?
210    {
211        return Ok(None);
212    }
213
214    // Use jsonb's built-in to_bool() method - strict conversion
215    raw_jsonb
216        .to_bool()
217        .map(Some)
218        .map_err(|e| common::execution_error(format!("Failed to convert to boolean: {}", e)))
219}
220
221/// Create the json_extract UDF for extracting JSONPath from JSON data
222///
223/// # Arguments
224/// * First parameter: JSONB binary data (LargeBinary)
225/// * Second parameter: JSONPath expression as string (Utf8)
226///
227/// # Returns
228/// String representation of the extracted value, or null if path not found
229pub fn json_extract_udf() -> ScalarUDF {
230    create_udf(
231        "json_extract",
232        vec![DataType::LargeBinary, DataType::Utf8],
233        DataType::Utf8,
234        Volatility::Immutable,
235        Arc::new(json_extract_columnar_impl),
236    )
237}
238
239/// Create the json_extract_with_type UDF that returns JSONB bytes with type information
240///
241/// # Arguments
242/// * First parameter: JSONB binary data (LargeBinary)
243/// * Second parameter: JSONPath expression as string (Utf8)
244///
245/// # Returns
246/// A struct with two fields:
247/// - value: LargeBinary (the extracted JSONB value)
248/// - type_tag: UInt8 (type information: 0=null, 1=bool, 2=int64, 3=float64, 4=string, 5=array, 6=object)
249pub fn json_extract_with_type_udf() -> ScalarUDF {
250    use arrow_schema::Fields;
251
252    let return_type = DataType::Struct(Fields::from(vec![
253        arrow_schema::Field::new("value", DataType::LargeBinary, true),
254        arrow_schema::Field::new("type_tag", DataType::UInt8, false),
255    ]));
256
257    create_udf(
258        "json_extract_with_type",
259        vec![DataType::LargeBinary, DataType::Utf8],
260        return_type,
261        Volatility::Immutable,
262        Arc::new(json_extract_with_type_columnar_impl),
263    )
264}
265
266/// Implementation of json_extract function with ColumnarValue
267fn json_extract_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
268    let arrays = common::columnar_to_arrays(args);
269    let result = json_extract_impl(&arrays)?;
270    Ok(ColumnarValue::Array(result))
271}
272
273/// Implementation of json_extract_with_type function with ColumnarValue
274fn json_extract_with_type_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
275    let arrays = common::columnar_to_arrays(args);
276    let result = json_extract_with_type_impl(&arrays)?;
277    Ok(ColumnarValue::Array(result))
278}
279
280/// Implementation of json_extract function
281fn json_extract_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
282    common::validate_arg_count(args, 2, "json_extract")?;
283
284    let jsonb_array = common::extract_jsonb_array(args)?;
285    let path_array = common::extract_string_array(args, 1)?;
286    let mut builder = StringBuilder::with_capacity(jsonb_array.len(), 1024);
287
288    for i in 0..jsonb_array.len() {
289        if jsonb_array.is_null(i) {
290            builder.append_null();
291        } else if let Some(path) = common::get_string_value_at(path_array, i) {
292            let jsonb_bytes = jsonb_array.value(i);
293            match extract_json_path(jsonb_bytes, path)? {
294                Some(value) => builder.append_value(&value),
295                None => builder.append_null(),
296            }
297        } else {
298            builder.append_null();
299        }
300    }
301
302    Ok(Arc::new(builder.finish()))
303}
304
305/// Implementation of json_extract_with_type function
306fn json_extract_with_type_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
307    use arrow_array::StructArray;
308    use arrow_array::builder::{LargeBinaryBuilder, UInt8Builder};
309
310    common::validate_arg_count(args, 2, "json_extract_with_type")?;
311
312    let jsonb_array = common::extract_jsonb_array(args)?;
313    let path_array = common::extract_string_array(args, 1)?;
314
315    let mut value_builder = LargeBinaryBuilder::with_capacity(jsonb_array.len(), 1024);
316    let mut type_builder = UInt8Builder::with_capacity(jsonb_array.len());
317
318    for i in 0..jsonb_array.len() {
319        if jsonb_array.is_null(i) {
320            value_builder.append_null();
321            type_builder.append_value(JsonbType::Null.as_u8());
322        } else if let Some(path) = common::get_string_value_at(path_array, i) {
323            let jsonb_bytes = jsonb_array.value(i);
324            match extract_json_path_with_type(jsonb_bytes, path)? {
325                Some((value_bytes, type_tag)) => {
326                    value_builder.append_value(&value_bytes);
327                    type_builder.append_value(type_tag);
328                }
329                None => {
330                    value_builder.append_null();
331                    type_builder.append_value(JsonbType::Null.as_u8());
332                }
333            }
334        } else {
335            value_builder.append_null();
336            type_builder.append_value(JsonbType::Null.as_u8());
337        }
338    }
339
340    // Create struct array with two fields
341    let value_array = Arc::new(value_builder.finish()) as ArrayRef;
342    let type_array = Arc::new(type_builder.finish()) as ArrayRef;
343
344    let struct_array = StructArray::from(vec![
345        (
346            Arc::new(arrow_schema::Field::new(
347                "value",
348                DataType::LargeBinary,
349                true,
350            )),
351            value_array,
352        ),
353        (
354            Arc::new(arrow_schema::Field::new("type_tag", DataType::UInt8, false)),
355            type_array,
356        ),
357    ]);
358
359    Ok(Arc::new(struct_array))
360}
361
362/// Extract value from JSONB using JSONPath and return with type information
363/// Returns (JSONB bytes, type_tag) where type_tag represents the JsonbType
364fn extract_json_path_with_type(jsonb_bytes: &[u8], path: &str) -> Result<Option<(Vec<u8>, u8)>> {
365    let json_path = common::parse_json_path(path)?;
366
367    match select_json_path(jsonb_bytes, &json_path) {
368        Ok(Some(owned_value)) => {
369            let raw = owned_value.as_raw();
370
371            // Determine type using JsonbType enum
372            let jsonb_type = if raw.is_null().unwrap_or(false) {
373                JsonbType::Null
374            } else if raw.is_boolean().unwrap_or(false) {
375                JsonbType::Boolean
376            } else if raw.is_number().unwrap_or(false) {
377                let is_float_storage =
378                    matches!(raw.as_number(), Ok(Some(lance_jsonb::Number::Float64(_))));
379                if !is_float_storage && raw.is_i64().unwrap_or(false) {
380                    JsonbType::Int64
381                } else {
382                    JsonbType::Float64
383                }
384            } else if raw.is_string().unwrap_or(false) {
385                JsonbType::String
386            } else if raw.is_array().unwrap_or(false) {
387                JsonbType::Array
388            } else if raw.is_object().unwrap_or(false) {
389                JsonbType::Object
390            } else {
391                JsonbType::String // default to string
392            };
393
394            // Return the JSONB bytes and type tag as u8
395            Ok(Some((owned_value.to_vec(), jsonb_type.as_u8())))
396        }
397        Ok(None) => Ok(None),
398        Err(e) => Err(common::execution_error(format!(
399            "Failed to select value from path '{}': {}",
400            path, e
401        ))),
402    }
403}
404
405/// Extract value from JSONB using JSONPath
406///
407/// Note: Uses `select_value` so JSONPath expressions matching multiple values
408/// return a JSON array instead of silently dropping all but the first match.
409fn extract_json_path(jsonb_bytes: &[u8], path: &str) -> Result<Option<String>> {
410    let json_path = common::parse_json_path(path)?;
411
412    match select_json_path(jsonb_bytes, &json_path) {
413        Ok(value) => Ok(value.map(|value| value.to_string())),
414        Err(e) => Err(common::execution_error(format!(
415            "Failed to select value from path '{}': {}",
416            path, e
417        ))),
418    }
419}
420
421/// Create the json_exists UDF for checking if a JSONPath exists
422///
423/// # Arguments
424/// * First parameter: JSONB binary data (LargeBinary)
425/// * Second parameter: JSONPath expression as string (Utf8)
426///
427/// # Returns
428/// Boolean indicating whether the path exists in the JSON data
429pub fn json_exists_udf() -> ScalarUDF {
430    create_udf(
431        "json_exists",
432        vec![DataType::LargeBinary, DataType::Utf8],
433        DataType::Boolean,
434        Volatility::Immutable,
435        Arc::new(json_exists_columnar_impl),
436    )
437}
438
439/// Implementation of json_exists function with ColumnarValue
440fn json_exists_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
441    let arrays = common::columnar_to_arrays(args);
442    let result = json_exists_impl(&arrays)?;
443    Ok(ColumnarValue::Array(result))
444}
445
446/// Implementation of json_exists function
447fn json_exists_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
448    common::validate_arg_count(args, 2, "json_exists")?;
449
450    let jsonb_array = common::extract_jsonb_array(args)?;
451    let path_array = common::extract_string_array(args, 1)?;
452
453    let mut builder = BooleanBuilder::with_capacity(jsonb_array.len());
454
455    for i in 0..jsonb_array.len() {
456        if jsonb_array.is_null(i) {
457            builder.append_null();
458        } else if let Some(path) = common::get_string_value_at(path_array, i) {
459            let jsonb_bytes = jsonb_array.value(i);
460            let exists = check_json_path_exists(jsonb_bytes, path)?;
461            builder.append_value(exists);
462        } else {
463            builder.append_null();
464        }
465    }
466
467    Ok(Arc::new(builder.finish()))
468}
469
470/// Check if a JSONPath exists in JSONB
471fn check_json_path_exists(jsonb_bytes: &[u8], path: &str) -> Result<bool> {
472    let json_path = common::parse_json_path(path)?;
473
474    match json_path_exists(jsonb_bytes, &json_path) {
475        Ok(exists) => Ok(exists),
476        Err(e) => Err(common::execution_error(format!(
477            "Failed to check existence of path '{}': {}",
478            path, e
479        ))),
480    }
481}
482
483/// Create the json_get UDF for getting a field value as JSON string
484///
485/// # Arguments
486/// * First parameter: JSONB binary data (LargeBinary)
487/// * Second parameter: Field name or array index as string (Utf8)
488///
489/// # Returns
490/// Raw JSONB bytes of the field value, or null if not found
491pub fn json_get_udf() -> ScalarUDF {
492    create_udf(
493        "json_get",
494        vec![DataType::LargeBinary, DataType::Utf8],
495        DataType::LargeBinary,
496        Volatility::Immutable,
497        Arc::new(json_get_columnar_impl),
498    )
499}
500
501/// Implementation of json_get function with ColumnarValue
502fn json_get_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
503    let arrays = common::columnar_to_arrays(args);
504    let result = json_get_impl(&arrays)?;
505    Ok(ColumnarValue::Array(result))
506}
507
508/// Implementation of json_get function
509fn json_get_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
510    common::validate_arg_count(args, 2, "json_get")?;
511
512    let jsonb_array = common::extract_jsonb_array(args)?;
513    let key_array = common::extract_string_array(args, 1)?;
514
515    let mut builder = LargeBinaryBuilder::with_capacity(jsonb_array.len(), 0);
516
517    for i in 0..jsonb_array.len() {
518        if jsonb_array.is_null(i) {
519            builder.append_null();
520        } else if let Some(key) = common::get_string_value_at(key_array, i) {
521            let jsonb_bytes = jsonb_array.value(i);
522            let raw_jsonb = lance_jsonb::RawJsonb::new(jsonb_bytes);
523
524            match common::get_json_value_by_key(&raw_jsonb, key)? {
525                Some(value) => builder.append_value(value.as_raw().as_ref()),
526                None => builder.append_null(),
527            }
528        } else {
529            builder.append_null();
530        }
531    }
532
533    Ok(Arc::new(builder.finish()))
534}
535
536/// Create the json_get_string UDF for getting a string value
537///
538/// # Arguments
539/// * First parameter: JSONB binary data (LargeBinary)
540/// * Second parameter: Field name or array index as string (Utf8)
541///
542/// # Returns
543/// String value with type coercion (numbers/booleans converted to strings)
544pub fn json_get_string_udf() -> ScalarUDF {
545    create_udf(
546        "json_get_string",
547        vec![DataType::LargeBinary, DataType::Utf8],
548        DataType::Utf8,
549        Volatility::Immutable,
550        Arc::new(json_get_string_columnar_impl),
551    )
552}
553
554/// Implementation of json_get_string function with ColumnarValue
555fn json_get_string_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
556    let arrays = common::columnar_to_arrays(args);
557    let result = json_get_string_impl(&arrays)?;
558    Ok(ColumnarValue::Array(result))
559}
560
561/// Implementation of json_get_string function
562fn json_get_string_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
563    common::validate_arg_count(args, 2, "json_get_string")?;
564
565    let jsonb_array = common::extract_jsonb_array(args)?;
566    let key_array = common::extract_string_array(args, 1)?;
567
568    let mut builder = StringBuilder::with_capacity(jsonb_array.len(), 1024);
569
570    for i in 0..jsonb_array.len() {
571        if jsonb_array.is_null(i) {
572            builder.append_null();
573        } else if let Some(key) = common::get_string_value_at(key_array, i) {
574            let jsonb_bytes = jsonb_array.value(i);
575            let raw_jsonb = lance_jsonb::RawJsonb::new(jsonb_bytes);
576
577            match common::get_json_value_by_key(&raw_jsonb, key)? {
578                Some(value) => match json_value_to_string(value)? {
579                    Some(string_val) => builder.append_value(&string_val),
580                    None => builder.append_null(),
581                },
582                None => builder.append_null(),
583            }
584        } else {
585            builder.append_null();
586        }
587    }
588
589    Ok(Arc::new(builder.finish()))
590}
591
592/// Create the json_get_int UDF for getting an integer value
593///
594/// # Arguments
595/// * First parameter: JSONB binary data (LargeBinary)
596/// * Second parameter: Field name or array index as string (Utf8)
597///
598/// # Returns
599/// Integer value with type coercion (strings/floats/booleans converted to int)
600pub fn json_get_int_udf() -> ScalarUDF {
601    create_udf(
602        "json_get_int",
603        vec![DataType::LargeBinary, DataType::Utf8],
604        DataType::Int64,
605        Volatility::Immutable,
606        Arc::new(json_get_int_columnar_impl),
607    )
608}
609
610/// Implementation of json_get_int function with ColumnarValue
611fn json_get_int_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
612    let arrays = common::columnar_to_arrays(args);
613    let result = json_get_int_impl(&arrays)?;
614    Ok(ColumnarValue::Array(result))
615}
616
617/// Implementation of json_get_int function
618fn json_get_int_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
619    common::validate_arg_count(args, 2, "json_get_int")?;
620
621    let jsonb_array = common::extract_jsonb_array(args)?;
622    let key_array = common::extract_string_array(args, 1)?;
623
624    let mut builder = Int64Builder::with_capacity(jsonb_array.len());
625
626    for i in 0..jsonb_array.len() {
627        if jsonb_array.is_null(i) {
628            builder.append_null();
629        } else if let Some(key) = common::get_string_value_at(key_array, i) {
630            let jsonb_bytes = jsonb_array.value(i);
631            let raw_jsonb = lance_jsonb::RawJsonb::new(jsonb_bytes);
632
633            match common::get_json_value_by_key(&raw_jsonb, key)? {
634                Some(value) => match json_value_to_int(value)? {
635                    Some(int_val) => builder.append_value(int_val),
636                    None => builder.append_null(),
637                },
638                None => builder.append_null(),
639            }
640        } else {
641            builder.append_null();
642        }
643    }
644
645    Ok(Arc::new(builder.finish()))
646}
647
648/// Create the json_get_float UDF for getting a float value
649///
650/// # Arguments
651/// * First parameter: JSONB binary data (LargeBinary)
652/// * Second parameter: Field name or array index as string (Utf8)
653///
654/// # Returns
655/// Float value with type coercion (strings/integers/booleans converted to float)
656pub fn json_get_float_udf() -> ScalarUDF {
657    create_udf(
658        "json_get_float",
659        vec![DataType::LargeBinary, DataType::Utf8],
660        DataType::Float64,
661        Volatility::Immutable,
662        Arc::new(json_get_float_columnar_impl),
663    )
664}
665
666/// Implementation of json_get_float function with ColumnarValue
667fn json_get_float_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
668    let arrays = common::columnar_to_arrays(args);
669    let result = json_get_float_impl(&arrays)?;
670    Ok(ColumnarValue::Array(result))
671}
672
673/// Implementation of json_get_float function
674fn json_get_float_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
675    common::validate_arg_count(args, 2, "json_get_float")?;
676
677    let jsonb_array = common::extract_jsonb_array(args)?;
678    let key_array = common::extract_string_array(args, 1)?;
679
680    let mut builder = Float64Builder::with_capacity(jsonb_array.len());
681
682    for i in 0..jsonb_array.len() {
683        if jsonb_array.is_null(i) {
684            builder.append_null();
685        } else if let Some(key) = common::get_string_value_at(key_array, i) {
686            let jsonb_bytes = jsonb_array.value(i);
687            let raw_jsonb = lance_jsonb::RawJsonb::new(jsonb_bytes);
688
689            match common::get_json_value_by_key(&raw_jsonb, key)? {
690                Some(value) => match json_value_to_float(value)? {
691                    Some(float_val) => builder.append_value(float_val),
692                    None => builder.append_null(),
693                },
694                None => builder.append_null(),
695            }
696        } else {
697            builder.append_null();
698        }
699    }
700
701    Ok(Arc::new(builder.finish()))
702}
703
704/// Create the json_get_bool UDF for getting a boolean value
705///
706/// # Arguments
707/// * First parameter: JSONB binary data (LargeBinary)
708/// * Second parameter: Field name or array index as string (Utf8)
709///
710/// # Returns
711/// Boolean value with flexible type coercion (strings like 'true'/'yes'/'1' become true)
712pub fn json_get_bool_udf() -> ScalarUDF {
713    create_udf(
714        "json_get_bool",
715        vec![DataType::LargeBinary, DataType::Utf8],
716        DataType::Boolean,
717        Volatility::Immutable,
718        Arc::new(json_get_bool_columnar_impl),
719    )
720}
721
722/// Implementation of json_get_bool function with ColumnarValue
723fn json_get_bool_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
724    let arrays = common::columnar_to_arrays(args);
725    let result = json_get_bool_impl(&arrays)?;
726    Ok(ColumnarValue::Array(result))
727}
728
729/// Implementation of json_get_bool function
730fn json_get_bool_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
731    common::validate_arg_count(args, 2, "json_get_bool")?;
732
733    let jsonb_array = common::extract_jsonb_array(args)?;
734    let key_array = common::extract_string_array(args, 1)?;
735
736    let mut builder = BooleanBuilder::with_capacity(jsonb_array.len());
737
738    for i in 0..jsonb_array.len() {
739        if jsonb_array.is_null(i) {
740            builder.append_null();
741        } else if let Some(key) = common::get_string_value_at(key_array, i) {
742            let jsonb_bytes = jsonb_array.value(i);
743            let raw_jsonb = lance_jsonb::RawJsonb::new(jsonb_bytes);
744
745            match common::get_json_value_by_key(&raw_jsonb, key)? {
746                Some(value) => match json_value_to_bool(value)? {
747                    Some(bool_val) => builder.append_value(bool_val),
748                    None => builder.append_null(),
749                },
750                None => builder.append_null(),
751            }
752        } else {
753            builder.append_null();
754        }
755    }
756
757    Ok(Arc::new(builder.finish()))
758}
759
760/// Create the json_array_contains UDF for checking if array contains a value
761///
762/// # Arguments
763/// * First parameter: JSONB binary data (LargeBinary)
764/// * Second parameter: JSONPath to array location (Utf8)
765/// * Third parameter: Value to search for as string (Utf8)
766///
767/// # Returns
768/// Boolean indicating whether the array contains the specified value
769pub fn json_array_contains_udf() -> ScalarUDF {
770    create_udf(
771        "json_array_contains",
772        vec![DataType::LargeBinary, DataType::Utf8, DataType::Utf8],
773        DataType::Boolean,
774        Volatility::Immutable,
775        Arc::new(json_array_contains_columnar_impl),
776    )
777}
778
779/// Implementation of json_array_contains function with ColumnarValue
780fn json_array_contains_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
781    let arrays = common::columnar_to_arrays(args);
782    let result = json_array_contains_impl(&arrays)?;
783    Ok(ColumnarValue::Array(result))
784}
785
786/// Implementation of json_array_contains function
787fn json_array_contains_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
788    common::validate_arg_count(args, 3, "json_array_contains")?;
789
790    let jsonb_array = common::extract_jsonb_array(args)?;
791    let path_array = common::extract_string_array(args, 1)?;
792    let value_array = common::extract_string_array(args, 2)?;
793
794    let mut builder = BooleanBuilder::with_capacity(jsonb_array.len());
795
796    for i in 0..jsonb_array.len() {
797        if jsonb_array.is_null(i) {
798            builder.append_null();
799        } else {
800            let path = common::get_string_value_at(path_array, i);
801            let value = common::get_string_value_at(value_array, i);
802
803            match (path, value) {
804                (Some(p), Some(v)) => {
805                    let jsonb_bytes = jsonb_array.value(i);
806                    let contains = check_array_contains(jsonb_bytes, p, v)?;
807                    builder.append_value(contains);
808                }
809                _ => builder.append_null(),
810            }
811        }
812    }
813
814    Ok(Arc::new(builder.finish()))
815}
816
817/// Check if a JSON array at path contains a value
818fn check_array_contains(jsonb_bytes: &[u8], path: &str, value: &str) -> Result<bool> {
819    let json_path = common::parse_json_path(path)?;
820
821    match select_json_path_values(jsonb_bytes, &json_path) {
822        Ok(values) => {
823            for v in values {
824                // Convert to raw JSONB for direct access
825                let raw = v.as_raw();
826                // Check if it's an array by trying to iterate
827                let mut index = 0;
828                loop {
829                    match raw.get_by_index(index) {
830                        Ok(Some(elem)) => {
831                            let elem_str = elem.to_string();
832                            // Compare as JSON strings (with quotes for strings)
833                            if elem_str == value || elem_str == format!("\"{}\"", value) {
834                                return Ok(true);
835                            }
836                            index += 1;
837                        }
838                        Ok(None) => break, // End of array
839                        Err(_) => break,   // Not an array or error
840                    }
841                }
842            }
843            Ok(false)
844        }
845        Err(e) => Err(common::execution_error(format!(
846            "Failed to check array contains at path '{}': {}",
847            path, e
848        ))),
849    }
850}
851
852/// Create the json_array_length UDF for getting array length
853///
854/// # Arguments
855/// * First parameter: JSONB binary data (LargeBinary)
856/// * Second parameter: JSONPath to array location (Utf8)
857///
858/// # Returns
859/// Integer length of the JSON array, or null if path doesn't point to an array
860pub fn json_array_length_udf() -> ScalarUDF {
861    create_udf(
862        "json_array_length",
863        vec![DataType::LargeBinary, DataType::Utf8],
864        DataType::Int64,
865        Volatility::Immutable,
866        Arc::new(json_array_length_columnar_impl),
867    )
868}
869
870/// Implementation of json_array_length function with ColumnarValue
871fn json_array_length_columnar_impl(args: &[ColumnarValue]) -> Result<ColumnarValue> {
872    let arrays = common::columnar_to_arrays(args);
873    let result = json_array_length_impl(&arrays)?;
874    Ok(ColumnarValue::Array(result))
875}
876
877/// Implementation of json_array_length function
878fn json_array_length_impl(args: &[ArrayRef]) -> Result<ArrayRef> {
879    common::validate_arg_count(args, 2, "json_array_length")?;
880
881    let jsonb_array = common::extract_jsonb_array(args)?;
882    let path_array = common::extract_string_array(args, 1)?;
883
884    let mut builder = Int64Builder::with_capacity(jsonb_array.len());
885
886    for i in 0..jsonb_array.len() {
887        if jsonb_array.is_null(i) {
888            builder.append_null();
889        } else if let Some(path) = common::get_string_value_at(path_array, i) {
890            let jsonb_bytes = jsonb_array.value(i);
891            match get_array_length(jsonb_bytes, path)? {
892                Some(len) => builder.append_value(len),
893                None => builder.append_null(),
894            }
895        } else {
896            builder.append_null();
897        }
898    }
899
900    Ok(Arc::new(builder.finish()))
901}
902
903/// Get the length of a JSON array at path
904fn get_array_length(jsonb_bytes: &[u8], path: &str) -> Result<Option<i64>> {
905    let json_path = common::parse_json_path(path)?;
906
907    match select_json_path_values(jsonb_bytes, &json_path) {
908        Ok(values) => {
909            if values.is_empty() {
910                return Ok(None);
911            }
912            let first = &values[0];
913            let raw = first.as_raw();
914
915            // Count array elements by iterating
916            let mut count = 0;
917            loop {
918                match raw.get_by_index(count) {
919                    Ok(Some(_)) => count += 1,
920                    Ok(None) => break, // End of array
921                    Err(_) => {
922                        // Not an array
923                        if count == 0 {
924                            return Err(common::execution_error(format!(
925                                "Path '{}' does not point to an array",
926                                path
927                            )));
928                        }
929                        break;
930                    }
931                }
932            }
933            Ok(Some(count as i64))
934        }
935        Err(e) => Err(common::execution_error(format!(
936            "Failed to get array length at path '{}': {}",
937            path, e
938        ))),
939    }
940}
941
942#[cfg(test)]
943mod tests {
944    use super::*;
945    use arrow_array::builder::LargeBinaryBuilder;
946    use arrow_array::{BooleanArray, Float64Array, Int64Array};
947
948    fn create_test_jsonb(json_str: &str) -> Vec<u8> {
949        lance_jsonb::parse_value(json_str.as_bytes())
950            .unwrap()
951            .to_vec()
952    }
953
954    #[test]
955    fn test_jsonb_type_enum() {
956        // Test enum conversion to/from u8
957        assert_eq!(JsonbType::Null.as_u8(), 0);
958        assert_eq!(JsonbType::Boolean.as_u8(), 1);
959        assert_eq!(JsonbType::Int64.as_u8(), 2);
960        assert_eq!(JsonbType::Float64.as_u8(), 3);
961        assert_eq!(JsonbType::String.as_u8(), 4);
962        assert_eq!(JsonbType::Array.as_u8(), 5);
963        assert_eq!(JsonbType::Object.as_u8(), 6);
964
965        // Test from_u8 conversion
966        assert_eq!(JsonbType::from_u8(0), Some(JsonbType::Null));
967        assert_eq!(JsonbType::from_u8(1), Some(JsonbType::Boolean));
968        assert_eq!(JsonbType::from_u8(2), Some(JsonbType::Int64));
969        assert_eq!(JsonbType::from_u8(3), Some(JsonbType::Float64));
970        assert_eq!(JsonbType::from_u8(4), Some(JsonbType::String));
971        assert_eq!(JsonbType::from_u8(5), Some(JsonbType::Array));
972        assert_eq!(JsonbType::from_u8(6), Some(JsonbType::Object));
973        assert_eq!(JsonbType::from_u8(7), None); // Invalid value
974    }
975
976    #[tokio::test]
977    async fn test_json_extract_udf() -> Result<()> {
978        let json = r#"{"user": {"name": "Alice", "age": 30}, "tags": ["python", "ml"]}"#;
979        let jsonb_bytes = create_test_jsonb(json);
980
981        let mut binary_builder = LargeBinaryBuilder::new();
982        binary_builder.append_value(&jsonb_bytes);
983        binary_builder.append_value(&jsonb_bytes);
984        binary_builder.append_value(&jsonb_bytes);
985        binary_builder.append_null();
986
987        let jsonb_array = Arc::new(binary_builder.finish());
988        let path_array = Arc::new(StringArray::from(vec![
989            Some("$.user.name"),
990            Some("$.user.age"),
991            Some("$.tags[*]"),
992            Some("$.user.name"),
993        ]));
994
995        let result = json_extract_impl(&[jsonb_array, path_array])?;
996        let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
997
998        assert_eq!(string_array.len(), 4);
999        assert_eq!(string_array.value(0), "\"Alice\"");
1000        assert_eq!(string_array.value(1), "30");
1001        assert_eq!(string_array.value(2), "[\"python\",\"ml\"]");
1002        assert!(string_array.is_null(3));
1003
1004        Ok(())
1005    }
1006
1007    #[tokio::test]
1008    async fn test_json_exists_udf() -> Result<()> {
1009        let json = r#"{"user": {"name": "Alice", "age": 30}, "tags": ["rust", "json"]}"#;
1010        let jsonb_bytes = create_test_jsonb(json);
1011
1012        let mut binary_builder = LargeBinaryBuilder::new();
1013        binary_builder.append_value(&jsonb_bytes);
1014        binary_builder.append_value(&jsonb_bytes);
1015        binary_builder.append_value(&jsonb_bytes);
1016        binary_builder.append_null();
1017
1018        let jsonb_array = Arc::new(binary_builder.finish());
1019        let path_array = Arc::new(StringArray::from(vec![
1020            Some("$.user.name"),
1021            Some("$.user.email"),
1022            Some("$.tags"),
1023            Some("$.any"),
1024        ]));
1025
1026        let result = json_exists_impl(&[jsonb_array, path_array])?;
1027        let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
1028
1029        assert_eq!(bool_array.len(), 4);
1030        assert!(bool_array.value(0));
1031        assert!(!bool_array.value(1));
1032        assert!(bool_array.value(2));
1033        assert!(bool_array.is_null(3));
1034
1035        Ok(())
1036    }
1037
1038    #[tokio::test]
1039    async fn test_json_get_string_udf() -> Result<()> {
1040        // Test valid string conversions
1041        let json = r#"{"str": "hello", "num": 123, "bool": true, "null": null}"#;
1042        let jsonb_bytes = create_test_jsonb(json);
1043
1044        let mut binary_builder = LargeBinaryBuilder::new();
1045        binary_builder.append_value(&jsonb_bytes);
1046        binary_builder.append_value(&jsonb_bytes);
1047        binary_builder.append_value(&jsonb_bytes);
1048        binary_builder.append_value(&jsonb_bytes);
1049
1050        let jsonb_array = Arc::new(binary_builder.finish());
1051        let key_array = Arc::new(StringArray::from(vec![
1052            Some("str"),
1053            Some("num"),
1054            Some("bool"),
1055            Some("null"),
1056        ]));
1057
1058        let result = json_get_string_impl(&[jsonb_array, key_array])?;
1059        let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
1060
1061        assert_eq!(string_array.len(), 4);
1062        assert_eq!(string_array.value(0), "hello");
1063        assert_eq!(string_array.value(1), "123");
1064        assert_eq!(string_array.value(2), "true");
1065        assert!(string_array.is_null(3));
1066
1067        Ok(())
1068    }
1069
1070    #[tokio::test]
1071    async fn test_json_get_int_udf() -> Result<()> {
1072        let json = r#"{"int": 42, "str_num": "99", "bool": true, "0": 7}"#;
1073        let jsonb_bytes = create_test_jsonb(json);
1074
1075        let mut binary_builder = LargeBinaryBuilder::new();
1076        for _ in 0..4 {
1077            binary_builder.append_value(&jsonb_bytes);
1078        }
1079
1080        let jsonb_array = Arc::new(binary_builder.finish());
1081        let key_array = Arc::new(StringArray::from(vec![
1082            Some("int"),
1083            Some("str_num"),
1084            Some("bool"),
1085            Some("0"),
1086        ]));
1087
1088        let result = json_get_int_impl(&[jsonb_array, key_array])?;
1089        let int_array = result.as_any().downcast_ref::<Int64Array>().unwrap();
1090
1091        assert_eq!(int_array.len(), 4);
1092        assert_eq!(int_array.value(0), 42);
1093        assert_eq!(int_array.value(1), 99);
1094        assert_eq!(int_array.value(2), 1); // jsonb converts true to 1
1095        assert_eq!(int_array.value(3), 7);
1096
1097        Ok(())
1098    }
1099
1100    #[tokio::test]
1101    async fn test_json_get_float_udf() -> Result<()> {
1102        let json = r#"{
1103            "float_decimal": 1.5,
1104            "float_neg": -2.5,
1105            "float_int_value": 1.0,
1106            "float_exp": 1e2,
1107            "int_pos": 42,
1108            "int_neg": -7,
1109            "big_int": 9223372036854775808,
1110            "str_num": "3.5",
1111            "bool_true": true,
1112            "null_val": null
1113        }"#;
1114        let jsonb_bytes = create_test_jsonb(json);
1115
1116        let mut binary_builder = LargeBinaryBuilder::new();
1117        for _ in 0..11 {
1118            binary_builder.append_value(&jsonb_bytes);
1119        }
1120        let jsonb_array = Arc::new(binary_builder.finish());
1121        let key_array = Arc::new(StringArray::from(vec![
1122            Some("float_decimal"),
1123            Some("float_neg"),
1124            Some("float_int_value"),
1125            Some("float_exp"),
1126            Some("int_pos"),
1127            Some("int_neg"),
1128            Some("big_int"),
1129            Some("str_num"),
1130            Some("bool_true"),
1131            Some("null_val"),
1132            Some("missing"),
1133        ]));
1134
1135        let result = json_get_float_impl(&[jsonb_array, key_array])?;
1136        let float_array = result.as_any().downcast_ref::<Float64Array>().unwrap();
1137
1138        assert_eq!(float_array.len(), 11);
1139        assert_eq!(float_array.value(0), 1.5);
1140        assert_eq!(float_array.value(1), -2.5);
1141        assert_eq!(float_array.value(2), 1.0);
1142        assert_eq!(float_array.value(3), 100.0);
1143        assert_eq!(float_array.value(4), 42.0);
1144        assert_eq!(float_array.value(5), -7.0);
1145        // 2^63 is exactly representable in f64.
1146        assert_eq!(float_array.value(6), 9223372036854775808.0);
1147        assert_eq!(float_array.value(7), 3.5);
1148        assert_eq!(float_array.value(8), 1.0); // jsonb converts true to 1.0
1149        assert!(float_array.is_null(9));
1150        assert!(float_array.is_null(10));
1151
1152        Ok(())
1153    }
1154
1155    #[tokio::test]
1156    async fn test_json_get_bool_udf() -> Result<()> {
1157        let json =
1158            r#"{"bool_true": true, "bool_false": false, "str_true": "true", "str_false": "false"}"#;
1159        let jsonb_bytes = create_test_jsonb(json);
1160
1161        let mut binary_builder = LargeBinaryBuilder::new();
1162        binary_builder.append_value(&jsonb_bytes);
1163        binary_builder.append_value(&jsonb_bytes);
1164        binary_builder.append_value(&jsonb_bytes);
1165        binary_builder.append_value(&jsonb_bytes);
1166
1167        let jsonb_array = Arc::new(binary_builder.finish());
1168        let key_array = Arc::new(StringArray::from(vec![
1169            Some("bool_true"),
1170            Some("bool_false"),
1171            Some("str_true"),
1172            Some("str_false"),
1173        ]));
1174
1175        let result = json_get_bool_impl(&[jsonb_array, key_array])?;
1176        let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
1177
1178        assert_eq!(bool_array.len(), 4);
1179        assert!(bool_array.value(0));
1180        assert!(!bool_array.value(1));
1181        assert!(bool_array.value(2)); // "true" string converts to true
1182        assert!(!bool_array.value(3)); // "false" string converts to false
1183
1184        Ok(())
1185    }
1186
1187    #[tokio::test]
1188    async fn test_json_array_contains_udf() -> Result<()> {
1189        let json = r#"{"tags": ["rust", "json", "database"], "nums": [1, 2, 3]}"#;
1190        let jsonb_bytes = create_test_jsonb(json);
1191
1192        let mut binary_builder = LargeBinaryBuilder::new();
1193        binary_builder.append_value(&jsonb_bytes);
1194        binary_builder.append_value(&jsonb_bytes);
1195        binary_builder.append_value(&jsonb_bytes);
1196        binary_builder.append_null();
1197
1198        let jsonb_array = Arc::new(binary_builder.finish());
1199        let path_array = Arc::new(StringArray::from(vec![
1200            Some("$.tags"),
1201            Some("$.tags"),
1202            Some("$.nums"),
1203            Some("$.tags"),
1204        ]));
1205        let value_array = Arc::new(StringArray::from(vec![
1206            Some("rust"),
1207            Some("python"),
1208            Some("2"),
1209            Some("any"),
1210        ]));
1211
1212        let result = json_array_contains_impl(&[jsonb_array, path_array, value_array])?;
1213        let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
1214
1215        assert_eq!(bool_array.len(), 4);
1216        assert!(bool_array.value(0));
1217        assert!(!bool_array.value(1));
1218        assert!(bool_array.value(2));
1219        assert!(bool_array.is_null(3));
1220
1221        Ok(())
1222    }
1223
1224    #[tokio::test]
1225    async fn test_json_array_length_udf() -> Result<()> {
1226        let json = r#"{"empty": [], "tags": ["a", "b", "c"], "nested": {"arr": [1, 2]}}"#;
1227        let jsonb_bytes = create_test_jsonb(json);
1228
1229        let mut binary_builder = LargeBinaryBuilder::new();
1230        binary_builder.append_value(&jsonb_bytes);
1231        binary_builder.append_value(&jsonb_bytes);
1232        binary_builder.append_value(&jsonb_bytes);
1233        binary_builder.append_null();
1234
1235        let jsonb_array = Arc::new(binary_builder.finish());
1236        let path_array = Arc::new(StringArray::from(vec![
1237            Some("$.empty"),
1238            Some("$.tags"),
1239            Some("$.nested.arr"),
1240            Some("$.any"),
1241        ]));
1242
1243        let result = json_array_length_impl(&[jsonb_array, path_array])?;
1244        let int_array = result.as_any().downcast_ref::<Int64Array>().unwrap();
1245
1246        assert_eq!(int_array.len(), 4);
1247        assert_eq!(int_array.value(0), 0);
1248        assert_eq!(int_array.value(1), 3);
1249        assert_eq!(int_array.value(2), 2);
1250        assert!(int_array.is_null(3));
1251
1252        Ok(())
1253    }
1254
1255    #[tokio::test]
1256    async fn test_json_extract_with_type() -> Result<()> {
1257        use arrow_array::StructArray;
1258        use arrow_array::UInt8Array;
1259
1260        let cases: &[(&str, JsonbType)] = &[
1261            (r#"{"v": 1}"#, JsonbType::Int64),
1262            (r#"{"v": 0}"#, JsonbType::Int64),
1263            (r#"{"v": -42}"#, JsonbType::Int64),
1264            (r#"{"v": 9223372036854775807}"#, JsonbType::Int64), // i64::MAX
1265            (r#"{"v": 9223372036854775808}"#, JsonbType::Float64), // i64::MAX + 1
1266            (r#"{"v": 1.0}"#, JsonbType::Float64),
1267            (r#"{"v": 2.7}"#, JsonbType::Float64),
1268            (r#"{"v": 1.5}"#, JsonbType::Float64),
1269            (r#"{"v": -1.5}"#, JsonbType::Float64),
1270            (r#"{"v": 1e2}"#, JsonbType::Float64),
1271        ];
1272
1273        for (json, expected) in cases {
1274            let bytes = create_test_jsonb(json);
1275            let mut binary_builder = LargeBinaryBuilder::new();
1276            binary_builder.append_value(&bytes);
1277            let jsonb_array: ArrayRef = Arc::new(binary_builder.finish());
1278            let path_array: ArrayRef = Arc::new(StringArray::from(vec![Some("$.v")]));
1279
1280            let result = json_extract_with_type_impl(&[jsonb_array, path_array])?;
1281            let struct_array = result.as_any().downcast_ref::<StructArray>().unwrap();
1282            let type_tags = struct_array
1283                .column_by_name("type_tag")
1284                .unwrap()
1285                .as_any()
1286                .downcast_ref::<UInt8Array>()
1287                .unwrap();
1288            assert_eq!(type_tags.value(0), expected.as_u8());
1289        }
1290
1291        Ok(())
1292    }
1293
1294    #[tokio::test]
1295    async fn test_json_extract_with_wildcard() -> Result<()> {
1296        use arrow_array::StructArray;
1297        use arrow_array::UInt8Array;
1298
1299        let json = r#"{"items": [{"price": 1}, {"price": 2}]}"#;
1300        let jsonb_bytes = create_test_jsonb(json);
1301
1302        let mut binary_builder = LargeBinaryBuilder::new();
1303        binary_builder.append_value(&jsonb_bytes);
1304
1305        let jsonb_array: ArrayRef = Arc::new(binary_builder.finish());
1306        let path_array: ArrayRef = Arc::new(StringArray::from(vec![Some("$.items[*].price")]));
1307
1308        let result = json_extract_with_type_impl(&[jsonb_array, path_array])?;
1309        let struct_array = result.as_any().downcast_ref::<StructArray>().unwrap();
1310        let values = struct_array
1311            .column_by_name("value")
1312            .unwrap()
1313            .as_any()
1314            .downcast_ref::<LargeBinaryArray>()
1315            .unwrap();
1316        let type_tags = struct_array
1317            .column_by_name("type_tag")
1318            .unwrap()
1319            .as_any()
1320            .downcast_ref::<UInt8Array>()
1321            .unwrap();
1322
1323        assert_eq!(
1324            lance_jsonb::RawJsonb::new(values.value(0)).to_string(),
1325            "[1,2]"
1326        );
1327        assert_eq!(type_tags.value(0), JsonbType::Array.as_u8());
1328
1329        Ok(())
1330    }
1331
1332    #[tokio::test]
1333    async fn test_json_array_access() -> Result<()> {
1334        let json = r#"["first", "second", "third"]"#;
1335        let jsonb_bytes = create_test_jsonb(json);
1336
1337        let mut binary_builder = LargeBinaryBuilder::new();
1338        binary_builder.append_value(&jsonb_bytes);
1339        binary_builder.append_value(&jsonb_bytes);
1340        binary_builder.append_value(&jsonb_bytes);
1341
1342        let jsonb_array = Arc::new(binary_builder.finish());
1343        let key_array = Arc::new(StringArray::from(vec![
1344            Some("0"),
1345            Some("1"),
1346            Some("10"), // Out of bounds
1347        ]));
1348
1349        let result = json_get_string_impl(&[jsonb_array, key_array])?;
1350        let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
1351
1352        assert_eq!(string_array.len(), 3);
1353        assert_eq!(string_array.value(0), "first");
1354        assert_eq!(string_array.value(1), "second");
1355        assert!(string_array.is_null(2));
1356
1357        Ok(())
1358    }
1359
1360    #[tokio::test]
1361    async fn test_json_get_numeric_object_key() -> Result<()> {
1362        let obj_bytes = create_test_jsonb(r#"{"0": "from_object", "1": 42}"#);
1363        let arr_bytes = create_test_jsonb(r#"["zero", "one", "two"]"#);
1364
1365        let mut binary_builder = LargeBinaryBuilder::new();
1366        binary_builder.append_value(&obj_bytes);
1367        binary_builder.append_value(&arr_bytes);
1368        binary_builder.append_value(&arr_bytes);
1369
1370        let jsonb_array = Arc::new(binary_builder.finish());
1371        let key_array = Arc::new(StringArray::from(vec![
1372            Some("0"),   // Numeric key on object: looks up the "0" field.
1373            Some("0"),   // Numeric key on array: looks up index 0.
1374            Some("foo"), // Non-numeric key on array: no match.
1375        ]));
1376
1377        let result = json_get_string_impl(&[jsonb_array, key_array])?;
1378        let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
1379
1380        assert_eq!(string_array.len(), 3);
1381        assert_eq!(string_array.value(0), "from_object");
1382        assert_eq!(string_array.value(1), "zero");
1383        assert!(string_array.is_null(2));
1384
1385        Ok(())
1386    }
1387}
lance_datafusion/udf/json.rs

lance_datafusion/udf/
json.rs
