diff --git a/crates/paimon/src/arrow/mod.rs b/crates/paimon/src/arrow/mod.rs index 34140b974..7273e1751 100644 --- a/crates/paimon/src/arrow/mod.rs +++ b/crates/paimon/src/arrow/mod.rs @@ -26,6 +26,8 @@ pub(crate) mod partition; mod read_budget; pub(crate) mod residual; mod row_filter; +mod string_transform; +pub(crate) use string_transform::evaluate_string_transform; pub(crate) mod schema_evolution; pub(crate) mod shredding; mod variant; diff --git a/crates/paimon/src/arrow/string_transform.rs b/crates/paimon/src/arrow/string_transform.rs new file mode 100644 index 000000000..71d5dbab7 --- /dev/null +++ b/crates/paimon/src/arrow/string_transform.rs @@ -0,0 +1,447 @@ +// 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. + +//! Arrow execution of Java's character and position transforms. + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::spec::{PadDirection, Transform, TransformInput, TrimFlag}; +use crate::Result; +use arrow_array::{cast::AsArray, Array, ArrayRef, Int32Array, Int64Array, StringArray}; +use arrow_schema::DataType as ArrowType; + +enum Input<'a> { + Text(Option<&'a str>), + Number(Option), + IntegerText(&'a str), + Strings(StringArray), + Integers(Int64Array), +} + +impl Input<'_> { + fn is_null(&self, row: usize) -> bool { + match self { + Self::Text(value) => value.is_none(), + Self::Number(value) => value.is_none(), + Self::IntegerText(_) => false, + Self::Strings(values) => values.is_null(row), + Self::Integers(values) => values.is_null(row), + } + } + fn text(&self, row: usize) -> Result> { + match self { + Self::Text(value) => Ok(*value), + Self::Strings(values) => Ok(values.is_valid(row).then(|| values.value(row))), + _ => Err(invalid("string transform requires a character input")), + } + } + + fn integer(&self, row: usize) -> Result> { + let value = match self { + Self::IntegerText(text) => return parse_integer(text).map(Some), + Self::Number(value) => *value, + Self::Integers(values) => values.is_valid(row).then(|| values.value(row)), + _ => return Err(invalid("position transform requires an integer input")), + }; + value + .map(|v| { + i32::try_from(v).map_err(|_| invalid("transform position exceeds Java INT range")) + }) + .transpose() + } +} + +pub(crate) fn evaluate_string_transform( + transform: &Transform, + rows: usize, + column: &dyn Fn(usize) -> Result, +) -> Result { + let inputs = transform + .inputs() + .iter() + .map(|input| match input { + TransformInput::Literal(value) => Ok(Input::Text(value.as_deref())), + TransformInput::Integer(value) => Ok(Input::Number(*value)), + TransformInput::IntegerText(text) => Ok(Input::IntegerText(text)), + TransformInput::Field(index) => { + let array = column(*index)?; + match array.data_type() { + ArrowType::Utf8 => Ok(Input::Strings(array.as_string::().clone())), + ArrowType::Int8 | ArrowType::Int16 | ArrowType::Int32 | ArrowType::Int64 => { + let array = arrow_cast::cast(&array, &ArrowType::Int64) + .map_err(|e| invalid(&e.to_string()))?; + Ok(Input::Integers( + array + .as_primitive::() + .clone(), + )) + } + other => Err(invalid(&format!("invalid transform input type {other:?}"))), + } + } + }) + .collect::>>()?; + if matches!(transform, Transform::Length(_) | Transform::BitLength(_)) { + let values = (0..rows) + .map(|row| { + Ok(inputs[0].text(row)?.map(|text| { + if matches!(transform, Transform::Length(_)) { + text.chars().count() as i32 + } else { + (text.len() as i32).wrapping_mul(8) + } + })) + }) + .collect::>>()?; + return Ok(Arc::new(Int32Array::from(values))); + } + let values = (0..rows) + .map(|row| evaluate_row(transform, &inputs, row)) + .collect::>>()?; + Ok(Arc::new(StringArray::from(values))) +} + +fn evaluate_row(transform: &Transform, inputs: &[Input<'_>], row: usize) -> Result> { + let source = inputs[0].text(row)?; + let value = match transform { + Transform::Substring(_) => { + // SubstringTransform checks source and null positions before + // evaluating positions; a hidden NULL source cannot raise overflow. + let Some(source) = source else { + return Ok(None); + }; + if inputs[1..].iter().any(|input| input.is_null(row)) { + return Ok(None); + } + let Some(position) = inputs[1].integer(row)? else { + return Ok(None); + }; + let length = if inputs.len() == 3 { + let Some(length) = inputs[2].integer(row)? else { + return Ok(None); + }; + length + } else { + i32::MAX + }; + substring_sql(source, position, length).to_owned() + } + Transform::Trim(_, flag) => { + let Some(source) = source else { + return Ok(None); + }; + let chars = if inputs.len() == 2 { + let Some(chars) = inputs[1].text(row)? else { + return Ok(None); + }; + chars + } else { + " " + }; + let contains = |c| chars.contains(c); + match flag { + TrimFlag::Leading => source.trim_start_matches(contains), + TrimFlag::Trailing => source.trim_end_matches(contains), + TrimFlag::Both => source.trim_matches(contains), + } + .to_owned() + } + Transform::Pad(_, direction) => { + let length = inputs[1].integer(row)?; + let padding = inputs[2].text(row)?; + let (Some(source), Some(length), Some(padding)) = (source, length, padding) else { + return Ok(None); + }; + let count = source.chars().count() as i32; + let needed = length.wrapping_sub(count); + if needed <= 0 || padding.is_empty() { + substring(source, 0, length).to_owned() + } else { + let pad_chars = padding.chars().count(); + let repeats = needed as usize / pad_chars; + let tail = substring(padding, 0, needed % pad_chars as i32); + let size = padding + .len() + .checked_mul(repeats) + .and_then(|n| n.checked_add(tail.len())) + .and_then(|n| n.checked_add(source.len())) + .filter(|n| *n <= i32::MAX as usize) + .ok_or_else(|| invalid("PAD result exceeds Arrow UTF-8 storage"))?; + let mut result = String::new(); + result + .try_reserve(size) + .map_err(|e| invalid(&format!("cannot reserve PAD result: {e}")))?; + if *direction == PadDirection::Right { + result.push_str(source); + } + for _ in 0..repeats { + result.push_str(padding); + } + result.push_str(tail); + if *direction == PadDirection::Left { + result.push_str(source); + } + result + } + } + Transform::Translate(_) => { + let matching = inputs[1].text(row)?; + let replacement = inputs[2].text(row)?; + let (Some(source), Some(matching), Some(replacement)) = (source, matching, replacement) + else { + return Ok(None); + }; + let replacements: Vec = replacement.chars().collect(); + let mut mapping = HashMap::new(); + for (index, ch) in matching.chars().enumerate() { + mapping + .entry(ch) + .or_insert(replacements.get(index).copied().filter(|c| *c != '\0')); + } + source + .chars() + .filter_map(|ch| mapping.get(&ch).copied().unwrap_or(Some(ch))) + .collect() + } + Transform::Overlay(_) => { + let replacement = inputs[1].text(row)?; + let position = inputs[2].integer(row)?; + let length = if inputs.len() == 4 { + inputs[3].integer(row)? + } else { + None + }; + let (Some(source), Some(replacement), Some(position)) = (source, replacement, position) + else { + return Ok(None); + }; + if inputs.len() == 4 && length.is_none() { + return Ok(None); + } + let replaced = length + .filter(|length| *length >= 0) + .unwrap_or_else(|| replacement.chars().count() as i32); + format!( + "{}{}{}", + substring_sql(source, 1, position.wrapping_sub(1)), + replacement, + substring_sql(source, position.wrapping_add(replaced), i32::MAX) + ) + } + _ => return Err(invalid("unsupported character transform")), + }; + Ok(Some(value)) +} + +fn substring_sql(source: &str, position: i32, length: i32) -> &str { + let count = source.chars().count() as i32; + let start = if position > 0 { + position - 1 + } else if position < 0 { + count + position + } else { + 0 + }; + let end = (i64::from(start) + i64::from(length)).clamp(i64::from(i32::MIN), i64::from(i32::MAX)) + as i32; + substring(source, start, end) +} + +fn substring(source: &str, start: i32, end: i32) -> &str { + if end <= start || end <= 0 { + return ""; + } + let start = source + .char_indices() + .nth(start.max(0) as usize) + .map_or(source.len(), |(index, _)| index); + let end = source + .char_indices() + .nth(end as usize) + .map_or(source.len(), |(index, _)| index); + &source[start..end.max(start)] +} + +fn invalid(message: &str) -> crate::Error { + crate::Error::DataInvalid { + message: message.to_owned(), + source: None, + } +} + +/// Integer.parseInt uses Character.digit on UTF-16 chars: decimal digits in +/// BMP scripts are accepted, supplementary digits and whitespace are rejected. +fn parse_integer(text: &str) -> Result { + const ZEROES: &[u32] = &[ + 0x0030, 0x0660, 0x06f0, 0x07c0, 0x0966, 0x09e6, 0x0a66, 0x0ae6, 0x0b66, 0x0be6, 0x0c66, + 0x0ce6, 0x0d66, 0x0de6, 0x0e50, 0x0ed0, 0x0f20, 0x1040, 0x1090, 0x17e0, 0x1810, 0x1946, + 0x19d0, 0x1a80, 0x1a90, 0x1b50, 0x1bb0, 0x1c40, 0x1c50, 0xa620, 0xa8d0, 0xa900, 0xa9d0, + 0xa9f0, 0xaa50, 0xabf0, 0xff10, + ]; + let (negative, digits) = match text.strip_prefix('-') { + Some(digits) => (true, digits), + None => (false, text.strip_prefix('+').unwrap_or(text)), + }; + if digits.is_empty() { + return Err(invalid("empty transform position")); + } + // Accumulate negatively, so i32::MIN remains representable. + let mut value = 0i32; + for ch in digits.chars() { + let digit = ZEROES + .iter() + .find_map(|zero| { + (u32::from(ch) >= *zero && u32::from(ch) < zero + 10) + .then(|| (u32::from(ch) - zero) as i32) + }) + .ok_or_else(|| invalid("invalid Java integer position"))?; + value = value + .checked_mul(10) + .and_then(|v| v.checked_sub(digit)) + .ok_or_else(|| invalid("transform position exceeds Java INT range"))?; + } + if negative { + Ok(value) + } else { + value + .checked_neg() + .ok_or_else(|| invalid("transform position exceeds Java INT range")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{BigIntType, DataField, DataType, VarCharType}; + + fn apply(rule: &str, source: Vec>) -> Result { + let fields = vec![DataField::new( + 0, + "s".into(), + DataType::VarChar(VarCharType::string_type()), + )]; + let transform = Transform::from_rest_json(rule, &fields)?; + let rows = source.len(); + let source = Arc::new(StringArray::from(source)) as ArrayRef; + evaluate_string_transform(&transform, rows, &|_| Ok(source.clone())) + } + + #[test] + fn substring_positions_are_parsed_after_null_checks() { + for position in [ + r#""bad""#, + r#""999999999999999999999""#, + "18446744073709551617", + ] { + let rule = format!(r#"{{"name":"SUBSTRING","inputs":[{{"name":"s"}}, {position}]}}"#); + assert!(apply(&rule, vec![None]).unwrap().is_null(0)); + assert!(apply(&rule, vec![Some("abc")]).is_err()); + let rule = + format!(r#"{{"name":"SUBSTRING","inputs":[{{"name":"s"}}, {position}, null]}}"#); + assert!(apply(&rule, vec![Some("abc")]).unwrap().is_null(0)); + } + let result = apply( + r#"{"name":"SUBSTRING","inputs":[{"name":"s"},"+٢",1]}"#, + vec![Some("abc")], + ) + .unwrap(); + assert_eq!(result.as_string::().value(0), "b"); + for position in ["1.5", "1e2", "true"] { + assert!(apply( + &format!(r#"{{"name":"SUBSTRING","inputs":[null,{position}]}}"#), + vec![None] + ) + .is_err()); + } + } + + #[test] + fn integer_syntax_and_range_match_java() { + for (text, value) in [ + ("-2147483648", i32::MIN), + ("+2147483647", i32::MAX), + ("٠٠٢", 2), + ("2", 2), + ] { + assert_eq!(parse_integer(text).unwrap(), value); + } + for text in ["", "+", " 2", "2 ", "1_0", "2147483648", "-2147483649", "𝟚"] { + assert!(parse_integer(text).is_err(), "{text}"); + } + } + + #[test] + fn pad_and_overlay_convert_integers_before_null_strings() { + for rule in [ + r#"{"name":"PAD","direction":"LEFT","inputs":[null,2147483648,"x"]}"#, + r#"{"name":"OVERLAY","inputs":[null,"x",2147483648,null]}"#, + r#"{"name":"OVERLAY","inputs":[null,"x",null,2147483648]}"#, + ] { + assert!(apply(rule, vec![None]).is_err()); + } + } + + #[test] + fn mixed_fields_are_bound_by_name_and_checked_before_execution() { + let fields = vec![ + DataField::new(0, "s".into(), DataType::VarChar(VarCharType::string_type())), + DataField::new(1, "n".into(), DataType::BigInt(BigIntType::new())), + ]; + for rule in [ + r#"{"name":"PAD","direction":"RIGHT","inputs":[{"name":"s","type":"INT"},3,"x"]}"#, + r#"{"name":"PAD","direction":"RIGHT","inputs":["a",{"name":"n","type":"STRING"},"x"]}"#, + r#"{"name":"OVERLAY","inputs":["a","x",{"name":"n","type":"STRING"}]}"#, + r#"{"name":"TRIM","trimFlag":"BOTH","inputs":[{"name":"s","type":"INT"}]}"#, + r#"{"name":"SUBSTRING","inputs":["a",{"name":"s","type":"INT"}]}"#, + ] { + assert!(Transform::from_rest_json(rule, &fields).is_err(), "{rule}"); + } + let transform = Transform::from_rest_json(r#"{"name":"SUBSTRING","inputs":[{"index":99,"name":"s","type":"INT"},{"index":99,"name":"n","type":"STRING"},1]}"#, &fields).unwrap(); + let columns: Vec = vec![ + Arc::new(StringArray::from(vec![Some("甲🙂乙"), None])), + Arc::new(Int64Array::from(vec![2, i64::MAX])), + ]; + let result = + evaluate_string_transform(&transform, 2, &|index| Ok(columns[index].clone())).unwrap(); + assert_eq!( + result.as_string::().iter().collect::>(), + vec![Some("🙂"), None] + ); + } + + #[test] + fn lengths_count_unicode_characters_and_utf8_bits_as_int() { + for (name, expected) in [ + ("LENGTH", vec![Some(3), Some(2), Some(0), None]), + ("BIT_LENGTH", vec![Some(80), Some(24), Some(0), None]), + ] { + let result = apply( + &format!(r#"{{"name":"{name}","inputs":[{{"name":"s","type":"STRING"}}]}}"#), + vec![Some("甲🙂乙"), Some("e\u{301}"), Some(""), None], + ) + .unwrap(); + assert_eq!( + result + .as_primitive::() + .iter() + .collect::>(), + expected + ); + } + } +} diff --git a/crates/paimon/src/spec/mod.rs b/crates/paimon/src/spec/mod.rs index 7969f9b3f..a108f89ea 100644 --- a/crates/paimon/src/spec/mod.rs +++ b/crates/paimon/src/spec/mod.rs @@ -133,8 +133,8 @@ pub(crate) use predicate::like_match; pub(crate) use predicate::row_id_leaf; pub(crate) use predicate::validate_datum_matches_type; pub use predicate::{ - field_idx_to_partition_idx, Datum, Predicate, PredicateBuilder, PredicateOperator, Transform, - TransformInput, + field_idx_to_partition_idx, Datum, PadDirection, Predicate, PredicateBuilder, + PredicateOperator, Transform, TransformInput, TrimFlag, }; pub(crate) mod murmur_hash; mod partition_statistics; diff --git a/crates/paimon/src/spec/predicate.rs b/crates/paimon/src/spec/predicate.rs index 21c5ec81a..79b4b25bd 100644 --- a/crates/paimon/src/spec/predicate.rs +++ b/crates/paimon/src/spec/predicate.rs @@ -1206,13 +1206,51 @@ pub enum Transform { Concat(Vec), /// Concatenation where the first input is the separator. ConcatWs(Vec), + /// SQL substring with one-based, zero or negative positions. + Substring(Vec), + /// Remove whole characters at the requested ends. + Trim(Vec, TrimFlag), + /// Number of Unicode characters, as INT. + Length(TransformInput), + /// Number of UTF-8 bits, as INT. + BitLength(TransformInput), + /// SQL LPAD/RPAD. + Pad(Vec, PadDirection), + /// Map or delete Unicode characters. + Translate(Vec), + /// SQL overlay with optional replacement length. + Overlay(Vec), } -/// One input of a string [`Transform`]. +/// The ends selected by Java's TRIM transform. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TrimFlag { + /// Remove leading characters. + Leading, + /// Remove trailing characters. + Trailing, + /// Remove characters at both ends. + Both, +} + +/// The side selected by Java's PAD transform. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PadDirection { + /// Prefix the padding. + Left, + /// Append the padding. + Right, +} + +/// One input of a [`Transform`]. #[derive(Debug, Clone, PartialEq)] pub enum TransformInput { /// A string literal (`None` = a JSON `null` input, i.e. SQL NULL). Literal(Option), + /// An integer position or length, checked as a Java int when evaluated. + Integer(Option), + /// SUBSTRING parses string and large integral positions only after NULL checks. + IntegerText(String), /// A field reference. Field(usize), } @@ -1233,11 +1271,8 @@ impl Transform { Transform::FieldRef(index) | Transform::Cast(index, _) => { out.insert(*index); } - Transform::Upper(inputs) - | Transform::Lower(inputs) - | Transform::Concat(inputs) - | Transform::ConcatWs(inputs) => { - for input in inputs { + _ => { + for input in self.inputs() { if let TransformInput::Field(index) = input { out.insert(*index); } @@ -1245,6 +1280,22 @@ impl Transform { } } } + + pub(crate) fn inputs(&self) -> &[TransformInput] { + match self { + Self::Upper(inputs) + | Self::Lower(inputs) + | Self::Concat(inputs) + | Self::ConcatWs(inputs) + | Self::Substring(inputs) + | Self::Trim(inputs, _) + | Self::Pad(inputs, _) + | Self::Translate(inputs) + | Self::Overlay(inputs) => inputs, + Self::Length(input) | Self::BitLength(input) => std::slice::from_ref(input), + _ => &[], + } + } } fn parse_rest_transform(value: &serde_json::Value, fields: &[DataField]) -> Result { @@ -1273,6 +1324,7 @@ fn parse_rest_transform(value: &serde_json::Value, fields: &[DataField]) -> Resu serde_json::Value::Null => Ok(TransformInput::Literal(None)), serde_json::Value::Object(_) => { let index = field_ref_index(input)?; + check_input_type(input, false)?; match fields[index].data_type() { DataType::Char(_) | DataType::VarChar(_) => { Ok(TransformInput::Field(index)) @@ -1319,10 +1371,171 @@ fn parse_rest_transform(value: &serde_json::Value, fields: &[DataField]) -> Resu "LOWER" => Ok(Transform::Lower(string_inputs(1, 1)?)), "CONCAT" => Ok(Transform::Concat(string_inputs(0, usize::MAX)?)), "CONCAT_WS" => Ok(Transform::ConcatWs(string_inputs(2, usize::MAX)?)), + "SUBSTRING" => Ok(Transform::Substring(transform_inputs( + value, + fields, + 2, + 3, + &[1, 2], + true, + )?)), + "TRIM" => { + let flag = match str_prop(value, "trimFlag")? { + "LEADING" => TrimFlag::Leading, + "TRAILING" => TrimFlag::Trailing, + "BOTH" => TrimFlag::Both, + flag => return Err(rest_json_err(format!("invalid trimFlag `{flag}`"))), + }; + Ok(Transform::Trim(string_inputs(1, 2)?, flag)) + } + "LENGTH" => Ok(Transform::Length(string_inputs(1, 1)?.remove(0))), + "BIT_LENGTH" => Ok(Transform::BitLength(string_inputs(1, 1)?.remove(0))), + "PAD" => { + let direction = match str_prop(value, "direction")? { + "LEFT" => PadDirection::Left, + "RIGHT" => PadDirection::Right, + direction => { + return Err(rest_json_err(format!( + "invalid PAD direction `{direction}`" + ))) + } + }; + Ok(Transform::Pad( + transform_inputs(value, fields, 3, 3, &[1], false)?, + direction, + )) + } + "TRANSLATE" => Ok(Transform::Translate(string_inputs(3, 3)?)), + "OVERLAY" => Ok(Transform::Overlay(transform_inputs( + value, + fields, + 3, + 4, + &[2, 3], + false, + )?)), other => Err(rest_json_err(format!("transform `{other}`"))), } } +/// Bind mixed string/integer inputs by name to the actual table schema. +fn transform_inputs( + value: &serde_json::Value, + fields: &[DataField], + min: usize, + max: usize, + integers: &[usize], + string_positions: bool, +) -> Result> { + let inputs = value + .get("inputs") + .and_then(serde_json::Value::as_array) + .ok_or_else(|| rest_json_err("transform without inputs"))?; + if !(min..=max).contains(&inputs.len()) { + return Err(rest_json_err(format!( + "transform expects {min}..={max} inputs, got {}", + inputs.len() + ))); + } + inputs + .iter() + .enumerate() + .map(|(slot, input)| { + let integer = integers.contains(&slot); + match input { + serde_json::Value::Null => Ok(if integer { + TransformInput::Integer(None) + } else { + TransformInput::Literal(None) + }), + serde_json::Value::String(text) if !integer => { + Ok(TransformInput::Literal(Some(text.clone()))) + } + serde_json::Value::String(text) if string_positions => { + Ok(TransformInput::IntegerText(text.clone())) + } + serde_json::Value::Number(number) if integer => { + // arbitrary_precision retains the lexical integer even beyond + // i64. SUBSTRING must not convert it before row NULL checks. + let text = number.to_string(); + if text + .strip_prefix('-') + .unwrap_or(&text) + .bytes() + .all(|c| c.is_ascii_digit()) + { + match number.as_i64() { + Some(number) => Ok(TransformInput::Integer(Some(number))), + None if string_positions => Ok(TransformInput::IntegerText(text)), + None => Err(rest_json_err("position exceeds BIGINT range")), + } + } else { + Err(rest_json_err("position must be an integral number")) + } + } + serde_json::Value::Object(_) => { + let name = str_prop(input, "name")?; + if !string_positions { + check_input_type(input, integer)?; + } else if let Some(stored) = input.get("type").filter(|value| !value.is_null()) + { + serde_json::from_value::(stored.clone()) + .map_err(rest_json_err)?; + } + let index = fields + .iter() + .position(|field| field.name() == name) + .ok_or_else(|| rest_json_err(format!("unknown field `{name}`")))?; + let valid = if integer { + matches!( + fields[index].data_type(), + DataType::TinyInt(_) + | DataType::SmallInt(_) + | DataType::Int(_) + | DataType::BigInt(_) + ) + } else { + matches!( + fields[index].data_type(), + DataType::Char(_) | DataType::VarChar(_) + ) + }; + if !valid { + return Err(rest_json_err(format!( + "invalid type for transform input `{name}`" + ))); + } + Ok(TransformInput::Field(index)) + } + _ => Err(rest_json_err(format!("invalid transform input {input}"))), + } + }) + .collect() +} + +/// Java checks the wire FieldRef family before replacing it from the schema. +/// SUBSTRING deliberately allows arbitrary stored families and validates the +/// rebound field at execution instead. +fn check_input_type(input: &serde_json::Value, integer: bool) -> Result<()> { + let stored: DataType = input + .get("type") + .and_then(|value| serde_json::from_value(value.clone()).ok()) + .ok_or_else(|| rest_json_err("transform field without a parseable type"))?; + let valid = if integer { + matches!( + stored, + DataType::TinyInt(_) | DataType::SmallInt(_) | DataType::Int(_) | DataType::BigInt(_) + ) + } else { + matches!(stored, DataType::Char(_) | DataType::VarChar(_)) + }; + if valid { + Ok(()) + } else { + Err(rest_json_err("invalid stored type for transform field")) + } +} + // --------------------------------------------------------------------------- // PredicateBuilder // --------------------------------------------------------------------------- diff --git a/crates/paimon/src/spec/schema.rs b/crates/paimon/src/spec/schema.rs index 2377b4be2..69b13bbe3 100644 --- a/crates/paimon/src/spec/schema.rs +++ b/crates/paimon/src/spec/schema.rs @@ -178,6 +178,36 @@ impl TableSchema { /// fail, so the stored value wins instead, and an absent one leaves the /// default standing. pub fn copy_with_options(&self, mut extra: HashMap) -> Self { + self.retain_immutable_options(&mut extra); + let mut new_schema = self.clone(); + new_schema.options.extend(extra); + new_schema + } + + /// REST returns a Schema without its creation timestamp. Independently + /// loaded copies can have different local timestamps while describing the + /// same schema. Options are compared separately when copying a handle. + pub(crate) fn has_same_definition(&self, other: &Self) -> bool { + self.version == other.version + && self.id == other.id + && self.fields == other.fields + && self.highest_field_id == other.highest_field_id + && self.partition_keys == other.partition_keys + && self.primary_keys == other.primary_keys + && self.comment == other.comment + } + + /// Replace dynamic options without allowing an override of stored immutable + /// options. Unlike merging, absence in this map removes a mutable option. + pub(crate) fn copy_with_replaced_dynamic_options( + &self, + mut options: HashMap, + ) -> Self { + self.retain_immutable_options(&mut options); + self.copy_with_replaced_options(options) + } + + fn retain_immutable_options(&self, extra: &mut HashMap) { if self.core_options().query_auth_enabled() { extra.insert(QUERY_AUTH_ENABLED_OPTION.to_string(), "true".to_string()); } @@ -194,9 +224,6 @@ impl TableSchema { None => extra.remove(key), }; } - let mut new_schema = self.clone(); - new_schema.options.extend(extra); - new_schema } /// Create a copy of this schema with the options replaced entirely, diff --git a/crates/paimon/src/table/goldens/split_v1_query_auth.bin b/crates/paimon/src/table/goldens/split_v1_query_auth.bin new file mode 100644 index 000000000..3978a4398 Binary files /dev/null and b/crates/paimon/src/table/goldens/split_v1_query_auth.bin differ diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index ff1db75b0..bd8f2cd58 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -645,11 +645,33 @@ impl Table { /// and does not perform time travel. Unlike that constructor, it retains the /// FileIO provider, REST environment and identity of this table. The branch /// selects its metadata namespace without loading a different schema. + /// An options-only replacement of the catalog-loaded schema keeps its REST + /// authorization session and the stored immutable options. /// Any cached time-travel resolution is discarded: subsequent scans resolve /// the supplied options without replacing the supplied fields. pub fn copy_with_resolved_schema(&self, schema: TableSchema, branch: &str) -> Result { schema.validate_resolved_structure()?; validate_branch_name(branch)?; + // An options-only copy still reads the exact schema loaded by the + // catalog. Preserve its session through the existing dynamic-option + // path; supplied fields, IDs, keys, comments and branch must all match. + // REST synthesizes creation times independently in different clients. + if self.query_auth_session.is_some() + && !self.time_traveled + && self.travel_snapshot.is_none() + && branch == self.branch + && schema.has_same_definition(&self.schema) + { + return Ok(Self { + schema: self + .schema + .copy_with_replaced_dynamic_options(schema.options().clone()), + time_traveled: false, + travel_snapshot: None, + applied_dynamic_option_keys: HashSet::new(), + ..self.clone() + }); + } let schema_manager = SchemaManager::new(self.file_io.clone(), self.location.clone()); let schema_manager = if branch == DEFAULT_MAIN_BRANCH { schema_manager diff --git a/crates/paimon/src/table/query_auth.rs b/crates/paimon/src/table/query_auth.rs index 3ae7d3497..3de5625a5 100644 --- a/crates/paimon/src/table/query_auth.rs +++ b/crates/paimon/src/table/query_auth.rs @@ -24,12 +24,20 @@ pub(crate) use rules::{filter_batch, mask_batch, mask_inputs, ColumnMask, Rules} /// The server's answer for one user on one table; `session` ties it to the /// handle that asked, as the response names no table or user. -#[derive(Debug, PartialEq)] +#[derive(Debug)] pub(crate) struct QueryAuthGrant { session: u64, /// The columns the request asked about; `None` asked about the whole table. select: Option>, rules: Rules, + /// Original wire rules for Java QueryAuthSplit metadata transport. + response: Option, +} + +impl PartialEq for QueryAuthGrant { + fn eq(&self, other: &Self) -> bool { + self.session == other.session && self.select == other.select && self.rules == other.rules + } } impl QueryAuthGrant { @@ -40,7 +48,9 @@ impl QueryAuthGrant { select: Option>, fields: &[crate::spec::DataField], ) -> crate::Result { - Ok(Self::new(session, select, Rules::parse(response, fields)?)) + let mut grant = Self::new(session, select, Rules::parse(response, fields)?); + grant.response = Some(response.clone()); + Ok(grant) } pub(crate) fn new(session: u64, select: Option>, rules: Rules) -> Self { @@ -48,9 +58,14 @@ impl QueryAuthGrant { session, select, rules, + response: None, } } + pub(crate) fn response(&self) -> Option<&AuthTableQueryResponse> { + self.response.as_ref() + } + /// No row filter and no masking. pub(crate) fn is_unrestricted(&self) -> bool { self.rules.is_empty() @@ -405,13 +420,83 @@ mod tests { } #[tokio::test] - async fn test_a_schema_replaced_copy_loses_its_session() { + async fn test_options_only_schema_replacement_keeps_its_session() { let table = crate::table::rest_query_auth_table().await; assert!(table.query_auth_session().is_some()); + let mut json = serde_json::to_value(table.schema()).unwrap(); + json["timeMillis"] = serde_json::json!(1); let copy = table - .copy_with_resolved_schema(table.schema().clone(), "main") + .copy_with_resolved_schema(serde_json::from_value(json).unwrap(), "main") .unwrap(); - assert!(copy.query_auth_session().is_none()); + assert_eq!(copy.query_auth_session(), table.query_auth_session()); + let table = table.copy_with_options(std::collections::HashMap::from([ + ("scan.snapshot-id".into(), "7".into()), + ("read.batch-size".into(), "99".into()), + ])); + let schema = table + .schema() + .copy_with_replaced_options(std::collections::HashMap::from([ + ("query-auth.enabled".into(), "false".into()), + ("type".into(), "format".into()), + ("data-file.path-directory".into(), "/tmp/other".into()), + ])); + let copy = table.copy_with_resolved_schema(schema, "main").unwrap(); + assert_eq!(copy.query_auth_session(), table.query_auth_session()); + assert!(copy.schema().core_options().query_auth_enabled()); + for key in [ + "scan.snapshot-id", + "read.batch-size", + "type", + "data-file.path-directory", + ] { + assert!(!copy.schema().options().contains_key(key), "{key}"); + } + assert!(!copy.reads_another_schema().unwrap()); + } + + #[tokio::test] + async fn test_changed_schema_or_branch_loses_its_session() { + let table = crate::table::rest_query_auth_table().await; + let mut json = serde_json::to_value(table.schema()).unwrap(); + json["id"] = serde_json::json!(table.schema().id() + 1); + let changed = serde_json::from_value(json).unwrap(); + assert!(table + .copy_with_resolved_schema(changed, "main") + .unwrap() + .query_auth_session() + .is_none()); + assert!(table + .copy_with_resolved_schema(table.schema().clone(), "test") + .unwrap() + .query_auth_session() + .is_none()); + let assembled = crate::table::query_auth_table(); + assert!(assembled + .copy_with_resolved_schema(assembled.schema().clone(), "main") + .unwrap() + .query_auth_session() + .is_none()); + } + + #[tokio::test] + async fn test_resolved_schema_copy_cannot_reauthorize_a_historical_handle() { + let table = crate::table::rest_query_auth_table().await; + let grant = super::QueryAuthGrant::new( + table.query_auth_session().unwrap(), + None, + Default::default(), + ); + let mut travelled = table.clone(); + travelled.time_traveled = true; + let pinned = table.copy_with_pinned_snapshot(None); + for view in [travelled, pinned] { + let schema = view + .schema() + .copy_with_replaced_options(table.schema().options().clone()); + let copy = view.copy_with_resolved_schema(schema, "main").unwrap(); + assert!(copy.query_auth_session().is_none()); + assert!(!grant.matches_table(©)); + } } #[tokio::test] diff --git a/crates/paimon/src/table/query_auth/rules.rs b/crates/paimon/src/table/query_auth/rules.rs index 086500117..d1e71ac01 100644 --- a/crates/paimon/src/table/query_auth/rules.rs +++ b/crates/paimon/src/table/query_auth/rules.rs @@ -391,6 +391,9 @@ pub(crate) fn mask_batch( | Transform::Lower(_) | Transform::Concat(_) | Transform::ConcatWs(_) => string_mask(batch, &mask.transform, &input)?, + _ => { + crate::arrow::evaluate_string_transform(&mask.transform, batch.num_rows(), &input)? + } }; let masked = cast_to(&masked, &target_type)?; for position in targets { @@ -408,6 +411,8 @@ pub(crate) fn mask_batch( fn mask_can_be_null(transform: &Transform, fields: &[DataField]) -> bool { let nullable = |input: &TransformInput| match input { TransformInput::Literal(literal) => literal.is_none(), + TransformInput::Integer(value) => value.is_none(), + TransformInput::IntegerText(_) => false, TransformInput::Field(index) => fields[*index].data_type().is_nullable(), }; match transform { @@ -415,11 +420,9 @@ fn mask_can_be_null(transform: &Transform, fields: &[DataField]) -> bool { Transform::FieldRef(index) | Transform::Cast(index, _) => { fields[*index].data_type().is_nullable() } - Transform::Upper(inputs) | Transform::Lower(inputs) | Transform::Concat(inputs) => { - inputs.iter().any(nullable) - } // Only a null separator makes CONCAT_WS null. Transform::ConcatWs(inputs) => inputs.first().is_some_and(nullable), + _ => transform.inputs().iter().any(nullable), } } @@ -431,10 +434,8 @@ fn mask_output_type(transform: &Transform, fields: &[DataField]) -> Option crate::arrow::paimon_type_to_arrow(to).ok(), - Transform::Upper(_) - | Transform::Lower(_) - | Transform::Concat(_) - | Transform::ConcatWs(_) => Some(arrow_schema::DataType::Utf8), + Transform::Length(_) | Transform::BitLength(_) => Some(arrow_schema::DataType::Int32), + _ => Some(arrow_schema::DataType::Utf8), } } @@ -536,6 +537,9 @@ fn string_mask( .iter() .map(|i| match i { TransformInput::Literal(literal) => Ok(StringInput::Literal(literal.as_deref())), + TransformInput::Integer(_) | TransformInput::IntegerText(_) => { + Err(unsupported("string mask has an integer input")) + } TransformInput::Field(index) => Ok(StringInput::Column( cast_to(&input(*index)?, &arrow_schema::DataType::Utf8)? .as_string::() @@ -651,6 +655,50 @@ mod tests { DataType::VarChar(VarCharType::string_type()) } + #[test] + fn extended_string_masks_match_java_character_semantics() { + let fields = vec![DataField::new(0, "value".into(), string())]; + let source = batch( + &fields, + vec![Arc::new(StringArray::from(vec![ + Some("甲🙂乙"), + Some("abcabc"), + Some(""), + None, + ]))], + ); + for (rule, expected) in [ + ( + r#"{"name":"SUBSTRING","inputs":[{"name":"value","type":"STRING"},2,2]}"#, + vec![Some("🙂乙"), Some("bc"), Some(""), None], + ), + ( + r#"{"name":"TRIM","trimFlag":"BOTH","inputs":[{"name":"value","type":"STRING"},"甲乙ac"]}"#, + vec![Some("🙂"), Some("bcab"), Some(""), None], + ), + ( + r#"{"name":"PAD","direction":"LEFT","inputs":[{"name":"value","type":"STRING"},5,"🙂x"]}"#, + vec![Some("🙂x甲🙂乙"), Some("abcab"), Some("🙂x🙂x🙂"), None], + ), + ( + r#"{"name":"TRANSLATE","inputs":[{"name":"value","type":"STRING"},"甲🙂乙aba","XY"]}"#, + vec![Some("XY"), Some("cc"), Some(""), None], + ), + ( + r#"{"name":"OVERLAY","inputs":[{"name":"value","type":"STRING"},"x",2,1]}"#, + vec![Some("甲x乙"), Some("axcabc"), Some("x"), None], + ), + ] { + let rules = parse(&[], &[("value", rule)], &fields).unwrap(); + let result = mask_batch(&source, &rules.masks, &fields, &fields).unwrap(); + assert_eq!( + result.column(0).as_string::(), + &StringArray::from(expected), + "{rule}" + ); + } + } + fn batch(fields: &[DataField], columns: Vec) -> RecordBatch { let schema: Vec = fields .iter() diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index 11d6bc3fa..593f432c1 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -771,8 +771,8 @@ impl DataSplit { DataSplitBuilder::new() } - /// The Java-compatible frames have no field for the marker, so a reader would - /// rebuild the split without it. Serde keeps it; these two must refuse. + /// A bare DataSplit frame has no authorization marker and must refuse. + /// SplitSerializer can instead transport the Java QueryAuthSplit wrapper. fn ensure_serializable_without_grant(&self) -> crate::Result<()> { if self.query_auth_required { return Err(crate::table::query_auth::unsupported( @@ -940,7 +940,8 @@ impl DataSplit { /// Java `SplitSerializer#serialize` frame: magic + version + type id + body. The cross-language /// entry point. A plain split serializes as `DataSplit` (type 1); a split carrying row ranges as - /// `IndexedSplit` (type 3) wrapping the DataSplit body plus the ranges. Byte-compatible with + /// `IndexedSplit` (type 3) wrapping the DataSplit body plus the ranges. Authorized metadata + /// uses `QueryAuthSplit` (type 5), retaining the complete captured REST response. Byte-compatible with /// `compatibility/split-v1-data` / `split-v1-indexed`. pub fn serialize_split_v1(&self) -> crate::Result> { self.serialize_split_v1_with_scores(None) @@ -949,6 +950,29 @@ impl DataSplit { /// Serialize aligned vector scores using Java's IndexedSplit wire format. /// Scores belong to the selection, in ascending expanded-range order. pub fn serialize_split_v1_with_scores(&self, scores: Option<&[f32]>) -> crate::Result> { + if self.query_auth_required { + let response = self + .query_auth_grant + .as_ref() + .and_then(|grant| grant.response()) + .ok_or_else(|| { + crate::table::query_auth::unsupported( + "the split has no live wire authorization", + ) + })?; + // The nested DataSplit contains metadata only. The outer frame + // carries the rules and cannot decode as an ordinary DataSplit. + let mut metadata = self.clone(); + metadata.query_auth_required = false; + metadata.query_auth_grant = None; + let mut out = Vec::new(); + out.extend_from_slice(&SPLIT_SER_MAGIC.to_be_bytes()); + out.extend_from_slice(&SPLIT_SER_VERSION.to_be_bytes()); + out.extend_from_slice(&SPLIT_SER_TYPE_QUERY_AUTH_SPLIT.to_be_bytes()); + out.extend_from_slice(&metadata.serialize_split_v1_with_scores(scores)?); + write_auth_response(&mut out, response)?; + return Ok(out); + } self.ensure_serializable_without_grant()?; if let Some(scores) = scores { let count = self.row_ranges.as_ref().and_then(|ranges| { @@ -999,42 +1023,68 @@ impl DataSplit { } /// Reverse of [`DataSplit::serialize_split_v1`]: parse a Java `SplitSerializer` frame into a - /// `DataSplit` (type 1) or an `IndexedSplit` (type 3, which carries `row_ranges`). Other split - /// type ids, and IndexedSplit vector scores, are unsupported. Consumes the entire buffer. + /// `DataSplit` (type 1), an `IndexedSplit` (type 3, which carries `row_ranges`), or a + /// `QueryAuthSplit` (type 5, retaining the authorization-required marker without minting + /// a grant). Other type ids, and IndexedSplit vector scores, are unsupported. + /// Consumes the entire buffer. pub fn deserialize_split_v1(data: &[u8]) -> crate::Result { let mut cur = data; - let magic = read_i64(&mut cur)?; + let split = Self::read_split_v1(&mut cur, 0)?; + if !cur.is_empty() { + return Err(crate::Error::DataInvalid { + message: format!("{} trailing bytes after SPLIT_V1 frame", cur.len()), + source: None, + }); + } + Ok(split) + } + + fn read_split_v1(cur: &mut &[u8], depth: usize) -> crate::Result { + if depth >= 32 { + return Err(crate::Error::DataInvalid { + message: "QueryAuthSplit nesting exceeds 32".into(), + source: None, + }); + } + let magic = read_i64(cur)?; if magic != SPLIT_SER_MAGIC { return Err(crate::Error::DataInvalid { message: format!("invalid SPLIT_V1 magic: {magic:#018x}"), source: None, }); } - let version = read_i32(&mut cur)?; + let version = read_i32(cur)?; if version != SPLIT_SER_VERSION { return Err(crate::Error::Unsupported { message: format!("SplitSerializer version {version} not supported"), }); } - let type_id = read_i32(&mut cur)?; + let type_id = read_i32(cur)?; let split = match type_id { - SPLIT_SER_TYPE_DATA_SPLIT => Self::read_body(&mut cur)?, + SPLIT_SER_TYPE_QUERY_AUTH_SPLIT => { + let split = Self::read_split_v1(cur, depth + 1)?; + read_auth_response(cur)?; + // Rules transported in a frame are not a grant minted by this + // catalog session. Reading it still requires the original split. + split.planned(None) + } + SPLIT_SER_TYPE_DATA_SPLIT => Self::read_body(cur)?, SPLIT_SER_TYPE_INDEXED_SPLIT => { - let im = read_i64(&mut cur)?; + let im = read_i64(cur)?; if im != INDEXED_SPLIT_MAGIC { return Err(crate::Error::DataInvalid { message: format!("invalid IndexedSplit magic: {im:#018x}"), source: None, }); } - let iv = read_i32(&mut cur)?; + let iv = read_i32(cur)?; if iv != INDEXED_SPLIT_VERSION { return Err(crate::Error::Unsupported { message: format!("IndexedSplit version {iv} not supported"), }); } - let body = Self::read_body(&mut cur)?; - let ranges_n = read_i32(&mut cur)?; + let body = Self::read_body(cur)?; + let ranges_n = read_i32(cur)?; if ranges_n < 0 { return Err(crate::Error::DataInvalid { message: "negative row_ranges count".into(), @@ -1045,8 +1095,8 @@ impl DataSplit { // 16 bytes (two i64), so a crafted huge count cannot over-reserve. let mut ranges = Vec::with_capacity((ranges_n as usize).min(cur.len() / 16)); for _ in 0..ranges_n { - let from = read_i64(&mut cur)?; - let to = read_i64(&mut cur)?; + let from = read_i64(cur)?; + let to = read_i64(cur)?; if from > to { return Err(crate::Error::DataInvalid { message: format!("row range from {from} > to {to}"), @@ -1055,7 +1105,7 @@ impl DataSplit { } ranges.push(RowRange::new(from, to)); } - let scores_flag = read_u8(&mut cur)?; + let scores_flag = read_u8(cur)?; if scores_flag != 0 { return Err(crate::Error::Unsupported { message: "IndexedSplit scores are not supported".into(), @@ -1071,12 +1121,6 @@ impl DataSplit { }); } }; - if !cur.is_empty() { - return Err(crate::Error::DataInvalid { - message: format!("{} trailing bytes after SPLIT_V1 frame", cur.len()), - source: None, - }); - } Ok(split) } } @@ -1094,6 +1138,87 @@ const SPLIT_SER_MAGIC: i64 = 0x53504C49545F5631; // "SPLIT_V1" const SPLIT_SER_VERSION: i32 = 1; const SPLIT_SER_TYPE_DATA_SPLIT: i32 = 1; const SPLIT_SER_TYPE_INDEXED_SPLIT: i32 = 3; +const SPLIT_SER_TYPE_QUERY_AUTH_SPLIT: i32 = 5; + +fn write_auth_response( + out: &mut Vec, + response: &crate::api::AuthTableQueryResponse, +) -> crate::Result<()> { + fn string(out: &mut Vec, value: &str) -> crate::Result<()> { + let length = i32::try_from(value.len()).map_err(|e| crate::Error::DataInvalid { + message: e.to_string(), + source: None, + })?; + out.extend_from_slice(&length.to_be_bytes()); + out.extend_from_slice(value.as_bytes()); + Ok(()) + } + out.push(1); + match &response.filter { + None => out.push(0), + Some(filters) => { + out.push(1); + out.extend_from_slice(&(filters.len() as i32).to_be_bytes()); + for filter in filters { + string(out, filter)?; + } + } + } + match &response.column_masking { + None => out.push(0), + Some(masks) => { + out.push(1); + out.extend_from_slice(&(masks.len() as i32).to_be_bytes()); + let mut masks: Vec<_> = masks.iter().collect(); + // Java's TreeMap compares UTF-16 units, including supplementary keys. + masks.sort_by(|(a, _), (b, _)| a.encode_utf16().cmp(b.encode_utf16())); + for (name, mask) in masks { + string(out, name)?; + string(out, mask)?; + } + } + } + Ok(()) +} + +/// Consume wire rules without creating a trusted grant from untrusted bytes. +fn read_auth_response(cur: &mut &[u8]) -> crate::Result<()> { + fn string(cur: &mut &[u8]) -> crate::Result<()> { + let length = read_i32(cur)?; + if length < 0 { + return Err(crate::Error::DataInvalid { + message: "NULL authorization rule in QueryAuthSplit".into(), + source: None, + }); + } + let bytes = take(cur, length as usize)?; + std::str::from_utf8(bytes).map_err(|e| crate::Error::DataInvalid { + message: e.to_string(), + source: None, + })?; + Ok(()) + } + if read_u8(cur)? == 0 { + return Ok(()); + } + for multiplier in [1, 2] { + if read_u8(cur)? != 0 { + let count = read_i32(cur)?; + if count < 0 || count as usize > cur.len() / (4 * multiplier) { + return Err(crate::Error::DataInvalid { + message: "invalid QueryAuthSplit rule count".into(), + source: None, + }); + } + for _ in 0..count { + for _ in 0..multiplier { + string(cur)?; + } + } + } + } + Ok(()) +} /// Java `IndexedSplit#MAGIC` / `VERSION`. const INDEXED_SPLIT_MAGIC: i64 = -938472394838495695; @@ -2194,6 +2319,105 @@ mod tests { } } + #[test] + fn query_auth_frame_matches_java_and_does_not_restore_a_live_grant() { + let response = crate::api::AuthTableQueryResponse { + filter: Some(vec!["row甲🙂\0".into()]), + column_masking: Some(std::collections::HashMap::from([ + ("\u{e000}".into(), "bmp".into()), + ("𐀀".into(), "supplementary".into()), + ])), + }; + // Captured from Java QueryAuthSplit.serialize: UTF-8 with int length, + // NUL as one byte, and keys sorted by UTF-16 code units. + let expected = hex::decode("0101000000010000000b726f77e794b2f09f998200010000000200000004f09080800000000d737570706c656d656e7461727900000003ee808000000003626d70").unwrap(); + let mut tail = Vec::new(); + write_auth_response(&mut tail, &response).unwrap(); + assert_eq!(tail, expected); + let grant = + QueryAuthGrant::parse(&crate::api::AuthTableQueryResponse::default(), 1, None, &[]) + .unwrap(); + // Use a valid parsed policy for the live frame. The arbitrary strings + // above test transport independently of grant validation. + let mut frame_tail = Vec::new(); + write_auth_response(&mut frame_tail, grant.response().unwrap()).unwrap(); + let ordinary = v1_data_split_builder().build().unwrap(); + let marked = ordinary.clone().planned(Some(Arc::new(grant))); + assert!(marked.serialize().is_err()); + let frame = marked.serialize_split_v1().unwrap(); + let mut expected_frame = Vec::new(); + expected_frame.extend_from_slice(&SPLIT_SER_MAGIC.to_be_bytes()); + expected_frame.extend_from_slice(&SPLIT_SER_VERSION.to_be_bytes()); + expected_frame.extend_from_slice(&SPLIT_SER_TYPE_QUERY_AUTH_SPLIT.to_be_bytes()); + expected_frame.extend_from_slice(&ordinary.serialize_split_v1().unwrap()); + expected_frame.extend_from_slice(&frame_tail); + assert_eq!(frame, expected_frame); + let copy = DataSplit::deserialize_split_v1(&frame).unwrap(); + assert!(copy.query_auth_required()); + assert!(copy.query_auth_grant().is_none()); + assert_eq!(copy.data_files(), ordinary.data_files()); + assert!(copy.serialize().is_err()); + assert!(copy.serialize_split_v1().is_err()); + } + + #[test] + fn deserialize_java_query_auth_golden_keeps_the_required_marker() { + // Java SplitSerializer(QueryAuthSplit(DataSplit v9, policy)) with BMP, + // supplementary and NUL text in its policy strings. + let split = + DataSplit::deserialize_split_v1(include_bytes!("goldens/split_v1_query_auth.bin")) + .unwrap(); + assert_eq!(split.bucket(), 3); + assert_eq!(split.data_files().len(), 2); + assert!(split.query_auth_required()); + assert!(split.query_auth_grant().is_none()); + assert!(split.serialize_split_v1().is_err()); + } + + #[test] + fn query_auth_frame_bounds_policy_lengths_and_nesting() { + let ordinary = v1_data_split_builder() + .build() + .unwrap() + .serialize_split_v1() + .unwrap(); + let mut prefix = Vec::new(); + prefix.extend_from_slice(&SPLIT_SER_MAGIC.to_be_bytes()); + prefix.extend_from_slice(&SPLIT_SER_VERSION.to_be_bytes()); + prefix.extend_from_slice(&SPLIT_SER_TYPE_QUERY_AUTH_SPLIT.to_be_bytes()); + let mut frame = prefix.clone(); + frame.extend_from_slice(&ordinary); + let policy_start = frame.len(); + let response = crate::api::AuthTableQueryResponse { + filter: Some(vec!["甲".repeat(30000)]), + column_masking: Some(Default::default()), + }; + write_auth_response(&mut frame, &response).unwrap(); + assert!(DataSplit::deserialize_split_v1(&frame) + .unwrap() + .query_auth_required()); + for tail in [ + vec![1, 0], + vec![1, 1, 0x7f, 0xff, 0xff, 0xff], + vec![1, 1, 0xff, 0xff, 0xff, 0xff], + vec![1, 1, 0, 0, 0, 1, 0xff, 0xff, 0xff, 0xff, 0], + ] { + frame.truncate(policy_start); + frame.extend_from_slice(&tail); + assert!(DataSplit::deserialize_split_v1(&frame).is_err()); + } + let mut nested = prefix.repeat(33); + nested.extend_from_slice(&ordinary); + nested.extend_from_slice(&[0; 33]); + assert!(DataSplit::deserialize_split_v1(&nested).is_err()); + frame.truncate(policy_start); + frame.extend_from_slice(&[0, 0]); + assert!( + DataSplit::deserialize_split_v1(&frame).is_err(), + "trailing bytes" + ); + } + #[test] fn deserialize_split_v1_rejects_huge_ranges_count_without_aborting() { let split = DataSplitBuilder::new()