From 9c2c0652cb9ccf15b3f99e35a3d51f215970c87c Mon Sep 17 00:00:00 2001 From: liuzhuo Date: Wed, 28 Oct 2026 11:33:53 +0800 Subject: [PATCH] [rust] Decode only selected columns in remote Arrow logs Push column selection into Arrow IPC decoding and remap historical schema alignment to the selected output. Preserve row and batch scan semantics, missing columns, and full-column decoding. Generated-by: OpenAI Codex (GPT-6) --- .../src/client/table/log_fetch_buffer.rs | 9 +- .../crates/fluss/src/client/table/scanner.rs | 90 ++++-- fluss-rust/crates/fluss/src/record/arrow.rs | 273 ++++++++++++++++-- 3 files changed, 327 insertions(+), 45 deletions(-) diff --git a/fluss-rust/crates/fluss/src/client/table/log_fetch_buffer.rs b/fluss-rust/crates/fluss/src/client/table/log_fetch_buffer.rs index 6cfd31846c9..5a1126e1b48 100644 --- a/fluss-rust/crates/fluss/src/client/table/log_fetch_buffer.rs +++ b/fluss-rust/crates/fluss/src/client/table/log_fetch_buffer.rs @@ -1551,7 +1551,14 @@ mod tests { tables.push(table); } let target = &tables[2]; - for fields in [vec![0], vec![2], vec![2, 0], vec![1, 0], vec![2, 1, 0]] { + for fields in [ + vec![0], + vec![2], + vec![2, 0], + vec![1, 0], + vec![2, 1, 0], + vec![2, 0, 2], + ] { let full_schema = to_arrow_schema(target.get_row_type())?; let output_schema = ReadContext::project_schema(full_schema.clone(), &fields)?; let row_type = Arc::new(target.get_row_type().project(&fields)?); diff --git a/fluss-rust/crates/fluss/src/client/table/scanner.rs b/fluss-rust/crates/fluss/src/client/table/scanner.rs index 81bf043cf20..1fc3dc60321 100644 --- a/fluss-rust/crates/fluss/src/client/table/scanner.rs +++ b/fluss-rust/crates/fluss/src/client/table/scanner.rs @@ -2745,6 +2745,26 @@ mod tests { } async fn check_poll_remote_projection(batch_mode: bool) -> Result<()> { + for compression in ["NONE", "LZ4_FRAME", "ZSTD"] { + for fields in [ + vec![2, 1, 0], + vec![0], + vec![1], + vec![2], + vec![2, 0, 2], + vec![0, 1, 2], + ] { + check_poll_remote_projection_case(batch_mode, &fields, compression).await?; + } + } + Ok(()) + } + + async fn check_poll_remote_projection_case( + batch_mode: bool, + fields: &[usize], + compression: &str, + ) -> Result<()> { use crate::metadata::DataField; use crate::row::DataGetters; use crate::test_utils::build_table_info_with_columns; @@ -2771,10 +2791,14 @@ mod tests { ], ); target.schema_id = 2; - let mut builder = MemoryLogRecordsArrowBuilder::new( - uncompressed_arrow_batch_config(1, source.get_row_type(), usize::MAX), - false, + let mut config = uncompressed_arrow_batch_config(1, source.get_row_type(), usize::MAX); + config.compression = crate::compression::ArrowCompressionInfo::from_conf( + &std::collections::HashMap::from([( + "table.log.arrow.compression.type".to_string(), + compression.to_string(), + )]), )?; + let mut builder = MemoryLogRecordsArrowBuilder::new(config, false)?; for name in [Some("alice"), None] { let mut row = GenericRow::new(2); row.set_field(0, 7_i32); @@ -2796,7 +2820,7 @@ mod tests { Arc::clone(&metadata), Arc::new(RpcClient::new()), &Config::default(), - Some(vec![2, 1, 0]), + Some(fields.to_vec()), fixed_schema, None, test_admin(&metadata), @@ -2809,7 +2833,7 @@ mod tests { Arc::clone(&metadata), Arc::clone(&inner.log_scanner_status), &Config::default(), - Some(vec![2, 1, 0]), + Some(fields.to_vec()), fixed_schema, None, test_scanner_metrics(&table_path), @@ -2846,23 +2870,32 @@ mod tests { let batch = batches[0].batch(); assert_eq!( batch.schema(), - to_arrow_schema(&target.get_row_type().project(&[2, 1, 0])?)? + to_arrow_schema(&target.get_row_type().project(fields)?)? ); assert_eq!(batch.num_rows(), 2); - assert_eq!(batch.column(0).null_count(), 2); - let names = batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(names.value(0), "alice"); - assert!(names.is_null(1)); - let ids = batch - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(ids.values().as_ref(), &[7, 7]); + for (column, field) in fields.iter().enumerate() { + match field { + 0 => { + let ids = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values().as_ref(), &[7, 7]); + } + 1 => { + let names = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(names.value(0), "alice"); + assert!(names.is_null(1)); + } + 2 => assert_eq!(batch.column(column).null_count(), 2), + _ => unreachable!(), + } + } assert!(scanner.poll(Duration::ZERO).await?.is_empty()); } else { let scanner = LogScanner { @@ -2873,11 +2906,20 @@ mod tests { assert_eq!(records.len(), 2); for (index, record) in records.iter().enumerate() { assert_eq!(record.offset(), index as i64); - assert!(record.row.is_null_at(0)?); - assert_eq!(record.row.is_null_at(1)?, index == 1); - assert_eq!(record.row.get_int(2)?, 7); + for (column, field) in fields.iter().enumerate() { + match field { + 0 => assert_eq!(record.row.get_int(column)?, 7), + 1 => { + assert_eq!(record.row.is_null_at(column)?, index == 1); + if index == 0 { + assert_eq!(record.row.get_string(column)?, "alice"); + } + } + 2 => assert!(record.row.is_null_at(column)?), + _ => unreachable!(), + } + } } - assert_eq!(records[0].row.get_string(1)?, "alice"); assert!(scanner.poll(Duration::ZERO).await?.is_empty()); } assert!(fetcher.resolver.resolve(1, true).is_some()); diff --git a/fluss-rust/crates/fluss/src/record/arrow.rs b/fluss-rust/crates/fluss/src/record/arrow.rs index 06597411b46..c9d437d7a3d 100644 --- a/fluss-rust/crates/fluss/src/record/arrow.rs +++ b/fluss-rust/crates/fluss/src/record/arrow.rs @@ -1233,6 +1233,7 @@ pub struct ReadContext { is_from_remote: bool, fluss_row_type: Option>, schema_alignment: Option>, + source_projection: Option>, } #[derive(Clone)] @@ -1259,6 +1260,7 @@ impl ReadContext { is_from_remote, fluss_row_type: None, schema_alignment: None, + source_projection: None, } } @@ -1313,6 +1315,39 @@ impl ReadContext { }), "schema alignment source and target types must match" ); + let schema_alignment = if self.is_from_remote { + let source_projection: Vec<_> = schema_alignment + .iter() + .filter(|&&index| index != UNEXIST_MAPPING) + .map(|&index| index as usize) + .collect(); + if source_projection + .iter() + .copied() + .eq(0..self.full_schema.fields().len()) + { + schema_alignment + } else { + // Alignment indexes address the projected decoder output, not the source schema. + let mut projected_index = 0; + let alignment = schema_alignment + .iter() + .map(|&index| { + if index == UNEXIST_MAPPING { + index + } else { + let index = projected_index; + projected_index += 1; + index + } + }) + .collect::>(); + self.source_projection = Some(source_projection); + Arc::from(alignment) + } + } else { + schema_alignment + }; self.target_schema = target_schema; self.schema_alignment = Some(schema_alignment); self @@ -1330,8 +1365,7 @@ impl ReadContext { // the logic is little bit of hard to understand, to refactor it to follow // java side let (need_do_reorder, sorted_fields) = { - // currently, for remote read, arrow log doesn't support projection pushdown, - // so, only need to do reordering when is not from remote + // Remote decoding preserves projection order; server responses use field order. if !is_from_remote { let mut sorted_fields = projected_fields.clone(); sorted_fields.sort_unstable(); @@ -1382,6 +1416,17 @@ impl ReadContext { } }; + let source_projection = if is_from_remote + && !project + .projected_fields + .iter() + .copied() + .eq(0..arrow_schema.fields().len()) + { + Some(project.projected_fields.clone()) + } else { + None + }; Ok(ReadContext { target_schema, full_schema: arrow_schema, @@ -1390,6 +1435,7 @@ impl ReadContext { is_from_remote, fluss_row_type: None, schema_alignment: None, + source_projection, }) } @@ -1427,11 +1473,15 @@ impl ReadContext { .map(|p| p.ordered_fields.as_slice()) } + fn remote_decode_projection(&self) -> Option<&[usize]> { + self.source_projection.as_deref() + } + pub(crate) fn record_batch(&self, data: Bytes) -> Result { let (batch_metadata, body_buffer, version) = parse_ipc_message(&data)?; let resolve_schema = { - // if from remote, no projection, need to use full schema + // Remote IPC metadata describes all source fields, even with decoder projection. if self.is_from_remote || self.schema_alignment.is_some() { self.full_schema.clone() } else { @@ -1456,7 +1506,11 @@ impl ReadContext { batch_metadata, resolve_schema, &HashMap::new(), - None, + if self.is_from_remote { + self.remote_decode_projection() + } else { + None + }, &version, )?; @@ -1464,9 +1518,7 @@ impl ReadContext { Some(projection) => { let reordered_columns = { // need to do reorder - if self.is_from_remote { - Some(&projection.projected_fields) - } else if projection.reordering_needed { + if !self.is_from_remote && projection.reordering_needed { Some(&projection.reordering_indexes) } else { None @@ -1504,21 +1556,10 @@ impl ReadContext { batch_metadata, self.full_schema.clone(), &HashMap::new(), - None, + self.remote_decode_projection(), &version, )?; - let record_batch = match &self.projection { - Some(projection) => { - let projected_columns: Vec<_> = projection - .projected_fields - .iter() - .map(|&idx| record_batch.column(idx).clone()) - .collect(); - RecordBatch::try_new(self.target_schema.clone(), projected_columns)? - } - None => record_batch, - }; let record_batch = match &self.schema_alignment { Some(schema_alignment) => align_record_batch_to_schema( record_batch, @@ -1547,7 +1588,11 @@ fn align_record_batch_to_schema( } } - Ok(RecordBatch::try_new(target_schema, columns)?) + Ok(RecordBatch::try_new_with_options( + target_schema, + columns, + &arrow::array::RecordBatchOptions::new().with_row_count(Some(row_count)), + )?) } pub struct ArrowLogRecordIterator { @@ -2227,6 +2272,194 @@ mod tests { ); } + fn encode_projection_batch( + batch: &RecordBatch, + compression: Option, + ) -> Result { + let options = IpcWriteOptions::default().try_with_compression(compression)?; + let mut writer = + StreamWriter::try_new_with_options(Vec::new(), batch.schema().as_ref(), options)?; + let header = writer.get_ref().len(); + writer.write(batch)?; + Ok(Bytes::copy_from_slice(&writer.get_ref()[header..])) + } + + #[test] + fn remote_decode_projection_preserves_columns_and_rows() -> Result<()> { + let row_type = Arc::new(RowType::new(vec![ + DataField::new("id", DataTypes::int(), None), + DataField::new("tags", DataTypes::array(DataTypes::int()), None), + DataField::new("name", DataTypes::string(), None), + DataField::new( + "attributes", + DataTypes::map(DataTypes::string(), DataTypes::int()), + None, + ), + DataField::new( + "nested", + DataTypes::row(vec![ + DataField::new("seq", DataTypes::int(), None), + DataField::new("label", DataTypes::string(), None), + ]), + None, + ), + ])); + let schema = to_arrow_schema(&row_type)?; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(arrow::array::Int32Array::from(vec![Some(1), None])), + list_int_array(vec![Some(3), None, Some(5)], vec![0, 2, 3], None), + Arc::new(arrow::array::StringArray::from(vec![Some("hello"), None])), + map_string_int_array(vec!["a", "b"], vec![Some(1), None], vec![0, 1, 2], None), + struct_int_string_array( + vec![Some(1), None], + vec![Some("x"), None], + Some(vec![true, false]), + ), + ], + )?; + for compression in [ + None, + Some(CompressionType::LZ4_FRAME), + Some(CompressionType::ZSTD), + ] { + for rows in [0, 2] { + let batch = batch.slice(0, rows); + let data = encode_projection_batch(&batch, compression)?; + for fields in [ + vec![], + vec![0], + vec![2, 0], + vec![1], + vec![2, 1, 2], + vec![0, 1, 2, 3, 4], + vec![4, 3], + vec![3, 0, 4, 3], + ] { + let context = ReadContext::with_projection_pushdown( + schema.clone(), + Arc::new(row_type.project(&fields)?), + fields.clone(), + true, + )?; + let expected = batch.project(&fields)?; + assert_eq!(context.record_batch(data.clone())?, expected); + assert_eq!( + context.record_batch_for_remote_log(data.clone())?.unwrap(), + expected + ); + } + } + } + Ok(()) + } + + #[test] + fn remote_decode_projection_remaps_historical_schema() -> Result<()> { + let row_type = Arc::new(RowType::new(vec![ + DataField::new("first", DataTypes::int(), None), + DataField::new("unused", DataTypes::string(), None), + DataField::new("last", DataTypes::int(), None), + ])); + let schema = to_arrow_schema(&row_type)?; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(arrow::array::Int32Array::from(vec![1, 2])), + Arc::new(arrow::array::StringArray::from(vec!["unused", "unused"])), + Arc::new(arrow::array::Int32Array::from(vec![3, 4])), + ], + )?; + for mapping in [ + vec![], + vec![-1], + vec![-1, -1], + vec![2, -1, 0, 2], + vec![0, -1], + vec![2, 0], + vec![0, -1, 1, 2], + vec![-1, 0, 1, 2], + vec![0, 1, 2, -1], + ] { + let target = Arc::new(arrow_schema::Schema::new_with_metadata( + mapping + .iter() + .enumerate() + .map(|(i, &source)| { + let data_type = if source == UNEXIST_MAPPING { + ArrowDataType::Int32 + } else { + schema.field(source as usize).data_type().clone() + }; + Field::new(format!("target_{i}"), data_type, true) + }) + .collect::>(), + HashMap::from([("version".into(), "target".into())]), + )); + let context = ReadContext::new(schema.clone(), row_type.clone(), true) + .with_target_schema_alignment(target.clone(), Arc::from(mapping.clone())); + let expected = align_record_batch_to_schema(batch.clone(), target, &mapping)?; + for compression in [ + None, + Some(CompressionType::LZ4_FRAME), + Some(CompressionType::ZSTD), + ] { + let data = encode_projection_batch(&batch, compression)?; + assert_eq!(context.record_batch(data.clone())?, expected); + assert_eq!( + context.record_batch_for_remote_log(data)?.unwrap(), + expected + ); + } + } + Ok(()) + } + + #[test] + fn remote_decode_projection_skips_unselected_compressed_buffers() -> Result<()> { + let row_type = Arc::new(RowType::new(vec![ + DataField::new("id", DataTypes::int(), None), + DataField::new("unused", DataTypes::string(), None), + ])); + let schema = to_arrow_schema(&row_type)?; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(arrow::array::Int32Array::from(vec![7; 100])), + Arc::new(arrow::array::StringArray::from(vec![ + "repeated payload"; + 100 + ])), + ], + )?; + let data = encode_projection_batch(&batch, Some(CompressionType::ZSTD))?; + let (metadata, body, _) = parse_ipc_message(&data)?; + let buffers = metadata.buffers().unwrap(); + let buffer = buffers.get(buffers.len() - 1); + let start = data.len() - body.len() + buffer.offset() as usize; + let end = start + buffer.length() as usize; + let mut data = data.to_vec(); + // Keep IPC framing and the uncompressed length; invalidate only the unused payload. + data[start + 8..end].fill(0); + let data = Bytes::from(data); + let full = ReadContext::new(schema.clone(), row_type.clone(), true); + assert!(full.record_batch(data.clone()).is_err()); + let selected = ReadContext::with_projection_pushdown( + schema, + Arc::new(row_type.project(&[0])?), + vec![0], + true, + )?; + let expected = batch.project(&[0])?; + assert_eq!(selected.record_batch(data.clone())?, expected); + assert_eq!( + selected.record_batch_for_remote_log(data)?.unwrap(), + expected + ); + Ok(()) + } + fn single_int_read_context() -> (ReadContext, SchemaRef) { let row_type = Arc::new(RowType::new(vec![DataField::new( "id",