Skip to content

Commit f71b883

Browse files
adriangbadriangb
committed
implement suggestion, fix test
1 parent 2783193 commit f71b883

2 files changed

Lines changed: 71 additions & 148 deletions

File tree

‎src/json_from_scalar.rs‎

Lines changed: 67 additions & 147 deletions
Original file line numberDiff line numberDiff line change
@@ -150,101 +150,50 @@ fn scalar_to_json_union_field(scalar: ScalarValue) -> DataFusionResult<JsonUnion
150150
}
151151
}
152152

153-
#[expect(clippy::too_many_lines)]
154153
fn array_to_json_union(array: &ArrayRef) -> DataFusionResult<JsonUnion> {
155-
let mut union = JsonUnion::new(array.len());
156-
157-
match array.data_type() {
158-
DataType::Null => {
159-
for _ in 0..array.len() {
160-
union.push(JsonUnionField::JsonNull);
161-
}
162-
}
163-
164-
DataType::Boolean => {
165-
let arr = array.as_boolean();
166-
for i in 0..arr.len() {
167-
if arr.is_null(i) {
168-
union.push_none();
169-
} else {
170-
union.push(JsonUnionField::Bool(arr.value(i)));
171-
}
172-
}
173-
}
174-
154+
Ok(match array.data_type() {
155+
DataType::Null => (0..array.len()).map(|_| Some(JsonUnionField::JsonNull)).collect(),
156+
DataType::Boolean => array.as_boolean().iter().map(|v| v.map(JsonUnionField::Bool)).collect(),
175157
// Integer types - coerce to i64
176-
DataType::Int8 => {
177-
let arr = array.as_primitive::<Int8Type>();
178-
for i in 0..arr.len() {
179-
if arr.is_null(i) {
180-
union.push_none();
181-
} else {
182-
union.push(JsonUnionField::Int(i64::from(arr.value(i))));
183-
}
184-
}
185-
}
186-
DataType::Int16 => {
187-
let arr = array.as_primitive::<Int16Type>();
188-
for i in 0..arr.len() {
189-
if arr.is_null(i) {
190-
union.push_none();
191-
} else {
192-
union.push(JsonUnionField::Int(i64::from(arr.value(i))));
193-
}
194-
}
195-
}
196-
DataType::Int32 => {
197-
let arr = array.as_primitive::<Int32Type>();
198-
for i in 0..arr.len() {
199-
if arr.is_null(i) {
200-
union.push_none();
201-
} else {
202-
union.push(JsonUnionField::Int(i64::from(arr.value(i))));
203-
}
204-
}
205-
}
206-
DataType::Int64 => {
207-
let arr = array.as_primitive::<Int64Type>();
208-
for i in 0..arr.len() {
209-
if arr.is_null(i) {
210-
union.push_none();
211-
} else {
212-
union.push(JsonUnionField::Int(arr.value(i)));
213-
}
214-
}
215-
}
216-
DataType::UInt8 => {
217-
let arr = array.as_primitive::<UInt8Type>();
218-
for i in 0..arr.len() {
219-
if arr.is_null(i) {
220-
union.push_none();
221-
} else {
222-
union.push(JsonUnionField::Int(i64::from(arr.value(i))));
223-
}
224-
}
225-
}
226-
DataType::UInt16 => {
227-
let arr = array.as_primitive::<UInt16Type>();
228-
for i in 0..arr.len() {
229-
if arr.is_null(i) {
230-
union.push_none();
231-
} else {
232-
union.push(JsonUnionField::Int(i64::from(arr.value(i))));
233-
}
234-
}
235-
}
236-
DataType::UInt32 => {
237-
let arr = array.as_primitive::<UInt32Type>();
238-
for i in 0..arr.len() {
239-
if arr.is_null(i) {
240-
union.push_none();
241-
} else {
242-
union.push(JsonUnionField::Int(i64::from(arr.value(i))));
243-
}
244-
}
245-
}
158+
DataType::Int8 => array
159+
.as_primitive::<Int8Type>()
160+
.iter()
161+
.map(|v| v.map(|x| JsonUnionField::Int(i64::from(x))))
162+
.collect(),
163+
DataType::Int16 => array
164+
.as_primitive::<Int16Type>()
165+
.iter()
166+
.map(|v| v.map(|x| JsonUnionField::Int(i64::from(x))))
167+
.collect(),
168+
DataType::Int32 => array
169+
.as_primitive::<Int32Type>()
170+
.iter()
171+
.map(|v| v.map(|x| JsonUnionField::Int(i64::from(x))))
172+
.collect(),
173+
DataType::Int64 => array
174+
.as_primitive::<Int64Type>()
175+
.iter()
176+
.map(|v| v.map(JsonUnionField::Int))
177+
.collect(),
178+
DataType::UInt8 => array
179+
.as_primitive::<UInt8Type>()
180+
.iter()
181+
.map(|v| v.map(|x| JsonUnionField::Int(i64::from(x))))
182+
.collect(),
183+
DataType::UInt16 => array
184+
.as_primitive::<UInt16Type>()
185+
.iter()
186+
.map(|v| v.map(|x| JsonUnionField::Int(i64::from(x))))
187+
.collect(),
188+
DataType::UInt32 => array
189+
.as_primitive::<UInt32Type>()
190+
.iter()
191+
.map(|v| v.map(|x| JsonUnionField::Int(i64::from(x))))
192+
.collect(),
246193
DataType::UInt64 => {
194+
// UInt64 requires explicit loop for fallible conversion
247195
let arr = array.as_primitive::<UInt64Type>();
196+
let mut union = JsonUnion::new(arr.len());
248197
for i in 0..arr.len() {
249198
if arr.is_null(i) {
250199
union.push_none();
@@ -254,66 +203,37 @@ fn array_to_json_union(array: &ArrayRef) -> DataFusionResult<JsonUnion> {
254203
})?));
255204
}
256205
}
206+
return Ok(union);
257207
}
258-
259208
// Float types - coerce to f64
260-
DataType::Float32 => {
261-
let arr = array.as_primitive::<Float32Type>();
262-
for i in 0..arr.len() {
263-
if arr.is_null(i) {
264-
union.push_none();
265-
} else {
266-
union.push(JsonUnionField::Float(f64::from(arr.value(i))));
267-
}
268-
}
269-
}
270-
DataType::Float64 => {
271-
let arr = array.as_primitive::<Float64Type>();
272-
for i in 0..arr.len() {
273-
if arr.is_null(i) {
274-
union.push_none();
275-
} else {
276-
union.push(JsonUnionField::Float(arr.value(i)));
277-
}
278-
}
279-
}
280-
209+
DataType::Float32 => array
210+
.as_primitive::<Float32Type>()
211+
.iter()
212+
.map(|v| v.map(|x| JsonUnionField::Float(f64::from(x))))
213+
.collect(),
214+
DataType::Float64 => array
215+
.as_primitive::<Float64Type>()
216+
.iter()
217+
.map(|v| v.map(JsonUnionField::Float))
218+
.collect(),
281219
// String types
282-
DataType::Utf8 => {
283-
let arr = array.as_string::<i32>();
284-
for i in 0..arr.len() {
285-
if arr.is_null(i) {
286-
union.push_none();
287-
} else {
288-
union.push(JsonUnionField::Str(arr.value(i).to_string()));
289-
}
290-
}
291-
}
292-
DataType::LargeUtf8 => {
293-
let arr = array.as_string::<i64>();
294-
for i in 0..arr.len() {
295-
if arr.is_null(i) {
296-
union.push_none();
297-
} else {
298-
union.push(JsonUnionField::Str(arr.value(i).to_string()));
299-
}
300-
}
301-
}
302-
DataType::Utf8View => {
303-
let arr = array.as_string_view();
304-
for i in 0..arr.len() {
305-
if arr.is_null(i) {
306-
union.push_none();
307-
} else {
308-
union.push(JsonUnionField::Str(arr.value(i).to_string()));
309-
}
310-
}
311-
}
312-
220+
DataType::Utf8 => array
221+
.as_string::<i32>()
222+
.iter()
223+
.map(|v| v.map(|s| JsonUnionField::Str(s.to_string())))
224+
.collect(),
225+
DataType::LargeUtf8 => array
226+
.as_string::<i64>()
227+
.iter()
228+
.map(|v| v.map(|s| JsonUnionField::Str(s.to_string())))
229+
.collect(),
230+
DataType::Utf8View => array
231+
.as_string_view()
232+
.iter()
233+
.map(|v| v.map(|s| JsonUnionField::Str(s.to_string())))
234+
.collect(),
313235
dt => {
314236
return exec_err!("Unsupported array type for json_from_scalar: {:?}", dt);
315237
}
316-
}
317-
318-
Ok(union)
238+
})
319239
}

‎tests/main.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2426,7 +2426,10 @@ async fn test_json_from_scalar_too_many_args() {
24262426
let result = run_query(sql).await;
24272427
assert!(result.is_err());
24282428
let err = result.unwrap_err().to_string();
2429-
assert!(err.contains("requires exactly one argument"));
2429+
assert!(
2430+
err.contains("The function 'json_from_scalar' expected 1 arguments but received 2"),
2431+
"Err: {err}"
2432+
);
24302433
}
24312434

24322435
#[tokio::test]

0 commit comments

Comments
 (0)