alter column's json type on demand

Signed-off-by: luofucong <[email protected]>

Author: MichaelScofield

Cargo.lock

@@ -9,6 +9,7 @@ workspace = true
 
 [dependencies]
 arc-swap = "1.6"
+arrow-schema.workspace = true
 common-base.workspace = true
 common-error.workspace = true
 common-macro.workspace = true

src/common/recordbatch/Cargo.toml

@@ -18,7 +18,7 @@ pub mod adapter;
 pub mod cursor;
 pub mod error;
 pub mod filter;
-mod recordbatch;
+pub mod recordbatch;
 pub mod util;
 
 use std::pin::Pin;

src/common/recordbatch/src/lib.rs

@@ -12,16 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::borrow::Cow;
 use std::collections::HashMap;
 use std::slice;
 use std::sync::Arc;
 
+use arrow_schema::extension::ExtensionType;
 use datafusion::arrow::util::pretty::pretty_format_batches;
 use datafusion_common::arrow::array::ArrayRef;
 use datafusion_common::arrow::compute;
 use datafusion_common::arrow::datatypes::{DataType as ArrowDataType, SchemaRef as ArrowSchemaRef};
 use datatypes::arrow::array::{Array, RecordBatchOptions};
 use datatypes::arrow::datatypes::{Field, Schema};
+use datatypes::extension::json::JsonExtensionType;
 use datatypes::prelude::DataType;
 use datatypes::schema::SchemaRef;
 use datatypes::vectors::{Helper, VectorRef};
@@ -32,7 +35,7 @@ use snafu::{OptionExt, ResultExt, ensure};
 use crate::DfRecordBatch;
 use crate::error::{
     self, ArrowComputeSnafu, CastVectorSnafu, ColumnNotExistsSnafu, DataTypesSnafu,
-    ProjectArrowRecordBatchSnafu, Result, SchemaConversionSnafu,
+    ProjectArrowRecordBatchSnafu, Result,
 };
 
 /// A two-dimensional batch of column-oriented data with a defined schema.
@@ -61,8 +64,6 @@ impl RecordBatch {
         // TODO(LFC): Remove the casting here once `Batch` is no longer used.
         let arrow_arrays = Self::cast_view_arrays(schema.arrow_schema(), arrow_arrays)?;
 
-        let schema = maybe_amend_with_struct_arrays(schema, &arrow_arrays)?;
-
         let df_record_batch = DfRecordBatch::try_new(schema.arrow_schema().clone(), arrow_arrays)
             .context(error::NewDfRecordBatchSnafu)?;
 
@@ -252,30 +253,44 @@ impl RecordBatch {
     }
 }
 
-fn maybe_amend_with_struct_arrays(schema: SchemaRef, arrays: &[ArrayRef]) -> Result<SchemaRef> {
-    if arrays
+/// Align the schema's datatypes with the actual arrays', only if there are json arrays.
+/// Because the actual datatype of json array is determined by its values, not by the pre-defined
+/// schema. So the datatype in schema could be lagged behind the actual one in the array.
+pub fn maybe_align_schema_with_json_arrays<'a>(
+    schema: &'a ArrowSchemaRef,
+    arrays: &[ArrayRef],
+) -> Cow<'a, ArrowSchemaRef> {
+    if schema
+        .fields()
         .iter()
-        .any(|x| matches!(x.data_type(), ArrowDataType::Struct(_)))
+        .any(|f| f.extension_type_name() == Some(JsonExtensionType::NAME))
     {
-        let schema = schema.arrow_schema();
         let mut fields = Vec::with_capacity(schema.fields().len());
         for (f, a) in schema.fields().iter().zip(arrays.iter()) {
-            let field_type = f.data_type();
-            let array_type = a.data_type();
-            match (field_type, array_type) {
-                (ArrowDataType::Struct(_), ArrowDataType::Struct(x)) => {
-                    let field = Field::new_struct(f.name(), x.clone(), f.is_nullable());
-                    fields.push(Arc::new(field));
+            if f.extension_type_name() == Some(JsonExtensionType::NAME) {
+                let field_type = f.data_type();
+                let array_type = a.data_type();
+                match (field_type, array_type) {
+                    (ArrowDataType::Struct(x), ArrowDataType::Struct(y)) if x != y => {
+                        common_telemetry::debug!(
+                            "Align field {} datatype: {field_type} => {array_type}",
+                            f.name(),
+                        );
+                        let field = Field::new_struct(f.name(), y.clone(), f.is_nullable())
+                            .with_metadata(f.metadata().clone());
+                        fields.push(Arc::new(field));
+                    }
+                    _ => fields.push(f.clone()),
                 }
-                _ => fields.push(f.clone()),
+            } else {
+                fields.push(f.clone())
             }
         }
         let schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()));
-        let schema = schema.try_into().context(SchemaConversionSnafu)?;
-        Ok(Arc::new(schema))
+        Cow::Owned(schema)
     } else {
-        // Fast path: must not need to amend if there are no struct arrays.
-        Ok(schema)
+        // Fast path: must not need to align if there are no json types.
+        Cow::Borrowed(schema)
     }
 }
 

