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 30c171d3 feat: expand Python native reads with nested projection and
row kinds (#887)
30c171d3 is described below
commit 30c171d3d5c86837ecfc05f9d32d039162407c39
Author: Jingsong Lee <[email protected]>
AuthorDate: Sun Sep 20 19:00:11 2026 +0800
feat: expand Python native reads with nested projection and row kinds (#887)
---
bindings/python/src/read.rs | 267 +++++++++++++++++++++++++++-
crates/paimon/src/arrow/format/parquet.rs | 251 ++++++++++++++++++++++++--
crates/paimon/src/arrow/shredding/map.rs | 166 ++++++++++++++---
crates/paimon/src/table/table_read.rs | 120 ++++++++++++-
crates/paimon/tests/audit_log_table_test.rs | 31 ++++
5 files changed, 795 insertions(+), 40 deletions(-)
diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs
index df706c0b..f73a8ded 100644
--- a/bindings/python/src/read.rs
+++ b/bindings/python/src/read.rs
@@ -22,7 +22,7 @@ use std::sync::{Arc, Mutex};
use arrow::pyarrow::ToPyArrow;
use arrow::record_batch::RecordBatch;
use futures::TryStreamExt;
-use paimon::spec::Predicate;
+use paimon::spec::{DataField, DataType, Predicate, RowType};
use paimon::table::{ArrowRecordBatchStream, DataSplit, IncrementalScanMode,
RowRange, Table};
use paimon_datafusion::runtime::runtime;
use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError};
@@ -33,6 +33,9 @@ use tokio::sync::Notify;
use crate::error::to_py_err;
use crate::predicate::dict_to_predicate;
+const MAP_SELECTED_KEYS_PREFIX: &str = "__PAIMON_MAP_SELECTED_KEYS:";
+const MAP_SELECTED_KEYS_DELIMITER: char = ';';
+
/// Time-travel selector option names, in the core's resolution priority order.
const TIME_TRAVEL_SELECTORS: [&str; 5] = [
"scan.timestamp-millis",
@@ -70,12 +73,15 @@ fn find_time_travel_selector(opts: &HashMap<String,
String>) -> Option<(&str, &s
fn apply_read_config(
builder: &mut paimon::table::ReadBuilder<'_>,
projection: &Option<Vec<String>>,
+ read_type: &Option<Vec<DataField>>,
limit: Option<usize>,
filter: &Option<Predicate>,
case_sensitive: bool,
) -> PyResult<()> {
builder.with_case_sensitive(case_sensitive);
- if let Some(projection) = projection {
+ if let Some(read_type) = read_type {
+ builder.with_read_type(read_type.clone());
+ } else if let Some(projection) = projection {
let cols: Vec<&str> = projection.iter().map(String::as_str).collect();
builder.with_projection(&cols).map_err(to_py_err)?;
}
@@ -88,6 +94,106 @@ fn apply_read_config(
Ok(())
}
+/// Resolve flat output paths to the authoritative nested read type used by the
+/// core reader. ROW children are pruned recursively. MAP values remain MAPs,
+/// but their requested string keys are carried in the temporary field
+/// description so shared-shredding Parquet files can prune physical columns.
+fn project_nested_read_type(
+ fields: &[DataField],
+ paths: &[Vec<String>],
+) -> PyResult<Vec<DataField>> {
+ if paths.iter().any(Vec::is_empty) {
+ return Err(PyValueError::new_err(
+ "nested projection paths must not be empty",
+ ));
+ }
+
+ let mut result = Vec::new();
+ let mut seen = std::collections::HashSet::new();
+ for path in paths {
+ let name = &path[0];
+ if !seen.insert(name.clone()) {
+ continue;
+ }
+ let field = fields
+ .iter()
+ .find(|field| field.name() == name)
+ .ok_or_else(|| {
+ PyValueError::new_err(format!(
+ "nested projection field '{}' does not exist",
+ path.join(".")
+ ))
+ })?;
+ let matching: Vec<&[String]> = paths
+ .iter()
+ .filter(|candidate| candidate.first() == Some(name))
+ .map(|candidate| candidate[1..].as_ref())
+ .collect();
+ result.push(project_nested_field(field, &matching, name)?);
+ }
+ Ok(result)
+}
+
+fn project_nested_field(
+ field: &DataField,
+ tails: &[&[String]],
+ full_name: &str,
+) -> PyResult<DataField> {
+ if tails.iter().any(|tail| tail.is_empty()) {
+ return Ok(field.clone());
+ }
+ match field.data_type() {
+ DataType::Row(row) => {
+ let child_paths: Vec<Vec<String>> = tails.iter().map(|tail|
tail.to_vec()).collect();
+ let children = project_nested_read_type(row.fields(),
&child_paths).map_err(|_| {
+ PyValueError::new_err(format!(
+ "nested projection field '{}' does not exist",
+ tails
+ .first()
+ .map(|tail| format!("{full_name}.{}", tail.join(".")))
+ .unwrap_or_else(|| full_name.to_string())
+ ))
+ })?;
+ Ok(DataField::new(
+ field.id(),
+ field.name().to_string(),
+ DataType::Row(RowType::with_nullable(
+ field.data_type().is_nullable(),
+ children,
+ )),
+ )
+ .with_description(field.description().map(str::to_string)))
+ }
+ DataType::Map(_) if tails.iter().all(|tail| tail.len() == 1) => {
+ let mut keys = Vec::new();
+ for tail in tails {
+ let key = &tail[0];
+ if key.contains(MAP_SELECTED_KEYS_DELIMITER)
+ || key.starts_with(MAP_SELECTED_KEYS_PREFIX)
+ {
+ // Keep the complete MAP for keys which cannot be encoded
+ // by the cross-language selected-key convention.
+ return Ok(field.clone());
+ }
+ if !keys.contains(key) {
+ keys.push(key.clone());
+ }
+ }
+ Ok(field.clone().with_description(Some(format!(
+ "{MAP_SELECTED_KEYS_PREFIX}{}",
+ keys.join(&MAP_SELECTED_KEYS_DELIMITER.to_string())
+ ))))
+ }
+ _ => Err(PyValueError::new_err(format!(
+ "nested projection field '{}' is not a ROW or MAP",
+ tails
+ .first()
+ .map(|tail| format!("{full_name}.{}", tail.join(".")))
+ .unwrap_or_else(|| full_name.to_string())
+ ))),
+ }
+}
+
/// Extract a sequence of Python `Split` objects into core `DataSplit`s.
Accepts
/// any iterable (list/tuple/generator). Runs under the GIL since it touches
/// Python objects. A non-iterable argument or a non-`Split` element raises
@@ -111,11 +217,13 @@ fn extract_splits(splits: &Bound<'_, PyAny>) ->
PyResult<Vec<DataSplit>> {
pub struct PyReadBuilder {
table: Arc<Table>,
projection: Option<Vec<String>>,
+ read_type: Option<Vec<DataField>>,
limit: Option<usize>,
filter: Option<Predicate>,
row_ranges: Option<Vec<RowRange>>,
case_sensitive: bool,
blob_parallelism: Option<usize>,
+ include_row_kind: bool,
}
impl PyReadBuilder {
@@ -123,11 +231,13 @@ impl PyReadBuilder {
Self {
table,
projection: None,
+ read_type: None,
limit: None,
filter: None,
row_ranges: None,
case_sensitive: true,
blob_parallelism: None,
+ include_row_kind: false,
}
}
@@ -169,11 +279,13 @@ impl PyReadBuilder {
Ok(Self {
table: Arc::new(traveled),
projection: None,
+ read_type: None,
limit: None,
filter: None,
row_ranges: None,
case_sensitive: true,
blob_parallelism: None,
+ include_row_kind: false,
})
}
}
@@ -182,14 +294,55 @@ impl PyReadBuilder {
impl PyReadBuilder {
fn with_projection(mut slf: PyRefMut<'_, Self>, columns: Vec<String>) ->
PyRefMut<'_, Self> {
slf.projection = Some(columns);
+ slf.read_type = None;
slf
}
+ /// Project top-level fields or nested ROW leaves by their exact name
paths.
+ /// MAP paths keep the complete MAP so the caller can extract literal keys.
+ fn with_nested_projection(
+ mut slf: PyRefMut<'_, Self>,
+ paths: Vec<Vec<String>>,
+ ) -> PyResult<PyRefMut<'_, Self>> {
+ let schema = slf.table.schema();
+ let mut read_type = project_nested_read_type(schema.fields(), &paths)?;
+ let options = schema.options();
+ let has_default_aggregation =
options.contains_key("fields.default-aggregate-function");
+ for projected in &mut read_type {
+ let selected_map = matches!(projected.data_type(),
DataType::Map(_))
+ && projected
+ .description()
+ .is_some_and(|description|
description.starts_with(MAP_SELECTED_KEYS_PREFIX));
+ let has_field_aggregation =
+ options.contains_key(&format!("fields.{}.aggregate-function",
projected.name()));
+ if selected_map && (has_default_aggregation ||
has_field_aggregation) {
+ // Merge engines may need the complete MAP to calculate the
+ // projected key. Match PyPaimon's complete-MAP fallback.
+ *projected = schema
+ .fields()
+ .iter()
+ .find(|field| field.name() == projected.name())
+ .expect("nested projection came from this schema")
+ .clone();
+ }
+ }
+ slf.read_type = Some(read_type);
+ slf.projection = None;
+ Ok(slf)
+ }
+
fn with_limit(mut slf: PyRefMut<'_, Self>, limit: usize) -> PyRefMut<'_,
Self> {
slf.limit = Some(limit);
slf
}
+ /// Include a leading `rowkind` string column. Streaming primary-key
+ /// splits preserve their physical change kinds; snapshot rows are `+I`.
+ fn with_include_row_kind(mut slf: PyRefMut<'_, Self>, include: bool) ->
PyRefMut<'_, Self> {
+ slf.include_row_kind = include;
+ slf
+ }
+
/// Set the maximum number of concurrent BLOB range reads for this read.
fn with_blob_parallelism(
mut slf: PyRefMut<'_, Self>,
@@ -251,6 +404,7 @@ impl PyReadBuilder {
PyTableScan {
table: Arc::clone(&self.table),
projection: self.projection.clone(),
+ read_type: self.read_type.clone(),
limit: self.limit,
filter: self.filter.clone(),
row_ranges: self.row_ranges.clone(),
@@ -273,10 +427,12 @@ impl PyReadBuilder {
PyTableRead {
table: Arc::clone(&self.table),
projection: self.projection.clone(),
+ read_type: self.read_type.clone(),
limit: self.limit,
filter: self.filter.clone(),
case_sensitive: self.case_sensitive,
blob_parallelism: self.blob_parallelism,
+ include_row_kind: self.include_row_kind,
}
}
}
@@ -285,6 +441,7 @@ impl PyReadBuilder {
pub struct PyTableScan {
table: Arc<Table>,
projection: Option<Vec<String>>,
+ read_type: Option<Vec<DataField>>,
limit: Option<usize>,
filter: Option<Predicate>,
row_ranges: Option<Vec<RowRange>>,
@@ -336,6 +493,7 @@ impl PyTableScan {
apply_read_config(
&mut builder,
&self.projection,
+ &self.read_type,
self.limit,
&self.filter,
self.case_sensitive,
@@ -396,10 +554,12 @@ impl PyTableScan {
pub struct PyTableRead {
table: Arc<Table>,
projection: Option<Vec<String>>,
+ read_type: Option<Vec<DataField>>,
limit: Option<usize>,
filter: Option<Predicate>,
case_sensitive: bool,
blob_parallelism: Option<usize>,
+ include_row_kind: bool,
}
impl PyTableRead {
@@ -414,6 +574,7 @@ impl PyTableRead {
apply_read_config(
&mut builder,
&self.projection,
+ &self.read_type,
self.limit,
&self.filter,
self.case_sensitive,
@@ -427,7 +588,11 @@ impl PyTableRead {
// path so an invalid projection fails consistently regardless of
// how many splits are passed.
let read = builder.new_read().map_err(to_py_err)?;
- read.to_arrow(&splits).map_err(to_py_err)
+ if self.include_row_kind {
+ read.to_arrow_with_row_kind(&splits).map_err(to_py_err)
+ } else {
+ read.to_arrow(&splits).map_err(to_py_err)
+ }
})
}
}
@@ -639,6 +804,102 @@ mod tests {
use super::*;
+ use paimon::spec::{DataField, DataType, IntType, MapType, RowType,
VarCharType};
+
+ #[test]
+ fn nested_projection_prunes_rows_and_marks_selected_map_keys() {
+ let fields = vec![
+ DataField::new(0, "id".to_string(), DataType::Int(IntType::new())),
+ DataField::new(
+ 1,
+ "payload".to_string(),
+ DataType::Row(RowType::new(vec![
+ DataField::new(2, "version".to_string(),
DataType::Int(IntType::new())),
+ DataField::new(
+ 3,
+ "details".to_string(),
+ DataType::Row(RowType::new(vec![
+ DataField::new(
+ 4,
+ "name".to_string(),
+ DataType::VarChar(VarCharType::string_type()),
+ ),
+ DataField::new(5, "score".to_string(),
DataType::Int(IntType::new())),
+ ])),
+ ),
+ ])),
+ ),
+ DataField::new(
+ 6,
+ "attrs".to_string(),
+ DataType::Map(MapType::new(
+ DataType::VarChar(VarCharType::string_type()),
+ DataType::Int(IntType::new()),
+ )),
+ ),
+ ];
+
+ let projected = project_nested_read_type(
+ &fields,
+ &[
+ vec![
+ "payload".to_string(),
+ "details".to_string(),
+ "score".to_string(),
+ ],
+ vec!["id".to_string()],
+ vec!["attrs".to_string(), "selected".to_string()],
+ ],
+ )
+ .unwrap();
+
+ assert_eq!(
+ projected.iter().map(DataField::name).collect::<Vec<_>>(),
+ vec!["payload", "id", "attrs"]
+ );
+ let DataType::Row(payload) = projected[0].data_type() else {
+ panic!("payload must remain a row");
+ };
+ assert_eq!(payload.fields().len(), 1);
+ assert_eq!(payload.fields()[0].name(), "details");
+ let DataType::Row(details) = payload.fields()[0].data_type() else {
+ panic!("details must remain a row");
+ };
+ assert_eq!(
+ details
+ .fields()
+ .iter()
+ .map(DataField::name)
+ .collect::<Vec<_>>(),
+ vec!["score"]
+ );
+ assert_eq!(projected[2].data_type(), fields[2].data_type());
+ assert_eq!(
+ projected[2].description(),
+ Some("__PAIMON_MAP_SELECTED_KEYS:selected")
+ );
+ }
+
+ #[test]
+ fn nested_projection_rejects_an_invalid_struct_path() {
+ let fields = vec![DataField::new(
+ 1,
+ "payload".to_string(),
+ DataType::Row(RowType::new(vec![DataField::new(
+ 2,
+ "version".to_string(),
+ DataType::Int(IntType::new()),
+ )])),
+ )];
+
+ let error = project_nested_read_type(
+ &fields,
+ &[vec!["payload".to_string(), "missing".to_string()]],
+ )
+ .unwrap_err();
+ assert!(error.to_string().contains("payload.missing"));
+ }
+
#[test]
fn record_batch_reader_close_interrupts_pending_next() {
let (polled_tx, polled_rx) = mpsc::channel();
diff --git a/crates/paimon/src/arrow/format/parquet.rs
b/crates/paimon/src/arrow/format/parquet.rs
index a8ddca4b..bc7403a3 100644
--- a/crates/paimon/src/arrow/format/parquet.rs
+++ b/crates/paimon/src/arrow/format/parquet.rs
@@ -519,7 +519,17 @@ impl FormatFileReader for ParquetFormatReader {
})
.collect();
- let mask = ProjectionMask::roots(&parquet_schema, root_indices);
+ // Build the MAP plan before the projection mask: a selected-key MAP
+ // can then decode only __field_mapping plus the direct/overflow
+ // children which may contain the requested keys.
+ let map_read_plan =
+ MapShreddingReadPlan::create(&scan_fields,
batch_stream_builder.schema())?
+ .map(Arc::new);
+ let mask = if let Some(plan) = map_read_plan.as_deref() {
+ map_shredding_projection_mask(&parquet_schema, &root_indices, plan)
+ } else {
+ ProjectionMask::roots(&parquet_schema, root_indices)
+ };
batch_stream_builder =
batch_stream_builder.with_projection(mask.clone());
let mut decoder_predicates = build_parquet_row_filter(&parquet_schema,
preds, file_fields)?
@@ -614,14 +624,6 @@ impl FormatFileReader for ParquetFormatReader {
batch_stream_builder = batch_stream_builder.with_batch_size(size);
}
- // MAP shared-shredding read plan, built from the per-field metadata
- // committed into the file footer at write time. `None` when no scanned
- // field is shared-shredded. Assembly must happen before any residual
- // predicate evaluation so predicates see logical MAP columns.
- let map_read_plan =
- MapShreddingReadPlan::create(&scan_fields,
batch_stream_builder.schema())?
- .map(Arc::new);
-
// A normal Parquet stream fetches and decodes row groups one by one.
// For remote object stores, a full scan of a compacted file can
// therefore serialize dozens of independent range requests behind one
@@ -784,6 +786,32 @@ impl FormatFileReader for ParquetFormatReader {
}
}
+fn map_shredding_projection_mask(
+ parquet_schema: &parquet::schema::types::SchemaDescriptor,
+ selected_roots: &[usize],
+ plan: &MapShreddingReadPlan,
+) -> ProjectionMask {
+ let selected_roots: std::collections::HashSet<usize> =
selected_roots.iter().copied().collect();
+ let root_fields = parquet_schema.root_schema().get_fields();
+ let leaves = (0..parquet_schema.num_columns()).filter(|leaf_index| {
+ let root_index = parquet_schema.get_column_root_idx(*leaf_index);
+ if !selected_roots.contains(&root_index) {
+ return false;
+ }
+ let root_name = root_fields[root_index].name();
+ let Some(children) = plan.projected_physical_children(root_name) else {
+ return true;
+ };
+ parquet_schema
+ .column(*leaf_index)
+ .path()
+ .parts()
+ .get(1)
+ .is_some_and(|child| children.contains(child))
+ });
+ ProjectionMask::leaves(parquet_schema, leaves)
+}
+
fn projected_row_group_bytes(row_group: &RowGroupMetaData, projection:
&ProjectionMask) -> u64 {
row_group
.columns()
@@ -4931,6 +4959,15 @@ mod tests {
rows: &[Option<Vec<(&str, Option<i64>)>>],
ids: &[i32],
max_columns: usize,
+ ) -> (String, crate::io::FileIO, u64, Vec<DataField>) {
+ write_map_shredding_file_with_compression(rows, ids, max_columns,
"zstd").await
+ }
+
+ async fn write_map_shredding_file_with_compression(
+ rows: &[Option<Vec<(&str, Option<i64>)>>],
+ ids: &[i32],
+ max_columns: usize,
+ compression: &str,
) -> (String, crate::io::FileIO, u64, Vec<DataField>) {
let fields = map_shredding_fields();
let options = HashMap::from([
@@ -4959,7 +4996,7 @@ mod tests {
let mut writer = create_format_writer(
&output,
logical_schema,
- "zstd",
+ compression,
1,
None,
Some(&fields),
@@ -5050,6 +5087,200 @@ mod tests {
assert_int64_map_rows(tags, &rows);
}
+ #[tokio::test]
+ async fn test_parquet_map_shredding_selected_key_prunes_and_assembles() {
+ let rows: Vec<Option<Vec<(&str, Option<i64>)>>> = vec![
+ Some(vec![("a", Some(10)), ("b", None), ("c", Some(30))]),
+ None,
+ Some(vec![]),
+ Some(vec![("b", Some(40)), ("a", Some(50))]),
+ ];
+ let ids = vec![1, 2, 3, 4];
+ let (path, file_io, file_size, fields) =
write_map_shredding_file(&rows, &ids, 2).await;
+ let selected_fields = vec![
+ fields[0].clone(),
+ fields[1].clone().with_description(Some(format!(
+ "{}c",
+ crate::arrow::shredding::map::SELECTED_KEYS_PREFIX
+ ))),
+ ];
+
+ let input = file_io.new_input(&path).unwrap();
+ let file_reader = input.reader().await.unwrap();
+ let reader = create_format_reader(&path, false,
&selected_fields).unwrap();
+ let batches = reader
+ .read_batch_stream(
+ Box::new(file_reader),
+ file_size,
+ &selected_fields,
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ let tags = batches[0]
+ .column_by_name("tags")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<MapArray>()
+ .unwrap();
+ assert_int64_map_rows(
+ tags,
+ &[
+ Some(vec![("c", Some(30))]),
+ None,
+ Some(vec![]),
+ Some(vec![]),
+ ],
+ );
+
+ let missing_fields =
vec![fields[1].clone().with_description(Some(format!(
+ "{}missing",
+ crate::arrow::shredding::map::SELECTED_KEYS_PREFIX
+ )))];
+ let missing_reader =
file_io.new_input(&path).unwrap().reader().await.unwrap();
+ let missing = create_format_reader(&path, false, &missing_fields)
+ .unwrap()
+ .read_batch_stream(
+ Box::new(missing_reader),
+ file_size,
+ &missing_fields,
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_int64_map_rows(
+ missing[0]
+ .column_by_name("tags")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<MapArray>()
+ .unwrap(),
+ &[Some(vec![]), None, Some(vec![]), Some(vec![])],
+ );
+
+ // An empty encoded suffix is the selected-key protocol's encoding for
+ // the valid empty-string key, not an empty key set.
+ let empty_key_rows = vec![Some(vec![("", Some(10)), ("other",
Some(11))])];
+ let (empty_path, empty_file_io, empty_file_size, empty_fields) =
+ write_map_shredding_file(&empty_key_rows, &[1], 2).await;
+ let selected_empty_fields =
vec![empty_fields[1].clone().with_description(Some(
+ crate::arrow::shredding::map::SELECTED_KEYS_PREFIX.to_string(),
+ ))];
+ let empty_key_reader = empty_file_io
+ .new_input(&empty_path)
+ .unwrap()
+ .reader()
+ .await
+ .unwrap();
+ let empty_key = create_format_reader(&empty_path, false,
&selected_empty_fields)
+ .unwrap()
+ .read_batch_stream(
+ Box::new(empty_key_reader),
+ empty_file_size,
+ &selected_empty_fields,
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+ assert_int64_map_rows(
+ empty_key[0]
+ .column_by_name("tags")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<MapArray>()
+ .unwrap(),
+ &[Some(vec![("", Some(10))])],
+ );
+ }
+
+ #[tokio::test]
+ async fn test_parquet_map_shredding_selected_key_reduces_data_io() {
+ const ROWS: usize = 100_000;
+ let rows = (0..ROWS)
+ .map(|row| {
+ Some(vec![
+ ("a", Some((row as i64).wrapping_mul(982_451_653))),
+ ("b", Some((row as i64).wrapping_mul(961_748_941))),
+ ])
+ })
+ .collect::<Vec<_>>();
+ let ids = (0..ROWS as i32).collect::<Vec<_>>();
+ let (path, file_io, file_size, fields) =
+ write_map_shredding_file_with_compression(&rows, &ids, 2,
"none").await;
+ let data = file_io.new_input(&path).unwrap().read().await.unwrap();
+
+ let selected_fields =
vec![fields[1].clone().with_description(Some(format!(
+ "{}a",
+ crate::arrow::shredding::map::SELECTED_KEYS_PREFIX
+ )))];
+ let selected_reader = TrackingFileRead::new(data.clone());
+ let selected_tracker = selected_reader.clone();
+ let selected_stream = ParquetFormatReader::default()
+ .read_batch_stream(
+ Box::new(selected_reader),
+ file_size,
+ &selected_fields,
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap();
+ selected_tracker.reset();
+ let selected_rows = selected_stream
+ .try_fold(0usize, |count, batch| async move {
+ Ok(count + batch.num_rows())
+ })
+ .await
+ .unwrap();
+ let selected_bytes = selected_tracker.bytes_read();
+
+ let full_reader = TrackingFileRead::new(data);
+ let full_tracker = full_reader.clone();
+ let full_stream = ParquetFormatReader::default()
+ .read_batch_stream(
+ Box::new(full_reader),
+ file_size,
+ &fields[1..],
+ None,
+ None,
+ None,
+ )
+ .await
+ .unwrap();
+ full_tracker.reset();
+ let full_rows = full_stream
+ .try_fold(0usize, |count, batch| async move {
+ Ok(count + batch.num_rows())
+ })
+ .await
+ .unwrap();
+ let full_bytes = full_tracker.bytes_read();
+
+ assert_eq!(selected_rows, ROWS);
+ assert_eq!(full_rows, ROWS);
+ assert!(
+ selected_bytes * 5 < full_bytes * 4,
+ "selected-key read used {selected_bytes} bytes; full MAP read used
{full_bytes} bytes"
+ );
+ }
+
#[tokio::test]
async fn
test_parquet_map_shredding_residual_path_assembles_before_filter() {
// An `Or` predicate is not fully enforced by the parquet row filter,
so
diff --git a/crates/paimon/src/arrow/shredding/map.rs
b/crates/paimon/src/arrow/shredding/map.rs
index b887324c..b4e9bdac 100644
--- a/crates/paimon/src/arrow/shredding/map.rs
+++ b/crates/paimon/src/arrow/shredding/map.rs
@@ -37,7 +37,7 @@ use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
use arrow_schema::{DataType as ArrowDataType, Fields, Schema as ArrowSchema};
use arrow_select::interleave::interleave;
use arrow_select::take::take;
-use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::sync::Arc;
// ---------------------------------------------------------------------------
@@ -56,6 +56,11 @@ const OVERFLOW_SET_KEY: &str =
"paimon.map.shared-shredding.overflow-set";
const NUM_COLUMNS_KEY: &str = "paimon.map.shared-shredding.num-columns";
const MAX_ROW_WIDTH_KEY: &str = "paimon.map.shared-shredding.max-row-width";
+/// Temporary read-field description used by Java/Python to request literal
+/// MAP keys without changing the logical MAP output type.
+pub(crate) const SELECTED_KEYS_PREFIX: &str = "__PAIMON_MAP_SELECTED_KEYS:";
+const SELECTED_KEYS_DELIMITER: char = ';';
+
const FIELD_MAPPING_NAME: &str = "__field_mapping";
const OVERFLOW_NAME: &str = "__overflow";
@@ -1017,6 +1022,9 @@ pub(crate) struct MapShreddingReadPlan {
struct MapReadContext {
num_columns: usize,
name_by_id: HashMap<i32, String>,
+ selected_names: Option<HashSet<String>>,
+ selected_columns: BTreeSet<usize>,
+ include_overflow: bool,
}
impl MapShreddingReadPlan {
@@ -1052,6 +1060,21 @@ impl MapShreddingReadPlan {
DataField::new(field.id(), field.name().to_string(),
physical_type)
.with_description(field.description().map(ToString::to_string)),
);
+ let selected_names = selected_map_keys(field)?;
+ let mut selected_columns = BTreeSet::new();
+ let mut include_overflow = true;
+ if let Some(names) = selected_names.as_ref() {
+ include_overflow = false;
+ for name in names {
+ let Some(field_id) =
field_meta.name_to_id.get(name) else {
+ continue;
+ };
+ if let Some(columns) =
field_meta.field_to_columns.get(field_id) {
+
selected_columns.extend(columns.iter().copied());
+ }
+ include_overflow |=
field_meta.overflow_set.contains(field_id);
+ }
+ }
contexts.insert(
index,
MapReadContext {
@@ -1061,6 +1084,9 @@ impl MapShreddingReadPlan {
.iter()
.map(|(name, &id)| (id, name.clone()))
.collect(),
+ selected_names,
+ selected_columns,
+ include_overflow,
},
);
converted = true;
@@ -1078,6 +1104,51 @@ impl MapShreddingReadPlan {
contexts,
}))
}
+
+ /// Return the physical struct children needed for one logical field.
+ /// `None` means the whole root field must be decoded.
+ pub(crate) fn projected_physical_children(&self, field_name: &str) ->
Option<HashSet<String>> {
+ let field_index = self
+ .logical_fields
+ .iter()
+ .position(|field| field.name() == field_name)?;
+ let ctx = self.contexts.get(&field_index)?;
+ ctx.selected_names.as_ref()?;
+
+ let mut children = HashSet::from([FIELD_MAPPING_NAME.to_string()]);
+ children.extend(
+ ctx.selected_columns
+ .iter()
+ .map(|index| physical_column_name(*index)),
+ );
+ if ctx.include_overflow {
+ children.insert(OVERFLOW_NAME.to_string());
+ }
+ Some(children)
+ }
+}
+
+fn selected_map_keys(field: &DataField) -> Result<Option<HashSet<String>>> {
+ let Some(description) = field.description() else {
+ return Ok(None);
+ };
+ let Some(encoded) = description.strip_prefix(SELECTED_KEYS_PREFIX) else {
+ return Ok(None);
+ };
+ let keys: HashSet<String> = encoded
+ .split(SELECTED_KEYS_DELIMITER)
+ .map(ToString::to_string)
+ .collect();
+ if keys.len() != encoded.split(SELECTED_KEYS_DELIMITER).count() {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "Selected-key MAP field '{}' contains duplicate keys",
+ field.name()
+ ),
+ source: None,
+ });
+ }
+ Ok(Some(keys))
}
fn validate_physical_struct_children(
@@ -1345,19 +1416,18 @@ fn struct_array_to_logical_map(
) -> Result<ArrayRef> {
let num_rows = array.len();
let k = ctx.num_columns;
- if array.num_columns() < k + 1 {
- return Err(Error::DataInvalid {
+ let field_mapping_index = array
+ .fields()
+ .iter()
+ .position(|field| field.name() == FIELD_MAPPING_NAME)
+ .ok_or_else(|| Error::DataInvalid {
message: format!(
- "Shared-shredding physical column '{column_name}' has {}
children, expected at least {}",
- array.num_columns(),
- k + 1
+ "Shared-shredding physical column '{column_name}' is missing
{FIELD_MAPPING_NAME}"
),
source: None,
- });
- }
-
+ })?;
let field_mapping = array
- .column(0)
+ .column(field_mapping_index)
.as_any()
.downcast_ref::<ListArray>()
.ok_or_else(|| Error::DataInvalid {
@@ -1397,11 +1467,13 @@ fn struct_array_to_logical_map(
};
let logical_value_type = entry_fields[1].data_type();
- let has_overflow_column = array.num_columns() > k + 1;
- let overflow = if has_overflow_column {
- Some(
+ let overflow = array
+ .fields()
+ .iter()
+ .position(|field| field.name() == OVERFLOW_NAME)
+ .map(|index| {
array
- .column(k + 1)
+ .column(index)
.as_any()
.downcast_ref::<MapArray>()
.ok_or_else(|| Error::DataInvalid {
@@ -1409,11 +1481,9 @@ fn struct_array_to_logical_map(
"Shared-shredding __overflow of '{column_name}' must
be MapArray"
),
source: None,
- })?,
- )
- } else {
- None
- };
+ })
+ })
+ .transpose()?;
let normalize_source = |value: &dyn Array| {
arrow_cast::cast(value, logical_value_type).map_err(|e|
Error::UnexpectedError {
@@ -1427,9 +1497,31 @@ fn struct_array_to_logical_map(
// Normalize each physical source once. For metadata-only differences this
// rebuilds nested Arrow wrappers while reusing their value buffers. The
// final interleave then copies only entries referenced by the mapping.
- let mut value_sources = Vec::with_capacity(k +
usize::from(overflow.is_some()));
- for i in 0..k {
- value_sources.push(normalize_source(array.column(i + 1).as_ref())?);
+ let mut value_sources = Vec::with_capacity(array.num_columns());
+ let mut direct_source_indices = vec![None; k];
+ for (position, physical_field) in array.fields().iter().enumerate() {
+ let Some(suffix) = physical_field.name().strip_prefix("__col_") else {
+ continue;
+ };
+ let physical_index = suffix.parse::<usize>().map_err(|_|
Error::DataInvalid {
+ message: format!(
+ "Shared-shredding physical column '{column_name}' has invalid
child '{}'",
+ physical_field.name()
+ ),
+ source: None,
+ })?;
+ if physical_index >= k {
+ return Err(Error::DataInvalid {
+ message: format!(
+ "Shared-shredding physical column '{column_name}' child
'{}' exceeds metadata column count {k}",
+ physical_field.name()
+ ),
+ source: None,
+ });
+ }
+ let source_index = value_sources.len();
+ value_sources.push(normalize_source(array.column(position).as_ref())?);
+ direct_source_indices[physical_index] = Some(source_index);
}
let overflow_source_index = if let Some(overflow) = overflow {
let index = value_sources.len();
@@ -1438,6 +1530,9 @@ fn struct_array_to_logical_map(
} else {
None
};
+ if value_sources.is_empty() {
+ value_sources.push(arrow_array::new_empty_array(logical_value_type));
+ }
let mut key_builder = arrow_array::builder::StringBuilder::new();
let mut value_indices: Vec<(usize, usize)> = Vec::new();
@@ -1466,7 +1561,7 @@ fn struct_array_to_logical_map(
source: None,
});
}
- for i in 0..k {
+ for (i, direct_source_index) in
direct_source_indices.iter().enumerate() {
let field_id = if mapping_values.is_null(fm_start + i) {
-1
} else {
@@ -1478,8 +1573,21 @@ fn struct_array_to_logical_map(
let Some(name) = ctx.name_by_id.get(&field_id) else {
continue;
};
+ if ctx
+ .selected_names
+ .as_ref()
+ .is_some_and(|selected| !selected.contains(name))
+ {
+ continue;
+ }
+ let source_index = (*direct_source_index).ok_or_else(||
Error::DataInvalid {
+ message: format!(
+ "Shared-shredding physical column '{column_name}' is
missing __col_{i} required for key '{name}'"
+ ),
+ source: None,
+ })?;
key_builder.append_value(name);
- value_indices.push((i, row));
+ value_indices.push((source_index, row));
count += 1;
}
@@ -1502,6 +1610,13 @@ fn struct_array_to_logical_map(
let Some(name) = ctx.name_by_id.get(&field_id) else {
continue;
};
+ if ctx
+ .selected_names
+ .as_ref()
+ .is_some_and(|selected| !selected.contains(name))
+ {
+ continue;
+ }
key_builder.append_value(name);
value_indices.push((overflow_source_index.expect("overflow
source exists"), j));
count += 1;
@@ -1732,6 +1847,9 @@ mod tests {
&MapReadContext {
num_columns,
name_by_id,
+ selected_names: None,
+ selected_columns: BTreeSet::new(),
+ include_overflow: true,
},
&MapType::new(string_type(), value_type),
"nested",
diff --git a/crates/paimon/src/table/table_read.rs
b/crates/paimon/src/table/table_read.rs
index d232614b..4881fc82 100644
--- a/crates/paimon/src/table/table_read.rs
+++ b/crates/paimon/src/table/table_read.rs
@@ -205,6 +205,25 @@ impl<'a> TableRead<'a> {
}
}
+ /// Returns rows with a leading `rowkind` column.
+ ///
+ /// Materialized snapshot rows are inserts (`+I`). Streaming primary-key
+ /// splits retain their physical `_VALUE_KIND` values so change events are
+ /// exposed as `+I`, `-U`, `+U`, or `-D` without merge reconciliation.
+ pub fn to_arrow_with_row_kind(
+ &self,
+ data_splits: &[DataSplit],
+ ) -> crate::Result<ArrowRecordBatchStream> {
+ self.ensure_query_auth_allowed()?;
+ match &self.0 {
+ TableReadKind::Paimon(read) =>
read.to_arrow_with_row_kind(data_splits),
+ TableReadKind::Format(read) => {
+ let schema = audit_schema_for_read_type(read.read_type(),
false)?;
+ prepend_insert_row_kind_stream(read.to_arrow(data_splits)?,
schema)
+ }
+ }
+ }
+
/// Returns an [`ArrowRecordBatchStream`] for an incremental scan plan.
///
/// Delta/Changelog use [`IncrementalSplit::Data`]. Diff uses
@@ -415,16 +434,57 @@ impl<'a> PaimonTableRead<'a> {
}
}
+ fn to_arrow_with_row_kind(
+ &self,
+ data_splits: &[DataSplit],
+ ) -> crate::Result<ArrowRecordBatchStream> {
+ let schema = audit_schema_for_read_type(&self.read_type, false)?;
+ let (streaming, materialized): (Vec<_>, Vec<_>) = data_splits
+ .iter()
+ .cloned()
+ .partition(DataSplit::is_streaming);
+ let mut streams = Vec::with_capacity(2);
+ if !materialized.is_empty() {
+ streams.push(prepend_insert_row_kind_stream(
+ self.to_arrow(&materialized)?,
+ Arc::clone(&schema),
+ )?);
+ }
+ if !streaming.is_empty() {
+ if self.table.schema().primary_keys().is_empty() {
+ streams.push(prepend_insert_row_kind_stream(
+ self.to_arrow(&streaming)?,
+ Arc::clone(&schema),
+ )?);
+ } else {
+ streams.push(self.audit_raw_data_splits(&streaming, true,
false)?);
+ }
+ }
+ Ok(Box::pin(stream::select_all(streams)))
+ }
+
fn audit_raw_stream(
&self,
plan: &IncrementalPlan,
has_value_kind: bool,
) -> crate::Result<ArrowRecordBatchStream> {
plan.validate()?;
- let core_options = self.table.schema().core_options();
let data_splits = plan.data_splits();
+ self.audit_raw_data_splits(
+ &data_splits,
+ has_value_kind,
+ audit_sequence_number_enabled(self.table),
+ )
+ }
+
+ fn audit_raw_data_splits(
+ &self,
+ data_splits: &[DataSplit],
+ has_value_kind: bool,
+ include_sequence: bool,
+ ) -> crate::Result<ArrowRecordBatchStream> {
+ let core_options = self.table.schema().core_options();
let user_read_type = self.read_type.clone();
- let include_sequence = audit_sequence_number_enabled(self.table);
let audit_schema = audit_schema_for_read_type(&user_read_type,
include_sequence)?;
let mut read_type = user_read_type.clone();
@@ -460,7 +520,7 @@ impl<'a> PaimonTableRead<'a> {
.with_parquet_read_budget(Some(self.parquet_read_budget()?))
.with_table_options(self.table.schema().options().clone())
.with_mosaic_prefetch(configured_mosaic_prefetch(self.table)?);
- let raw_stream = reader.read(&data_splits)?;
+ let raw_stream = reader.read(data_splits)?;
Ok(Box::pin(async_stream::try_stream! {
futures::pin_mut!(raw_stream);
@@ -982,6 +1042,34 @@ impl<'a> PaimonTableRead<'a> {
}
}
+fn prepend_insert_row_kind_stream(
+ stream: ArrowRecordBatchStream,
+ schema: Arc<ArrowSchema>,
+) -> crate::Result<ArrowRecordBatchStream> {
+ Ok(Box::pin(async_stream::try_stream! {
+ futures::pin_mut!(stream);
+ while let Some(batch) = stream.next().await {
+ yield prepend_insert_row_kind(batch?, Arc::clone(&schema))?;
+ }
+ }))
+}
+
+fn prepend_insert_row_kind(
+ batch: RecordBatch,
+ schema: Arc<ArrowSchema>,
+) -> crate::Result<RecordBatch> {
+ let mut columns = Vec::with_capacity(batch.num_columns() + 1);
+ columns.push(Arc::new(StringArray::from(vec!["+I"; batch.num_rows()])) as
ArrayRef);
+ columns.extend(batch.columns().iter().cloned());
+ let options =
RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+ RecordBatch::try_new_with_options(schema, columns,
&options).map_err(|error| {
+ crate::Error::UnexpectedError {
+ message: format!("Failed to prepend row kind column: {error}"),
+ source: Some(Box::new(error)),
+ }
+ })
+}
+
fn audit_schema_for_read_type(
read_type: &[DataField],
include_sequence: bool,
@@ -1569,6 +1657,32 @@ mod tests {
use crate::table::source::DataSplitBuilder;
use futures::TryStreamExt;
+ #[test]
+ fn test_prepend_insert_row_kind_preserves_zero_column_row_count() {
+ let batch = RecordBatch::try_new_with_options(
+ Arc::new(ArrowSchema::empty()),
+ vec![],
+ &RecordBatchOptions::new().with_row_count(Some(3)),
+ )
+ .unwrap();
+
+ let actual =
+ prepend_insert_row_kind(batch, audit_schema_for_read_type(&[],
false).unwrap())
+ .unwrap();
+
+ assert_eq!(actual.num_rows(), 3);
+ assert_eq!(actual.schema().fields()[0].name(), ROW_KIND_FIELD_NAME);
+ let kinds = actual
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ assert_eq!(
+ kinds.iter().collect::<Vec<_>>(),
+ vec![Some("+I"), Some("+I"), Some("+I")]
+ );
+ }
+
pub(super) fn file(name: &str, level: i32, delete_row_count: Option<i64>)
-> DataFileMeta {
DataFileMeta {
file_name: name.to_string(),
diff --git a/crates/paimon/tests/audit_log_table_test.rs
b/crates/paimon/tests/audit_log_table_test.rs
index 94a31d43..75327aa8 100644
--- a/crates/paimon/tests/audit_log_table_test.rs
+++ b/crates/paimon/tests/audit_log_table_test.rs
@@ -237,6 +237,17 @@ async fn audit_log_delta_scan_preserves_pk_row_kinds() {
("-U".to_string(), 2, 20),
]
);
+
+ let splits = plan.data_splits();
+ assert!(splits.iter().all(|split| split.is_streaming()));
+ let read = table.new_read_builder().new_read().unwrap();
+ let direct: Vec<RecordBatch> = read
+ .to_arrow_with_row_kind(&splits)
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ assert_eq!(collect_audit_rows(&direct), collect_audit_rows(&batches));
}
#[test]
@@ -318,6 +329,26 @@ async fn audit_log_exposes_sequence_number_when_enabled() {
let rows = collect_audit_rows_with_sequence(&batches);
assert_eq!(rows.len(), 2);
assert!(rows.iter().all(|(_, seq, _, _)| *seq >= 0));
+
+ let direct: Vec<RecordBatch> = table
+ .new_read_builder()
+ .new_read()
+ .unwrap()
+ .to_arrow_with_row_kind(&plan.data_splits())
+ .unwrap()
+ .try_collect()
+ .await
+ .unwrap();
+ assert_eq!(
+ direct[0]
+ .schema()
+ .fields()
+ .iter()
+ .map(|field| field.name().clone())
+ .collect::<Vec<_>>(),
+ vec!["rowkind".to_string(), "id".to_string(), "value".to_string()],
+ "the Python-compatible row-kind read must not leak the audit-only
sequence column",
+ );
}
#[tokio::test]
