diff --git a/src/datatypes/src/extension/json.rs b/src/datatypes/src/extension/json.rs index 18aee22a96..cb821df71d 100644 --- a/src/datatypes/src/extension/json.rs +++ b/src/datatypes/src/extension/json.rs @@ -15,17 +15,12 @@ use std::collections::HashMap; use std::sync::Arc; -#[cfg(test)] -use arrow_schema::FieldRef; use arrow_schema::extension::{ EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType, }; -use arrow_schema::{ArrowError, DataType, Field}; -#[cfg(test)] +use arrow_schema::{ArrowError, DataType, Field, FieldRef}; use parquet_variant_compute::VariantType; use serde::{Deserialize, Serialize}; -#[cfg(test)] -use snafu::OptionExt; use snafu::{ResultExt, ensure}; use crate::error::InvalidJson2LayoutSnafu; @@ -43,7 +38,17 @@ pub struct Json2PhysicalLayout { } impl Json2PhysicalLayout { - /// Parses and validates the JSON2 layout version from a root field. + /// Parses the JSON2 layout version from a root field. + /// + /// The version is read from the extension metadata; missing metadata + /// defaults to V1. + /// + /// Supported versions: + /// - V1 + /// - V2 + /// + /// Errors when the field is not a JSON2 extension or the version is not + /// one of the supported versions. pub fn try_from_root(field: &Field) -> crate::error::Result { ensure!( is_json2_extension_type(field), @@ -89,36 +94,28 @@ fn parse_version(metadata: &str) -> crate::error::Result> { } /// Returns the remainder field of a JSON2 v2 root. -#[cfg(test)] -fn json2_remainder_field(field: &Field) -> crate::error::Result<&FieldRef> { - if !Json2PhysicalLayout::try_from_root(field)?.is_version_2() { - return InvalidJson2LayoutSnafu { - reason: format!("JSON2 root '{}' is not layout v2", field.name()), - } - .fail(); +pub(crate) fn json2_remainder_field(field: &Field) -> crate::error::Result> { + let layout = Json2PhysicalLayout::try_from_root(field)?; + if !layout.is_version_2() { + return Ok(None); } - let DataType::Struct(fields) = field.data_type() else { - return InvalidJson2LayoutSnafu { - reason: format!( - "expecting the Struct datatype, actual: '{}'", - field.data_type(), - ), - } - .fail(); + let remainder = if let DataType::Struct(fields) = field.data_type() { + fields + .iter() + .find(|x| x.name() == JSON2_REMAINDER_FIELD_NAME) + } else { + None }; - let remainder = fields - .iter() - .find(|x| x.name() == JSON2_REMAINDER_FIELD_NAME) - .context(InvalidJson2LayoutSnafu { - reason: "remainder field not found", + + if let Some(remainder) = remainder { + let _ = remainder.try_extension_type::().map_err(|e| { + InvalidJson2LayoutSnafu { + reason: e.to_string(), + } + .build() })?; - let _ = remainder.try_extension_type::().map_err(|e| { - InvalidJson2LayoutSnafu { - reason: e.to_string(), - } - .build() - })?; + } Ok(remainder) } @@ -439,7 +436,7 @@ mod tests { JsonSettings::default(), )))); assert!(Json2PhysicalLayout::try_from_root(&v2)?.is_version_2()); - let remainder = json2_remainder_field(&v2)?; + let remainder = json2_remainder_field(&v2)?.unwrap(); assert_eq!(JSON2_REMAINDER_FIELD_NAME, remainder.name()); Ok(()) } diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index 9654f84c6d..0426334c3d 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -32,7 +32,7 @@ use crate::data_type::ConcreteDataType; use crate::error::{ AlignJsonArraySnafu, ArrowComputeSnafu, InvalidJsonSnafu, InvalidJsonbSnafu, Result, }; -use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, Json2PhysicalLayout}; +use crate::extension::json::{JSON2_REMAINDER_FIELD_NAME, json2_remainder_field}; use crate::json::value::{decode_json_variant, encode_serde_json_as_jsonb}; use crate::prelude::{DataType as _, Value as GreptimeValue}; use crate::value::{ListValue, StructValue}; @@ -105,23 +105,14 @@ impl JsonArray<'_> { } /// Projects a physical JSON2 array to a logical query type. - pub fn project_json2(&self, field: &Field, target: &DataType) -> Result { - let layout = Json2PhysicalLayout::try_from_root(field)?; - if !layout.is_version_2() { - return self.project_to(target); + /// + /// TODO(LFC) Supersede `project_to_v2` to `project_to`. + pub fn project_to_v2(&self, field: &Field, target: &DataType) -> Result { + if json2_remainder_field(field)?.is_some() { + project_json_values(self.json2_values()?, target) + } else { + self.project_to(target) } - - let has_remainder = match field.data_type() { - DataType::Struct(fields) => fields - .iter() - .any(|x| x.name() == JSON2_REMAINDER_FIELD_NAME), - _ => false, - }; - if !has_remainder { - return self.project_to(target); - } - - project_json_values(self.json2_values()?, target) } fn json2_values(&self) -> Result> { @@ -1085,7 +1076,7 @@ mod test { ] .into(), ); - let projected = JsonArray::from(&array).project_json2(&field, &target)?; + let projected = JsonArray::from(&array).project_to_v2(&field, &target)?; assert_eq!( json!({"cold": 1, "count": 42}), JsonArray::from(&projected).try_get_value(0)? @@ -1109,7 +1100,7 @@ mod test { Json2ExtensionType::new(Arc::new(JsonMetadata::new_v2(JsonSettings::default()))), ); - let projected = JsonArray::from(&array).project_json2(&field, field.data_type())?; + let projected = JsonArray::from(&array).project_to_v2(&field, field.data_type())?; assert!(Arc::ptr_eq(&array, &projected)); Ok(()) }