brgr-s commented on code in PR #2961:
URL: https://github.com/apache/iceberg-rust/pull/2961#discussion_r3728352412
##########
crates/iceberg/src/arrow/reader/row_filter.rs:
##########
@@ -62,8 +67,131 @@ impl ArrowReader {
// creates the projection mask for the Arrow predicates.
let projection_mask = ProjectionMask::leaves(parquet_schema,
column_indices.clone());
let predicate_func = visit(&mut converter, predicates)?;
- let arrow_predicate = ArrowPredicateFn::new(projection_mask,
predicate_func);
- Ok(RowFilter::new(vec![Box::new(arrow_predicate)]))
+ Ok(Box::new(ArrowPredicateFn::new(
+ projection_mask,
+ predicate_func,
+ )))
+ }
+
+ /// Builds one Arrow row-filter predicate per equality-delete set. The
predicate is based
+ /// on a hash-set lookup (see `EqDeleteSet`). It keeps a row unless its
key tuple is present
+ /// in that set. A row is deleted when it matches any set (the predicates
are AND-ed by the `RowFilter`).
+ pub(super) fn build_equality_delete_predicates(
+ sets: &[Arc<EqDeleteSet>],
+ parquet_schema: &SchemaDescriptor,
+ arrow_schema: &ArrowSchemaRef,
+ use_position_fallback: bool,
+ ) -> Result<Vec<Box<dyn ArrowPredicate>>> {
+ let field_id_map =
+ Self::resolve_field_id_map(parquet_schema, arrow_schema,
use_position_fallback)?;
+
+ let mut predicates: Vec<Box<dyn ArrowPredicate>> = Vec::new();
+ for set in sets {
+ if set.is_empty() {
+ continue;
+ }
+
+ // Parquet leaf index for each key column, in `fields` order; a
column dropped
+ // from this file by schema evolution has no entry.
+ let leaf_indices: Vec<Option<usize>> = set
+ .fields
+ .iter()
+ .map(|(_, id, _)| field_id_map.get(id).copied())
+ .collect();
+
+ let mut column_indices: Vec<usize> =
leaf_indices.iter().flatten().copied().collect();
+ column_indices.sort_unstable();
+ column_indices.dedup();
+ let projection_mask = ProjectionMask::leaves(parquet_schema,
column_indices.clone());
+
+ // Position of each key column within the projected batch
(parquet-rs presents the
+ // masked leaves in ascending leaf-index order).
+ let batch_positions: Vec<Option<usize>> = leaf_indices
+ .iter()
+ .map(|leaf| leaf.and_then(|idx|
column_indices.binary_search(&idx).ok()))
+ .collect();
+
+ let target_types: Vec<Type> = set.fields.iter().map(|(_, _, ty)|
ty.clone()).collect();
+ let num_cols = set.fields.len();
+ let set = set.clone();
+
+ let predicate_func =
+ move |batch: RecordBatch| -> std::result::Result<BooleanArray,
ArrowError> {
+ let num_rows = batch.num_rows();
+
+ // Change each key column into `Datum`s once, promoting to
the
+ // table type so the keys match the parsed delete keys
under schema
+ // evolution. A column absent from this file reads as
all-null.
+ let mut columns: Vec<Vec<Option<Datum>>> =
Vec::with_capacity(num_cols);
+ for (i, target_type) in target_types.iter().enumerate() {
+ let Some(pos) = batch_positions[i] else {
+ columns.push(vec![None; num_rows]);
+ continue;
+ };
+ let array = batch.column(pos);
+ let source_type = arrow_type_to_type(array.data_type())
+ .map_err(|e|
ArrowError::ComputeError(e.to_string()))?;
+ let source_primitive = source_type
+ .as_primitive_type()
+ .ok_or_else(|| {
+ ArrowError::ComputeError(
+ "equality delete key column is not a
primitive type"
+ .to_string(),
+ )
+ })?
+ .clone();
+ let needs_promotion = source_type != *target_type;
+ let literals = arrow_primitive_to_literal(array,
&source_type)
+ .map_err(|e|
ArrowError::ComputeError(e.to_string()))?;
+
+ let mut column = Vec::with_capacity(num_rows);
+ for literal in literals {
+ let datum = match literal {
+ Some(literal) => {
+ let primitive =
+
literal.as_primitive_literal().ok_or_else(|| {
+ ArrowError::ComputeError(
+ "failed to convert to
primitive literal"
+ .to_string(),
+ )
+ })?;
+ let datum =
Datum::new(source_primitive.clone(), primitive);
+ let datum = if needs_promotion {
+ datum
+ .to(target_type)
+ .map_err(|e|
ArrowError::ComputeError(e.to_string()))?
+ } else {
+ datum
+ };
+ Some(datum)
+ }
+ None => None,
+ };
+ column.push(datum);
+ }
+ columns.push(column);
+ }
+
+ // One hash lookup per row.
+ let mut keep = Vec::with_capacity(num_rows);
+ let mut probe = EqDeleteKey(vec![None; num_cols]);
+ for row in 0..num_rows {
+ for (i, column) in columns.iter_mut().enumerate() {
+ // we can `take` because each cell is probed once.
+ probe.0[i] = std::mem::take(&mut column[row]);
+ }
+ keep.push(!set.keys.contains(&probe));
+ }
+ Ok(BooleanArray::from(keep))
+ };
Review Comment:
Regarding `RowConverter`, I am hesitant to put this in this PR. It is not a
small change and I think it should propably go in a follow-up, with benchmarks
that prove that it improves performance beyond what this PR does. Same
reasoning for not trying to remove `arrow_primitive_to_literal`.
For double pass and the double String allocation, I have an idea that I will
push today.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