src/common/recordbatch/src/recordbatch.rs

@@ -15,7 +15,6 @@
 use std::fmt;
 use std::sync::Arc;
 
-use arrow::compute::cast as arrow_array_cast;
 use arrow::datatypes::{
     DataType as ArrowDataType, IntervalUnit as ArrowIntervalUnit, TimeUnit as ArrowTimeUnit,
 };
@@ -368,8 +367,10 @@ impl ConcreteDataType {
 
     /// Checks if the data type can cast to another data type.
     pub fn can_arrow_type_cast_to(&self, to_type: &ConcreteDataType) -> bool {
-        let array = arrow_array::new_empty_array(&self.as_arrow_type());
-        arrow_array_cast(array.as_ref(), &to_type.as_arrow_type()).is_ok()
+        match (self, to_type) {
+            (ConcreteDataType::Json(this), ConcreteDataType::Json(that)) => that.is_include(this),
+            _ => arrow::compute::can_cast_types(&self.as_arrow_type(), &to_type.as_arrow_type()),
+        }
     }
 
     /// Try to cast data type as a [`DurationType`].

src/datatypes/src/data_type.rs

@@ -273,8 +273,9 @@ fn collect_fields(column_schemas: &[ColumnSchema]) -> Result<FieldsAndIndices> {
             _ => None,
         };
         if let Some(extype) = extype {
-            let metadata = HashMap::from([(TYPE_KEY.to_string(), extype.to_string())]);
-            field = field.with_metadata(metadata);
+            field
+                .metadata_mut()
+                .insert(TYPE_KEY.to_string(), extype.to_string());
         }
         fields.push(field);
         ensure!(

src/datatypes/src/schema.rs

@@ -13,11 +13,11 @@
 // limitations under the License.
 
 use std::collections::{BTreeMap, HashMap};
+use std::fmt::{Display, Formatter};
 use std::str::FromStr;
 use std::sync::Arc;
 
 use arrow::datatypes::DataType as ArrowDataType;
-use arrow_schema::Fields;
 use common_base::bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use snafu::ResultExt;
@@ -68,7 +68,7 @@ impl JsonType {
     /// - if not, the json is one of bool, number, string or array, make it a special field called
     ///   [JSON_PLAIN_FIELD_NAME] with metadata [JSON_PLAIN_FIELD_METADATA_KEY] = `"true"` in a
     ///   struct with only that field.
-    pub(crate) fn as_struct_type(&self) -> StructType {
+    pub fn as_struct_type(&self) -> StructType {
         match &self.format {
             JsonFormat::Jsonb => StructType::default(),
             JsonFormat::Native(inner) => match inner.as_ref() {
@@ -98,7 +98,7 @@ impl JsonType {
     }
 
     /// Try to merge this json type with others, error on datatype conflict.
-    pub(crate) fn merge(&mut self, other: &JsonType) -> Result<()> {
+    pub fn merge(&mut self, other: &JsonType) -> Result<()> {
         match (&self.format, &other.format) {
             (JsonFormat::Jsonb, JsonFormat::Jsonb) => Ok(()),
             (JsonFormat::Native(this), JsonFormat::Native(that)) => {
@@ -113,7 +113,8 @@ impl JsonType {
         }
     }
 
-    pub(crate) fn is_mergeable(&self, other: &JsonType) -> bool {
+    /// Check if it can merge with `other` json type.
+    pub fn is_mergeable(&self, other: &JsonType) -> bool {
         match (&self.format, &other.format) {
             (JsonFormat::Jsonb, JsonFormat::Jsonb) => true,
             (JsonFormat::Native(this), JsonFormat::Native(that)) => {
@@ -122,6 +123,49 @@ impl JsonType {
             _ => false,
         }
     }
+
+    /// Check if it includes all fields in `other` json type.
+    pub fn is_include(&self, other: &JsonType) -> bool {
+        match (&self.format, &other.format) {
+            (JsonFormat::Jsonb, JsonFormat::Jsonb) => true,
+            (JsonFormat::Native(this), JsonFormat::Native(that)) => {
+                is_include(this.as_ref(), that.as_ref())
+            }
+            _ => false,
+        }
+    }
+}
+
+fn is_include(this: &ConcreteDataType, that: &ConcreteDataType) -> bool {
+    fn is_include_struct(this: &StructType, that: &StructType) -> bool {
+        let this_fields = this.fields();
+        let this_fields = this_fields
+            .iter()
+            .map(|x| (x.name(), x))
+            .collect::<HashMap<_, _>>();
+
+        for that_field in that.fields().iter() {
+            let Some(this_field) = this_fields.get(that_field.name()) else {
+                return false;
+            };
+            if !is_include(this_field.data_type(), that_field.data_type()) {
+                return false;
+            }
+        }
+        true
+    }
+
+    match (this, that) {
+        (this, that) if this == that => true,
+        (ConcreteDataType::List(this), ConcreteDataType::List(that)) => {
+            is_include(this.item_type(), that.item_type())
+        }
+        (ConcreteDataType::Struct(this), ConcreteDataType::Struct(that)) => {
+            is_include_struct(this, that)
+        }
+        (_, ConcreteDataType::Null(_)) => true,
+        _ => false,
+    }
 }
 
 fn is_mergeable(this: &ConcreteDataType, that: &ConcreteDataType) -> bool {
@@ -223,7 +267,7 @@ impl DataType for JsonType {
     fn as_arrow_type(&self) -> ArrowDataType {
         match self.format {
             JsonFormat::Jsonb => ArrowDataType::Binary,
-            JsonFormat::Native(_) => ArrowDataType::Struct(Fields::empty()),
+            JsonFormat::Native(_) => self.as_struct_type().as_arrow_type(),
         }
     }
 
@@ -242,6 +286,12 @@ impl DataType for JsonType {
     }
 }
 
+impl Display for JsonType {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name())
+    }
+}
+
 /// Converts a json type value to string
 pub fn jsonb_to_string(val: &[u8]) -> Result<String> {
     match jsonb::from_slice(val) {

src/datatypes/src/types/json_type.rs

@@ -379,10 +379,8 @@ impl MutableVector for StructVectorBuilder {
                 },
                 StructValueRef::Ref(val) => self.push_struct_value(val)?,
                 StructValueRef::RefList { val, fields } => {
-                    let struct_value = StructValue::try_new(
-                        val.iter().map(|v| Value::from(v.clone())).collect(),
-                        fields.clone(),
-                    )?;
+                    let struct_value =
+                        StructValue::try_new(val.into_iter().map(Value::from).collect(), fields)?;
                     self.push_struct_value(&struct_value)?;
                 }
             }
@@ -429,12 +427,17 @@ impl ScalarVectorBuilder for StructVectorBuilder {
             .value_builders
             .iter_mut()
             .map(|b| b.to_vector().to_arrow_array())
-            .collect();
-        let struct_array = StructArray::new(
-            self.fields.as_arrow_fields(),
-            arrays,
-            self.null_buffer.finish(),
-        );
+            .collect::<Vec<_>>();
+
+        let struct_array = if arrays.is_empty() {
+            StructArray::new_empty_fields(self.len(), self.null_buffer.finish())
+        } else {
+            StructArray::new(
+                self.fields.as_arrow_fields(),
+                arrays,
+                self.null_buffer.finish(),
+            )
+        };
 
         StructVector::try_new(self.fields.clone(), struct_array).unwrap()
     }

src/datatypes/src/vectors/struct_vector.rs

@@ -96,6 +96,7 @@ impl PrimaryKeyWriteFormat {
     }
 
     /// Gets the arrow schema to store in parquet.
+    #[cfg(test)]
     pub(crate) fn arrow_schema(&self) -> &SchemaRef {
         &self.arrow_schema
     }
@@ -135,7 +136,11 @@ impl PrimaryKeyWriteFormat {
         }
         columns.push(batch.op_types().to_arrow_array());
 
-        RecordBatch::try_new(self.arrow_schema.clone(), columns).context(NewRecordBatchSnafu)
+        let schema = common_recordbatch::recordbatch::maybe_align_schema_with_json_arrays(
+            &self.arrow_schema,
+            &columns,
+        );
+        RecordBatch::try_new(schema.as_ref().clone(), columns).context(NewRecordBatchSnafu)
     }
 }
 

src/mito2/src/sst/parquet/format.rs

@@ -411,7 +411,7 @@ where
         let arrow_batch = write_format.convert_batch(&batch)?;
 
         let start = Instant::now();
-        self.maybe_init_writer(write_format.arrow_schema(), opts)
+        self.maybe_init_writer(arrow_batch.schema_ref(), opts)
             .await?
             .write(&arrow_batch)
             .await