leaves12138 commented on code in PR #458: URL: https://github.com/apache/paimon-rust/pull/458#discussion_r3525106537
########## crates/paimon/src/variant.rs: ########## @@ -0,0 +1,1465 @@ +// 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_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<u8>, + metadata: Vec<u8>, +} + +impl GenericVariant { + /// Parse a JSON string into the same binary Variant representation produced by Java Paimon. + pub fn parse_json(json: &str) -> Result<Self> { + JsonVariantParser::new(json).parse() + } + + /// Build an owned Variant from existing buffers after validating the metadata header. + pub fn from_parts(value: Vec<u8>, metadata: Vec<u8>) -> Result<Self> { + 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<'_>> { + VariantRef::new(&self.value, &self.metadata, 0) + } + + pub fn is_variant_null(&self) -> Result<bool> { + self.as_ref()?.is_null() + } + + pub fn get_path(&self, path: &str) -> Result<Option<VariantRef<'_>>> { + self.as_ref()?.get_path(path) + } + + pub fn to_json(&self) -> Result<String> { + 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<Self> { + validate_payload(value, metadata)?; + check_index(value, pos)?; + Ok(Self { + value, + metadata, + pos, + }) + } + + pub fn kind(&self) -> Result<VariantKind> { + value_kind(self.value, self.pos) + } + + pub fn is_null(&self) -> Result<bool> { + 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<GenericVariant> { + Ok(GenericVariant { + value: self.value_slice()?.to_vec(), + metadata: self.metadata.to_vec(), + }) + } + + pub fn get_path(&self, path: &str) -> Result<Option<VariantRef<'a>>> { + 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<bool> { + get_boolean(self.value, self.pos) + } + + pub fn get_long(&self) -> Result<i64> { + get_long(self.value, self.pos) + } + + pub fn get_double(&self) -> Result<f64> { + get_double(self.value, self.pos) + } + + pub fn get_float(&self) -> Result<f32> { + get_float(self.value, self.pos) + } + + pub fn get_string(&self) -> Result<String> { + get_string(self.value, self.pos) + } + + pub fn get_decimal(&self) -> Result<VariantDecimal> { + get_decimal(self.value, self.pos) + } + + pub fn to_json(&self) -> Result<String> { + 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<Option<VariantRef<'a>>> { + 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(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( + self.value, + self.metadata, + layout.data_start + offset, + ) + .map(Some); + } + } + } + } + Ok(None) + } + + fn get_element_at_index(&self, index: usize) -> Result<Option<VariantRef<'a>>> { + 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(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<()> { Review Comment: This validation is used by the row/binary-row writers before accepting external Arrow `value`/`metadata` buffers, but it only checks the metadata version and total sizes. For example, `value = []` and `metadata = [0x01]` passes `VariantType::validate_payload`, so a non-null Variant with an empty value buffer can be persisted and only fail later when `VariantRef`/`to_json` reads it. Please validate at least the root value buffer here (and ideally recursively validate object/array offsets and metadata ids), or avoid using this helper at ingestion boundaries where malformed Variant bytes must be rejected. ########## crates/integrations/datafusion/src/variant_functions.rs: ########## @@ -0,0 +1,851 @@ +// 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<ArrowDataType> { + Ok(variant_arrow_type()) + } + + fn return_field_from_args(&self, _args: ReturnFieldArgs) -> DFResult<FieldRef> { + Ok(Arc::new(Field::new( + self.name(), + variant_arrow_type(), + true, + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult<ColumnarValue> { + 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<ArrowDataType> { + Ok(ArrowDataType::Boolean) + } + + fn return_field_from_args(&self, _args: ReturnFieldArgs) -> DFResult<FieldRef> { + Ok(Arc::new(Field::new( + self.name(), + ArrowDataType::Boolean, + false, + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult<ColumnarValue> { + 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<ArrowDataType> { + internal_err("return_field_from_args should be used for variant_get") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> DFResult<FieldRef> { + if args.arg_fields.len() != 2 && args.arg_fields.len() != 3 { + return plan_err(format!("{} expects 2 or 3 arguments", self.name())); + } + let output = variant_get_output_type(args.scalar_arguments.get(2).and_then(|v| *v))?; Review Comment: If the third argument is not a scalar literal, `scalar_arguments.get(2)` is `None`, so the planner silently chooses the default `Variant` return type. I reproduced this with `SELECT variant_get(parse_json('{"age":26}'), '$.age', ty) FROM t` where `ty = 'int'`: the result schema is still the Variant struct instead of an integer (or a planning error). Since DataFusion needs a fixed output type, please reject non-literal third arguments during planning, rather than treating them as omitted. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
