Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions src/variant_get.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<FieldRef> {
fn try_parse_type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result<FieldRef> {
let type_name = match scalar {
ScalarValue::Utf8(Some(value))
| ScalarValue::Utf8View(Some(value))
Expand All @@ -42,9 +42,9 @@ fn type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result<Field
Ok(Arc::new(Field::new(field_name, casted_type, true)))
}

fn type_hint_from_value(field_name: &str, arg: &ColumnarValue) -> Result<FieldRef> {
fn try_parse_type_hint(field_name: &str, arg: &ColumnarValue) -> Result<FieldRef> {
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")
}
Expand Down Expand Up @@ -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![
Expand Down Expand Up @@ -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) {
Expand Down
127 changes: 126 additions & 1 deletion tests/test_files/variant_get.slt
Original file line number Diff line number Diff line change
Expand Up @@ -17,4 +17,129 @@ NULL
NULL
NULL
NULL
NULL
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
Loading