dwsmith1983 commented on code in PR #5365: URL: https://github.com/apache/datafusion-comet/pull/5365#discussion_r3826290502
########## native/core/src/execution/delta_dv.rs: ########## @@ -0,0 +1,587 @@ +// 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. + +//! Delta Lake deletion-vector decoding and translation into DataFusion +//! [`ParquetAccessPlan`]s (feature = "delta"). +//! +//! Wire formats implemented here (from delta-spark's `DeletionVectorStore` / +//! `RoaringBitmapArray`, v3.3.2): +//! - On-disk DV file: 1 version byte at the start of the file; at +//! `descriptor.offset`: `[i32 BE size][data: size bytes][i32 BE CRC32(data)]`. +//! - `data`: `[i32 LE magic]` then either +//! - magic 1681511376 ("native"): `[i32 LE count]`, then per bitmap +//! `[i32 LE size][standard 32-bit RoaringBitmap]`, keys implicit (index); +//! - magic 1681511377 ("portable", the spec's 64-bit extension): `[i64 LE +//! count]`, then per bitmap `[i32 LE key][standard 32-bit RoaringBitmap]` +//! with keys ascending -- exactly [`RoaringTreemap`]'s serialized form. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata; +use datafusion::datasource::physical_plan::parquet::ParquetAccessPlan; +use datafusion::execution::runtime_env::RuntimeEnv; +use futures::{StreamExt, TryStreamExt}; +use object_store::ObjectStoreExt; +use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; +use parquet::file::metadata::PageIndexPolicy; +use roaring::{RoaringBitmap, RoaringTreemap}; + +use crate::execution::operators::ExecutionError; +use crate::execution::operators::ExecutionError::GeneralError; +use crate::parquet::parquet_support::prepare_object_store_with_configs; +use datafusion_comet_proto::spark_operator::DeltaSparkDvDescriptor; + +const NATIVE_MAGIC: i32 = 1681511376; +const PORTABLE_MAGIC: i32 = 1681511377; + +/// Unframe a DV blob read from `descriptor.offset` of a DV file: +/// `[i32 BE size][data][i32 BE crc]`. Verifies both the size against the +/// descriptor's `size_in_bytes` and the CRC32 checksum. +pub fn unframe_dv_blob(blob: &[u8], expected_size: usize) -> Result<&[u8], ExecutionError> { + if blob.len() < 8 { + return Err(GeneralError(format!( + "Deletion vector blob too short: {} bytes", + blob.len() + ))); + } + let size = i32::from_be_bytes(blob[0..4].try_into().unwrap()); + if size < 0 || size as usize != expected_size { + return Err(GeneralError(format!( + "Deletion vector size mismatch: file says {size}, descriptor says {expected_size}" + ))); + } + let end = 4 + size as usize; + if blob.len() < end + 4 { + return Err(GeneralError(format!( + "Deletion vector blob truncated: need {} bytes, have {}", + end + 4, + blob.len() + ))); + } + let data = &blob[4..end]; + let expected_crc = i32::from_be_bytes(blob[end..end + 4].try_into().unwrap()); + let actual_crc = crc32fast::hash(data) as i32; + if expected_crc != actual_crc { + return Err(GeneralError( + "Deletion vector checksum mismatch".to_string(), + )); + } + Ok(data) +} + +/// Deserialize the magic-prefixed RoaringBitmapArray into a 64-bit treemap of +/// deleted row indexes. +pub fn deserialize_dv_bitmap(data: &[u8]) -> Result<RoaringTreemap, ExecutionError> { + if data.len() < 4 { + return Err(GeneralError( + "Deletion vector bitmap too short for magic number".to_string(), + )); + } + let magic = i32::from_le_bytes(data[0..4].try_into().unwrap()); + let rest = &data[4..]; + match magic { + PORTABLE_MAGIC => RoaringTreemap::deserialize_from(rest) + .map_err(|e| GeneralError(format!("Invalid portable deletion vector bitmap: {e}"))), + NATIVE_MAGIC => { + if rest.len() < 4 { + return Err(GeneralError( + "Native deletion vector bitmap missing count".to_string(), + )); + } + let count = i32::from_le_bytes(rest[0..4].try_into().unwrap()); + if count < 0 { + return Err(GeneralError(format!( + "Invalid RoaringBitmapArray length ({count} < 0)" + ))); + } + let mut pos = 4usize; + let mut treemap = RoaringTreemap::new(); + for key in 0..count as u64 { + if rest.len() < pos + 4 { + return Err(GeneralError( + "Native deletion vector bitmap truncated".to_string(), + )); + } + let size = i32::from_le_bytes(rest[pos..pos + 4].try_into().unwrap()); + pos += 4; + if size < 0 || rest.len() < pos + size as usize { + return Err(GeneralError( + "Native deletion vector bitmap truncated".to_string(), + )); + } + let bitmap = RoaringBitmap::deserialize_from(&rest[pos..pos + size as usize]) + .map_err(|e| { + GeneralError(format!("Invalid deletion vector sub-bitmap: {e}")) + })?; + pos += size as usize; + for value in bitmap { + treemap.insert((key << 32) | value as u64); + } + } + Ok(treemap) + } + other => Err(GeneralError(format!( + "Unexpected RoaringBitmapArray magic number {other}" + ))), + } +} + +/// Translate deleted row indexes into a [`ParquetAccessPlan`]: fully-deleted +/// row groups become `Skip`, untouched groups stay `Scan`, and partially +/// deleted groups get a `RowSelection` selecting the complement of the deleted +/// rows. Page-index pruning later INTERSECTS with these selections, so DV +/// skips and page skips compose. +pub fn build_access_plan( + row_group_row_counts: &[i64], + deleted: &RoaringTreemap, +) -> Result<ParquetAccessPlan, ExecutionError> { + let mut plan = ParquetAccessPlan::new_all(row_group_row_counts.len()); + // Single sweep over the (sorted) deleted row indexes, bucketing by row group. + let mut deleted_iter = deleted.iter().peekable(); + let mut group_start = 0u64; + for (idx, &num_rows) in row_group_row_counts.iter().enumerate() { + let num_rows = num_rows as u64; + let group_end = group_start + num_rows; + let mut selectors: Vec<RowSelector> = Vec::new(); + let mut cursor = group_start; + let mut deleted_in_group = 0u64; + while let Some(&row) = deleted_iter.peek() { + if row >= group_end { + break; + } + deleted_iter.next(); + deleted_in_group += 1; + if row > cursor { + selectors.push(RowSelector::select((row - cursor) as usize)); + } + // Merge runs of consecutive deleted rows into one skip. + match selectors.last_mut() { + Some(last) if last.skip => last.row_count += 1, + _ => selectors.push(RowSelector::skip(1)), Review Comment: Two layers now. At planning, `spark.comet.scan.deltaNative.dv.maxDeletedRowsPerFile` (default 1M) declines pathological DVs before they reach native; the bound is cardinality-based so it's deliberately pessimistic, a contiguous delete declines the same as an alternating one, but run structure isn't knowable without decoding bitmaps on the driver, and the knob makes it recoverable. Natively, the exact selector bytes (derived from `size_of::<RowSelector>()`) are reserved against the memory pool, released with the plan, with multi-file accumulation and release-on-failure covered by tests. I looked at building the access plan at file-open time as you suggested, but DataFusion 54.1's ParquetOpener reads the plan out of the file extensions in `create_initial_plan` with no hook to compute it lazily, so that needs an upstream API first. -- 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]
