diff --git a/src/common.rs b/src/common.rs index 66505cb..b4c17ce 100644 --- a/src/common.rs +++ b/src/common.rs @@ -68,6 +68,7 @@ fn dict_key_type(d: &DataType) -> Option { None } +/// Parsed representation of a JSON path element. #[derive(Debug)] pub enum JsonPath<'s> { Key(&'s str), @@ -97,13 +98,17 @@ impl From for JsonPath<'_> { } #[derive(Debug)] -enum JsonPathArgs<'a> { +pub enum JsonPathArgs<'a> { Array(&'a ArrayRef), Scalars(Vec>), } impl<'s> JsonPathArgs<'s> { - fn extract_path(path_args: &'s [ColumnarValue]) -> DataFusionResult { + /// Extract the JSON path from the arguments to a json function. + /// The path can be: + /// - A single array argument (`json_get(json_data, keys_column)`) + /// - One or more scalar arguments (`json_get(json_data, 'foo', 0)`) + pub fn extract_path(path_args: &'s [ColumnarValue]) -> DataFusionResult { // If there is a single argument as an array, we know how to handle it if let Some((ColumnarValue::Array(array), &[])) = path_args.split_first() { return Ok(Self::Array(array)); diff --git a/src/common_union.rs b/src/common_union.rs index 36fcb79..cd88389 100644 --- a/src/common_union.rs +++ b/src/common_union.rs @@ -101,6 +101,13 @@ impl JsonUnion { } } +/// Get the `DataType` for a `JsonUnion`. +/// This will be the data type of the result of `json_array->'field'`. +#[must_use] +pub fn json_union_data_type() -> DataType { + JsonUnion::data_type() +} + /// So we can do `collect::()` impl FromIterator> for JsonUnion { fn from_iter>>(iter: I) -> Self { diff --git a/src/lib.rs b/src/lib.rs index cb0f25a..a11dab7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,7 +20,9 @@ mod json_length; mod json_object_keys; mod rewrite; -pub use common_union::{JsonUnionEncoder, JsonUnionValue}; +pub use common_union::{json_union_data_type, JsonUnionEncoder, JsonUnionValue}; + +pub use common::JsonPath; pub mod functions { pub use crate::json_as_text::json_as_text;