diff --git a/src/variant_get.rs b/src/variant_get.rs index 43e06e9..0d3fe19 100644 --- a/src/variant_get.rs +++ b/src/variant_get.rs @@ -20,7 +20,7 @@ use crate::shared::{ try_field_as_variant_array, try_parse_string_columnar, try_parse_string_scalar, }; -fn type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result { +fn try_parse_type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result { let type_name = match scalar { ScalarValue::Utf8(Some(value)) | ScalarValue::Utf8View(Some(value)) @@ -42,9 +42,9 @@ fn type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result Result { +fn try_parse_type_hint(field_name: &str, arg: &ColumnarValue) -> Result { match arg { - ColumnarValue::Scalar(value) => type_hint_from_scalar(field_name, value), + ColumnarValue::Scalar(value) => try_parse_type_hint_from_scalar(field_name, value), ColumnarValue::Array(_) => { exec_err!("type hint argument must be a scalar UTF8 literal") } @@ -98,7 +98,7 @@ impl ScalarUDFImpl for VariantGetUdf { let scalar = maybe_scalar.ok_or_else(|| { exec_datafusion_err!("type hint argument to variant_get must be a literal") })?; - return type_hint_from_scalar(self.name(), scalar); + return try_parse_type_hint_from_scalar(self.name(), scalar); } let data_type = DataType::Struct(Fields::from(vec![ @@ -129,7 +129,7 @@ impl ScalarUDFImpl for VariantGetUdf { try_field_as_variant_array(variant_field.as_ref())?; let type_field = type_arg - .map(|arg| type_hint_from_value(self.name(), arg)) + .map(|arg| try_parse_type_hint(self.name(), arg)) .transpose()?; let out = match (variant_arg, variant_path) { diff --git a/tests/test_files/variant_get.slt b/tests/test_files/variant_get.slt index 18d7e05..ef3ba69 100644 --- a/tests/test_files/variant_get.slt +++ b/tests/test_files/variant_get.slt @@ -17,4 +17,129 @@ NULL NULL NULL NULL -NULL \ No newline at end of file +NULL + +# test variant_get with type hint (third argument) +# test 1: get string field with utf8 type hint +query T +SELECT variant_get(json_to_variant('{"name": "Alice", "age": 30}'), 'name', 'Utf8'); +---- +Alice + +# test 2: get numeric field with int64 type hint +query I +SELECT variant_get(json_to_variant('{"name": "Bob", "age": 25}'), 'age', 'Int64'); +---- +25 + +# test 3: get numeric field with int32 type hint +query I +SELECT variant_get(json_to_variant('{"score": 100}'), 'score', 'Int32'); +---- +100 + +# test 4: get numeric field with float64 type hint +query R +SELECT variant_get(json_to_variant('{"price": 19.99}'), 'price', 'Float64'); +---- +19.99 + +# test 5: get boolean field with boolean type hint +query B +SELECT variant_get(json_to_variant('{"active": true}'), 'active', 'Boolean'); +---- +true + +# test 6: get null field with type hint (should return null) +query T +SELECT variant_get(json_to_variant('{"name": null}'), 'name', 'Utf8'); +---- +NULL + +# test 7: get non-existent field with type hint (should return null) +query T +SELECT variant_get(json_to_variant('{"name": "Charlie"}'), 'age', 'Int64'); +---- +NULL + +# test 8: get nested field with type hint +query T +SELECT variant_get(json_to_variant('{"user": {"name": "David"}}'), 'user.name', 'Utf8'); +---- +David + +# test 9: get array element with type hint +query I +SELECT variant_get(json_to_variant('{"items": [10, 20, 30]}'), 'items[1]', 'Int64'); +---- +20 + +# test 10: get multiple rows with type hint +query I +SELECT variant_get(json_to_variant(json_str), 'age', 'Int64') FROM json_data WHERE id <= 2; +---- +30 +25 + +# test 11: get string field from various rows with utf8 type hint +query T +SELECT variant_get(json_to_variant(json_str), 'name', 'Utf8') FROM json_data WHERE id <= 2; +---- +Alice +Bob + +# test 12: test with int16 type hint +query I +SELECT variant_get(json_to_variant('{"count": 5}'), 'count', 'Int16'); +---- +5 + +# test 13: test with int8 type hint +query I +SELECT variant_get(json_to_variant('{"level": 3}'), 'level', 'Int8'); +---- +3 + +# test 14: test with float32 type hint +query R +SELECT variant_get(json_to_variant('{"ratio": 0.5}'), 'ratio', 'Float32'); +---- +0.5 + +# test 15: test with large number and int64 type hint +query I +SELECT variant_get(json_to_variant('{"big_number": 9223372036854775807}'), 'big_number', 'Int64'); +---- +9223372036854775807 + +# test 16: test mixed types in same query - some with type hint, some without +query TT +SELECT + variant_get(json_to_variant('{"name": "Eve", "age": 28}'), 'name', 'Utf8'), + variant_pretty(variant_get(json_to_variant('{"name": "Eve", "age": 28}'), 'age')); +---- +Eve Int8(28) + +# test 17: test with nested path and int64 type hint +query I +SELECT variant_get(json_to_variant('{"data": {"count": 42}}'), 'data.count', 'Int64'); +---- +42 + +# test 18: test boolean false with type hint +query B +SELECT variant_get(json_to_variant('{"enabled": false}'), 'enabled', 'Boolean'); +---- +false + +# test 19: test with decimal-like number and float64 type hint +query R +SELECT variant_get(json_to_variant('{"temperature": -5.5}'), 'temperature', 'Float64'); +---- +-5.5 + +# test 20: test with zero values +query I +SELECT variant_get(json_to_variant('{"zero": 0}'), 'zero', 'Int64'); +---- +0 \ No newline at end of file