diff --git a/src/json_get_array.rs b/src/json_get_array.rs index 0a37b6a..9bbb5d5 100644 --- a/src/json_get_array.rs +++ b/src/json_get_array.rs @@ -11,33 +11,6 @@ use jiter::Peek; use crate::common::{check_args, get_err, invoke, jiter_json_find, GetError, JsonPath}; use crate::common_macros::make_udf_function; -struct StrArrayColumn { - rows: GenericListArray, -} - -impl FromIterator>> for StrArrayColumn { - fn from_iter>>>(iter: T) -> Self { - let string_builder = StringBuilder::new(); - let mut list_builder = ListBuilder::new(string_builder); - - for row in iter { - if let Some(row) = row { - for elem in row { - list_builder.values().append_value(elem); - } - - list_builder.append(true); - } else { - list_builder.append(false); - } - } - - Self { - rows: list_builder.finish(), - } - } -} - make_udf_function!( JsonGetArray, json_get_array, @@ -78,7 +51,7 @@ impl ScalarUDFImpl for JsonGetArray { } fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult { - invoke::>( + invoke::( args, jiter_json_get_array, |c| Ok(Arc::new(c.rows) as ArrayRef), @@ -87,7 +60,7 @@ impl ScalarUDFImpl for JsonGetArray { let mut list_builder = ListBuilder::new(string_builder); if let Some(row) = i { - for elem in row { + for elem in row.elements { list_builder.values().append_value(elem); } } @@ -102,12 +75,43 @@ impl ScalarUDFImpl for JsonGetArray { } } -fn jiter_json_get_array(json_data: Option<&str>, path: &[JsonPath]) -> Result, GetError> { +struct JsonArray { + rows: GenericListArray, +} + +struct JsonArrayField { + elements: Vec, +} + +impl FromIterator> for JsonArray { + fn from_iter>>(iter: T) -> Self { + let string_builder = StringBuilder::new(); + let mut list_builder = ListBuilder::new(string_builder); + + for row in iter { + if let Some(row) = row { + for elem in row.elements { + list_builder.values().append_value(elem); + } + + list_builder.append(true); + } else { + list_builder.append(false); + } + } + + Self { + rows: list_builder.finish(), + } + } +} + +fn jiter_json_get_array(json_data: Option<&str>, path: &[JsonPath]) -> Result { if let Some((mut jiter, peek)) = jiter_json_find(json_data, path) { match peek { Peek::Array => { let mut peek_opt = jiter.known_array()?; - let mut array_values = Vec::new(); + let mut elements = Vec::new(); while let Some(peek) = peek_opt { let start = jiter.current_index(); @@ -115,12 +119,12 @@ fn jiter_json_get_array(json_data: Option<&str>, path: &[JsonPath]) -> Result get_err!(), }