rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-10 07:22:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: separate a json2 extension type (#8745)

Browse Source 
Signed-off-by: luofucong <luofc@foxmail.com>
This commit is contained in:
LFC
2026-08-05 12:15:05 +00:00
committed by GitHub
parent ec09d8809a
commit 95d9d92e42
30 changed files with 388 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/query/src/datafusion/json_expr_planner.rs
+5 -10
View File
@@ -15,7 +15,6 @@
use std::sync::{Arc, LazyLock};
use arrow_schema::Field;
use arrow_schema::extension::ExtensionType;
use common_function::scalars::json::json_get::JsonGetWithType;
use common_function::scalars::udf::create_udf;
use datafusion_common::arrow::datatypes::DataType;
@@ -23,7 +22,7 @@ use datafusion_common::{Column, DFSchema, Result, ScalarValue, TableReference};
use datafusion_expr::expr::{BinaryExpr, ScalarFunction};
use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawBinaryExpr};
use datafusion_expr::{Expr, ExprSchemable, Operator, ScalarUDF};
use datatypes::extension::json::JsonExtensionType;
use datatypes::extension::json::is_json2_extension_type;
use either::Either;
use sqlparser::ast::BinaryOperator;
@@ -83,7 +82,7 @@ impl ExprPlanner for JsonExprPlanner {
qualifier: Option<&TableReference>,
nested_names: &[String],
) -> Result<PlannerResult<Vec<Expr>>> {
if field.extension_type_name() != Some(JsonExtensionType::NAME) {
if !is_json2_extension_type(field) {
return Ok(PlannerResult::Original(Vec::new()));
}
@@ -154,7 +153,7 @@ fn parse_sql_op(op: &BinaryOperator) -> Option<Operator> {
#[cfg(test)]
mod tests {
use arrow_schema::Fields;
use datatypes::extension::json::JsonMetadata;
use datatypes::extension::json::Json2ExtensionType;
use super::*;
@@ -169,11 +168,6 @@ mod tests {
))
}
fn json_field(name: &str) -> Field {
Field::new(name, DataType::Binary, true)
.with_extension_type(JsonExtensionType::new(Arc::new(JsonMetadata::default())))
}
#[test]
fn test_plan_binary_op() -> Result<()> {
let planner = JsonExprPlanner;
@@ -249,7 +243,8 @@ mod tests {
let nested_names = vec!["payload".to_string(), "cpu".to_string()];
let planned = planner.plan_compound_identifier(
&json_field("labels"),
&Field::new("labels", DataType::Struct(Fields::empty()), true)
.with_extension_type(Json2ExtensionType::default()),
Some(&qualifier),
&nested_names,
)?;
src/query/src/dist_plan/merge_scan.rs
+43 -4
View File
@@ -46,7 +46,10 @@ use datafusion_common::{Column as ColumnExpr, DataFusionError, Result};
use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore};
use datafusion_physical_expr::expressions::Column;
use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalSortExpr};
use datatypes::extension::json::is_json_extension_type;
use datatypes::extension::json::{
Json2ExtensionType, is_any_json_extension_type, is_json2_extension_type,
is_legacy_json2_extension_type,
};
use futures_util::StreamExt;
use greptime_proto::v1::region::RegionRequestHeader;
use meter_core::data::ReadItem;
@@ -122,7 +125,7 @@ fn merge_scan_schema_error_count_for_test() -> u64 {
/// parts of the field and must match.
fn json_fields_compatible(expected_field: &Field, actual_field: &Field) -> bool {
let is_json = |field: &Field| {
is_json_extension_type(field)
is_any_json_extension_type(field)
|| field
.metadata()
.get(datatypes::schema::TYPE_KEY)
@@ -913,11 +916,18 @@ fn maybe_amend_json2_field(schema: &ArrowSchema) -> ArrowSchemaRef {
let schema = schema.clone();
let mut new_fields = Vec::with_capacity(schema.fields().len());
for field in schema.fields().iter() {
let new_field = if is_json_extension_type(field)
let new_field = if is_json2_extension_type(field)
&& matches!(field.data_type(), DataType::Struct(fields) if fields.is_empty())
{
let is_legacy_json2 = is_legacy_json2_extension_type(field);
let mut new_field = field.as_ref().clone();
new_field.set_data_type(DataType::Binary);
if is_legacy_json2 {
// Pre-type-hint JSON2 is identified partly by its Struct data type. Promote the
// ephemeral field before rewriting it to Binary so later checks retain its JSON2
// identity.
new_field = new_field.with_extension_type(Json2ExtensionType::default());
}
Arc::new(new_field)
} else {
field.clone()
@@ -1291,7 +1301,10 @@ mod tests {
use std::task::{Context, Poll};
use arrow::array::{Int64Array, TimestampMillisecondArray};
use arrow_schema::{DataType as TestArrowDataType, Field, TimeUnit};
use arrow_schema::extension::{
EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
};
use arrow_schema::{DataType as TestArrowDataType, Field, Fields, TimeUnit};
use async_trait::async_trait;
use common_base::Plugins;
use common_meta::peer::Peer;
@@ -1311,6 +1324,7 @@ mod tests {
use datafusion_physical_expr::expressions::{
Column, DynamicFilterPhysicalExpr, lit as physical_lit,
};
use datatypes::extension::json::JsonExtensionType;
use datatypes::prelude::{ConcreteDataType, VectorRef};
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::vectors::{Int64Vector, StringVector, TimestampMillisecondVector};
@@ -1340,6 +1354,31 @@ mod tests {
QueryId::from(Uuid::from_u128(value))
}
#[test]
fn test_amend_legacy_json2_field_preserves_json2_identity() {
let field = Field::new("j", DataType::Struct(Fields::empty()), true).with_metadata(
StdHashMap::from([
(
EXTENSION_TYPE_NAME_KEY.to_string(),
JsonExtensionType::NAME.to_string(),
),
(
EXTENSION_TYPE_METADATA_KEY.to_string(),
serde_json::json!({
"json_structure_settings": { "Structured": null }
})
.to_string(),
),
]),
);
let schema = maybe_amend_json2_field(&ArrowSchema::new(vec![field]));
let field = schema.field(0);
assert_eq!(&DataType::Binary, field.data_type());
assert_eq!(Some(Json2ExtensionType::NAME), field.extension_type_name());
assert!(is_json2_extension_type(field));
}
fn merge_scan_exec_with_sorted_input(
region_count: u64,
target_partition: usize,
src/query/src/optimizer/json_type_concretize.rs
+2 -4
View File
@@ -156,7 +156,7 @@ mod tests {
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::{LogicalPlanBuilder, col, lit};
use datafusion_optimizer::OptimizerContext;
use datatypes::extension::json::{JsonExtensionType, JsonMetadata};
use datatypes::extension::json::Json2ExtensionType;
use datatypes::schema::ColumnSchema;
use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
use store_api::storage::{ConcreteDataType, RegionId};
@@ -196,9 +196,7 @@ mod tests {
ConcreteDataType::json2(JsonNativeType::Object(JsonObjectType::new())),
true,
);
json_column
.with_extension_type(&JsonExtensionType::new(Arc::new(JsonMetadata::default())))
.unwrap();
json_column.with_extension_type(&Json2ExtensionType::default());
builder
.push_column_metadata(ColumnMetadata {
column_schema: json_column,
src/query/src/sql/show_create_table.rs
+9 -12
View File
@@ -18,7 +18,7 @@ use std::collections::HashMap;
use arrow_schema::extension::ExtensionType;
use common_meta::SchemaOptions;
use datatypes::extension::json::JsonExtensionType;
use datatypes::extension::json::{Json2ExtensionType, parse_legacy_json2_settings};
use datatypes::schema::{
COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND,
COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE,
@@ -211,12 +211,12 @@ fn create_column(column_schema: &ColumnSchema, quote_style: char) -> Result<Colu
data_type = DataType::Custom(ObjectName::from(vec![Ident::new("JSON2")]), vec![]);
}
if let Some(json_extension) = column_schema.extension_type::<JsonExtensionType>()? {
let settings = json_extension
.metadata()
.json_settings
.clone()
.unwrap_or_default();
let settings = if let Some(extension) = column_schema.extension_type::<Json2ExtensionType>()? {
Some(extension.metadata().json_settings().clone())
} else {
parse_legacy_json2_settings(column_schema.metadata())?
};
if let Some(settings) = settings {
extensions.set_json_settings(settings).context(SqlSnafu)?;
}
@@ -321,6 +321,7 @@ mod tests {
use std::time::Duration;
use common_time::timestamp::TimeUnit;
use datatypes::extension::json::JsonExtensionType;
use datatypes::prelude::ConcreteDataType;
use datatypes::schema::{
FulltextOptions, Schema, SchemaRef, SkippingIndexOptions, VectorIndexOptions,
@@ -433,11 +434,7 @@ WITH(
#[test]
fn test_show_create_legacy_json_with_json_extension() {
let mut json_column = ColumnSchema::new("j", ConcreteDataType::json_datatype(), true);
json_column
.with_extension_type(&JsonExtensionType::new(Arc::new(
datatypes::extension::json::JsonMetadata::default(),
)))
.unwrap();
json_column.with_extension_type(&JsonExtensionType);
let table_schema = SchemaRef::new(Schema::new(vec![
json_column,