diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index de6e9328..8ede2a05 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -54,6 +54,7 @@ mod system_tables; mod table; mod table_function_args; mod update; +mod variant_functions; mod vector_search; use std::collections::HashMap; @@ -74,4 +75,5 @@ pub use physical_plan::PaimonTableScan; pub use relation_planner::PaimonRelationPlanner; pub use sql_context::SQLContext; pub use table::PaimonTableProvider; +pub use variant_functions::register_variant_functions; pub use vector_search::{register_vector_search, VectorSearchFunction}; diff --git a/crates/integrations/datafusion/src/sql_context.rs b/crates/integrations/datafusion/src/sql_context.rs index 6170d6bc..c550effe 100644 --- a/crates/integrations/datafusion/src/sql_context.rs +++ b/crates/integrations/datafusion/src/sql_context.rs @@ -62,7 +62,7 @@ use paimon::spec::{ ArrayType as PaimonArrayType, BigIntType, BlobType, BooleanType, DataField as PaimonDataField, DataType as PaimonDataType, DateType, Datum, DecimalType, DoubleType, FloatType, IntType, LocalZonedTimestampType, MapType as PaimonMapType, RowType as PaimonRowType, SchemaChange, - SmallIntType, TimestampType, TinyIntType, VarBinaryType, VarCharType, + SmallIntType, TimestampType, TinyIntType, VarBinaryType, VarCharType, VariantType, }; use crate::error::to_datafusion_error; @@ -106,6 +106,7 @@ impl SQLContext { )) .build(); let ctx = SessionContext::new_with_state(state); + crate::variant_functions::register_variant_functions(&ctx); Self { ctx, catalogs: HashMap::new(), @@ -1683,6 +1684,13 @@ fn sql_data_type_to_paimon_type( )) } SqlType::Blob(_) => Ok(PaimonDataType::Blob(BlobType::with_nullable(nullable))), + SqlType::Custom(name, modifiers) + if name.to_string().eq_ignore_ascii_case("VARIANT") && modifiers.is_empty() => + { + Ok(PaimonDataType::Variant(VariantType::with_nullable( + nullable, + ))) + } SqlType::Date => Ok(PaimonDataType::Date(DateType::with_nullable(nullable))), SqlType::Timestamp(precision, tz_info) => { let precision = match precision { @@ -2122,7 +2130,8 @@ fn datum_to_constant_array( | Datum::Timestamp { .. } | Datum::LocalZonedTimestamp { .. } | Datum::Decimal { .. } - | Datum::Bytes(_) => Err(DataFusionError::Plan(format!( + | Datum::Bytes(_) + | Datum::Variant { .. } => Err(DataFusionError::Plan(format!( "Unsupported datum type for partition column: {d}" ))), }, @@ -2835,6 +2844,15 @@ mod tests { ); } + #[test] + fn test_sql_type_variant() { + use datafusion::sql::sqlparser::ast::{DataType as SqlType, Ident, ObjectName}; + assert_sql_type_to_paimon( + SqlType::Custom(ObjectName::from(Ident::new("VARIANT")), vec![]), + PaimonDataType::Variant(VariantType::new()), + ); + } + #[test] fn test_sql_type_date() { use datafusion::sql::sqlparser::ast::DataType as SqlType; diff --git a/crates/integrations/datafusion/src/variant_functions.rs b/crates/integrations/datafusion/src/variant_functions.rs new file mode 100644 index 00000000..6eb8ab85 --- /dev/null +++ b/crates/integrations/datafusion/src/variant_functions.rs @@ -0,0 +1,877 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use datafusion::arrow::array::{ + Array, ArrayRef, BinaryArray, BinaryBuilder, BooleanBuilder, LargeStringArray, StringArray, + StringViewArray, StructArray, +}; +use datafusion::arrow::buffer::{BooleanBuffer, NullBuffer}; +use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, FieldRef, Fields}; +use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue}; +use datafusion::logical_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion::prelude::SessionContext; +use paimon::variant::{GenericVariant, VariantDecimal, VariantKind, VariantRef}; + +pub fn register_variant_functions(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(ParseJsonFunc::new(false))); + ctx.register_udf(ScalarUDF::from(ParseJsonFunc::new(true))); + ctx.register_udf(ScalarUDF::from(IsVariantNullFunc::new())); + ctx.register_udf(ScalarUDF::from(VariantGetFunc::new(false))); + ctx.register_udf(ScalarUDF::from(VariantGetFunc::new(true))); +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct ParseJsonFunc { + try_parse: bool, + signature: Signature, +} + +impl ParseJsonFunc { + fn new(try_parse: bool) -> Self { + Self { + try_parse, + signature: Signature::string(1, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for ParseJsonFunc { + fn name(&self) -> &str { + if self.try_parse { + "try_parse_json" + } else { + "parse_json" + } + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[ArrowDataType]) -> DFResult { + Ok(variant_arrow_type()) + } + + fn return_field_from_args(&self, _args: ReturnFieldArgs) -> DFResult { + Ok(Arc::new(Field::new( + self.name(), + variant_arrow_type(), + true, + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult { + if args.args.len() != 1 { + return plan_err(format!("{} expects 1 argument", self.name())); + } + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let input = arrays[0].as_ref(); + let mut values = Vec::with_capacity(input.len()); + for row in 0..input.len() { + let Some(json) = string_at(input, row)? else { + values.push(None); + continue; + }; + match GenericVariant::parse_json(&json) { + Ok(variant) => values.push(Some(variant)), + Err(e) if self.try_parse => { + let _ = e; + values.push(None); + } + Err(e) => return Err(to_df_error(e)), + } + } + Ok(ColumnarValue::Array(variant_array(values)?)) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct IsVariantNullFunc { + signature: Signature, +} + +impl IsVariantNullFunc { + fn new() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for IsVariantNullFunc { + fn name(&self) -> &str { + "is_variant_null" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[ArrowDataType]) -> DFResult { + Ok(ArrowDataType::Boolean) + } + + fn return_field_from_args(&self, _args: ReturnFieldArgs) -> DFResult { + Ok(Arc::new(Field::new( + self.name(), + ArrowDataType::Boolean, + false, + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult { + if args.args.len() != 1 { + return plan_err("is_variant_null expects 1 argument"); + } + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let input = arrays[0].as_ref(); + let mut builder = BooleanBuilder::new(); + let Some((values, metadata)) = variant_children(input)? else { + for _ in 0..input.len() { + builder.append_value(false); + } + return Ok(ColumnarValue::Array(Arc::new(builder.finish()))); + }; + + for row in 0..input.len() { + if input.is_null(row) { + builder.append_value(false); + } else { + let variant = VariantRef::new(values.value(row), metadata.value(row), 0) + .map_err(to_df_error)?; + builder.append_value(variant.is_null().map_err(to_df_error)?); + } + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct VariantGetFunc { + try_get: bool, + signature: Signature, +} + +impl VariantGetFunc { + fn new(try_get: bool) -> Self { + Self { + try_get, + signature: Signature::variadic_any(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for VariantGetFunc { + fn name(&self) -> &str { + if self.try_get { + "try_variant_get" + } else { + "variant_get" + } + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[ArrowDataType]) -> DFResult { + internal_err("return_field_from_args should be used for variant_get") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> DFResult { + if args.arg_fields.len() != 2 && args.arg_fields.len() != 3 { + return plan_err(format!("{} expects 2 or 3 arguments", self.name())); + } + let output = match args.arg_fields.len() { + 2 => variant_get_output_type(None)?, + 3 => { + let Some(type_arg) = args.scalar_arguments.get(2).and_then(|v| *v) else { + return plan_err("variant_get type argument must be a string literal"); + }; + variant_get_output_type(Some(type_arg))? + } + _ => unreachable!("argument count checked above"), + }; + Ok(Arc::new(Field::new( + self.name(), + output.arrow_type().clone(), + true, + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult { + if args.args.len() != 2 && args.args.len() != 3 { + return plan_err(format!("{} expects 2 or 3 arguments", self.name())); + } + let output = if args.return_type() == &variant_arrow_type() { + VariantGetOutput::Variant + } else { + VariantGetOutput::Scalar(args.return_type().clone()) + }; + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let variants = arrays[0].as_ref(); + let paths = arrays[1].as_ref(); + let Some((values, metadata)) = variant_children(variants)? else { + return Ok(ColumnarValue::Array(null_array( + output.arrow_type(), + variants.len(), + ))); + }; + + match output { + VariantGetOutput::Variant => { + let mut result = Vec::with_capacity(variants.len()); + for row in 0..variants.len() { + result.push(self.variant_at_path(variants, values, metadata, paths, row)?); + } + Ok(ColumnarValue::Array(variant_array(result)?)) + } + VariantGetOutput::Scalar(data_type) => { + let mut scalars = Vec::with_capacity(variants.len()); + for row in 0..variants.len() { + match self.variant_at_path_ref(variants, values, metadata, paths, row)? { + Some(variant) => scalars.push(cast_variant_to_scalar( + variant, + &data_type, + !self.try_get, + )?), + None => scalars.push(ScalarValue::try_from(&data_type)?), + } + } + if scalars.is_empty() { + return Ok(ColumnarValue::Array(null_array(&data_type, 0))); + } + Ok(ColumnarValue::Array(ScalarValue::iter_to_array(scalars)?)) + } + } + } +} + +impl VariantGetFunc { + fn variant_at_path( + &self, + variants: &dyn Array, + values: &BinaryArray, + metadata: &BinaryArray, + paths: &dyn Array, + row: usize, + ) -> DFResult> { + self.variant_at_path_ref(variants, values, metadata, paths, row)? + .map(|variant| variant.to_owned_variant().map_err(to_df_error)) + .transpose() + } + + fn variant_at_path_ref<'a>( + &self, + variants: &dyn Array, + values: &'a BinaryArray, + metadata: &'a BinaryArray, + paths: &dyn Array, + row: usize, + ) -> DFResult>> { + if variants.is_null(row) || paths.is_null(row) { + return Ok(None); + } + let path = string_at(paths, row)?; + let Some(path) = path else { + return Ok(None); + }; + let variant = + VariantRef::new(values.value(row), metadata.value(row), 0).map_err(to_df_error)?; + match variant.get_path(&path) { + Ok(value) => Ok(value), + Err(e) if self.try_get => { + let _ = e; + Ok(None) + } + Err(e) => Err(to_df_error(e)), + } + } +} + +#[derive(Clone, Debug)] +enum VariantGetOutput { + Variant, + Scalar(ArrowDataType), +} + +impl VariantGetOutput { + fn arrow_type(&self) -> &ArrowDataType { + match self { + Self::Variant => { + static VARIANT_TYPE: std::sync::LazyLock = + std::sync::LazyLock::new(variant_arrow_type); + &VARIANT_TYPE + } + Self::Scalar(data_type) => data_type, + } + } +} + +fn variant_get_output_type(type_arg: Option<&ScalarValue>) -> DFResult { + let Some(type_arg) = type_arg else { + return Ok(VariantGetOutput::Variant); + }; + let type_name = match type_arg { + ScalarValue::Utf8(Some(v)) + | ScalarValue::LargeUtf8(Some(v)) + | ScalarValue::Utf8View(Some(v)) => v, + ScalarValue::Utf8(None) | ScalarValue::LargeUtf8(None) | ScalarValue::Utf8View(None) => { + return plan_err("variant_get type argument must not be NULL"); + } + _ => return plan_err("variant_get type argument must be a string literal"), + }; + parse_variant_get_type(type_name) +} + +fn parse_variant_get_type(type_name: &str) -> DFResult { + let normalized = type_name.trim().to_ascii_lowercase(); + match normalized.as_str() { + "variant" => Ok(VariantGetOutput::Variant), + "boolean" | "bool" => Ok(VariantGetOutput::Scalar(ArrowDataType::Boolean)), + "byte" | "tinyint" => Ok(VariantGetOutput::Scalar(ArrowDataType::Int8)), + "short" | "smallint" => Ok(VariantGetOutput::Scalar(ArrowDataType::Int16)), + "int" | "integer" => Ok(VariantGetOutput::Scalar(ArrowDataType::Int32)), + "long" | "bigint" => Ok(VariantGetOutput::Scalar(ArrowDataType::Int64)), + "float" | "real" => Ok(VariantGetOutput::Scalar(ArrowDataType::Float32)), + "double" => Ok(VariantGetOutput::Scalar(ArrowDataType::Float64)), + "string" | "varchar" | "text" => Ok(VariantGetOutput::Scalar(ArrowDataType::Utf8)), + "decimal" => Ok(VariantGetOutput::Scalar(ArrowDataType::Decimal128(10, 0))), + _ if normalized.starts_with("decimal(") && normalized.ends_with(')') => { + let inner = &normalized["decimal(".len()..normalized.len() - 1]; + let Some((precision, scale)) = inner.split_once(',') else { + return plan_err(format!("Invalid decimal type for variant_get: {type_name}")); + }; + let precision = precision + .trim() + .parse::() + .map_err(|e| DataFusionError::Plan(format!("Invalid decimal precision: {e}")))?; + let scale = scale + .trim() + .parse::() + .map_err(|e| DataFusionError::Plan(format!("Invalid decimal scale: {e}")))?; + Ok(VariantGetOutput::Scalar(ArrowDataType::Decimal128( + precision, scale, + ))) + } + _ => plan_err(format!("Unsupported variant_get type: {type_name}")), + } +} + +fn cast_variant_to_scalar( + variant: VariantRef<'_>, + target: &ArrowDataType, + fail_on_error: bool, +) -> DFResult { + if variant.is_null().map_err(to_df_error)? { + return ScalarValue::try_from(target); + } + let result = match target { + ArrowDataType::Boolean => cast_to_boolean(variant), + ArrowDataType::Int8 => cast_to_i64(variant).and_then(|v| { + i8::try_from(v) + .map(ScalarValue::from) + .map_err(|_| invalid_cast()) + }), + ArrowDataType::Int16 => cast_to_i64(variant).and_then(|v| { + i16::try_from(v) + .map(ScalarValue::from) + .map_err(|_| invalid_cast()) + }), + ArrowDataType::Int32 => cast_to_i64(variant).and_then(|v| { + i32::try_from(v) + .map(ScalarValue::from) + .map_err(|_| invalid_cast()) + }), + ArrowDataType::Int64 => cast_to_i64(variant).map(ScalarValue::from), + ArrowDataType::Float32 => { + cast_to_f64(variant).map(|v| ScalarValue::Float32(Some(v as f32))) + } + ArrowDataType::Float64 => cast_to_f64(variant).map(ScalarValue::from), + ArrowDataType::Utf8 => cast_to_string(variant).map(ScalarValue::from), + ArrowDataType::Decimal128(precision, scale) => cast_to_decimal(variant, *precision, *scale), + _ => Err(invalid_cast()), + }; + + match result { + Ok(value) => Ok(value), + Err(e) if !fail_on_error => { + let _ = e; + ScalarValue::try_from(target) + } + Err(e) => Err(e), + } +} + +fn cast_to_boolean(variant: VariantRef<'_>) -> DFResult { + match variant.kind().map_err(to_df_error)? { + VariantKind::Boolean => Ok(ScalarValue::Boolean(Some( + variant.get_boolean().map_err(to_df_error)?, + ))), + VariantKind::String => match variant + .get_string() + .map_err(to_df_error)? + .to_ascii_lowercase() + .as_str() + { + "true" => Ok(ScalarValue::Boolean(Some(true))), + "false" => Ok(ScalarValue::Boolean(Some(false))), + _ => Err(invalid_cast()), + }, + _ => Err(invalid_cast()), + } +} + +fn cast_to_i64(variant: VariantRef<'_>) -> DFResult { + match variant.kind().map_err(to_df_error)? { + VariantKind::Long + | VariantKind::Date + | VariantKind::Timestamp + | VariantKind::TimestampNtz => variant.get_long().map_err(to_df_error), + VariantKind::String => variant + .get_string() + .map_err(to_df_error)? + .parse::() + .map_err(|_| invalid_cast()), + VariantKind::Decimal => { + let decimal = variant.get_decimal().map_err(to_df_error)?; + rescale_decimal(decimal.unscaled, decimal.scale, 0) + .and_then(|v| i64::try_from(v).map_err(|_| invalid_cast())) + } + _ => Err(invalid_cast()), + } +} + +fn cast_to_f64(variant: VariantRef<'_>) -> DFResult { + match variant.kind().map_err(to_df_error)? { + VariantKind::Long + | VariantKind::Date + | VariantKind::Timestamp + | VariantKind::TimestampNtz => Ok(variant.get_long().map_err(to_df_error)? as f64), + VariantKind::Double => variant.get_double().map_err(to_df_error), + VariantKind::Float => Ok(variant.get_float().map_err(to_df_error)? as f64), + VariantKind::Decimal => { + let decimal = variant.get_decimal().map_err(to_df_error)?; + Ok(decimal.unscaled as f64 / 10f64.powi(decimal.scale as i32)) + } + VariantKind::String => variant + .get_string() + .map_err(to_df_error)? + .parse::() + .map_err(|_| invalid_cast()), + _ => Err(invalid_cast()), + } +} + +fn cast_to_string(variant: VariantRef<'_>) -> DFResult { + match variant.kind().map_err(to_df_error)? { + VariantKind::Object | VariantKind::Array => variant.to_json().map_err(to_df_error), + VariantKind::Boolean => Ok(variant.get_boolean().map_err(to_df_error)?.to_string()), + VariantKind::Long + | VariantKind::Date + | VariantKind::Timestamp + | VariantKind::TimestampNtz => Ok(variant.get_long().map_err(to_df_error)?.to_string()), + VariantKind::String => variant.get_string().map_err(to_df_error), + VariantKind::Double => Ok(variant.get_double().map_err(to_df_error)?.to_string()), + VariantKind::Decimal => Ok(variant + .get_decimal() + .map_err(to_df_error)? + .to_plain_string()), + VariantKind::Float => Ok(variant.get_float().map_err(to_df_error)?.to_string()), + _ => variant.to_json().map_err(to_df_error), + } +} + +fn cast_to_decimal(variant: VariantRef<'_>, precision: u8, scale: i8) -> DFResult { + let unscaled = match variant.kind().map_err(to_df_error)? { + VariantKind::Long + | VariantKind::Date + | VariantKind::Timestamp + | VariantKind::TimestampNtz => { + rescale_decimal(variant.get_long().map_err(to_df_error)? as i128, 0, scale)? + } + VariantKind::Decimal => { + let decimal = variant.get_decimal().map_err(to_df_error)?; + rescale_decimal(decimal.unscaled, decimal.scale, scale)? + } + VariantKind::String => { + let parsed = parse_decimal_string(&variant.get_string().map_err(to_df_error)?) + .ok_or_else(invalid_cast)?; + rescale_decimal(parsed.unscaled, parsed.scale, scale)? + } + _ => return Err(invalid_cast()), + }; + if decimal_precision(unscaled) > precision { + return Err(invalid_cast()); + } + Ok(ScalarValue::Decimal128(Some(unscaled), precision, scale)) +} + +fn rescale_decimal(unscaled: i128, from_scale: i8, to_scale: i8) -> DFResult { + match to_scale.cmp(&from_scale) { + std::cmp::Ordering::Equal => Ok(unscaled), + std::cmp::Ordering::Greater => { + let factor = 10_i128 + .checked_pow((to_scale - from_scale) as u32) + .ok_or_else(invalid_cast)?; + unscaled.checked_mul(factor).ok_or_else(invalid_cast) + } + std::cmp::Ordering::Less => { + let factor = 10_i128 + .checked_pow((from_scale - to_scale) as u32) + .ok_or_else(invalid_cast)?; + if unscaled % factor == 0 { + Ok(unscaled / factor) + } else { + Err(invalid_cast()) + } + } + } +} + +fn parse_decimal_string(input: &str) -> Option { + let input = input.trim(); + if input.is_empty() || input.contains(['e', 'E']) { + return None; + } + let negative = input.starts_with('-'); + let unsigned = input.strip_prefix('-').unwrap_or(input); + if unsigned.is_empty() + || unsigned.matches('.').count() > 1 + || !unsigned.bytes().all(|ch| ch == b'.' || ch.is_ascii_digit()) + { + return None; + } + let scale = unsigned + .split_once('.') + .map(|(_, fraction)| fraction.len()) + .unwrap_or(0); + let digits: String = unsigned + .bytes() + .filter(|ch| *ch != b'.') + .map(char::from) + .collect(); + let significant = digits.trim_start_matches('0'); + let precision = if significant.is_empty() { + 1 + } else { + significant.len() + }; + if precision > 38 || scale > 38 { + return None; + } + let mut unscaled = digits.parse::().ok()?; + if negative { + unscaled = -unscaled; + } + Some(VariantDecimal { + unscaled, + precision: precision as u8, + scale: scale as i8, + }) +} + +fn decimal_precision(unscaled: i128) -> u8 { + let mut value = unscaled.unsigned_abs(); + if value == 0 { + return 1; + } + let mut precision = 0; + while value > 0 { + precision += 1; + value /= 10; + } + precision +} + +fn string_at(array: &dyn Array, row: usize) -> DFResult> { + if array.is_null(row) { + return Ok(None); + } + match array.data_type() { + ArrowDataType::Utf8 => Ok(Some( + array + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal("Expected Utf8 array".to_string()))? + .value(row) + .to_string(), + )), + ArrowDataType::LargeUtf8 => Ok(Some( + array + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal("Expected LargeUtf8 array".to_string()))? + .value(row) + .to_string(), + )), + ArrowDataType::Utf8View => Ok(Some( + array + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal("Expected Utf8View array".to_string()))? + .value(row) + .to_string(), + )), + other => plan_err(format!("Expected string array, got {other:?}")), + } +} + +fn variant_children(array: &dyn Array) -> DFResult> { + let ArrowDataType::Struct(fields) = array.data_type() else { + return Ok(None); + }; + if fields.len() != 2 + || fields[0].name() != "value" + || fields[0].data_type() != &ArrowDataType::Binary + || fields[1].name() != "metadata" + || fields[1].data_type() != &ArrowDataType::Binary + { + return Ok(None); + } + let array = array + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal("Expected Variant StructArray".to_string()))?; + let values = array + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal("Expected Variant.value BinaryArray".to_string()) + })?; + let metadata = array + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal("Expected Variant.metadata BinaryArray".to_string()) + })?; + Ok(Some((values, metadata))) +} + +fn variant_array(values: Vec>) -> DFResult { + let len = values.len(); + let mut value_builder = BinaryBuilder::new(); + let mut metadata_builder = BinaryBuilder::new(); + let mut validities = Vec::with_capacity(len); + for value in values { + match value { + Some(variant) => { + value_builder.append_value(variant.value()); + metadata_builder.append_value(variant.metadata()); + validities.push(true); + } + None => { + value_builder.append_value(&[] as &[u8]); + metadata_builder.append_value(&[] as &[u8]); + validities.push(false); + } + } + } + let nulls = if validities.iter().all(|valid| *valid) { + None + } else { + Some(NullBuffer::new(BooleanBuffer::from(validities))) + }; + let array = StructArray::try_new( + variant_fields(), + vec![ + Arc::new(value_builder.finish()), + Arc::new(metadata_builder.finish()), + ], + nulls, + )?; + Ok(Arc::new(array)) +} + +fn variant_arrow_type() -> ArrowDataType { + paimon::arrow::variant_arrow_type() +} + +fn variant_fields() -> Fields { + match variant_arrow_type() { + ArrowDataType::Struct(fields) => fields, + _ => unreachable!("variant_arrow_type must be a struct"), + } +} + +fn null_array(data_type: &ArrowDataType, len: usize) -> ArrayRef { + datafusion::arrow::array::new_null_array(data_type, len) +} + +fn invalid_cast() -> DataFusionError { + DataFusionError::Execution("Invalid Variant cast".to_string()) +} + +fn to_df_error(error: paimon::Error) -> DataFusionError { + DataFusionError::External(Box::new(error)) +} + +fn plan_err(message: impl Into) -> DFResult { + Err(DataFusionError::Plan(message.into())) +} + +fn internal_err(message: impl Into) -> DFResult { + Err(DataFusionError::Internal(message.into())) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{BooleanArray, Int32Array, StringArray}; + + async fn collect_one(sql: &str) -> datafusion::arrow::record_batch::RecordBatch { + let ctx = SessionContext::new(); + register_variant_functions(&ctx); + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + assert_eq!(batches.len(), 1); + batches.into_iter().next().unwrap() + } + + #[tokio::test] + async fn parse_json_and_variant_get_scalars() { + let batch = collect_one( + r#" + SELECT + variant_get(parse_json('{"age":26,"city":"Beijing","nested":{"name":"Alice"},"arr":[1,2,3]}'), '$.age', 'int') AS age, + variant_get(parse_json('{"age":26,"city":"Beijing","nested":{"name":"Alice"},"arr":[1,2,3]}'), '$.city', 'string') AS city, + variant_get(parse_json('{"age":26,"city":"Beijing","nested":{"name":"Alice"},"arr":[1,2,3]}'), '$.nested.name', 'string') AS name, + variant_get(parse_json('{"age":26,"city":"Beijing","nested":{"name":"Alice"},"arr":[1,2,3]}'), '$.arr[1]', 'int') AS arr_value + "#, + ) + .await; + + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 26 + ); + assert_eq!( + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "Beijing" + ); + assert_eq!( + batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "Alice" + ); + assert_eq!( + batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 2 + ); + } + + #[tokio::test] + async fn variant_null_is_distinct_from_sql_null() { + let batch = collect_one( + "SELECT is_variant_null(parse_json('null')) AS variant_null, is_variant_null(NULL) AS sql_null", + ) + .await; + let variant_null = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let sql_null = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(variant_null.value(0)); + assert!(!sql_null.value(0)); + } + + #[tokio::test] + async fn try_functions_return_null_on_invalid_input() { + let batch = collect_one( + r#" + SELECT + try_parse_json('{bad json') AS bad_json, + try_variant_get(parse_json('{"age":"not an int"}'), '$.age', 'int') AS bad_cast, + variant_get(parse_json('{}'), '$.missing', 'int') AS missing_path + "#, + ) + .await; + assert!(batch.column(0).is_null(0)); + assert!(batch.column(1).is_null(0)); + assert!(batch.column(2).is_null(0)); + } + + #[tokio::test] + async fn strict_functions_surface_errors() { + let ctx = SessionContext::new(); + register_variant_functions(&ctx); + let err = ctx + .sql("SELECT parse_json('{bad json')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(err.to_string().contains("Expected")); + + let err = ctx + .sql("SELECT variant_get(parse_json('{\"age\":\"not an int\"}'), '$.age', 'int')") + .await + .unwrap() + .collect() + .await + .unwrap_err(); + assert!(err.to_string().contains("Invalid Variant cast")); + } + + #[tokio::test] + async fn variant_get_rejects_non_literal_type_argument() { + let ctx = SessionContext::new(); + register_variant_functions(&ctx); + let sql = r#" + SELECT variant_get(parse_json('{"age":26}'), '$.age', type_name) + FROM (VALUES ('int')) AS t(type_name) + "#; + let err = match ctx.sql(sql).await { + Ok(df) => df.collect().await.unwrap_err(), + Err(err) => err, + }; + assert!(err + .to_string() + .contains("variant_get type argument must be a string literal")); + } +} diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 982f518a..579cbeb9 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -1215,6 +1215,7 @@ fn literal_scalar_for_parquet_filter( DataType::Time(_) | DataType::Timestamp(_) | DataType::LocalZonedTimestamp(_) + | DataType::Variant(_) | DataType::Blob(_) | DataType::Array(_) | DataType::Map(_) diff --git a/crates/paimon/src/arrow/format/row.rs b/crates/paimon/src/arrow/format/row.rs index 0ffcfd9e..211bc4d0 100644 --- a/crates/paimon/src/arrow/format/row.rs +++ b/crates/paimon/src/arrow/format/row.rs @@ -21,9 +21,12 @@ //! `org.apache.paimon.format.row.RowFormatWriter` in paimon-java. use super::{FilePredicates, FormatFileReader, FormatFileWriter}; -use crate::arrow::{arrow_to_paimon_type, build_target_arrow_schema, paimon_type_to_arrow}; +use crate::arrow::{ + arrow_to_paimon_type, build_target_arrow_schema, is_variant_arrow_fields, paimon_type_to_arrow, + variant_arrow_type, +}; use crate::io::{FileRead, FileWrite, OutputFile}; -use crate::spec::{DataField, DataType, IntType}; +use crate::spec::{DataField, DataType, IntType, VarBinaryType, VariantType}; use crate::table::{ArrowRecordBatchStream, RowRange}; use crate::Error; use arrow_array::builder::{ @@ -403,6 +406,7 @@ fn validate_arrow_type_for_row_field( ) | (ArrowDataType::Date32, DataType::Date(_)) | (ArrowDataType::Time32(TimeUnit::Millisecond), DataType::Time(_)) => true, + (ArrowDataType::Struct(fields), DataType::Variant(_)) => is_variant_arrow_fields(fields), (ArrowDataType::Timestamp(unit, tz), DataType::Timestamp(t)) => { *unit == timestamp_time_unit_for_precision(t.precision()) && tz.is_none() } @@ -520,6 +524,7 @@ fn validate_supported_type(data_type: &DataType) -> crate::Result<()> { | DataType::VarChar(_) | DataType::Binary(_) | DataType::VarBinary(_) + | DataType::Variant(_) | DataType::Blob(_) | DataType::Date(_) | DataType::Time(_) @@ -659,6 +664,10 @@ fn write_field_value( out, downcast::(array, data_type)?.value(row_idx), ), + DataType::Variant(_) => { + let row = downcast::(array, data_type)?; + write_variant_struct(out, row, row_idx)?; + } DataType::Date(_) => out.extend_from_slice( &downcast::(array, data_type)? .value(row_idx) @@ -712,6 +721,50 @@ fn write_field_value( Ok(()) } +fn write_variant_struct(out: &mut Vec, row: &StructArray, row_idx: usize) -> crate::Result<()> { + if row.num_columns() != 2 { + return Err(Error::DataInvalid { + message: format!(".row variant expected 2 columns, got {}", row.num_columns()), + source: None, + }); + } + let value = row.column(0); + let metadata = row.column(1); + if value.is_null(row_idx) || metadata.is_null(row_idx) { + return Err(Error::DataInvalid { + message: ".row variant value/metadata children must be non-null".to_string(), + source: None, + }); + } + let value = value + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: format!( + ".row variant value must be BinaryArray, got {:?}", + value.data_type() + ), + source: None, + })?; + let metadata = metadata + .as_any() + .downcast_ref::() + .ok_or_else(|| Error::DataInvalid { + message: format!( + ".row variant metadata must be BinaryArray, got {:?}", + metadata.data_type() + ), + source: None, + })?; + + let value = value.value(row_idx); + let metadata = metadata.value(row_idx); + VariantType::validate_payload(value, metadata)?; + write_bytes(out, value); + write_bytes(out, metadata); + Ok(()) +} + fn write_array_slice( out: &mut Vec, values: &ArrayRef, @@ -1146,6 +1199,15 @@ impl ColumnBuilder { DataType::Binary(_) | DataType::VarBinary(_) | DataType::Blob(_) => { Self::Binary(BinaryBuilder::new()) } + DataType::Variant(_) => Self::Row { + fields: variant_arrow_fields(), + columns: vec![ + ColumnBuilder::new(&variant_binary_type(), capacity)?, + ColumnBuilder::new(&variant_binary_type(), capacity)?, + ], + validities: Vec::with_capacity(capacity), + len: 0, + }, DataType::Date(_) => Self::Date(Date32Builder::with_capacity(capacity)), DataType::Time(_) => Self::Time(Time32MillisecondBuilder::with_capacity(capacity)), DataType::Timestamp(t) => timestamp_builder(t.precision(), false, capacity), @@ -1372,6 +1434,19 @@ impl ColumnBuilder { push_nested_offset(offsets, key_count)?; validities.push(true); } + ( + Self::Row { + columns, + validities, + len, + .. + }, + DataType::Variant(_), + ) => { + read_variant_into(input, columns)?; + *len += 1; + validities.push(true); + } ( Self::Row { columns, @@ -1565,6 +1640,44 @@ fn read_struct_into( Ok(()) } +fn read_variant_into( + input: &mut BlockInput<'_>, + builders: &mut [ColumnBuilder], +) -> crate::Result<()> { + if builders.len() != 2 { + return Err(Error::DataInvalid { + message: format!( + ".row variant reader expected 2 builders, got {}", + builders.len() + ), + source: None, + }); + } + let value = input.read_bytes()?; + let metadata = input.read_bytes()?; + VariantType::validate_payload(&value, &metadata)?; + append_binary_value(&mut builders[0], value, "Variant.value")?; + append_binary_value(&mut builders[1], metadata, "Variant.metadata")?; + Ok(()) +} + +fn append_binary_value( + builder: &mut ColumnBuilder, + value: Vec, + label: &str, +) -> crate::Result<()> { + match builder { + ColumnBuilder::Binary(b) => { + b.append_value(value); + Ok(()) + } + _ => Err(Error::DataInvalid { + message: format!(".row {label} builder must be Binary"), + source: None, + }), + } +} + fn arrow_fields_for_row(fields: &[DataField]) -> crate::Result { let fields = fields .iter() @@ -1579,6 +1692,20 @@ fn arrow_fields_for_row(fields: &[DataField]) -> crate::Result { Ok(fields.into()) } +fn variant_arrow_fields() -> Fields { + match variant_arrow_type() { + ArrowDataType::Struct(fields) => fields, + _ => unreachable!("variant_arrow_type always returns a Struct"), + } +} + +fn variant_binary_type() -> DataType { + DataType::VarBinary( + VarBinaryType::try_new(false, VarBinaryType::MAX_LENGTH) + .expect("variant binary child length is valid"), + ) +} + fn map_entries_field( key_type: &DataType, value_type: &DataType, @@ -2187,8 +2314,9 @@ mod tests { use crate::spec::{ ArrayType, BigIntType, BooleanType, DataType, DateType, DecimalType, DoubleType, FloatType, IntType, MapType, MultisetType, RowType, TimeType, TimestampType, VarBinaryType, - VarCharType, + VarCharType, VariantType, }; + use crate::variant::GenericVariant; use futures::TryStreamExt; use std::ops::Range; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -2342,6 +2470,125 @@ mod tests { assert_eq!(names.value(1), "ccc"); } + #[tokio::test] + async fn row_writer_reader_roundtrip_variant() { + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let path = "memory:/row-variant/data.row"; + let output = file_io.new_output(path).unwrap(); + let fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new( + 1, + "payload".to_string(), + DataType::Variant(VariantType::new()), + ), + ]; + let schema = build_target_arrow_schema(&fields).unwrap(); + let variant_fields = match variant_arrow_type() { + ArrowDataType::Struct(fields) => fields, + other => panic!("expected variant Struct, got {other:?}"), + }; + let first_variant = GenericVariant::parse_json("2").unwrap(); + let second_variant = GenericVariant::parse_json(r#"{"a":3}"#).unwrap(); + let variant_array = StructArray::try_new( + variant_fields, + vec![ + Arc::new(BinaryArray::from(vec![ + Some(first_variant.value()), + Some(second_variant.value()), + ])) as ArrayRef, + Arc::new(BinaryArray::from(vec![ + Some(first_variant.metadata()), + Some(second_variant.metadata()), + ])) as ArrayRef, + ], + None, + ) + .unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(variant_array), + ], + ) + .unwrap(); + let mut writer = RowFormatWriter::new(&output, schema, fields.clone(), 1) + .await + .unwrap(); + writer.write(&batch).await.unwrap(); + Box::new(writer).close().await.unwrap(); + + let bytes = file_io.new_input(path).unwrap().read().await.unwrap(); + let batches = RowFormatReader + .read_batch_stream( + Box::new(BytesFileRead(bytes.clone())), + bytes.len() as u64, + &fields, + None, + Some(8), + None, + ) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(batches.len(), 1); + let payload = batches[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let value = payload + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let metadata = payload + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(value.value(0), first_variant.value()); + assert_eq!(metadata.value(0), first_variant.metadata()); + assert_eq!(value.value(1), second_variant.value()); + assert_eq!(metadata.value(1), second_variant.metadata()); + } + + #[test] + fn row_variant_field_encoding_matches_java() { + let fields = match variant_arrow_type() { + ArrowDataType::Struct(fields) => fields, + other => panic!("expected variant Struct, got {other:?}"), + }; + let variant = GenericVariant::parse_json("2").unwrap(); + let variant_array = StructArray::try_new( + fields, + vec![ + Arc::new(BinaryArray::from(vec![Some(variant.value())])) as ArrayRef, + Arc::new(BinaryArray::from(vec![Some(variant.metadata())])) as ArrayRef, + ], + None, + ) + .unwrap(); + let array = Arc::new(variant_array) as ArrayRef; + let mut encoded = Vec::new(); + write_field_value( + &mut encoded, + &array, + 0, + &DataType::Variant(VariantType::new()), + ) + .unwrap(); + + let mut expected = Vec::new(); + write_bytes(&mut expected, variant.value()); + write_bytes(&mut expected, variant.metadata()); + assert_eq!(encoded, expected); + } + #[tokio::test] async fn row_reader_selection_across_blocks() { let fields = vec![ diff --git a/crates/paimon/src/arrow/format/vortex.rs b/crates/paimon/src/arrow/format/vortex.rs index 0a9717da..285f21f6 100644 --- a/crates/paimon/src/arrow/format/vortex.rs +++ b/crates/paimon/src/arrow/format/vortex.rs @@ -767,6 +767,7 @@ fn literal_scalar_for_arrow_filter( | DataType::Map(_) | DataType::Multiset(_) | DataType::Row(_) + | DataType::Variant(_) | DataType::Vector(_) => { return Ok(None); } diff --git a/crates/paimon/src/arrow/mod.rs b/crates/paimon/src/arrow/mod.rs index b897c8ca..29473cf0 100644 --- a/crates/paimon/src/arrow/mod.rs +++ b/crates/paimon/src/arrow/mod.rs @@ -22,12 +22,16 @@ pub(crate) mod schema_evolution; use crate::spec::{ ArrayType, BigIntType, BooleanType, DataField, DataType as PaimonDataType, DateType, DecimalType, DoubleType, FloatType, IntType, LocalZonedTimestampType, MapType, RowType, - SmallIntType, TimeType, TimestampType, TinyIntType, VarBinaryType, VarCharType, VectorType, + SmallIntType, TimeType, TimestampType, TinyIntType, VarBinaryType, VarCharType, VariantType, + VectorType, }; use arrow_schema::DataType as ArrowDataType; use arrow_schema::{Field as ArrowField, Schema as ArrowSchema, TimeUnit}; +use std::collections::HashMap; use std::sync::Arc; +const PARQUET_FIELD_ID_META_KEY: &str = "PARQUET:field_id"; + /// Converts a Paimon [`DataType`](PaimonDataType) to an Arrow [`DataType`](ArrowDataType). pub fn paimon_type_to_arrow(dt: &PaimonDataType) -> crate::Result { Ok(match dt { @@ -42,6 +46,7 @@ pub fn paimon_type_to_arrow(dt: &PaimonDataType) -> crate::Result PaimonDataType::Binary(_) | PaimonDataType::VarBinary(_) | PaimonDataType::Blob(_) => { ArrowDataType::Binary } + PaimonDataType::Variant(_) => variant_arrow_type(), PaimonDataType::Date(_) => ArrowDataType::Date32, PaimonDataType::Time(_) => ArrowDataType::Time32(TimeUnit::Millisecond), PaimonDataType::Timestamp(t) => { @@ -228,6 +233,11 @@ pub fn arrow_to_paimon_type( }) } ArrowDataType::Struct(fields) => { + if is_variant_arrow_fields(fields) && has_variant_arrow_field_ids(fields) { + return Ok(PaimonDataType::Variant(VariantType::with_nullable( + nullable, + ))); + } let field_slice: Vec = fields.iter().map(|f| f.as_ref().clone()).collect(); let paimon_fields = arrow_fields_to_paimon(&field_slice)?; Ok(PaimonDataType::Row(RowType::with_nullable( @@ -252,6 +262,51 @@ pub fn arrow_to_paimon_type( } } +pub fn variant_arrow_type() -> ArrowDataType { + ArrowDataType::Struct( + vec![ + arrow_field_with_paimon_id("value", ArrowDataType::Binary, false, 0), + arrow_field_with_paimon_id("metadata", ArrowDataType::Binary, false, 1), + ] + .into(), + ) +} + +pub(crate) fn is_variant_arrow_fields(fields: &arrow_schema::Fields) -> bool { + fields.len() == 2 + && fields[0].name() == "value" + && fields[0].data_type() == &ArrowDataType::Binary + && !fields[0].is_nullable() + && fields[1].name() == "metadata" + && fields[1].data_type() == &ArrowDataType::Binary + && !fields[1].is_nullable() +} + +fn has_variant_arrow_field_ids(fields: &arrow_schema::Fields) -> bool { + fields.len() == 2 + && arrow_field_id(&fields[0]) == Some(0) + && arrow_field_id(&fields[1]) == Some(1) +} + +fn arrow_field_id(field: &ArrowField) -> Option { + field + .metadata() + .get(PARQUET_FIELD_ID_META_KEY)? + .parse() + .ok() +} + +fn arrow_field_with_paimon_id( + name: impl Into, + data_type: ArrowDataType, + nullable: bool, + id: i32, +) -> ArrowField { + let mut metadata = HashMap::new(); + metadata.insert(PARQUET_FIELD_ID_META_KEY.to_string(), id.to_string()); + ArrowField::new(name, data_type, nullable).with_metadata(metadata) +} + /// Convert Arrow fields to Paimon [`DataField`]s with auto-assigned IDs starting from 0. pub fn arrow_fields_to_paimon(fields: &[ArrowField]) -> crate::Result> { fields @@ -270,10 +325,11 @@ pub fn build_target_arrow_schema(fields: &[DataField]) -> crate::Result>>()?; @@ -379,6 +435,54 @@ mod tests { assert_arrow_to_paimon(&ArrowDataType::Binary, true, &varbinary); } + #[test] + fn test_variant_roundtrip() { + let variant = PaimonDataType::Variant(VariantType::new()); + let arrow = variant_arrow_type(); + assert_paimon_to_arrow(&variant, &arrow); + assert_arrow_to_paimon(&arrow, true, &variant); + } + + #[test] + fn test_plain_value_metadata_struct_stays_row() { + let arrow = ArrowDataType::Struct( + vec![ + ArrowField::new("value", ArrowDataType::Binary, false), + ArrowField::new("metadata", ArrowDataType::Binary, false), + ] + .into(), + ); + let paimon = arrow_to_paimon_type(&arrow, true).unwrap(); + + assert!(matches!(paimon, PaimonDataType::Row(_))); + } + + #[test] + fn test_variant_arrow_field_ids_for_parquet() { + let schema = build_target_arrow_schema(&[DataField::new( + 7, + "payload".to_string(), + PaimonDataType::Variant(VariantType::new()), + )]) + .unwrap(); + let field = schema.field(0); + assert_eq!( + field.metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"7".to_string()) + ); + let ArrowDataType::Struct(fields) = field.data_type() else { + panic!("expected variant Struct"); + }; + assert_eq!( + fields[0].metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"0".to_string()) + ); + assert_eq!( + fields[1].metadata().get(PARQUET_FIELD_ID_META_KEY), + Some(&"1".to_string()) + ); + } + #[test] fn test_timestamp_roundtrip() { // millisecond precision diff --git a/crates/paimon/src/btree/key_serde.rs b/crates/paimon/src/btree/key_serde.rs index 2d7b37e0..22d72543 100644 --- a/crates/paimon/src/btree/key_serde.rs +++ b/crates/paimon/src/btree/key_serde.rs @@ -20,7 +20,7 @@ //! Reference: [org.apache.paimon.globalindex.btree.KeySerializer](https://github.com/apache/paimon/blob/master/paimon-common/src/main/java/org/apache/paimon/globalindex/btree/KeySerializer.java) use crate::btree::var_len::{decode_var_int_from_slice, encode_var_int}; -use crate::spec::{DataType, Datum}; +use crate::spec::{DataType, Datum, VariantType}; use std::cmp::Ordering; /// Timestamp precision <= 3 is compact (millis only). @@ -149,5 +149,14 @@ pub fn serialize_datum(datum: &Datum, data_type: &DataType) -> Vec { } } Datum::Bytes(v) => v.clone(), + Datum::Variant { value, metadata } => { + VariantType::validate_payload(value, metadata) + .expect("invalid Variant payload for BTree key"); + let mut bytes = Vec::with_capacity(4 + value.len() + metadata.len()); + bytes.extend_from_slice(&(value.len() as u32).to_le_bytes()); + bytes.extend_from_slice(value); + bytes.extend_from_slice(metadata); + bytes + } } } diff --git a/crates/paimon/src/lib.rs b/crates/paimon/src/lib.rs index 8c26fe65..da34c7d4 100644 --- a/crates/paimon/src/lib.rs +++ b/crates/paimon/src/lib.rs @@ -37,6 +37,7 @@ pub mod spec; pub mod table; #[cfg(feature = "fulltext")] pub mod tantivy; +pub mod variant; pub mod vector_search; pub mod vindex; diff --git a/crates/paimon/src/spec/binary_row.rs b/crates/paimon/src/spec/binary_row.rs index 145f64ca..cc6601b7 100644 --- a/crates/paimon/src/spec/binary_row.rs +++ b/crates/paimon/src/spec/binary_row.rs @@ -19,7 +19,7 @@ //! and BinaryRowBuilder for constructing BinaryRow instances. use crate::spec::murmur_hash::hash_by_words; -use crate::spec::{DataType, Datum}; +use crate::spec::{DataType, Datum, VariantType}; use arrow_array::RecordBatch; use serde::{Deserialize, Serialize}; use std::sync::LazyLock; @@ -343,6 +343,10 @@ impl BinaryRow { DataType::Binary(_) | DataType::VarBinary(_) => { Datum::Bytes(self.get_binary(pos)?.to_vec()) } + DataType::Variant(_) => { + let (value, metadata) = decode_variant_bytes(self.get_binary(pos)?)?; + Datum::Variant { value, metadata } + } DataType::Decimal(dt) => { let unscaled = self.get_decimal_unscaled(pos, dt.precision())?; Datum::Decimal { @@ -647,6 +651,11 @@ impl BinaryRowBuilder { self.write_binary(pos, b); } } + Datum::Variant { value, metadata } => { + let bytes = encode_variant_bytes(value, metadata) + .expect("invalid Variant payload for BinaryRow"); + self.write_binary(pos, &bytes); + } } } } @@ -772,6 +781,7 @@ pub fn extract_datum_from_arrow( .ok_or_else(|| type_mismatch_err("Binary", col_idx))?; Datum::Bytes(arr.value(row_idx).to_vec()) } + DataType::Variant(_) => extract_variant_datum_from_arrow(col, row_idx, col_idx)?, DataType::Timestamp(ts) => { if ts.precision() <= 3 { let arr = col @@ -829,6 +839,104 @@ pub fn extract_datum_from_arrow( Ok(Some(datum)) } +fn encode_variant_bytes(value: &[u8], metadata: &[u8]) -> crate::Result> { + VariantType::validate_payload(value, metadata)?; + let mut bytes = Vec::with_capacity(4 + value.len() + metadata.len()); + bytes.extend_from_slice(&(value.len() as u32).to_le_bytes()); + bytes.extend_from_slice(value); + bytes.extend_from_slice(metadata); + Ok(bytes) +} + +fn decode_variant_bytes(bytes: &[u8]) -> crate::Result<(Vec, Vec)> { + if bytes.len() < 4 { + return Err(crate::Error::DataInvalid { + message: format!("Variant bytes too short: {} bytes", bytes.len()), + source: None, + }); + } + let value_len = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as usize; + let value_end = 4usize + .checked_add(value_len) + .ok_or_else(|| crate::Error::DataInvalid { + message: "Variant value length overflows".to_string(), + source: None, + })?; + if value_end > bytes.len() { + return Err(crate::Error::DataInvalid { + message: format!( + "Variant value length {value_len} exceeds payload length {}", + bytes.len() + ), + source: None, + }); + } + let value = bytes[4..value_end].to_vec(); + let metadata = bytes[value_end..].to_vec(); + VariantType::validate_payload(&value, &metadata)?; + Ok((value, metadata)) +} + +fn extract_variant_datum_from_arrow( + col: &std::sync::Arc, + row_idx: usize, + col_idx: usize, +) -> crate::Result { + use arrow_array::Array; + + let arr = col + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch_err("Variant", col_idx))?; + validate_variant_struct_array(arr, col_idx)?; + let value = arr.column(0); + let metadata = arr.column(1); + if value.is_null(row_idx) || metadata.is_null(row_idx) { + return Err(crate::Error::DataInvalid { + message: format!("Variant Arrow struct at column {col_idx} has null child value"), + source: None, + }); + } + let value = value + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch_err("Variant.value", col_idx))? + .value(row_idx) + .to_vec(); + let metadata = metadata + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch_err("Variant.metadata", col_idx))? + .value(row_idx) + .to_vec(); + VariantType::validate_payload(&value, &metadata)?; + Ok(Datum::Variant { value, metadata }) +} + +fn validate_variant_struct_array( + arr: &arrow_array::StructArray, + col_idx: usize, +) -> crate::Result<()> { + if arr.num_columns() != 2 { + return Err(crate::Error::DataInvalid { + message: format!( + "Variant Arrow struct at column {col_idx} must have 2 fields, got {}", + arr.num_columns() + ), + source: None, + }); + } + arr.column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch_err("Variant.value", col_idx))?; + arr.column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch_err("Variant.metadata", col_idx))?; + Ok(()) +} + fn type_mismatch_err(expected: &str, col_idx: usize) -> crate::Error { crate::Error::DataInvalid { message: format!( @@ -858,6 +966,7 @@ enum TypedColumn<'a> { Date32(&'a arrow_array::Date32Array), Decimal128(&'a arrow_array::Decimal128Array, u32, u32), // (array, precision, scale) Binary(&'a arrow_array::BinaryArray), + Variant(&'a arrow_array::StructArray), TimestampMs(&'a arrow_array::TimestampMillisecondArray), TimestampUs(&'a arrow_array::TimestampMicrosecondArray), } @@ -943,6 +1052,14 @@ fn downcast_columns<'a>( .downcast_ref() .ok_or_else(|| type_mismatch_err("Binary", col_idx))?, ), + DataType::Variant(_) => { + let arr = col + .as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Variant", col_idx))?; + validate_variant_struct_array(arr, col_idx)?; + TypedColumn::Variant(arr) + } DataType::Timestamp(ts) => { if ts.precision() <= 3 { TypedColumn::TimestampMs( @@ -990,7 +1107,7 @@ fn write_typed_value( row_idx: usize, typed_col: &TypedColumn, _data_type: &DataType, -) { +) -> crate::Result<()> { use arrow_array::Array; match typed_col { TypedColumn::Boolean(arr) => { @@ -1109,6 +1226,32 @@ fn write_typed_value( } } } + TypedColumn::Variant(arr) => { + if arr.is_null(row_idx) { + builder.set_null_at(pos); + } else { + let value = arr.column(0); + let metadata = arr.column(1); + if value.is_null(row_idx) || metadata.is_null(row_idx) { + return Err(crate::Error::DataInvalid { + message: "Variant Arrow struct has null child value".to_string(), + source: None, + }); + } else if let (Some(value), Some(metadata)) = ( + value.as_any().downcast_ref::(), + metadata.as_any().downcast_ref::(), + ) { + let bytes = + encode_variant_bytes(value.value(row_idx), metadata.value(row_idx))?; + builder.write_binary(pos, &bytes); + } else { + return Err(crate::Error::DataInvalid { + message: "Variant Arrow struct children must be BinaryArray".to_string(), + source: None, + }); + } + } + } TypedColumn::TimestampMs(arr) => { if arr.is_null(row_idx) { builder.set_null_at(pos); @@ -1127,6 +1270,7 @@ fn write_typed_value( } } } + Ok(()) } /// Build BinaryRows for all rows in the batch for the given field indices. @@ -1144,7 +1288,7 @@ pub(crate) fn batch_build_binary_rows( for row_idx in 0..num_rows { let mut builder = BinaryRowBuilder::new(arity); for (pos, (typed_col, field)) in typed_columns.iter().enumerate() { - write_typed_value(&mut builder, pos, row_idx, typed_col, field.data_type()); + write_typed_value(&mut builder, pos, row_idx, typed_col, field.data_type())?; } rows.push(builder.build()); } @@ -1180,6 +1324,7 @@ pub fn batch_hash_codes( #[cfg(test)] mod tests { use super::*; + use crate::variant::GenericVariant; #[test] fn test_empty_binary_row() { @@ -1338,6 +1483,19 @@ mod tests { assert_eq!(row.get_binary(0).unwrap(), &[0x00, 0x01, 0x02, 0x03]); } + #[test] + fn test_variant_datum_roundtrip() { + let data_type = DataType::Variant(crate::spec::VariantType::new()); + let variant = GenericVariant::parse_json(r#"{"a":1}"#).unwrap(); + let datum = Datum::Variant { + value: variant.value().to_vec(), + metadata: variant.metadata().to_vec(), + }; + let row = BinaryRow::from_datums(&[(Some(&datum), &data_type)]); + + assert_eq!(row.get_datum(0, &data_type).unwrap(), Some(datum)); + } + #[test] fn test_mixed_types_partition_row() { let mut builder = BinaryRowBuilder::new(2); diff --git a/crates/paimon/src/spec/data_type_casts.rs b/crates/paimon/src/spec/data_type_casts.rs index c44c4a40..c4452044 100644 --- a/crates/paimon/src/spec/data_type_casts.rs +++ b/crates/paimon/src/spec/data_type_casts.rs @@ -162,6 +162,7 @@ fn explicit_cast_supported(source: &DataType, target: &DataType) -> bool { is_datetime(source) || is_character_string(source) || is_numeric(source) } DataType::Blob(_) + | DataType::Variant(_) | DataType::Array(_) | DataType::Map(_) | DataType::Multiset(_) diff --git a/crates/paimon/src/spec/partition_utils.rs b/crates/paimon/src/spec/partition_utils.rs index db0f2c2b..3db1884c 100644 --- a/crates/paimon/src/spec/partition_utils.rs +++ b/crates/paimon/src/spec/partition_utils.rs @@ -265,6 +265,7 @@ fn format_partition_value( | DataType::Double(_) | DataType::Binary(_) | DataType::VarBinary(_) + | DataType::Variant(_) | DataType::Blob(_) | DataType::Array(_) | DataType::Map(_) diff --git a/crates/paimon/src/spec/predicate.rs b/crates/paimon/src/spec/predicate.rs index a419f77f..5fb641d2 100644 --- a/crates/paimon/src/spec/predicate.rs +++ b/crates/paimon/src/spec/predicate.rs @@ -76,6 +76,10 @@ pub enum Datum { scale: u32, }, Bytes(Vec), + Variant { + value: Vec, + metadata: Vec, + }, } impl fmt::Display for Datum { @@ -97,6 +101,14 @@ impl fmt::Display for Datum { unscaled, scale, .. } => write!(f, "DEC({unscaled},s{scale})"), Self::Bytes(v) => write!(f, "BYTES(len={})", v.len()), + Self::Variant { value, metadata } => { + write!( + f, + "VARIANT(value_len={},metadata_len={})", + value.len(), + metadata.len() + ) + } } } } @@ -162,6 +174,22 @@ pub(crate) fn datum_cmp(lhs: &Datum, rhs: &Datum) -> Option { }, ) => decimal_cmp(*ua, *sa, *ub, *sb), (Datum::Bytes(a), Datum::Bytes(b)) => Some(java_bytes_cmp(a, b)), + ( + Datum::Variant { + value: va, + metadata: ma, + }, + Datum::Variant { + value: vb, + metadata: mb, + }, + ) => { + if va == vb && ma == mb { + Some(Ordering::Equal) + } else { + None + } + } _ => None, } } @@ -864,6 +892,7 @@ fn validate_datum_matches_type(datum: &Datum, data_type: &DataType) -> Result<() | (Datum::Decimal { .. }, DataType::Decimal(_)) | (Datum::Bytes(_), DataType::Binary(_)) | (Datum::Bytes(_), DataType::VarBinary(_)) + | (Datum::Variant { .. }, DataType::Variant(_)) ); if !ok { return Err(Error::ConfigInvalid { diff --git a/crates/paimon/src/spec/types.rs b/crates/paimon/src/spec/types.rs index 7269affc..b1e5ebbb 100644 --- a/crates/paimon/src/spec/types.rs +++ b/crates/paimon/src/spec/types.rs @@ -76,6 +76,8 @@ pub enum DataType { Binary(BinaryType), /// Data type of a variable-length binary string (=a sequence of bytes). VarBinary(VarBinaryType), + /// Data type of semi-structured values encoded as Variant value + metadata. + Variant(VariantType), /// Data type of binary large object. Blob(BlobType), /// Data type of a fixed-length character string. @@ -139,6 +141,7 @@ impl DataType { DataType::Float(v) => v.nullable, DataType::Binary(v) => v.nullable, DataType::VarBinary(v) => v.nullable, + DataType::Variant(v) => v.nullable, DataType::Blob(v) => v.nullable, DataType::Char(v) => v.nullable, DataType::VarChar(v) => v.nullable, @@ -176,6 +179,7 @@ impl DataType { DataType::VarBinary(v) => { DataType::VarBinary(VarBinaryType::try_new(nullable, v.length())?) } + DataType::Variant(_) => DataType::Variant(VariantType::with_nullable(nullable)), DataType::Blob(_) => DataType::Blob(BlobType::with_nullable(nullable)), DataType::Char(v) => DataType::Char(CharType::with_nullable(nullable, v.length())?), DataType::VarChar(v) => { @@ -632,6 +636,44 @@ impl BlobType { } } +/// VariantType for paimon. +/// +/// Data type of semi-structured values encoded as two binary buffers: `value` +/// and `metadata`. +/// +/// Impl Reference: . +#[serde_as] +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct VariantType { + #[serde_as(as = "FromInto>")] + nullable: bool, +} + +impl Default for VariantType { + fn default() -> Self { + Self::new() + } +} + +impl VariantType { + pub fn new() -> Self { + Self::with_nullable(true) + } + + pub fn with_nullable(nullable: bool) -> Self { + Self { nullable } + } + + pub fn family(&self) -> DataTypeFamily { + DataTypeFamily::PREDEFINED + } + + pub(crate) fn validate_payload(value: &[u8], metadata: &[u8]) -> crate::Result<()> { + crate::variant::validate_payload(value, metadata) + } +} + /// CharType for paimon. /// /// Data type of a fixed-length character string. @@ -1732,6 +1774,11 @@ mod serde_utils { const NAME: &'static str = "BLOB"; } + pub struct VARIANT; + impl DataTypeName for VARIANT { + const NAME: &'static str = "VARIANT"; + } + pub struct BINARY; impl DataTypeName for BINARY { const NAME: &'static str = "BINARY"; @@ -2171,6 +2218,14 @@ mod tests { length: 233, }), ), + ( + "variant_type", + DataType::Variant(VariantType { nullable: false }), + ), + ( + "variant_type_nullable", + DataType::Variant(VariantType { nullable: true }), + ), ( "varchar_type", DataType::VarChar(VarCharType { diff --git a/crates/paimon/src/variant.rs b/crates/paimon/src/variant.rs new file mode 100644 index 00000000..537a51a5 --- /dev/null +++ b/crates/paimon/src/variant.rs @@ -0,0 +1,1807 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Java-compatible helpers for Paimon Variant values. + +use std::collections::{HashMap, HashSet}; + +use crate::{Error, Result}; +use base64::{engine::general_purpose, Engine as _}; + +const BASIC_TYPE_BITS: u8 = 2; +const BASIC_TYPE_MASK: u8 = 0x3; +const TYPE_INFO_MASK: u8 = 0x3f; +const MAX_SHORT_STR_SIZE: usize = 0x3f; + +const PRIMITIVE: u8 = 0; +const SHORT_STR: u8 = 1; +const OBJECT: u8 = 2; +const ARRAY: u8 = 3; + +const NULL: u8 = 0; +const TRUE: u8 = 1; +const FALSE: u8 = 2; +const INT1: u8 = 3; +const INT2: u8 = 4; +const INT4: u8 = 5; +const INT8: u8 = 6; +const DOUBLE: u8 = 7; +const DECIMAL4: u8 = 8; +const DECIMAL8: u8 = 9; +const DECIMAL16: u8 = 10; +const DATE: u8 = 11; +const TIMESTAMP: u8 = 12; +const TIMESTAMP_NTZ: u8 = 13; +const FLOAT: u8 = 14; +const BINARY: u8 = 15; +const LONG_STR: u8 = 16; +const UUID: u8 = 20; + +const VERSION: u8 = 1; +const VERSION_MASK: u8 = 0x0f; + +const U8_MAX: usize = 0xff; +const U16_MAX: usize = 0xffff; +const U24_MAX: usize = 0xff_ffff; +const U32_SIZE: usize = 4; +const SIZE_LIMIT: usize = 128 * 1024 * 1024; +const MAX_NESTING_DEPTH: usize = 1000; + +const MAX_DECIMAL4_PRECISION: u8 = 9; +const MAX_DECIMAL8_PRECISION: u8 = 18; +const MAX_DECIMAL16_PRECISION: u8 = 38; +const BINARY_SEARCH_THRESHOLD: usize = 32; + +/// An owned Paimon Variant value encoded as Java-compatible `value` and `metadata` buffers. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GenericVariant { + value: Vec, + metadata: Vec, +} + +impl GenericVariant { + /// Parse a JSON string into the same binary Variant representation produced by Java Paimon. + pub fn parse_json(json: &str) -> Result { + JsonVariantParser::new(json).parse() + } + + /// Build an owned Variant from existing buffers after validating the metadata header. + pub fn from_parts(value: Vec, metadata: Vec) -> Result { + VariantRef::new(&value, &metadata, 0)?; + Ok(Self { value, metadata }) + } + + pub fn value(&self) -> &[u8] { + &self.value + } + + pub fn metadata(&self) -> &[u8] { + &self.metadata + } + + pub fn as_ref(&self) -> Result> { + VariantRef::new(&self.value, &self.metadata, 0) + } + + pub fn is_variant_null(&self) -> Result { + self.as_ref()?.is_null() + } + + pub fn get_path(&self, path: &str) -> Result>> { + self.as_ref()?.get_path(path) + } + + pub fn to_json(&self) -> Result { + self.as_ref()?.to_json() + } +} + +/// A borrowed view into a Variant buffer. Sub-values share the original metadata dictionary. +#[derive(Clone, Copy, Debug)] +pub struct VariantRef<'a> { + value: &'a [u8], + metadata: &'a [u8], + pos: usize, +} + +impl<'a> VariantRef<'a> { + pub fn new(value: &'a [u8], metadata: &'a [u8], pos: usize) -> Result { + validate_payload(value, metadata)?; + Self::new_at(value, metadata, pos) + } + + fn new_at(value: &'a [u8], metadata: &'a [u8], pos: usize) -> Result { + check_index(value, pos)?; + Ok(Self { + value, + metadata, + pos, + }) + } + + pub fn kind(&self) -> Result { + value_kind(self.value, self.pos) + } + + pub fn is_null(&self) -> Result { + Ok(self.kind()? == VariantKind::Null) + } + + pub fn value_slice(&self) -> Result<&'a [u8]> { + let size = value_size(self.value, self.pos)?; + check_index(self.value, self.pos + size - 1)?; + Ok(&self.value[self.pos..self.pos + size]) + } + + pub fn metadata(&self) -> &'a [u8] { + self.metadata + } + + pub fn to_owned_variant(&self) -> Result { + Ok(GenericVariant { + value: self.value_slice()?.to_vec(), + metadata: self.metadata.to_vec(), + }) + } + + pub fn get_path(&self, path: &str) -> Result>> { + let mut current = *self; + for segment in parse_path(path)? { + match (segment, current.kind()?) { + (PathSegment::Key(key), VariantKind::Object) => { + let Some(next) = current.get_field_by_key(&key)? else { + return Ok(None); + }; + current = next; + } + (PathSegment::Index(index), VariantKind::Array) => { + let Some(next) = current.get_element_at_index(index)? else { + return Ok(None); + }; + current = next; + } + _ => return Ok(None), + } + } + Ok(Some(current)) + } + + pub fn get_boolean(&self) -> Result { + get_boolean(self.value, self.pos) + } + + pub fn get_long(&self) -> Result { + get_long(self.value, self.pos) + } + + pub fn get_double(&self) -> Result { + get_double(self.value, self.pos) + } + + pub fn get_float(&self) -> Result { + get_float(self.value, self.pos) + } + + pub fn get_string(&self) -> Result { + get_string(self.value, self.pos) + } + + pub fn get_decimal(&self) -> Result { + get_decimal(self.value, self.pos) + } + + pub fn to_json(&self) -> Result { + let mut out = String::new(); + write_json(self.value, self.metadata, self.pos, &mut out)?; + Ok(out) + } + + fn get_field_by_key(&self, key: &str) -> Result>> { + let layout = object_layout(self.value, self.pos)?; + if layout.size < BINARY_SEARCH_THRESHOLD { + for i in 0..layout.size { + let id = read_unsigned( + self.value, + layout.id_start + layout.id_size * i, + layout.id_size, + )?; + if key == get_metadata_key(self.metadata, id)? { + let offset = read_unsigned( + self.value, + layout.offset_start + layout.offset_size * i, + layout.offset_size, + )?; + return VariantRef::new_at( + self.value, + self.metadata, + layout.data_start + offset, + ) + .map(Some); + } + } + } else { + let mut low = 0usize; + let mut high = layout.size; + while low < high { + let mid = low + (high - low) / 2; + let id = read_unsigned( + self.value, + layout.id_start + layout.id_size * mid, + layout.id_size, + )?; + match java_string_cmp(&get_metadata_key(self.metadata, id)?, key) { + std::cmp::Ordering::Less => low = mid + 1, + std::cmp::Ordering::Greater => high = mid, + std::cmp::Ordering::Equal => { + let offset = read_unsigned( + self.value, + layout.offset_start + layout.offset_size * mid, + layout.offset_size, + )?; + return VariantRef::new_at( + self.value, + self.metadata, + layout.data_start + offset, + ) + .map(Some); + } + } + } + } + Ok(None) + } + + fn get_element_at_index(&self, index: usize) -> Result>> { + let layout = array_layout(self.value, self.pos)?; + if index >= layout.size { + return Ok(None); + } + let offset = read_unsigned( + self.value, + layout.offset_start + layout.offset_size * index, + layout.offset_size, + )?; + VariantRef::new_at(self.value, self.metadata, layout.data_start + offset).map(Some) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum VariantKind { + Object, + Array, + Null, + Boolean, + Long, + String, + Double, + Decimal, + Date, + Timestamp, + TimestampNtz, + Float, + Binary, + Uuid, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct VariantDecimal { + pub unscaled: i128, + pub precision: u8, + pub scale: i8, +} + +impl VariantDecimal { + pub fn to_plain_string(self) -> String { + decimal_to_plain_string(self.unscaled, self.scale, true) + } +} + +pub fn validate_payload(value: &[u8], metadata: &[u8]) -> Result<()> { + if metadata.is_empty() || (metadata[0] & VERSION_MASK) != VERSION { + return data_invalid("Malformed Variant metadata version"); + } + if value.len() > SIZE_LIMIT || metadata.len() > SIZE_LIMIT { + return data_invalid("Variant value or metadata exceeds constructor size limit"); + } + validate_metadata(metadata)?; + let root_size = validate_value(value, metadata, 0, 0)?; + if root_size != value.len() { + return data_invalid("Malformed Variant root size"); + } + Ok(()) +} + +fn validate_metadata(metadata: &[u8]) -> Result<()> { + let offset_size = metadata_offset_size(metadata)?; + let dict_size = read_unsigned(metadata, 1, offset_size)?; + let offset_start = 1 + offset_size; + let string_start = offset_start + (dict_size + 1) * offset_size; + check_range(metadata, offset_start, (dict_size + 1) * offset_size)?; + let string_size = read_unsigned( + metadata, + offset_start + dict_size * offset_size, + offset_size, + )?; + if string_start.checked_add(string_size) != Some(metadata.len()) { + return data_invalid("Malformed Variant metadata size"); + } + + let mut previous_offset = 0usize; + let mut seen = HashSet::with_capacity(dict_size); + for id in 0..dict_size { + let offset = read_unsigned(metadata, offset_start + id * offset_size, offset_size)?; + let next_offset = + read_unsigned(metadata, offset_start + (id + 1) * offset_size, offset_size)?; + if offset != previous_offset || offset > next_offset || next_offset > string_size { + return data_invalid("Malformed Variant metadata offsets"); + } + previous_offset = next_offset; + let bytes = &metadata[string_start + offset..string_start + next_offset]; + let key = std::str::from_utf8(bytes).map_err(|e| Error::DataInvalid { + message: "Malformed Variant metadata UTF-8".to_string(), + source: Some(Box::new(e)), + })?; + if !seen.insert(key) { + return data_invalid("Malformed Variant metadata duplicate key"); + } + } + Ok(()) +} + +fn validate_value(value: &[u8], metadata: &[u8], pos: usize, depth: usize) -> Result { + if depth > MAX_NESTING_DEPTH { + return data_invalid("Malformed Variant nesting depth"); + } + + match value_kind(value, pos)? { + VariantKind::Object => validate_object(value, metadata, pos, depth), + VariantKind::Array => validate_array(value, metadata, pos, depth), + VariantKind::Null => validate_sized_value(value, pos, value_size(value, pos)?), + VariantKind::Boolean => { + get_boolean(value, pos)?; + validate_sized_value(value, pos, value_size(value, pos)?) + } + VariantKind::Long + | VariantKind::Date + | VariantKind::Timestamp + | VariantKind::TimestampNtz => { + get_long(value, pos)?; + validate_sized_value(value, pos, value_size(value, pos)?) + } + VariantKind::String => { + get_string(value, pos)?; + validate_sized_value(value, pos, value_size(value, pos)?) + } + VariantKind::Double => { + get_double(value, pos)?; + validate_sized_value(value, pos, value_size(value, pos)?) + } + VariantKind::Decimal => { + get_decimal(value, pos)?; + validate_sized_value(value, pos, value_size(value, pos)?) + } + VariantKind::Float => { + get_float(value, pos)?; + validate_sized_value(value, pos, value_size(value, pos)?) + } + VariantKind::Binary => { + get_binary(value, pos)?; + validate_sized_value(value, pos, value_size(value, pos)?) + } + VariantKind::Uuid => { + get_uuid(value, pos)?; + validate_sized_value(value, pos, value_size(value, pos)?) + } + } +} + +fn validate_sized_value(value: &[u8], pos: usize, size: usize) -> Result { + check_range(value, pos, size)?; + Ok(size) +} + +fn validate_object(value: &[u8], metadata: &[u8], pos: usize, depth: usize) -> Result { + let layout = object_layout(value, pos)?; + let data_size = read_unsigned( + value, + layout.offset_start + layout.size * layout.offset_size, + layout.offset_size, + )?; + let data_end = layout + .data_start + .checked_add(data_size) + .ok_or_else(|| Error::DataInvalid { + message: "Malformed Variant object offsets".to_string(), + source: None, + })?; + if data_end > value.len() { + return data_invalid("Malformed Variant object offsets"); + } + + let mut ranges = Vec::with_capacity(layout.size); + let mut previous_key: Option = None; + for i in 0..layout.size { + let id = read_unsigned(value, layout.id_start + layout.id_size * i, layout.id_size)?; + let key = get_metadata_key(metadata, id)?; + if previous_key + .as_deref() + .is_some_and(|previous| java_string_cmp(previous, &key) != std::cmp::Ordering::Less) + { + return data_invalid("Malformed Variant object key order"); + } + previous_key = Some(key); + + let offset = read_unsigned( + value, + layout.offset_start + layout.offset_size * i, + layout.offset_size, + )?; + ranges.push(validate_child_range( + value, + metadata, + layout.data_start, + data_size, + offset, + depth + 1, + )?); + } + + let final_offset = read_unsigned( + value, + layout.offset_start + layout.offset_size * layout.size, + layout.offset_size, + )?; + if final_offset != data_size { + return data_invalid("Malformed Variant object offsets"); + } + validate_ranges_cover_data(&mut ranges, data_size)?; + Ok(layout.data_start - pos + data_size) +} + +fn validate_array(value: &[u8], metadata: &[u8], pos: usize, depth: usize) -> Result { + let layout = array_layout(value, pos)?; + let data_size = read_unsigned( + value, + layout.offset_start + layout.size * layout.offset_size, + layout.offset_size, + )?; + let data_end = layout + .data_start + .checked_add(data_size) + .ok_or_else(|| Error::DataInvalid { + message: "Malformed Variant array offsets".to_string(), + source: None, + })?; + if data_end > value.len() { + return data_invalid("Malformed Variant array offsets"); + } + + let mut previous_offset = 0usize; + for i in 0..layout.size { + let offset = read_unsigned( + value, + layout.offset_start + layout.offset_size * i, + layout.offset_size, + )?; + let next_offset = read_unsigned( + value, + layout.offset_start + layout.offset_size * (i + 1), + layout.offset_size, + )?; + validate_child_value( + value, + metadata, + layout.data_start, + data_size, + offset, + next_offset, + depth + 1, + )?; + previous_offset = validate_offset_order(previous_offset, offset, next_offset, data_size)?; + } + + let final_offset = read_unsigned( + value, + layout.offset_start + layout.offset_size * layout.size, + layout.offset_size, + )?; + if final_offset != data_size || final_offset < previous_offset { + return data_invalid("Malformed Variant array offsets"); + } + Ok(layout.data_start - pos + data_size) +} + +fn validate_child_value( + value: &[u8], + metadata: &[u8], + data_start: usize, + data_size: usize, + offset: usize, + next_offset: usize, + depth: usize, +) -> Result<()> { + if offset > next_offset || next_offset > data_size { + return data_invalid("Malformed Variant offsets"); + } + let child_pos = data_start + .checked_add(offset) + .ok_or_else(|| Error::DataInvalid { + message: "Malformed Variant offsets".to_string(), + source: None, + })?; + let child_size = validate_value(value, metadata, child_pos, depth)?; + if child_size != next_offset - offset { + return data_invalid("Malformed Variant child size"); + } + Ok(()) +} + +fn validate_child_range( + value: &[u8], + metadata: &[u8], + data_start: usize, + data_size: usize, + offset: usize, + depth: usize, +) -> Result<(usize, usize)> { + if offset > data_size { + return data_invalid("Malformed Variant offsets"); + } + let child_pos = data_start + .checked_add(offset) + .ok_or_else(|| Error::DataInvalid { + message: "Malformed Variant offsets".to_string(), + source: None, + })?; + let child_size = validate_value(value, metadata, child_pos, depth)?; + let end = offset + .checked_add(child_size) + .ok_or_else(|| Error::DataInvalid { + message: "Malformed Variant offsets".to_string(), + source: None, + })?; + if end > data_size { + return data_invalid("Malformed Variant offsets"); + } + Ok((offset, end)) +} + +fn validate_ranges_cover_data(ranges: &mut [(usize, usize)], data_size: usize) -> Result<()> { + ranges.sort_unstable_by_key(|(start, _)| *start); + let mut expected_start = 0usize; + for (start, end) in ranges { + if *start != expected_start || *end < *start { + return data_invalid("Malformed Variant offsets"); + } + expected_start = *end; + } + if expected_start != data_size { + return data_invalid("Malformed Variant offsets"); + } + Ok(()) +} + +fn validate_offset_order( + previous_offset: usize, + offset: usize, + next_offset: usize, + data_size: usize, +) -> Result { + if offset != previous_offset || offset > next_offset || next_offset > data_size { + return data_invalid("Malformed Variant offsets"); + } + Ok(next_offset) +} + +fn data_invalid(message: impl Into) -> Result { + Err(Error::DataInvalid { + message: message.into(), + source: None, + }) +} + +fn primitive_header(type_info: u8) -> u8 { + (type_info << BASIC_TYPE_BITS) | PRIMITIVE +} + +fn short_str_header(size: usize) -> u8 { + ((size as u8) << BASIC_TYPE_BITS) | SHORT_STR +} + +fn object_header(large_size: bool, id_size: usize, offset_size: usize) -> u8 { + (((large_size as u8) << (BASIC_TYPE_BITS + 4)) + | (((id_size - 1) as u8) << (BASIC_TYPE_BITS + 2)) + | (((offset_size - 1) as u8) << BASIC_TYPE_BITS)) + | OBJECT +} + +fn array_header(large_size: bool, offset_size: usize) -> u8 { + (((large_size as u8) << (BASIC_TYPE_BITS + 2)) | (((offset_size - 1) as u8) << BASIC_TYPE_BITS)) + | ARRAY +} + +fn integer_size(value: usize) -> Result { + if value > SIZE_LIMIT { + return data_invalid("Variant value exceeds size limit"); + } + Ok(if value <= U8_MAX { + 1 + } else if value <= U16_MAX { + 2 + } else if value <= U24_MAX { + 3 + } else { + 4 + }) +} + +fn check_index(bytes: &[u8], pos: usize) -> Result<()> { + if pos >= bytes.len() { + return data_invalid("Malformed Variant"); + } + Ok(()) +} + +fn check_range(bytes: &[u8], pos: usize, len: usize) -> Result<()> { + if len == 0 { + return Ok(()); + } + let end = pos + .checked_add(len) + .and_then(|v| v.checked_sub(1)) + .ok_or_else(|| Error::DataInvalid { + message: "Malformed Variant".to_string(), + source: None, + })?; + check_index(bytes, end) +} + +fn write_le_at(bytes: &mut [u8], pos: usize, value: usize, num_bytes: usize) { + for i in 0..num_bytes { + bytes[pos + i] = ((value >> (8 * i)) & 0xff) as u8; + } +} + +fn push_signed_le(bytes: &mut Vec, value: i128, num_bytes: usize) { + for i in 0..num_bytes { + bytes.push(((value >> (8 * i)) & 0xff) as u8); + } +} + +fn read_unsigned(bytes: &[u8], pos: usize, num_bytes: usize) -> Result { + check_range(bytes, pos, num_bytes)?; + let mut result = 0usize; + for i in 0..num_bytes { + result |= (bytes[pos + i] as usize) << (8 * i); + } + Ok(result) +} + +fn read_signed(bytes: &[u8], pos: usize, num_bytes: usize) -> Result { + check_range(bytes, pos, num_bytes)?; + let mut result = 0i64; + for i in 0..num_bytes { + result |= (bytes[pos + i] as i64) << (8 * i); + } + let shift = 64 - num_bytes * 8; + Ok((result << shift) >> shift) +} + +fn value_kind(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + let basic_type = value[pos] & BASIC_TYPE_MASK; + let type_info = (value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + Ok(match basic_type { + SHORT_STR => VariantKind::String, + OBJECT => VariantKind::Object, + ARRAY => VariantKind::Array, + _ => match type_info { + NULL => VariantKind::Null, + TRUE | FALSE => VariantKind::Boolean, + INT1 | INT2 | INT4 | INT8 => VariantKind::Long, + DOUBLE => VariantKind::Double, + DECIMAL4 | DECIMAL8 | DECIMAL16 => VariantKind::Decimal, + DATE => VariantKind::Date, + TIMESTAMP => VariantKind::Timestamp, + TIMESTAMP_NTZ => VariantKind::TimestampNtz, + FLOAT => VariantKind::Float, + BINARY => VariantKind::Binary, + LONG_STR => VariantKind::String, + UUID => VariantKind::Uuid, + _ => return data_invalid(format!("Unknown primitive type in Variant: {type_info}")), + }, + }) +} + +fn value_size(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + let basic_type = value[pos] & BASIC_TYPE_MASK; + let type_info = (value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + match basic_type { + SHORT_STR => Ok(1 + type_info as usize), + OBJECT => { + let layout = object_layout(value, pos)?; + let data_size = read_unsigned( + value, + layout.offset_start + layout.size * layout.offset_size, + layout.offset_size, + )?; + Ok(layout.data_start - pos + data_size) + } + ARRAY => { + let layout = array_layout(value, pos)?; + let data_size = read_unsigned( + value, + layout.offset_start + layout.size * layout.offset_size, + layout.offset_size, + )?; + Ok(layout.data_start - pos + data_size) + } + _ => match type_info { + NULL | TRUE | FALSE => Ok(1), + INT1 => Ok(2), + INT2 => Ok(3), + INT4 | DATE | FLOAT => Ok(5), + INT8 | DOUBLE | TIMESTAMP | TIMESTAMP_NTZ => Ok(9), + DECIMAL4 => Ok(6), + DECIMAL8 => Ok(10), + DECIMAL16 => Ok(18), + BINARY | LONG_STR => { + let len = read_unsigned(value, pos + 1, U32_SIZE)?; + Ok(1 + U32_SIZE + len) + } + UUID => Ok(17), + _ => data_invalid(format!("Unknown primitive type in Variant: {type_info}")), + }, + } +} + +#[derive(Clone, Copy)] +struct ObjectLayout { + size: usize, + id_size: usize, + offset_size: usize, + id_start: usize, + offset_start: usize, + data_start: usize, +} + +fn object_layout(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + if (value[pos] & BASIC_TYPE_MASK) != OBJECT { + return data_invalid("Expected Variant object"); + } + let type_info = (value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + let large_size = ((type_info >> 4) & 0x1) != 0; + let size_bytes = if large_size { U32_SIZE } else { 1 }; + let size = read_unsigned(value, pos + 1, size_bytes)?; + let id_size = ((type_info >> 2) & 0x3) as usize + 1; + let offset_size = (type_info & 0x3) as usize + 1; + let id_start = pos + 1 + size_bytes; + let offset_start = id_start + size * id_size; + let data_start = offset_start + (size + 1) * offset_size; + check_range(value, id_start, size * id_size)?; + check_range(value, offset_start, (size + 1) * offset_size)?; + Ok(ObjectLayout { + size, + id_size, + offset_size, + id_start, + offset_start, + data_start, + }) +} + +#[derive(Clone, Copy)] +struct ArrayLayout { + size: usize, + offset_size: usize, + offset_start: usize, + data_start: usize, +} + +fn array_layout(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + if (value[pos] & BASIC_TYPE_MASK) != ARRAY { + return data_invalid("Expected Variant array"); + } + let type_info = (value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + let large_size = ((type_info >> 2) & 0x1) != 0; + let size_bytes = if large_size { U32_SIZE } else { 1 }; + let size = read_unsigned(value, pos + 1, size_bytes)?; + let offset_size = (type_info & 0x3) as usize + 1; + let offset_start = pos + 1 + size_bytes; + let data_start = offset_start + (size + 1) * offset_size; + check_range(value, offset_start, (size + 1) * offset_size)?; + Ok(ArrayLayout { + size, + offset_size, + offset_start, + data_start, + }) +} + +fn metadata_offset_size(metadata: &[u8]) -> Result { + check_index(metadata, 0)?; + if (metadata[0] & VERSION_MASK) != VERSION { + return data_invalid("Malformed Variant metadata version"); + } + Ok(((metadata[0] >> 6) & 0x3) as usize + 1) +} + +fn get_metadata_key(metadata: &[u8], id: usize) -> Result { + let offset_size = metadata_offset_size(metadata)?; + let dict_size = read_unsigned(metadata, 1, offset_size)?; + if id >= dict_size { + return data_invalid("Malformed Variant metadata dictionary id"); + } + let string_start = 1 + (dict_size + 2) * offset_size; + let offset = read_unsigned(metadata, 1 + (id + 1) * offset_size, offset_size)?; + let next_offset = read_unsigned(metadata, 1 + (id + 2) * offset_size, offset_size)?; + if offset > next_offset { + return data_invalid("Malformed Variant metadata offsets"); + } + check_range(metadata, string_start + offset, next_offset - offset)?; + let bytes = &metadata[string_start + offset..string_start + next_offset]; + std::str::from_utf8(bytes) + .map(|v| v.to_string()) + .map_err(|e| Error::DataInvalid { + message: "Malformed Variant metadata UTF-8".to_string(), + source: Some(Box::new(e)), + }) +} + +fn get_boolean(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + let basic_type = value[pos] & BASIC_TYPE_MASK; + let type_info = (value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + if basic_type != PRIMITIVE || (type_info != TRUE && type_info != FALSE) { + return data_invalid("Expected Variant boolean"); + } + Ok(type_info == TRUE) +} + +fn get_long(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + if (value[pos] & BASIC_TYPE_MASK) != PRIMITIVE { + return data_invalid("Expected Variant long/date/timestamp"); + } + match (value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK { + INT1 => read_signed(value, pos + 1, 1), + INT2 => read_signed(value, pos + 1, 2), + INT4 | DATE => read_signed(value, pos + 1, 4), + INT8 | TIMESTAMP | TIMESTAMP_NTZ => read_signed(value, pos + 1, 8), + _ => data_invalid("Expected Variant long/date/timestamp"), + } +} + +fn get_double(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + if (value[pos] & BASIC_TYPE_MASK) != PRIMITIVE + || ((value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK) != DOUBLE + { + return data_invalid("Expected Variant double"); + } + Ok(f64::from_bits(read_signed(value, pos + 1, 8)? as u64)) +} + +fn get_float(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + if (value[pos] & BASIC_TYPE_MASK) != PRIMITIVE + || ((value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK) != FLOAT + { + return data_invalid("Expected Variant float"); + } + Ok(f32::from_bits(read_signed(value, pos + 1, 4)? as u32)) +} + +fn get_string(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + let basic_type = value[pos] & BASIC_TYPE_MASK; + let type_info = (value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + let (start, len) = if basic_type == SHORT_STR { + (pos + 1, type_info as usize) + } else if basic_type == PRIMITIVE && type_info == LONG_STR { + (pos + 1 + U32_SIZE, read_unsigned(value, pos + 1, U32_SIZE)?) + } else { + return data_invalid("Expected Variant string"); + }; + check_range(value, start, len)?; + std::str::from_utf8(&value[start..start + len]) + .map(|v| v.to_string()) + .map_err(|e| Error::DataInvalid { + message: "Malformed Variant string UTF-8".to_string(), + source: Some(Box::new(e)), + }) +} + +fn get_binary(value: &[u8], pos: usize) -> Result<&[u8]> { + check_index(value, pos)?; + if (value[pos] & BASIC_TYPE_MASK) != PRIMITIVE + || ((value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK) != BINARY + { + return data_invalid("Expected Variant binary"); + } + let len = read_unsigned(value, pos + 1, U32_SIZE)?; + let start = pos + 1 + U32_SIZE; + check_range(value, start, len)?; + Ok(&value[start..start + len]) +} + +fn get_uuid(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + if (value[pos] & BASIC_TYPE_MASK) != PRIMITIVE + || ((value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK) != UUID + { + return data_invalid("Expected Variant UUID"); + } + check_range(value, pos + 1, 16)?; + let mut bytes = [0u8; 16]; + bytes.copy_from_slice(&value[pos + 1..pos + 17]); + Ok(uuid::Uuid::from_bytes(bytes)) +} + +fn get_decimal(value: &[u8], pos: usize) -> Result { + check_index(value, pos)?; + if (value[pos] & BASIC_TYPE_MASK) != PRIMITIVE { + return data_invalid("Expected Variant decimal"); + } + let type_info = (value[pos] >> BASIC_TYPE_BITS) & TYPE_INFO_MASK; + let scale = value + .get(pos + 1) + .copied() + .ok_or_else(|| Error::DataInvalid { + message: "Malformed Variant decimal".to_string(), + source: None, + })? as i8; + let (unscaled, max_precision) = match type_info { + DECIMAL4 => ( + read_signed(value, pos + 2, 4)? as i128, + MAX_DECIMAL4_PRECISION, + ), + DECIMAL8 => ( + read_signed(value, pos + 2, 8)? as i128, + MAX_DECIMAL8_PRECISION, + ), + DECIMAL16 => { + check_range(value, pos + 2, 16)?; + let mut bytes = [0u8; 16]; + bytes.copy_from_slice(&value[pos + 2..pos + 18]); + (i128::from_le_bytes(bytes), MAX_DECIMAL16_PRECISION) + } + _ => return data_invalid("Expected Variant decimal"), + }; + let precision = decimal_precision(unscaled); + if precision > max_precision || scale < 0 || scale as u8 > max_precision { + return data_invalid("Malformed Variant decimal precision or scale"); + } + Ok(VariantDecimal { + unscaled, + precision, + scale, + }) +} + +fn write_json(value: &[u8], metadata: &[u8], pos: usize, out: &mut String) -> Result<()> { + match value_kind(value, pos)? { + VariantKind::Object => { + let layout = object_layout(value, pos)?; + out.push('{'); + for i in 0..layout.size { + if i != 0 { + out.push(','); + } + let id = + read_unsigned(value, layout.id_start + layout.id_size * i, layout.id_size)?; + let key = get_metadata_key(metadata, id)?; + out.push_str( + &serde_json::to_string(&key).map_err(|e| Error::DataInvalid { + message: "Failed to escape Variant object key".to_string(), + source: Some(Box::new(e)), + })?, + ); + out.push(':'); + let offset = read_unsigned( + value, + layout.offset_start + layout.offset_size * i, + layout.offset_size, + )?; + write_json(value, metadata, layout.data_start + offset, out)?; + } + out.push('}'); + } + VariantKind::Array => { + let layout = array_layout(value, pos)?; + out.push('['); + for i in 0..layout.size { + if i != 0 { + out.push(','); + } + let offset = read_unsigned( + value, + layout.offset_start + layout.offset_size * i, + layout.offset_size, + )?; + write_json(value, metadata, layout.data_start + offset, out)?; + } + out.push(']'); + } + VariantKind::Null => out.push_str("null"), + VariantKind::Boolean => out.push_str(if get_boolean(value, pos)? { + "true" + } else { + "false" + }), + VariantKind::Long => out.push_str(&get_long(value, pos)?.to_string()), + VariantKind::String => out.push_str( + &serde_json::to_string(&get_string(value, pos)?).map_err(|e| Error::DataInvalid { + message: "Failed to escape Variant string".to_string(), + source: Some(Box::new(e)), + })?, + ), + VariantKind::Double => out.push_str(&get_double(value, pos)?.to_string()), + VariantKind::Decimal => out.push_str(&get_decimal(value, pos)?.to_plain_string()), + VariantKind::Float => out.push_str(&get_float(value, pos)?.to_string()), + VariantKind::Binary => { + let encoded = general_purpose::STANDARD.encode(get_binary(value, pos)?); + out.push_str( + &serde_json::to_string(&encoded).map_err(|e| Error::DataInvalid { + message: "Failed to encode Variant binary".to_string(), + source: Some(Box::new(e)), + })?, + ); + } + VariantKind::Uuid => out.push_str( + &serde_json::to_string(&get_uuid(value, pos)?.to_string()).map_err(|e| { + Error::DataInvalid { + message: "Failed to stringify Variant UUID".to_string(), + source: Some(Box::new(e)), + } + })?, + ), + VariantKind::Date | VariantKind::Timestamp | VariantKind::TimestampNtz => { + out.push_str( + &serde_json::to_string(&get_long(value, pos)?.to_string()).map_err(|e| { + Error::DataInvalid { + message: "Failed to stringify Variant value".to_string(), + source: Some(Box::new(e)), + } + })?, + ); + } + } + Ok(()) +} + +#[derive(Clone, Debug, PartialEq, Eq)] +enum PathSegment { + Key(String), + Index(usize), +} + +fn parse_path(path: &str) -> Result> { + let bytes = path.as_bytes(); + if !bytes.starts_with(b"$") { + return data_invalid(format!("Invalid Variant path: {path}")); + } + let mut pos = 1usize; + let mut segments = Vec::new(); + while pos < bytes.len() { + match bytes[pos] { + b'.' => { + pos += 1; + let start = pos; + while pos < bytes.len() && bytes[pos] != b'.' && bytes[pos] != b'[' { + pos += 1; + } + if start == pos { + return data_invalid(format!("Invalid Variant path: {path}")); + } + segments.push(PathSegment::Key(path[start..pos].to_string())); + } + b'[' => { + pos += 1; + if pos >= bytes.len() { + return data_invalid(format!("Invalid Variant path: {path}")); + } + if bytes[pos] == b'\'' || bytes[pos] == b'"' { + let quote = bytes[pos]; + pos += 1; + let start = pos; + while pos < bytes.len() && bytes[pos] != quote { + pos += 1; + } + if pos >= bytes.len() || pos + 1 >= bytes.len() || bytes[pos + 1] != b']' { + return data_invalid(format!("Invalid Variant path: {path}")); + } + segments.push(PathSegment::Key(path[start..pos].to_string())); + pos += 2; + } else { + let start = pos; + while pos < bytes.len() && bytes[pos].is_ascii_digit() { + pos += 1; + } + if start == pos || pos >= bytes.len() || bytes[pos] != b']' { + return data_invalid(format!("Invalid Variant path: {path}")); + } + let index = + path[start..pos] + .parse::() + .map_err(|e| Error::DataInvalid { + message: format!("Invalid Variant path index: {path}"), + source: Some(Box::new(e)), + })?; + segments.push(PathSegment::Index(index)); + pos += 1; + } + } + _ => return data_invalid(format!("Invalid Variant path: {path}")), + } + } + Ok(segments) +} + +#[derive(Clone)] +struct FieldEntry { + key: String, + id: usize, + offset: usize, +} + +struct VariantBuilder { + value: Vec, + dictionary: HashMap, + dictionary_keys: Vec>, +} + +impl VariantBuilder { + fn new() -> Self { + Self { + value: Vec::with_capacity(128), + dictionary: HashMap::new(), + dictionary_keys: Vec::new(), + } + } + + fn result(self) -> Result { + let num_keys = self.dictionary_keys.len(); + let dictionary_string_size = self.dictionary_keys.iter().map(Vec::len).sum::(); + let max_size = dictionary_string_size.max(num_keys); + let offset_size = integer_size(max_size)?; + let offset_start = 1 + offset_size; + let string_start = offset_start + (num_keys + 1) * offset_size; + let metadata_size = string_start + dictionary_string_size; + if metadata_size > SIZE_LIMIT { + return data_invalid("Variant metadata exceeds size limit"); + } + + let mut metadata = vec![0u8; metadata_size]; + metadata[0] = VERSION | (((offset_size - 1) as u8) << 6); + write_le_at(&mut metadata, 1, num_keys, offset_size); + let mut current_offset = 0usize; + for (i, key) in self.dictionary_keys.iter().enumerate() { + write_le_at( + &mut metadata, + offset_start + i * offset_size, + current_offset, + offset_size, + ); + metadata[string_start + current_offset..string_start + current_offset + key.len()] + .copy_from_slice(key); + current_offset += key.len(); + } + write_le_at( + &mut metadata, + offset_start + num_keys * offset_size, + current_offset, + offset_size, + ); + GenericVariant::from_parts(self.value, metadata) + } + + fn add_key(&mut self, key: &str) -> usize { + if let Some(id) = self.dictionary.get(key) { + *id + } else { + let id = self.dictionary_keys.len(); + self.dictionary.insert(key.to_string(), id); + self.dictionary_keys.push(key.as_bytes().to_vec()); + id + } + } + + fn append_null(&mut self) { + self.value.push(primitive_header(NULL)); + } + + fn append_bool(&mut self, value: bool) { + self.value + .push(primitive_header(if value { TRUE } else { FALSE })); + } + + fn append_string(&mut self, value: &str) { + let bytes = value.as_bytes(); + if bytes.len() > MAX_SHORT_STR_SIZE { + self.value.push(primitive_header(LONG_STR)); + let pos = self.value.len(); + self.value.resize(pos + U32_SIZE, 0); + write_le_at(&mut self.value, pos, bytes.len(), U32_SIZE); + } else { + self.value.push(short_str_header(bytes.len())); + } + self.value.extend_from_slice(bytes); + } + + fn append_long(&mut self, value: i64) { + if value == value as i8 as i64 { + self.value.push(primitive_header(INT1)); + push_signed_le(&mut self.value, value as i128, 1); + } else if value == value as i16 as i64 { + self.value.push(primitive_header(INT2)); + push_signed_le(&mut self.value, value as i128, 2); + } else if value == value as i32 as i64 { + self.value.push(primitive_header(INT4)); + push_signed_le(&mut self.value, value as i128, 4); + } else { + self.value.push(primitive_header(INT8)); + push_signed_le(&mut self.value, value as i128, 8); + } + } + + fn append_double(&mut self, value: f64) { + self.value.push(primitive_header(DOUBLE)); + self.value.extend_from_slice(&value.to_bits().to_le_bytes()); + } + + fn append_decimal(&mut self, decimal: VariantDecimal) { + if decimal.scale as u8 <= MAX_DECIMAL4_PRECISION + && decimal.precision <= MAX_DECIMAL4_PRECISION + { + self.value.push(primitive_header(DECIMAL4)); + self.value.push(decimal.scale as u8); + push_signed_le(&mut self.value, decimal.unscaled, 4); + } else if decimal.scale as u8 <= MAX_DECIMAL8_PRECISION + && decimal.precision <= MAX_DECIMAL8_PRECISION + { + self.value.push(primitive_header(DECIMAL8)); + self.value.push(decimal.scale as u8); + push_signed_le(&mut self.value, decimal.unscaled, 8); + } else { + self.value.push(primitive_header(DECIMAL16)); + self.value.push(decimal.scale as u8); + self.value + .extend_from_slice(&decimal.unscaled.to_le_bytes()); + } + } + + fn finish_object(&mut self, start: usize, mut fields: Vec) -> Result<()> { + fields.sort_by(|a, b| java_string_cmp(&a.key, &b.key)); + for pair in fields.windows(2) { + if pair[0].key == pair[1].key { + return data_invalid("VARIANT_DUPLICATE_KEY"); + } + } + + let data_size = self.value.len() - start; + let size = fields.len(); + let large_size = size > U8_MAX; + let size_bytes = if large_size { U32_SIZE } else { 1 }; + let max_id = fields.iter().map(|f| f.id).max().unwrap_or(0); + let id_size = integer_size(max_id)?; + let offset_size = integer_size(data_size)?; + let header_size = 1 + size_bytes + size * id_size + (size + 1) * offset_size; + self.value.splice(start..start, vec![0; header_size]); + self.value[start] = object_header(large_size, id_size, offset_size); + write_le_at(&mut self.value, start + 1, size, size_bytes); + let id_start = start + 1 + size_bytes; + let offset_start = id_start + size * id_size; + for (i, field) in fields.iter().enumerate() { + write_le_at(&mut self.value, id_start + i * id_size, field.id, id_size); + write_le_at( + &mut self.value, + offset_start + i * offset_size, + field.offset, + offset_size, + ); + } + write_le_at( + &mut self.value, + offset_start + size * offset_size, + data_size, + offset_size, + ); + Ok(()) + } + + fn finish_array(&mut self, start: usize, offsets: Vec) -> Result<()> { + let data_size = self.value.len() - start; + let size = offsets.len(); + let large_size = size > U8_MAX; + let size_bytes = if large_size { U32_SIZE } else { 1 }; + let offset_size = integer_size(data_size)?; + let header_size = 1 + size_bytes + (size + 1) * offset_size; + self.value.splice(start..start, vec![0; header_size]); + self.value[start] = array_header(large_size, offset_size); + write_le_at(&mut self.value, start + 1, size, size_bytes); + let offset_start = start + 1 + size_bytes; + for (i, offset) in offsets.iter().enumerate() { + write_le_at( + &mut self.value, + offset_start + i * offset_size, + *offset, + offset_size, + ); + } + write_le_at( + &mut self.value, + offset_start + size * offset_size, + data_size, + offset_size, + ); + Ok(()) + } +} + +struct JsonVariantParser<'a> { + input: &'a str, + bytes: &'a [u8], + pos: usize, + builder: VariantBuilder, +} + +impl<'a> JsonVariantParser<'a> { + fn new(input: &'a str) -> Self { + Self { + input, + bytes: input.as_bytes(), + pos: 0, + builder: VariantBuilder::new(), + } + } + + fn parse(mut self) -> Result { + self.skip_ws(); + self.parse_value()?; + self.skip_ws(); + if self.pos != self.bytes.len() { + return self.parse_error("Trailing characters after JSON value"); + } + self.builder.result() + } + + fn parse_value(&mut self) -> Result<()> { + self.skip_ws(); + let Some(ch) = self.peek() else { + return self.parse_error("Unexpected end of JSON input"); + }; + match ch { + b'{' => self.parse_object(), + b'[' => self.parse_array(), + b'"' => { + let value = self.parse_string()?; + self.builder.append_string(&value); + Ok(()) + } + b't' => { + self.expect_literal("true")?; + self.builder.append_bool(true); + Ok(()) + } + b'f' => { + self.expect_literal("false")?; + self.builder.append_bool(false); + Ok(()) + } + b'n' => { + self.expect_literal("null")?; + self.builder.append_null(); + Ok(()) + } + b'-' | b'0'..=b'9' => self.parse_number(), + _ => self.parse_error("Unexpected JSON token"), + } + } + + fn parse_object(&mut self) -> Result<()> { + self.consume(b'{')?; + let start = self.builder.value.len(); + let mut fields = Vec::new(); + let mut seen = HashSet::new(); + self.skip_ws(); + if self.peek() == Some(b'}') { + self.pos += 1; + return self.builder.finish_object(start, fields); + } + loop { + self.skip_ws(); + let key = self.parse_string()?; + if !seen.insert(key.clone()) { + return self.parse_error("VARIANT_DUPLICATE_KEY"); + } + self.skip_ws(); + self.consume(b':')?; + let id = self.builder.add_key(&key); + let offset = self.builder.value.len() - start; + fields.push(FieldEntry { key, id, offset }); + self.parse_value()?; + self.skip_ws(); + match self.peek() { + Some(b',') => { + self.pos += 1; + } + Some(b'}') => { + self.pos += 1; + break; + } + _ => return self.parse_error("Expected ',' or '}' in JSON object"), + } + } + self.builder.finish_object(start, fields) + } + + fn parse_array(&mut self) -> Result<()> { + self.consume(b'[')?; + let start = self.builder.value.len(); + let mut offsets = Vec::new(); + self.skip_ws(); + if self.peek() == Some(b']') { + self.pos += 1; + return self.builder.finish_array(start, offsets); + } + loop { + offsets.push(self.builder.value.len() - start); + self.parse_value()?; + self.skip_ws(); + match self.peek() { + Some(b',') => { + self.pos += 1; + } + Some(b']') => { + self.pos += 1; + break; + } + _ => return self.parse_error("Expected ',' or ']' in JSON array"), + } + } + self.builder.finish_array(start, offsets) + } + + fn parse_string(&mut self) -> Result { + self.skip_ws(); + let start = self.pos; + self.consume(b'"')?; + let mut escaped = false; + while let Some(ch) = self.peek() { + self.pos += 1; + if escaped { + escaped = false; + continue; + } + match ch { + b'\\' => escaped = true, + b'"' => { + return serde_json::from_str(&self.input[start..self.pos]).map_err(|e| { + Error::DataInvalid { + message: "Invalid JSON string".to_string(), + source: Some(Box::new(e)), + } + }); + } + _ if ch < 0x20 => return self.parse_error("Invalid control character in string"), + _ => {} + } + } + self.parse_error("Unterminated JSON string") + } + + fn parse_number(&mut self) -> Result<()> { + let start = self.pos; + if self.peek() == Some(b'-') { + self.pos += 1; + } + match self.peek() { + Some(b'0') => self.pos += 1, + Some(b'1'..=b'9') => { + self.pos += 1; + while matches!(self.peek(), Some(b'0'..=b'9')) { + self.pos += 1; + } + } + _ => return self.parse_error("Invalid JSON number"), + } + let mut has_fraction = false; + if self.peek() == Some(b'.') { + has_fraction = true; + self.pos += 1; + let fraction_start = self.pos; + while matches!(self.peek(), Some(b'0'..=b'9')) { + self.pos += 1; + } + if self.pos == fraction_start { + return self.parse_error("Invalid JSON number fraction"); + } + } + let mut has_exponent = false; + if matches!(self.peek(), Some(b'e' | b'E')) { + has_exponent = true; + self.pos += 1; + if matches!(self.peek(), Some(b'+' | b'-')) { + self.pos += 1; + } + let exponent_start = self.pos; + while matches!(self.peek(), Some(b'0'..=b'9')) { + self.pos += 1; + } + if self.pos == exponent_start { + return self.parse_error("Invalid JSON number exponent"); + } + } + let token = &self.input[start..self.pos]; + if !has_fraction && !has_exponent { + if let Ok(value) = token.parse::() { + self.builder.append_long(value); + return Ok(()); + } + } + if !has_exponent { + if let Some(decimal) = parse_decimal_token(token)? { + self.builder.append_decimal(decimal); + return Ok(()); + } + } + let value = token.parse::().map_err(|e| Error::DataInvalid { + message: format!("Invalid JSON number: {token}"), + source: Some(Box::new(e)), + })?; + self.builder.append_double(value); + Ok(()) + } + + fn expect_literal(&mut self, literal: &str) -> Result<()> { + if self.input[self.pos..].starts_with(literal) { + self.pos += literal.len(); + Ok(()) + } else { + self.parse_error(format!("Expected JSON literal {literal}")) + } + } + + fn consume(&mut self, expected: u8) -> Result<()> { + self.skip_ws(); + if self.peek() == Some(expected) { + self.pos += 1; + Ok(()) + } else { + self.parse_error(format!("Expected '{}'", expected as char)) + } + } + + fn skip_ws(&mut self) { + while matches!(self.peek(), Some(b' ' | b'\n' | b'\r' | b'\t')) { + self.pos += 1; + } + } + + fn peek(&self) -> Option { + self.bytes.get(self.pos).copied() + } + + fn parse_error(&self, message: impl Into) -> Result { + data_invalid(format!("{} at byte {}", message.into(), self.pos)) + } +} + +fn parse_decimal_token(token: &str) -> Result> { + if !token + .bytes() + .all(|ch| ch == b'-' || ch == b'.' || ch.is_ascii_digit()) + { + return Ok(None); + } + let negative = token.starts_with('-'); + let unsigned = token.strip_prefix('-').unwrap_or(token); + let scale = unsigned + .split_once('.') + .map(|(_, fraction)| fraction.len()) + .unwrap_or(0); + if scale > MAX_DECIMAL16_PRECISION as usize { + return Ok(None); + } + let mut digits = String::with_capacity(unsigned.len()); + for ch in unsigned.bytes() { + if ch != b'.' { + digits.push(ch as char); + } + } + let significant = digits.trim_start_matches('0'); + let precision = if significant.is_empty() { + 1 + } else { + significant.len() + }; + if precision > MAX_DECIMAL16_PRECISION as usize { + return Ok(None); + } + let mut unscaled = digits.parse::().map_err(|e| Error::DataInvalid { + message: format!("Invalid decimal Variant number: {token}"), + source: Some(Box::new(e)), + })?; + if negative { + unscaled = -unscaled; + } + Ok(Some(VariantDecimal { + unscaled, + precision: precision as u8, + scale: scale as i8, + })) +} + +fn decimal_precision(unscaled: i128) -> u8 { + let mut value = unscaled.unsigned_abs(); + if value == 0 { + return 1; + } + let mut precision = 0u8; + while value > 0 { + precision += 1; + value /= 10; + } + precision +} + +fn decimal_to_plain_string(unscaled: i128, scale: i8, strip_trailing_zeros: bool) -> String { + if scale <= 0 { + return unscaled.to_string(); + } + let negative = unscaled < 0; + let digits = unscaled.unsigned_abs().to_string(); + let scale = scale as usize; + let mut result = if digits.len() > scale { + let split = digits.len() - scale; + format!("{}.{}", &digits[..split], &digits[split..]) + } else { + format!("0.{}{}", "0".repeat(scale - digits.len()), digits) + }; + if strip_trailing_zeros && result.contains('.') { + while result.ends_with('0') { + result.pop(); + } + if result.ends_with('.') { + result.pop(); + } + } + if negative && result != "0" { + result.insert(0, '-'); + } + result +} + +fn java_string_cmp(left: &str, right: &str) -> std::cmp::Ordering { + left.encode_utf16().cmp(right.encode_utf16()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn empty_metadata() -> Vec { + vec![VERSION, 0, 0] + } + + #[test] + fn parse_json_matches_java_basic_object_layout() { + let variant = GenericVariant::parse_json(r#"{"age":27,"city":"Beijing"}"#).unwrap(); + assert_eq!( + variant.value(), + &[ + 0x02, 0x02, 0x00, 0x01, 0x00, 0x02, 0x0a, 0x0c, 0x1b, 0x1d, b'B', b'e', b'i', b'j', + b'i', b'n', b'g' + ] + ); + assert_eq!( + variant.metadata(), + &[0x01, 0x02, 0x00, 0x03, 0x07, b'a', b'g', b'e', b'c', b'i', b't', b'y'] + ); + } + + #[test] + fn variant_get_path_reads_objects_and_arrays() { + let variant = + GenericVariant::parse_json(r#"{"object":{"name":"Alice"},"array":[1,2,null]}"#) + .unwrap(); + assert_eq!( + variant + .get_path("$.object.name") + .unwrap() + .unwrap() + .get_string() + .unwrap(), + "Alice" + ); + assert_eq!( + variant + .get_path("$.array[1]") + .unwrap() + .unwrap() + .get_long() + .unwrap(), + 2 + ); + assert!(variant + .get_path("$.array[2]") + .unwrap() + .unwrap() + .is_null() + .unwrap()); + assert!(variant.get_path("$.array[9]").unwrap().is_none()); + } + + #[test] + fn parse_json_rejects_duplicate_object_keys() { + let err = GenericVariant::parse_json(r#"{"a":1,"a":2}"#).unwrap_err(); + assert!(err.to_string().contains("VARIANT_DUPLICATE_KEY")); + } + + #[test] + fn validate_payload_rejects_malformed_root_value() { + let metadata = empty_metadata(); + assert!(validate_payload(&[], &metadata).is_err()); + + let truncated_short_string = [short_str_header(3), b'a']; + assert!(validate_payload(&truncated_short_string, &metadata).is_err()); + } + + #[test] + fn validate_payload_rejects_bad_object_metadata_ids() { + let variant = GenericVariant::parse_json(r#"{"a":1}"#).unwrap(); + let mut value = variant.value().to_vec(); + let layout = object_layout(&value, 0).unwrap(); + write_le_at(&mut value, layout.id_start, 1, layout.id_size); + + assert!(validate_payload(&value, variant.metadata()).is_err()); + } + + #[test] + fn validate_payload_rejects_bad_object_offsets() { + let variant = GenericVariant::parse_json(r#"{"a":1}"#).unwrap(); + let mut value = variant.value().to_vec(); + let layout = object_layout(&value, 0).unwrap(); + let data_size = value.len() - layout.data_start; + write_le_at( + &mut value, + layout.offset_start + layout.size * layout.offset_size, + data_size + 1, + layout.offset_size, + ); + + assert!(validate_payload(&value, variant.metadata()).is_err()); + } + + #[test] + fn json_number_encoding_keeps_decimal_and_double_distinct() { + let decimal = GenericVariant::parse_json(r#"{"d":123.4500}"#).unwrap(); + let d = decimal.get_path("$.d").unwrap().unwrap(); + assert_eq!(d.kind().unwrap(), VariantKind::Decimal); + assert_eq!(d.to_json().unwrap(), "123.45"); + + let double = GenericVariant::parse_json(r#"{"d":1.23e10}"#).unwrap(); + assert_eq!( + double.get_path("$.d").unwrap().unwrap().kind().unwrap(), + VariantKind::Double + ); + } +} diff --git a/crates/paimon/tests/fixtures/variant_type.json b/crates/paimon/tests/fixtures/variant_type.json new file mode 100644 index 00000000..a4b1f54d --- /dev/null +++ b/crates/paimon/tests/fixtures/variant_type.json @@ -0,0 +1 @@ +"VARIANT NOT NULL" diff --git a/crates/paimon/tests/fixtures/variant_type_nullable.json b/crates/paimon/tests/fixtures/variant_type_nullable.json new file mode 100644 index 00000000..e96faaf4 --- /dev/null +++ b/crates/paimon/tests/fixtures/variant_type_nullable.json @@ -0,0 +1 @@ +"VARIANT" diff --git a/docs/src/sql.md b/docs/src/sql.md index b43f84cb..283e03fa 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -95,6 +95,7 @@ The following SQL data types are supported in CREATE TABLE and mapped to their c | `DOUBLE` / `DOUBLE PRECISION` | DoubleType | | | `VARCHAR` / `TEXT` / `STRING` / `CHAR` | VarCharType | | | `BINARY` / `VARBINARY` / `BYTEA` | VarBinaryType | | +| `VARIANT` | VariantType | Semi-structured value encoded as value + metadata binary buffers | | `BLOB` | BlobType | Binary large object | | `DATE` | DateType | | | `TIMESTAMP[(p)]` | TimestampType | Precision p: 0/3/6/9, default 3 | @@ -104,6 +105,123 @@ The following SQL data types are supported in CREATE TABLE and mapped to their c | `MAP(key, value)` | MapType | e.g. `MAP(STRING, INT)` | | `STRUCT` | RowType | e.g. `STRUCT` | +### Variant Usage + +`VARIANT` stores semi-structured data using the same logical value + metadata binary shape as Paimon Java. Use it for JSON-like fields whose schema may differ row by row. + +Create `VARIANT` columns like ordinary table columns: + +```sql +CREATE TABLE paimon.my_db.user_events ( + user_id BIGINT NOT NULL, + event_time TIMESTAMP, + payload VARIANT, + attributes VARIANT, + dt STRING, + PRIMARY KEY (user_id, dt) +) PARTITIONED BY (dt) +WITH ('bucket' = '4'); +``` + +`VARIANT` columns can be nullable or `NOT NULL`: + +```sql +CREATE TABLE paimon.my_db.variant_examples ( + id INT NOT NULL, + payload VARIANT NOT NULL, + optional_payload VARIANT +); +``` + +Do not use `VARIANT` as a partition column. Partition values must be scalar strings, numbers, dates, or timestamps that can be encoded as stable partition names. + +Use `parse_json` when inserting JSON text into a `VARIANT` column: + +```sql +INSERT INTO paimon.my_db.user_events VALUES +( + 1, + TIMESTAMP '2024-01-01 10:00:00', + parse_json('{"event":"login","device":{"os":"ios","version":17},"score":98.5}'), + parse_json('{"city":"Beijing","tags":["new","mobile"],"vip":true}'), + '2024-01-01' +); +``` + +`parse_json` rejects invalid JSON and duplicate object keys. Use `try_parse_json` when malformed JSON should become SQL `NULL` instead of failing the query: + +```sql +INSERT INTO paimon.my_db.user_events +SELECT + user_id, + event_time, + try_parse_json(raw_payload), + try_parse_json(raw_attributes), + dt +FROM staging_events; +``` + +`SQLContext::new` registers Spark-compatible scalar functions for common `VARIANT` workflows: + +```sql +SELECT + user_id, + variant_get(payload, '$.event', 'string') AS event_name, + variant_get(payload, '$.device.os', 'string') AS os, + variant_get(payload, '$.score', 'double') AS score, + variant_get(attributes, '$.tags[0]', 'string') AS first_tag +FROM paimon.my_db.user_events +WHERE variant_get(attributes, '$.vip', 'boolean') = true; +``` + +Supported functions: + +| Function | Notes | +|---|---| +| `parse_json(json)` | Parses a JSON string into `VARIANT`; invalid JSON returns an error | +| `try_parse_json(json)` | Parses a JSON string into `VARIANT`; invalid JSON returns `NULL` | +| `variant_get(v, path[, type])` | Extracts a path; missing paths return `NULL`; invalid casts return an error | +| `try_variant_get(v, path[, type])` | Extracts a path; missing paths, invalid paths, and invalid casts return `NULL` | +| `is_variant_null(v)` | Returns true for JSON `null` inside `VARIANT`; SQL `NULL` returns false | + +Path syntax supports the root path `$`, object access (`$.field`), quoted object access (`$["field"]` or `$['field']`), array indexes (`$[0]`), and nested combinations such as `$.items[0].price`. + +The optional `type` argument is a string literal. Supported result types are `variant` (or omitted), `boolean`, `byte` / `tinyint`, `short` / `smallint`, `int` / `integer`, `long` / `bigint`, `float`, `double`, `decimal(p, s)`, and `string`. + +When `type` is omitted or set to `variant`, `variant_get` returns a nested `VARIANT` value that can be passed to another `variant_get` call: + +```sql +SELECT + variant_get( + variant_get(payload, '$.device'), + '$.os', + 'string' + ) AS os +FROM paimon.my_db.user_events; +``` + +Missing paths return SQL `NULL`. JSON `null` is represented as a non-SQL-null Variant value, so use `is_variant_null` when you need to distinguish it: + +```sql +SELECT + is_variant_null(parse_json('null')) AS json_null, + is_variant_null(NULL) AS sql_null; +``` + +Current limitations: + +- `schema_of_variant`, `schema_of_variant_agg`, `to_variant_object`, `variant_explode`, and `variant_explode_outer` are not implemented yet. +- `variant_get` currently casts to scalar types and `VARIANT`. It does not yet cast directly to `ARRAY`, `MAP`, or `STRUCT`. +- Predicate pushdown is not applied through `variant_get`; DataFusion evaluates Variant filters after reading rows. + +With a raw DataFusion `SessionContext`, register these scalar functions explicitly: + +```rust +use paimon_datafusion::register_variant_functions; + +register_variant_functions(&ctx); +``` + ## DDL ### CREATE DATABASE / CREATE SCHEMA / DROP SCHEMA @@ -280,6 +398,19 @@ INSERT INTO paimon.my_db.users VALUES (1, 'alice'), (2, 'bob'), (3, 'carol'); INSERT INTO paimon.my_db.users SELECT * FROM source_table; ``` +For `VARIANT` columns, convert JSON text with `parse_json` or `try_parse_json`: + +```sql +INSERT INTO paimon.my_db.user_events (user_id, event_time, payload, attributes, dt) +VALUES ( + 1, + TIMESTAMP '2024-01-01 10:00:00', + parse_json('{"event":"login","device":{"os":"ios"}}'), + try_parse_json('{"vip":true,"tags":["mobile"]}'), + '2024-01-01' +); +``` + For primary-key tables, records with duplicate keys are deduplicated according to the merge engine (default: Deduplicate engine, where the last written value wins). ### Mosaic Read Scope @@ -492,6 +623,29 @@ All DataFusion query capabilities are supported (JOINs, aggregations, subqueries SELECT id, name FROM paimon.my_db.users WHERE id > 10 ORDER BY id LIMIT 100; ``` +### Variant Queries + +Use `variant_get` to extract fields from `VARIANT` columns. Provide a target type string when the query needs a scalar result: + +```sql +SELECT + user_id, + variant_get(payload, '$.event', 'string') AS event_name, + variant_get(payload, '$.device.os', 'string') AS device_os, + variant_get(attributes, '$.vip', 'boolean') AS is_vip +FROM paimon.my_db.user_events +WHERE variant_get(payload, '$.event', 'string') = 'login'; +``` + +Use `try_variant_get` when incompatible values should return `NULL`: + +```sql +SELECT + user_id, + try_variant_get(payload, '$.score', 'double') AS score +FROM paimon.my_db.user_events; +``` + ### Column Projection Only the required columns are read, reducing I/O: