Skip to content

Commit b35db3a

Browse files
friendlymatthewfriendlymatthew
committed
Update some nits, add slt test cases
1 parent 29c4c60 commit b35db3a

2 files changed

Lines changed: 131 additions & 6 deletions

File tree

‎src/variant_get.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use crate::shared::{
2020
try_field_as_variant_array, try_parse_string_columnar, try_parse_string_scalar,
2121
};
2222

23-
fn type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result<FieldRef> {
23+
fn try_parse_type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result<FieldRef> {
2424
let type_name = match scalar {
2525
ScalarValue::Utf8(Some(value))
2626
| ScalarValue::Utf8View(Some(value))
@@ -42,9 +42,9 @@ fn type_hint_from_scalar(field_name: &str, scalar: &ScalarValue) -> Result<Field
4242
Ok(Arc::new(Field::new(field_name, casted_type, true)))
4343
}
4444

45-
fn type_hint_from_value(field_name: &str, arg: &ColumnarValue) -> Result<FieldRef> {
45+
fn try_parse_type_hint(field_name: &str, arg: &ColumnarValue) -> Result<FieldRef> {
4646
match arg {
47-
ColumnarValue::Scalar(value) => type_hint_from_scalar(field_name, value),
47+
ColumnarValue::Scalar(value) => try_parse_type_hint_from_scalar(field_name, value),
4848
ColumnarValue::Array(_) => {
4949
exec_err!("type hint argument must be a scalar UTF8 literal")
5050
}
@@ -98,7 +98,7 @@ impl ScalarUDFImpl for VariantGetUdf {
9898
let scalar = maybe_scalar.ok_or_else(|| {
9999
exec_datafusion_err!("type hint argument to variant_get must be a literal")
100100
})?;
101-
return type_hint_from_scalar(self.name(), scalar);
101+
return try_parse_type_hint_from_scalar(self.name(), scalar);
102102
}
103103

104104
let data_type = DataType::Struct(Fields::from(vec![
@@ -129,7 +129,7 @@ impl ScalarUDFImpl for VariantGetUdf {
129129
try_field_as_variant_array(variant_field.as_ref())?;
130130

131131
let type_field = type_arg
132-
.map(|arg| type_hint_from_value(self.name(), arg))
132+
.map(|arg| try_parse_type_hint(self.name(), arg))
133133
.transpose()?;
134134

135135
let out = match (variant_arg, variant_path) {

‎tests/test_files/variant_get.slt‎

Lines changed: 126 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,4 +17,129 @@ NULL
1717
NULL
1818
NULL
1919
NULL
20-
NULL
20+
NULL
21+
22+
# test variant_get with type hint (third argument)
23+
# test 1: get string field with utf8 type hint
24+
query T
25+
SELECT variant_get(json_to_variant('{"name": "Alice", "age": 30}'), 'name', 'Utf8');
26+
----
27+
Alice
28+
29+
# test 2: get numeric field with int64 type hint
30+
query I
31+
SELECT variant_get(json_to_variant('{"name": "Bob", "age": 25}'), 'age', 'Int64');
32+
----
33+
25
34+
35+
# test 3: get numeric field with int32 type hint
36+
query I
37+
SELECT variant_get(json_to_variant('{"score": 100}'), 'score', 'Int32');
38+
----
39+
100
40+
41+
# test 4: get numeric field with float64 type hint
42+
query R
43+
SELECT variant_get(json_to_variant('{"price": 19.99}'), 'price', 'Float64');
44+
----
45+
19.99
46+
47+
# test 5: get boolean field with boolean type hint
48+
query B
49+
SELECT variant_get(json_to_variant('{"active": true}'), 'active', 'Boolean');
50+
----
51+
true
52+
53+
# test 6: get null field with type hint (should return null)
54+
query T
55+
SELECT variant_get(json_to_variant('{"name": null}'), 'name', 'Utf8');
56+
----
57+
NULL
58+
59+
# test 7: get non-existent field with type hint (should return null)
60+
query T
61+
SELECT variant_get(json_to_variant('{"name": "Charlie"}'), 'age', 'Int64');
62+
----
63+
NULL
64+
65+
# test 8: get nested field with type hint
66+
query T
67+
SELECT variant_get(json_to_variant('{"user": {"name": "David"}}'), 'user.name', 'Utf8');
68+
----
69+
David
70+
71+
# test 9: get array element with type hint
72+
query I
73+
SELECT variant_get(json_to_variant('{"items": [10, 20, 30]}'), 'items[1]', 'Int64');
74+
----
75+
20
76+
77+
# test 10: get multiple rows with type hint
78+
query I
79+
SELECT variant_get(json_to_variant(json_str), 'age', 'Int64') FROM json_data WHERE id <= 2;
80+
----
81+
30
82+
25
83+
84+
# test 11: get string field from various rows with utf8 type hint
85+
query T
86+
SELECT variant_get(json_to_variant(json_str), 'name', 'Utf8') FROM json_data WHERE id <= 2;
87+
----
88+
Alice
89+
Bob
90+
91+
# test 12: test with int16 type hint
92+
query I
93+
SELECT variant_get(json_to_variant('{"count": 5}'), 'count', 'Int16');
94+
----
95+
5
96+
97+
# test 13: test with int8 type hint
98+
query I
99+
SELECT variant_get(json_to_variant('{"level": 3}'), 'level', 'Int8');
100+
----
101+
3
102+
103+
# test 14: test with float32 type hint
104+
query R
105+
SELECT variant_get(json_to_variant('{"ratio": 0.5}'), 'ratio', 'Float32');
106+
----
107+
0.5
108+
109+
# test 15: test with large number and int64 type hint
110+
query I
111+
SELECT variant_get(json_to_variant('{"big_number": 9223372036854775807}'), 'big_number', 'Int64');
112+
----
113+
9223372036854775807
114+
115+
# test 16: test mixed types in same query - some with type hint, some without
116+
query TT
117+
SELECT
118+
variant_get(json_to_variant('{"name": "Eve", "age": 28}'), 'name', 'Utf8'),
119+
variant_pretty(variant_get(json_to_variant('{"name": "Eve", "age": 28}'), 'age'));
120+
----
121+
Eve Int8(28)
122+
123+
# test 17: test with nested path and int64 type hint
124+
query I
125+
SELECT variant_get(json_to_variant('{"data": {"count": 42}}'), 'data.count', 'Int64');
126+
----
127+
42
128+
129+
# test 18: test boolean false with type hint
130+
query B
131+
SELECT variant_get(json_to_variant('{"enabled": false}'), 'enabled', 'Boolean');
132+
----
133+
false
134+
135+
# test 19: test with decimal-like number and float64 type hint
136+
query R
137+
SELECT variant_get(json_to_variant('{"temperature": -5.5}'), 'temperature', 'Float64');
138+
----
139+
-5.5
140+
141+
# test 20: test with zero values
142+
query I
143+
SELECT variant_get(json_to_variant('{"zero": 0}'), 'zero', 'Int64');
144+
----
145+
0

0 commit comments

Comments
 (0)