sunchao commented on code in PR #5365: URL: https://github.com/apache/datafusion-comet/pull/5365#discussion_r3986708322
########## native/core/src/execution/planner/delta_spark_scan.rs: ########## @@ -0,0 +1,792 @@ +// 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. + +//! JVM-planned Delta handler for the generic `OpStruct::ContribScan` dispatcher, feature-gated +//! behind `delta`. +//! +//! delta-spark has already done log replay, snapshot resolution, and partition pruning by the +//! time the scan reaches Comet, so the envelope carries a concrete file list (plus deletion +//! vector descriptors) and the read path reuses the exact same shared parquet scan builder as +//! `NativeScan` -- inheriting row-group stats pruning, page-index pruning, and filter pushdown. +//! Sibling of the kernel-planned handler in `delta_scan.rs`; the two claim different +//! `type_url`s within the same `ContribScan` envelope. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::execution::object_store::ObjectStoreUrl; +use object_store::path::Path; +use object_store::ObjectStore; +use url::Url; + +use datafusion_comet_proto::spark_operator::{ + ContribScan, DeltaSparkScan, Operator, SparkFilePartition, SparkPartitionedFile, +}; +use prost::Message; + +use crate::execution::operators::ExecutionError; +use crate::execution::operators::ExecutionError::GeneralError; +use crate::execution::planner::PhysicalPlanner; +use crate::execution::planner::PlanCreationResult; +use crate::parquet::objectstore::s3_blob_fs_support::normalize_object_store_url; +use crate::parquet::parquet_support::{ + hash_object_store_configs, object_store_url_key, prepare_object_store_with_config_hash, +}; + +/// Type name the JVM-planned Delta contrib claims within the `ContribScan` envelope. The +/// contrib jar packs a `DeltaSparkScan` with a `type_url` of +/// `type.googleapis.com/comet.contrib.delta_spark.DeltaSparkScan`; dispatch keys on the +/// contrib-owned suffix, same convention as the kernel path's `delta_scan.rs`. +const DELTA_SPARK_SCAN_TYPE_NAME: &str = "comet.contrib.delta_spark.DeltaSparkScan"; + +/// Contrib entry point for the `OpStruct::ContribScan` dispatcher. Returns `Some(result)` when +/// the envelope carries a JVM-planned Delta scan, or `None` when the `type_url` belongs to some +/// other contrib. +pub(crate) fn try_plan_contrib_scan( + planner: &PhysicalPlanner, + spark_plan: &Operator, + contrib: &ContribScan, +) -> Option<PlanCreationResult> { + if !contrib.type_url.ends_with(DELTA_SPARK_SCAN_TYPE_NAME) { + return None; + } + Some( + DeltaSparkScan::decode(contrib.value.as_slice()) + .map_err(|e| { + GeneralError(format!( + "Failed to decode DeltaSparkScan from contrib_scan: {e}" + )) + }) + .and_then(|scan| plan_delta_spark_scan(planner, spark_plan, &scan)), + ) +} + +fn plan_delta_spark_scan( + planner: &PhysicalPlanner, + spark_plan: &Operator, + scan: &DeltaSparkScan, +) -> PlanCreationResult { + // Delta data files are plain parquet; the read path deliberately reuses + // the same shared parquet scan builder as NativeScan so Delta inherits + // row-group stats pruning, page-index pruning, and filter pushdown. Only + // the file list arrives in Delta-specific form. Note delta_common's + // column_mapping_mode is informational in M1: the actual field-id + // matching switch is common.use_field_id, same as the Iceberg path. + let common = scan + .common + .as_ref() + .ok_or_else(|| GeneralError("DeltaSparkScan missing common data".into()))?; + + let delta_partition = scan + .file_partition + .as_ref() + .ok_or_else(|| GeneralError("DeltaSparkScan missing file_partition".into()))?; + + let spark_partition = SparkFilePartition { + partitioned_file: delta_partition + .partitioned_file + .iter() + .map(|f| { + f.file.clone().ok_or_else(|| { + GeneralError("DeltaSparkPartitionedFile missing inner file".into()) + }) + }) + .collect::<Result<Vec<_>, _>>()?, + }; + + // Defense-in-depth against a stale or bypassed JVM gate: DeltaScanSupport.declineReason + // (multiStoreReason) already declines data files spanning multiple object-store authorities + // at planning time, but prepare_scan_store_and_files below resolves this whole partition's + // ObjectStoreUrl from the FIRST file only and then strips every other file down to its bare + // object-store path -- a file that actually lives under a different authority would + // silently read through the first file's store handle. Checked here rather than inside + // prepare_scan_store_and_files itself, which is shared with plain NativeScan and out of + // scope for this Delta-specific invariant. + check_same_object_store_authority(&spark_partition.partitioned_file)?; + + let (object_store_url, mut files) = + planner.prepare_scan_store_and_files(common, &spark_partition)?; + + // Translate deletion vectors into per-file ParquetAccessPlans so deleted + // rows are skipped inside the reader (composing, by intersection, with + // page-index pruning). Fetching the bitmaps and footers is async I/O; + // create_plan runs on the JNI task thread outside the tokio context, so + // block_on here is safe and keeps the scan a plain DataSourceExec. + if delta_partition + .partitioned_file + .iter() + .any(|f| f.dv.is_some()) + { + let object_store_options: HashMap<String, String> = common + .object_store_options + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + let runtime_env = planner.session_ctx.runtime_env(); + // `object_store_options` is the same map for every file this partition resolves a store + // for, so its hash is loop-invariant too: computed once here rather than once per file + // inside `prepare_object_store_with_configs`. + let object_store_config_hash = hash_object_store_configs(&object_store_options); + + // Resolve every object store this partition's files touch -- the data + // files' shared authority (check_same_object_store_authority above + // has already verified every data file in this partition resolves to + // the same authority, so prepare_scan_store_and_files's + // first-file-only resolution is safe here) plus any on-disk deletion + // vector's authority, which may legitimately differ from the data + // files' and carries its own resolved store -- before entering the + // async DV runtime below. This MUST happen here, on the JNI thread + // outside the tokio runtime: + // constructing a cold S3 store issues its own internal + // Handle::block_on calls (credential-provider / bucket-region + // resolution), which panics when nested inside the + // get_runtime().block_on(...) a few lines down. See + // delta_dv::attach_access_plans's doc comment for the invariant this + // maintains -- the async path never builds a store. + let mut resolved_stores: HashMap<ObjectStoreUrl, Arc<dyn ObjectStore>> = HashMap::new(); + // Tracks, per resolved ObjectStoreUrl, the userinfo (and raw URL, for the error + // message) of the first URL that resolved to it. This closure is the ONE place in this + // scan that sees both data-file and deletion-vector URLs together, so the + // store-identity collision check (see check_store_identity's doc comment) lives here + // rather than as an extension of check_same_object_store_authority above, which sees + // only data files and would hard-error the legitimate cross-bucket DV shape. + let mut store_identities: HashMap<ObjectStoreUrl, (String, String)> = HashMap::new(); + let mut resolve_store = + |url: String| -> Result<(Path, Arc<dyn ObjectStore>), ExecutionError> { + let parsed_url = Url::parse(&url).map_err(|e| { + GeneralError(format!( + "Error parsing URL {}: {e}", + redacted_url_display(&url) + )) + })?; + let user_info = url_user_info(&parsed_url); + // The same s3a/alias rewrite the shared resolution path applies first, so the + // cheap key and cached-hit path below agree with what + // prepare_object_store_with_config_hash registers under. Re-parses the already + // parsed string only so the parse error above stays redacted. + let normalized = normalize_object_store_url(&url, &object_store_options)?; + + // Cheap, I/O-free cache key: no config hashing, no global object-store-cache + // lock, no runtime_env registration. Checked against the LOCAL `resolved_stores` + // map below before ever paying for the expensive resolution path -- most files + // in a partition share the same authority as an already-resolved file. + let (url_key, _is_hdfs_scheme) = object_store_url_key(&normalized); + let parsed_url = normalized.url; + let store_url = ObjectStoreUrl::parse(url_key)?; + check_store_identity(&store_url, &user_info, &url, &mut store_identities)?; + if let Some(store) = resolved_stores.get(&store_url) { Review Comment: ### Performance **[P2] Preserve local store-cache hits with isolated registration URLs** For a partition containing deletion vectors, `object_store_url_key` still produces the physical key (for example, `s3://bucket`), but [`prepare_object_store_with_config_hash`](https://github.com/apache/datafusion-comet/blob/3f86d7ec4865ec5c59bd9560a2f37686932e034e/native/core/src/parquet/parquet_support.rs#L775-L797) now returns `s3+comet-<hash>-native://bucket`. Line 207 inserts that returned key, so the next data-file or external-DV reference to the same bucket always misses this lookup. Every reference then takes the global cache read lock, registers the store again in the runtime and looks it up again. The intended once-per-store local memoization is lost. Native `file://` is the exception because its registration key is unchanged. Please use the same backend-aware identity for both lookup and insertion, preserving the distinction between native and Hadoop stores, and add a repeated-resolution regression. A source-extracted closure probe with in-memory URI/runtime doubles confirmed 16 slow-path calls for 16 same-store references, compared with one under the old registration-key control. The global cache still reuses the backend, so this is repeated planning work rather than evidence of additional storage I/O. -- 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]
