This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 29600d99 perf(arrow): read Variant metadata as a dictionary for field
extraction (#1057)
29600d99 is described below
commit 29600d9924a3c78e9aa0665206907c04dcf2e29e
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Wed Oct 7 19:39:03 2026 +0800
perf(arrow): read Variant metadata as a dictionary for field extraction
(#1057)
---
crates/paimon/src/arrow/format/parquet.rs | 165 ++++++++++++++++++++++++++-
crates/paimon/src/arrow/shredding/variant.rs | 42 ++++++-
crates/paimon/src/arrow/variant.rs | 125 +++++++++++++++++++-
crates/paimon/src/variant/numeric.rs | 100 ++++++++++++----
4 files changed, 399 insertions(+), 33 deletions(-)
diff --git a/crates/paimon/src/arrow/format/parquet.rs
b/crates/paimon/src/arrow/format/parquet.rs
index 5f7a9eb8..605699d1 100644
--- a/crates/paimon/src/arrow/format/parquet.rs
+++ b/crates/paimon/src/arrow/format/parquet.rs
@@ -656,6 +656,63 @@ fn cast_record_batch_to_schema(
})
}
+/// Options for per-row-group readers; they keep the main reader's Variant
metadata schema.
+fn rebuilt_reader_options(variant_schema: &Option<arrow_schema::SchemaRef>) ->
ArrowReaderOptions {
+ match variant_schema {
+ Some(schema) =>
ArrowReaderOptions::new().with_schema(Arc::clone(schema)),
+ None => ArrowReaderOptions::new(),
+ }
+}
+
+/// File schema with `metadata: Dictionary(Int32, Binary)` for plain Variant
columns read for extraction.
+fn dictionary_variant_metadata_schema(
+ schema: &arrow_schema::SchemaRef,
+ read_fields: &[DataField],
+) -> Option<arrow_schema::SchemaRef> {
+ let mut changed = false;
+ let fields = schema
+ .fields()
+ .iter()
+ .map(|field| {
+ let extraction = read_fields.iter().any(|read| {
+ read.name() == field.name()
+ &&
crate::spec::is_variant_extraction_row_type(read.data_type())
+ });
+ match field.data_type() {
+ arrow_schema::DataType::Struct(children)
+ if extraction &&
crate::arrow::is_variant_arrow_fields(children) =>
+ {
+ changed = true;
+ let mut children = children
+ .iter()
+ .map(|child| child.as_ref().clone())
+ .collect::<Vec<_>>();
+ children[1] =
+ children[1]
+ .clone()
+ .with_data_type(arrow_schema::DataType::Dictionary(
+ Box::new(arrow_schema::DataType::Int32),
+ Box::new(arrow_schema::DataType::Binary),
+ ));
+ Arc::new(
+ field
+ .as_ref()
+ .clone()
+
.with_data_type(arrow_schema::DataType::Struct(children.into())),
+ )
+ }
+ _ => Arc::clone(field),
+ }
+ })
+ .collect::<Vec<_>>();
+ changed.then(|| {
+ Arc::new(arrow_schema::Schema::new_with_metadata(
+ fields,
+ schema.metadata().clone(),
+ ))
+ })
+}
+
#[async_trait]
impl FormatFileReader for ParquetFormatReader {
async fn read_batch_stream(
@@ -689,9 +746,24 @@ impl FormatFileReader for ParquetFormatReader {
arrow_options =
arrow_options.with_offset_index_policy(PageIndexPolicy::Optional);
}
}
- let mut batch_stream_builder =
-
ParquetRecordBatchStreamBuilder::new_with_options(arrow_file_reader,
arrow_options)
- .await?;
+ let mut batch_stream_builder =
ParquetRecordBatchStreamBuilder::new_with_options(
+ arrow_file_reader,
+ arrow_options.clone(),
+ )
+ .await?;
+ // Variant columns read only for extraction keep their metadata
dictionary-encoded.
+ let variant_schema =
+ dictionary_variant_metadata_schema(batch_stream_builder.schema(),
read_fields);
+ if let Some(schema) = variant_schema.clone() {
+ let reader_metadata = ArrowReaderMetadata::try_new(
+ Arc::clone(batch_stream_builder.metadata()),
+ arrow_options.with_schema(schema),
+ )?;
+ let reader = ArrowFileReader::new(file_size,
Arc::clone(&shared_reader))
+ .with_metadata_cache_enabled(self.metadata_cache_enabled);
+ batch_stream_builder =
+ ParquetRecordBatchStreamBuilder::new_with_metadata(reader,
reader_metadata);
+ }
let parquet_schema = batch_stream_builder.parquet_schema().clone();
@@ -906,7 +978,7 @@ impl FormatFileReader for ParquetFormatReader {
let row_group_count = selected_row_groups.len();
let reader_metadata = ArrowReaderMetadata::try_new(
batch_stream_builder.metadata().clone(),
- ArrowReaderOptions::new(),
+ rebuilt_reader_options(&variant_schema),
)?;
let read_budget = Arc::clone(read_budget.expect("checked above"));
let (row_group_tx, mut row_group_rx) =
mpsc::channel(row_group_parallelism);
@@ -1004,7 +1076,7 @@ impl FormatFileReader for ParquetFormatReader {
let selected = selected_row_groups.expect("resource-aware reads
need a selection plan");
let metadata = ArrowReaderMetadata::try_new(
batch_stream_builder.metadata().clone(),
- ArrowReaderOptions::new(),
+ rebuilt_reader_options(&variant_schema),
)?;
let residual = (!all_enforced).then(|| FilePredicates {
predicates: preds.to_vec(),
@@ -7406,4 +7478,87 @@ mod tests {
};
assert!(message.contains("file.compression"), "message: {message}");
}
+
+ #[tokio::test]
+ async fn variant_extraction_reads_metadata_as_dictionary() {
+ let rows = [r#"{"x":1,"y":2.5}"#, r#"{"y":-1,"x":4}"#, r#"{"x":6}"#];
+ let variants = rows
+ .iter()
+ .map(|json| GenericVariant::parse_json(json).unwrap())
+ .collect::<Vec<_>>();
+ let ArrowDataType::Struct(variant_fields) = variant_arrow_type() else {
+ unreachable!()
+ };
+ let column = StructArray::new(
+ variant_fields.clone(),
+ vec![
+ Arc::new(BinaryArray::from_iter_values(
+ variants.iter().map(|v| v.value()),
+ )),
+ Arc::new(BinaryArray::from_iter_values(
+ variants.iter().map(|v| v.metadata()),
+ )),
+ ],
+ None,
+ );
+ let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+ "v",
+ ArrowDataType::Struct(variant_fields),
+ true,
+ )]));
+ let batch =
+ RecordBatch::try_new(Arc::clone(&schema),
vec![Arc::new(column.clone())]).unwrap();
+ let mut buf = Vec::new();
+ let mut writer = AsyncArrowWriter::try_new(&mut buf, schema,
None).unwrap();
+ writer.write(&batch).await.unwrap();
+ writer.close().await.unwrap();
+
+ let float = || DataType::Float(crate::spec::FloatType::new());
+ let row_type = crate::spec::variant_extraction_row(
+ true,
+ vec![
+ (float(), "$.x".to_string(), false, "UTC".to_string()),
+ (float(), "$.y".to_string(), false, "UTC".to_string()),
+ ],
+ )
+ .unwrap();
+ let fields = vec![DataField::new(
+ 0,
+ "v".to_string(),
+ DataType::Row(row_type.clone()),
+ )];
+ let data = Bytes::from(buf);
+ let batches = ParquetFormatReader::default()
+ .read_batch_stream(
+ Box::new(TrackingFileRead::new(data.clone())),
+ data.len() as u64,
+ &fields,
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ let read = batches[0]
+ .column_by_name("v")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ assert!(matches!(
+ read.column(1).data_type(),
+ ArrowDataType::Dictionary(key, value)
+ if key.as_ref() == &ArrowDataType::Int32 && value.as_ref() ==
&ArrowDataType::Binary
+ ));
+ let expected =
+
crate::arrow::shredding::variant::assemble_variant_extraction_array(&column,
&row_type)
+ .unwrap();
+ let actual =
+
crate::arrow::shredding::variant::assemble_variant_extraction_array(read,
&row_type)
+ .unwrap();
+ assert_eq!(actual.as_ref(), expected.as_ref());
+ }
}
diff --git a/crates/paimon/src/arrow/shredding/variant.rs
b/crates/paimon/src/arrow/shredding/variant.rs
index 83e90ca4..4244d998 100644
--- a/crates/paimon/src/arrow/shredding/variant.rs
+++ b/crates/paimon/src/arrow/shredding/variant.rs
@@ -904,6 +904,8 @@ pub(crate) fn assemble_variant_extraction_array(
if let Some(projected) = assemble_plain_variant_projection(input, fields,
&metadata)? {
return Ok(projected);
}
+ let materialized = materialize_variant_metadata(input)?;
+ let input = &materialized;
let shredded_type = if is_shredded_variant_array(input) {
let DataType::Row(row) = arrow_to_paimon_type(input.data_type(),
true)? else {
@@ -1174,7 +1176,45 @@ fn is_plain_variant_array(array: &dyn Array) -> bool {
let ArrowDataType::Struct(fields) = array.data_type() else {
return false;
};
- is_variant_arrow_fields(fields)
+ is_variant_arrow_fields(fields) || has_dictionary_variant_metadata(fields)
+}
+
+/// Plain Variant whose metadata child the Parquet reader kept
dictionary-encoded.
+fn has_dictionary_variant_metadata(fields: &Fields) -> bool {
+ fields.len() == 2
+ && fields[0].name() == "value"
+ && fields[0].data_type() == &ArrowDataType::Binary
+ && fields[1].name() == "metadata"
+ && matches!(fields[1].data_type(), ArrowDataType::Dictionary(key,
value)
+ if key.as_ref() == &ArrowDataType::Int32 && value.as_ref() ==
&ArrowDataType::Binary)
+}
+
+/// Turns dictionary-encoded Variant metadata back into Binary for the generic
row path.
+fn materialize_variant_metadata(input: &StructArray) -> Result<StructArray> {
+ if !has_dictionary_variant_metadata(input.fields()) {
+ return Ok(input.clone());
+ }
+ let metadata = arrow_cast::cast(input.column(1),
&ArrowDataType::Binary).map_err(|e| {
+ Error::DataInvalid {
+ message: format!("Failed to decode dictionary Variant metadata:
{e}"),
+ source: Some(Box::new(e)),
+ }
+ })?;
+ let mut fields = input
+ .fields()
+ .iter()
+ .map(|f| f.as_ref().clone())
+ .collect::<Vec<_>>();
+ fields[1] = fields[1].clone().with_data_type(ArrowDataType::Binary);
+ StructArray::try_new(
+ fields.into(),
+ vec![input.column(0).clone(), metadata],
+ input.nulls().cloned(),
+ )
+ .map_err(|e| Error::DataInvalid {
+ message: format!("Failed to rebuild Variant struct: {e}"),
+ source: Some(Box::new(e)),
+ })
}
fn shredded_rows_to_struct_array(
diff --git a/crates/paimon/src/arrow/variant.rs
b/crates/paimon/src/arrow/variant.rs
index 07619aa0..3ac4d4f2 100644
--- a/crates/paimon/src/arrow/variant.rs
+++ b/crates/paimon/src/arrow/variant.rs
@@ -18,7 +18,8 @@
use std::sync::Arc;
use arrow_array::builder::{FixedSizeListBuilder, Float32Builder};
-use arrow_array::{Array, BinaryArray, FixedSizeListArray, StructArray};
+use arrow_array::types::Int32Type;
+use arrow_array::{Array, BinaryArray, DictionaryArray, FixedSizeListArray,
StructArray};
use arrow_schema::{DataType, Field};
use crate::variant::VariantFloat32Projection;
@@ -69,14 +70,14 @@ pub(crate) fn extract_numeric_fields(
})?;
let value_column = variant_binary_child(column, 0, "value")?;
- let metadata_column = variant_binary_child(column, 1, "metadata")?;
+ let metadata_column = VariantMetadataColumn::new(column)?;
let mut builder = FixedSizeListBuilder::with_capacity(
Float32Builder::with_capacity(value_capacity),
width,
column.len(),
)
.with_field(Arc::new(Field::new("item", DataType::Float32, true)));
- let mut projection_metadata = None;
+ let mut projection_metadata: Option<MetadataIdentity<'_>> = None;
let mut projection = None;
let mut offsets = Vec::new();
let mut extracted = vec![None; fields.len()];
@@ -92,10 +93,10 @@ pub(crate) fn extract_numeric_fields(
));
}
- let metadata = metadata_column.value(row);
- if projection_metadata != Some(metadata) {
+ let (metadata, identity) = metadata_column.value(row);
+ if !projection_metadata.is_some_and(|previous|
previous.same_as(&identity)) {
projection = Some(VariantFloat32Projection::new(metadata,
fields)?);
- projection_metadata = Some(metadata);
+ projection_metadata = Some(identity);
}
match mode {
NumericFieldMode::Cast { fail_on_error } => {
@@ -123,6 +124,90 @@ pub(crate) fn extract_numeric_fields(
Ok(builder.finish())
}
+/// Variant metadata read either as plain binary or, without per-row copies,
dictionary-encoded.
+enum VariantMetadataColumn<'a> {
+ Plain(&'a BinaryArray),
+ Dictionary {
+ keys: &'a DictionaryArray<Int32Type>,
+ values: &'a BinaryArray,
+ },
+}
+
+/// Identifies the metadata a projection was built for; dictionary keys avoid
comparing bytes.
+#[derive(Clone, Copy)]
+enum MetadataIdentity<'a> {
+ Bytes(&'a [u8]),
+ Key(i32),
+}
+
+impl MetadataIdentity<'_> {
+ fn same_as(&self, other: &Self) -> bool {
+ match (self, other) {
+ (Self::Key(a), Self::Key(b)) => a == b,
+ (Self::Bytes(a), Self::Bytes(b)) => {
+ (a.as_ptr() == b.as_ptr() && a.len() == b.len()) || a == b
+ }
+ _ => false,
+ }
+ }
+}
+
+impl<'a> VariantMetadataColumn<'a> {
+ fn new(column: &'a StructArray) -> Result<Self> {
+ let Some(field) = column.fields().get(1) else {
+ return data_invalid("Expected Variant struct fields value and
metadata");
+ };
+ if column.num_columns() != 2 || field.name() != "metadata" {
+ return data_invalid("Expected Variant struct fields value and
metadata");
+ }
+ match field.data_type() {
+ DataType::Binary => Ok(Self::Plain(variant_binary_child(column, 1,
"metadata")?)),
+ DataType::Dictionary(key, value)
+ if key.as_ref() == &DataType::Int32 && value.as_ref() ==
&DataType::Binary =>
+ {
+ let keys = column
+ .column(1)
+ .as_any()
+ .downcast_ref::<DictionaryArray<Int32Type>>()
+ .ok_or_else(|| Error::DataInvalid {
+ message: "Variant metadata dictionary must use Int32
keys".to_string(),
+ source: None,
+ })?;
+ let values = keys
+ .values()
+ .as_any()
+ .downcast_ref::<BinaryArray>()
+ .ok_or_else(|| Error::DataInvalid {
+ message: "Variant metadata dictionary values must be
Binary".to_string(),
+ source: None,
+ })?;
+ Ok(Self::Dictionary { keys, values })
+ }
+ _ => data_invalid("Expected Variant struct fields value and
metadata"),
+ }
+ }
+
+ fn is_null(&self, row: usize) -> bool {
+ match self {
+ Self::Plain(array) => array.is_null(row),
+ Self::Dictionary { keys, .. } => keys.is_null(row),
+ }
+ }
+
+ fn value(&self, row: usize) -> (&'a [u8], MetadataIdentity<'a>) {
+ match self {
+ Self::Plain(array) => {
+ let bytes = array.value(row);
+ (bytes, MetadataIdentity::Bytes(bytes))
+ }
+ Self::Dictionary { keys, values } => {
+ let key = keys.keys().value(row);
+ (values.value(key as usize), MetadataIdentity::Key(key))
+ }
+ }
+ }
+}
+
fn variant_binary_child<'a>(
column: &'a StructArray,
index: usize,
@@ -198,6 +283,34 @@ mod tests {
)
}
+ #[test]
+ fn dictionary_metadata_matches_plain_metadata() {
+ let plain = variant_column(&[
+ Some(r#"{"x":1,"y":2.5}"#),
+ None,
+ Some(r#"{"y":-1,"x":4}"#),
+ Some(r#"{"z":true}"#),
+ ]);
+ let dictionary_type =
+ DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Binary));
+ let metadata = arrow_cast::cast(plain.column(1),
&dictionary_type).unwrap();
+ let mut fields = plain
+ .fields()
+ .iter()
+ .map(|f| f.as_ref().clone())
+ .collect::<Vec<_>>();
+ fields[1] = fields[1].clone().with_data_type(dictionary_type);
+ let dictionary = StructArray::new(
+ fields.into(),
+ vec![plain.column(0).clone(), metadata],
+ plain.nulls().cloned(),
+ );
+ let names = vec!["x".to_string(), "y".to_string()];
+ let expected = variant_get_numeric_fields(&plain, &names).unwrap();
+ let actual = variant_get_numeric_fields(&dictionary, &names).unwrap();
+ assert_eq!(actual, expected);
+ }
+
#[test]
fn extracts_numeric_fields_in_requested_order() {
let input = variant_column(&[
diff --git a/crates/paimon/src/variant/numeric.rs
b/crates/paimon/src/variant/numeric.rs
index 8c797481..9aa66fba 100644
--- a/crates/paimon/src/variant/numeric.rs
+++ b/crates/paimon/src/variant/numeric.rs
@@ -21,6 +21,9 @@ pub(crate) struct VariantFloat32Projection {
fields: HashMap<usize, ProjectedField>,
output_width: usize,
last_layout: Option<ObjectLayoutProjection>,
+ // Offsets of non-monotonic objects in ascending order and each offset's
rank; reused while valid.
+ order: Vec<usize>,
+ rank: Vec<usize>,
}
struct ObjectLayoutProjection {
@@ -68,6 +71,8 @@ impl VariantFloat32Projection {
fields: projected,
output_width: fields.len(),
last_layout: None,
+ order: Vec::new(),
+ rank: Vec::new(),
})
}
@@ -141,33 +146,33 @@ impl VariantFloat32Projection {
let monotonic =
offsets.first() == Some(&0) && offsets.windows(2).all(|pair|
pair[0] < pair[1]);
- let sorted_offsets = if monotonic {
- None
- } else {
- let mut sorted = offsets.clone();
- sorted.sort_unstable();
- if sorted.first() != Some(&0)
- || sorted.last() != Some(&data_size)
- || sorted.windows(2).any(|pair| pair[0] == pair[1])
- {
+ if !monotonic && !self.order_fits(offsets, data_size) {
+ self.order.clear();
+ self.order.extend(0..offsets.len());
+ self.order.sort_unstable_by_key(|index| offsets[*index]);
+ if !self.order_fits(offsets, data_size) {
return data_invalid("Malformed Variant object offsets");
}
- Some(sorted)
- };
+ self.rank.clear();
+ self.rank.resize(offsets.len(), 0);
+ for (position, index) in self.order.iter().enumerate() {
+ self.rank[*index] = position;
+ }
+ }
self.prepare_layout(value, metadata, &layout)?;
- for (index, id) in &self.last_layout.as_ref().unwrap().selected {
+ let selected = &self.last_layout.as_ref().unwrap().selected;
+ for (index, id) in selected {
let projected = self.fields.get(id).unwrap();
let start = offsets[*index];
- let end = match &sorted_offsets {
- None => offsets[*index + 1],
- Some(sorted) => {
- let next = sorted.partition_point(|offset| *offset <=
start);
- *sorted.get(next).ok_or_else(|| Error::DataInvalid {
- message: "Malformed Variant object
offsets".to_string(),
- source: None,
- })?
- }
+ let end = if monotonic {
+ offsets[*index + 1]
+ } else {
+ let next = self.rank[*index] + 1;
+ offsets[*self.order.get(next).ok_or_else(|| Error::DataInvalid
{
+ message: "Malformed Variant object offsets".to_string(),
+ source: None,
+ })?]
};
let child_pos =
layout
@@ -210,6 +215,17 @@ impl VariantFloat32Projection {
Ok(())
}
+ /// True if the cached order lists `offsets` strictly ascending from 0 to
`data_size`.
+ fn order_fits(&self, offsets: &[usize], data_size: usize) -> bool {
+ self.order.len() == offsets.len()
+ && self.order.first().map(|index| offsets[*index]) == Some(0)
+ && self.order.last().map(|index| offsets[*index]) ==
Some(data_size)
+ && self
+ .order
+ .windows(2)
+ .all(|pair| offsets[pair[0]] < offsets[pair[1]])
+ }
+
fn prepare_layout(
&mut self,
value: &[u8],
@@ -306,4 +322,46 @@ mod tests {
.unwrap();
assert_eq!(output, vec![Some(2.0), Some(1.0), Some(2.0)]);
}
+
+ #[test]
+ fn cached_offset_order_still_rejects_malformed_rows() {
+ let variant = GenericVariant::parse_json(r#"{"a":1,"b":2}"#).unwrap();
+ let mut value = variant.value().to_vec();
+ let layout = object_layout(&value, 0).unwrap();
+ let first = read_unsigned(&value, layout.offset_start,
layout.offset_size).unwrap();
+ let second = read_unsigned(
+ &value,
+ layout.offset_start + layout.offset_size,
+ layout.offset_size,
+ )
+ .unwrap();
+ write_le_at(&mut value, layout.offset_start, second,
layout.offset_size);
+ write_le_at(
+ &mut value,
+ layout.offset_start + layout.offset_size,
+ first,
+ layout.offset_size,
+ );
+ let fields = vec!["a".to_string(), "b".to_string()];
+ let mut projection = VariantFloat32Projection::new(variant.metadata(),
&fields).unwrap();
+ let mut offsets = Vec::new();
+ let mut output = vec![None; fields.len()];
+ for _ in 0..2 {
+ projection
+ .extract_float32(&value, variant.metadata(), &mut offsets,
&mut output)
+ .unwrap();
+ assert_eq!(output, vec![Some(2.0), Some(1.0)]);
+ }
+ let mut duplicate = value.clone();
+ write_le_at(&mut duplicate, layout.offset_start, 0,
layout.offset_size);
+ write_le_at(
+ &mut duplicate,
+ layout.offset_start + layout.offset_size,
+ 0,
+ layout.offset_size,
+ );
+ assert!(projection
+ .extract_float32(&duplicate, variant.metadata(), &mut offsets,
&mut output)
+ .is_err());
+ }
}
