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: &jsonb::RawJsonb,
121        key: &str,
122    ) -> Result<Option<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<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: 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: 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: 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: 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(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 = 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 = 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 = 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 = 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 = 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        jsonb::parse_value(json_str.as_bytes()).unwrap().to_vec()
950    }
951
952    #[test]
953    fn test_jsonb_type_enum() {
954        // Test enum conversion to/from u8
955        assert_eq!(JsonbType::Null.as_u8(), 0);
956        assert_eq!(JsonbType::Boolean.as_u8(), 1);
957        assert_eq!(JsonbType::Int64.as_u8(), 2);
958        assert_eq!(JsonbType::Float64.as_u8(), 3);
959        assert_eq!(JsonbType::String.as_u8(), 4);
960        assert_eq!(JsonbType::Array.as_u8(), 5);
961        assert_eq!(JsonbType::Object.as_u8(), 6);
962
963        // Test from_u8 conversion
964        assert_eq!(JsonbType::from_u8(0), Some(JsonbType::Null));
965        assert_eq!(JsonbType::from_u8(1), Some(JsonbType::Boolean));
966        assert_eq!(JsonbType::from_u8(2), Some(JsonbType::Int64));
967        assert_eq!(JsonbType::from_u8(3), Some(JsonbType::Float64));
968        assert_eq!(JsonbType::from_u8(4), Some(JsonbType::String));
969        assert_eq!(JsonbType::from_u8(5), Some(JsonbType::Array));
970        assert_eq!(JsonbType::from_u8(6), Some(JsonbType::Object));
971        assert_eq!(JsonbType::from_u8(7), None); // Invalid value
972    }
973
974    #[tokio::test]
975    async fn test_json_extract_udf() -> Result<()> {
976        let json = r#"{"user": {"name": "Alice", "age": 30}, "tags": ["python", "ml"]}"#;
977        let jsonb_bytes = create_test_jsonb(json);
978
979        let mut binary_builder = LargeBinaryBuilder::new();
980        binary_builder.append_value(&jsonb_bytes);
981        binary_builder.append_value(&jsonb_bytes);
982        binary_builder.append_value(&jsonb_bytes);
983        binary_builder.append_null();
984
985        let jsonb_array = Arc::new(binary_builder.finish());
986        let path_array = Arc::new(StringArray::from(vec![
987            Some("$.user.name"),
988            Some("$.user.age"),
989            Some("$.tags[*]"),
990            Some("$.user.name"),
991        ]));
992
993        let result = json_extract_impl(&[jsonb_array, path_array])?;
994        let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
995
996        assert_eq!(string_array.len(), 4);
997        assert_eq!(string_array.value(0), "\"Alice\"");
998        assert_eq!(string_array.value(1), "30");
999        assert_eq!(string_array.value(2), "[\"python\",\"ml\"]");
1000        assert!(string_array.is_null(3));
1001
1002        Ok(())
1003    }
1004
1005    #[tokio::test]
1006    async fn test_json_exists_udf() -> Result<()> {
1007        let json = r#"{"user": {"name": "Alice", "age": 30}, "tags": ["rust", "json"]}"#;
1008        let jsonb_bytes = create_test_jsonb(json);
1009
1010        let mut binary_builder = LargeBinaryBuilder::new();
1011        binary_builder.append_value(&jsonb_bytes);
1012        binary_builder.append_value(&jsonb_bytes);
1013        binary_builder.append_value(&jsonb_bytes);
1014        binary_builder.append_null();
1015
1016        let jsonb_array = Arc::new(binary_builder.finish());
1017        let path_array = Arc::new(StringArray::from(vec![
1018            Some("$.user.name"),
1019            Some("$.user.email"),
1020            Some("$.tags"),
1021            Some("$.any"),
1022        ]));
1023
1024        let result = json_exists_impl(&[jsonb_array, path_array])?;
1025        let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
1026
1027        assert_eq!(bool_array.len(), 4);
1028        assert!(bool_array.value(0));
1029        assert!(!bool_array.value(1));
1030        assert!(bool_array.value(2));
1031        assert!(bool_array.is_null(3));
1032
1033        Ok(())
1034    }
1035
1036    #[tokio::test]
1037    async fn test_json_get_string_udf() -> Result<()> {
1038        // Test valid string conversions
1039        let json = r#"{"str": "hello", "num": 123, "bool": true, "null": null}"#;
1040        let jsonb_bytes = create_test_jsonb(json);
1041
1042        let mut binary_builder = LargeBinaryBuilder::new();
1043        binary_builder.append_value(&jsonb_bytes);
1044        binary_builder.append_value(&jsonb_bytes);
1045        binary_builder.append_value(&jsonb_bytes);
1046        binary_builder.append_value(&jsonb_bytes);
1047
1048        let jsonb_array = Arc::new(binary_builder.finish());
1049        let key_array = Arc::new(StringArray::from(vec![
1050            Some("str"),
1051            Some("num"),
1052            Some("bool"),
1053            Some("null"),
1054        ]));
1055
1056        let result = json_get_string_impl(&[jsonb_array, key_array])?;
1057        let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
1058
1059        assert_eq!(string_array.len(), 4);
1060        assert_eq!(string_array.value(0), "hello");
1061        assert_eq!(string_array.value(1), "123");
1062        assert_eq!(string_array.value(2), "true");
1063        assert!(string_array.is_null(3));
1064
1065        Ok(())
1066    }
1067
1068    #[tokio::test]
1069    async fn test_json_get_int_udf() -> Result<()> {
1070        let json = r#"{"int": 42, "str_num": "99", "bool": true, "0": 7}"#;
1071        let jsonb_bytes = create_test_jsonb(json);
1072
1073        let mut binary_builder = LargeBinaryBuilder::new();
1074        for _ in 0..4 {
1075            binary_builder.append_value(&jsonb_bytes);
1076        }
1077
1078        let jsonb_array = Arc::new(binary_builder.finish());
1079        let key_array = Arc::new(StringArray::from(vec![
1080            Some("int"),
1081            Some("str_num"),
1082            Some("bool"),
1083            Some("0"),
1084        ]));
1085
1086        let result = json_get_int_impl(&[jsonb_array, key_array])?;
1087        let int_array = result.as_any().downcast_ref::<Int64Array>().unwrap();
1088
1089        assert_eq!(int_array.len(), 4);
1090        assert_eq!(int_array.value(0), 42);
1091        assert_eq!(int_array.value(1), 99);
1092        assert_eq!(int_array.value(2), 1); // jsonb converts true to 1
1093        assert_eq!(int_array.value(3), 7);
1094
1095        Ok(())
1096    }
1097
1098    #[tokio::test]
1099    async fn test_json_get_float_udf() -> Result<()> {
1100        let json = r#"{
1101            "float_decimal": 1.5,
1102            "float_neg": -2.5,
1103            "float_int_value": 1.0,
1104            "float_exp": 1e2,
1105            "int_pos": 42,
1106            "int_neg": -7,
1107            "big_int": 9223372036854775808,
1108            "str_num": "3.5",
1109            "bool_true": true,
1110            "null_val": null
1111        }"#;
1112        let jsonb_bytes = create_test_jsonb(json);
1113
1114        let mut binary_builder = LargeBinaryBuilder::new();
1115        for _ in 0..11 {
1116            binary_builder.append_value(&jsonb_bytes);
1117        }
1118        let jsonb_array = Arc::new(binary_builder.finish());
1119        let key_array = Arc::new(StringArray::from(vec![
1120            Some("float_decimal"),
1121            Some("float_neg"),
1122            Some("float_int_value"),
1123            Some("float_exp"),
1124            Some("int_pos"),
1125            Some("int_neg"),
1126            Some("big_int"),
1127            Some("str_num"),
1128            Some("bool_true"),
1129            Some("null_val"),
1130            Some("missing"),
1131        ]));
1132
1133        let result = json_get_float_impl(&[jsonb_array, key_array])?;
1134        let float_array = result.as_any().downcast_ref::<Float64Array>().unwrap();
1135
1136        assert_eq!(float_array.len(), 11);
1137        assert_eq!(float_array.value(0), 1.5);
1138        assert_eq!(float_array.value(1), -2.5);
1139        assert_eq!(float_array.value(2), 1.0);
1140        assert_eq!(float_array.value(3), 100.0);
1141        assert_eq!(float_array.value(4), 42.0);
1142        assert_eq!(float_array.value(5), -7.0);
1143        // 2^63 is exactly representable in f64.
1144        assert_eq!(float_array.value(6), 9223372036854775808.0);
1145        assert_eq!(float_array.value(7), 3.5);
1146        assert_eq!(float_array.value(8), 1.0); // jsonb converts true to 1.0
1147        assert!(float_array.is_null(9));
1148        assert!(float_array.is_null(10));
1149
1150        Ok(())
1151    }
1152
1153    #[tokio::test]
1154    async fn test_json_get_bool_udf() -> Result<()> {
1155        let json =
1156            r#"{"bool_true": true, "bool_false": false, "str_true": "true", "str_false": "false"}"#;
1157        let jsonb_bytes = create_test_jsonb(json);
1158
1159        let mut binary_builder = LargeBinaryBuilder::new();
1160        binary_builder.append_value(&jsonb_bytes);
1161        binary_builder.append_value(&jsonb_bytes);
1162        binary_builder.append_value(&jsonb_bytes);
1163        binary_builder.append_value(&jsonb_bytes);
1164
1165        let jsonb_array = Arc::new(binary_builder.finish());
1166        let key_array = Arc::new(StringArray::from(vec![
1167            Some("bool_true"),
1168            Some("bool_false"),
1169            Some("str_true"),
1170            Some("str_false"),
1171        ]));
1172
1173        let result = json_get_bool_impl(&[jsonb_array, key_array])?;
1174        let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
1175
1176        assert_eq!(bool_array.len(), 4);
1177        assert!(bool_array.value(0));
1178        assert!(!bool_array.value(1));
1179        assert!(bool_array.value(2)); // "true" string converts to true
1180        assert!(!bool_array.value(3)); // "false" string converts to false
1181
1182        Ok(())
1183    }
1184
1185    #[tokio::test]
1186    async fn test_json_array_contains_udf() -> Result<()> {
1187        let json = r#"{"tags": ["rust", "json", "database"], "nums": [1, 2, 3]}"#;
1188        let jsonb_bytes = create_test_jsonb(json);
1189
1190        let mut binary_builder = LargeBinaryBuilder::new();
1191        binary_builder.append_value(&jsonb_bytes);
1192        binary_builder.append_value(&jsonb_bytes);
1193        binary_builder.append_value(&jsonb_bytes);
1194        binary_builder.append_null();
1195
1196        let jsonb_array = Arc::new(binary_builder.finish());
1197        let path_array = Arc::new(StringArray::from(vec![
1198            Some("$.tags"),
1199            Some("$.tags"),
1200            Some("$.nums"),
1201            Some("$.tags"),
1202        ]));
1203        let value_array = Arc::new(StringArray::from(vec![
1204            Some("rust"),
1205            Some("python"),
1206            Some("2"),
1207            Some("any"),
1208        ]));
1209
1210        let result = json_array_contains_impl(&[jsonb_array, path_array, value_array])?;
1211        let bool_array = result.as_any().downcast_ref::<BooleanArray>().unwrap();
1212
1213        assert_eq!(bool_array.len(), 4);
1214        assert!(bool_array.value(0));
1215        assert!(!bool_array.value(1));
1216        assert!(bool_array.value(2));
1217        assert!(bool_array.is_null(3));
1218
1219        Ok(())
1220    }
1221
1222    #[tokio::test]
1223    async fn test_json_array_length_udf() -> Result<()> {
1224        let json = r#"{"empty": [], "tags": ["a", "b", "c"], "nested": {"arr": [1, 2]}}"#;
1225        let jsonb_bytes = create_test_jsonb(json);
1226
1227        let mut binary_builder = LargeBinaryBuilder::new();
1228        binary_builder.append_value(&jsonb_bytes);
1229        binary_builder.append_value(&jsonb_bytes);
1230        binary_builder.append_value(&jsonb_bytes);
1231        binary_builder.append_null();
1232
1233        let jsonb_array = Arc::new(binary_builder.finish());
1234        let path_array = Arc::new(StringArray::from(vec![
1235            Some("$.empty"),
1236            Some("$.tags"),
1237            Some("$.nested.arr"),
1238            Some("$.any"),
1239        ]));
1240
1241        let result = json_array_length_impl(&[jsonb_array, path_array])?;
1242        let int_array = result.as_any().downcast_ref::<Int64Array>().unwrap();
1243
1244        assert_eq!(int_array.len(), 4);
1245        assert_eq!(int_array.value(0), 0);
1246        assert_eq!(int_array.value(1), 3);
1247        assert_eq!(int_array.value(2), 2);
1248        assert!(int_array.is_null(3));
1249
1250        Ok(())
1251    }
1252
1253    #[tokio::test]
1254    async fn test_json_extract_with_type() -> Result<()> {
1255        use arrow_array::StructArray;
1256        use arrow_array::UInt8Array;
1257
1258        let cases: &[(&str, JsonbType)] = &[
1259            (r#"{"v": 1}"#, JsonbType::Int64),
1260            (r#"{"v": 0}"#, JsonbType::Int64),
1261            (r#"{"v": -42}"#, JsonbType::Int64),
1262            (r#"{"v": 9223372036854775807}"#, JsonbType::Int64), // i64::MAX
1263            (r#"{"v": 9223372036854775808}"#, JsonbType::Float64), // i64::MAX + 1
1264            (r#"{"v": 1.0}"#, JsonbType::Float64),
1265            (r#"{"v": 2.7}"#, JsonbType::Float64),
1266            (r#"{"v": 1.5}"#, JsonbType::Float64),
1267            (r#"{"v": -1.5}"#, JsonbType::Float64),
1268            (r#"{"v": 1e2}"#, JsonbType::Float64),
1269        ];
1270
1271        for (json, expected) in cases {
1272            let bytes = create_test_jsonb(json);
1273            let mut binary_builder = LargeBinaryBuilder::new();
1274            binary_builder.append_value(&bytes);
1275            let jsonb_array: ArrayRef = Arc::new(binary_builder.finish());
1276            let path_array: ArrayRef = Arc::new(StringArray::from(vec![Some("$.v")]));
1277
1278            let result = json_extract_with_type_impl(&[jsonb_array, path_array])?;
1279            let struct_array = result.as_any().downcast_ref::<StructArray>().unwrap();
1280            let type_tags = struct_array
1281                .column_by_name("type_tag")
1282                .unwrap()
1283                .as_any()
1284                .downcast_ref::<UInt8Array>()
1285                .unwrap();
1286            assert_eq!(type_tags.value(0), expected.as_u8());
1287        }
1288
1289        Ok(())
1290    }
1291
1292    #[tokio::test]
1293    async fn test_json_extract_with_wildcard() -> Result<()> {
1294        use arrow_array::StructArray;
1295        use arrow_array::UInt8Array;
1296
1297        let json = r#"{"items": [{"price": 1}, {"price": 2}]}"#;
1298        let jsonb_bytes = create_test_jsonb(json);
1299
1300        let mut binary_builder = LargeBinaryBuilder::new();
1301        binary_builder.append_value(&jsonb_bytes);
1302
1303        let jsonb_array: ArrayRef = Arc::new(binary_builder.finish());
1304        let path_array: ArrayRef = Arc::new(StringArray::from(vec![Some("$.items[*].price")]));
1305
1306        let result = json_extract_with_type_impl(&[jsonb_array, path_array])?;
1307        let struct_array = result.as_any().downcast_ref::<StructArray>().unwrap();
1308        let values = struct_array
1309            .column_by_name("value")
1310            .unwrap()
1311            .as_any()
1312            .downcast_ref::<LargeBinaryArray>()
1313            .unwrap();
1314        let type_tags = struct_array
1315            .column_by_name("type_tag")
1316            .unwrap()
1317            .as_any()
1318            .downcast_ref::<UInt8Array>()
1319            .unwrap();
1320
1321        assert_eq!(jsonb::RawJsonb::new(values.value(0)).to_string(), "[1,2]");
1322        assert_eq!(type_tags.value(0), JsonbType::Array.as_u8());
1323
1324        Ok(())
1325    }
1326
1327    #[tokio::test]
1328    async fn test_json_array_access() -> Result<()> {
1329        let json = r#"["first", "second", "third"]"#;
1330        let jsonb_bytes = create_test_jsonb(json);
1331
1332        let mut binary_builder = LargeBinaryBuilder::new();
1333        binary_builder.append_value(&jsonb_bytes);
1334        binary_builder.append_value(&jsonb_bytes);
1335        binary_builder.append_value(&jsonb_bytes);
1336
1337        let jsonb_array = Arc::new(binary_builder.finish());
1338        let key_array = Arc::new(StringArray::from(vec![
1339            Some("0"),
1340            Some("1"),
1341            Some("10"), // Out of bounds
1342        ]));
1343
1344        let result = json_get_string_impl(&[jsonb_array, key_array])?;
1345        let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
1346
1347        assert_eq!(string_array.len(), 3);
1348        assert_eq!(string_array.value(0), "first");
1349        assert_eq!(string_array.value(1), "second");
1350        assert!(string_array.is_null(2));
1351
1352        Ok(())
1353    }
1354
1355    #[tokio::test]
1356    async fn test_json_get_numeric_object_key() -> Result<()> {
1357        let obj_bytes = create_test_jsonb(r#"{"0": "from_object", "1": 42}"#);
1358        let arr_bytes = create_test_jsonb(r#"["zero", "one", "two"]"#);
1359
1360        let mut binary_builder = LargeBinaryBuilder::new();
1361        binary_builder.append_value(&obj_bytes);
1362        binary_builder.append_value(&arr_bytes);
1363        binary_builder.append_value(&arr_bytes);
1364
1365        let jsonb_array = Arc::new(binary_builder.finish());
1366        let key_array = Arc::new(StringArray::from(vec![
1367            Some("0"),   // Numeric key on object: looks up the "0" field.
1368            Some("0"),   // Numeric key on array: looks up index 0.
1369            Some("foo"), // Non-numeric key on array: no match.
1370        ]));
1371
1372        let result = json_get_string_impl(&[jsonb_array, key_array])?;
1373        let string_array = result.as_any().downcast_ref::<StringArray>().unwrap();
1374
1375        assert_eq!(string_array.len(), 3);
1376        assert_eq!(string_array.value(0), "from_object");
1377        assert_eq!(string_array.value(1), "zero");
1378        assert!(string_array.is_null(2));
1379
1380        Ok(())
1381    }
1382}
lance_datafusion/udf/json.rs

lance_datafusion/udf/
json.rs
