diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 85a2a0fa6a..7b6b9557aa 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -354,6 +354,7 @@ jobs: org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite org.apache.spark.sql.comet.CometTaskMetricsSuite + org.apache.spark.sql.comet.CometFileLocalityManagerSuite org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite org.apache.spark.sql.comet.PlanDataInjectorSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index abde1554f6..50b6c03859 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -170,6 +170,7 @@ jobs: org.apache.spark.sql.comet.CometTPCDSV1_4_PlanStabilitySuite org.apache.spark.sql.comet.CometTPCDSV2_7_PlanStabilitySuite org.apache.spark.sql.comet.CometTaskMetricsSuite + org.apache.spark.sql.comet.CometFileLocalityManagerSuite org.apache.spark.sql.comet.CometDppFallbackRepro3949Suite org.apache.spark.sql.comet.CometShuffleFallbackStickinessSuite org.apache.spark.sql.comet.PlanDataInjectorSuite diff --git a/native/Cargo.lock b/native/Cargo.lock index 1658c480a8..6568480927 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -1990,8 +1990,10 @@ dependencies = [ "aws-credential-types", "criterion", "datafusion", + "datafusion-comet-block-cache", "datafusion-comet-common", "datafusion-comet-jni-bridge", + "datafusion-comet-object-store-cache", "datafusion-comet-objectstore-hdfs", "datafusion-comet-proto", "datafusion-comet-shuffle", @@ -2035,6 +2037,16 @@ dependencies = [ "uuid", ] +[[package]] +name = "datafusion-comet-block-cache" +version = "1.0.0" +dependencies = [ + "async-trait", + "bytes", + "futures", + "tokio", +] + [[package]] name = "datafusion-comet-common" version = "1.0.0" @@ -2079,6 +2091,19 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "datafusion-comet-object-store-cache" +version = "1.0.0" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "datafusion-comet-block-cache", + "futures", + "object_store", + "tokio", +] + [[package]] name = "datafusion-comet-objectstore-hdfs" version = "1.0.0" diff --git a/native/Cargo.toml b/native/Cargo.toml index 3e797eb968..cea8183a2a 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -16,8 +16,8 @@ # under the License. [workspace] -default-members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle"] -members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "hdfs", "fs-hdfs"] +default-members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "block-cache", "object-store-cache"] +members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "hdfs", "fs-hdfs", "block-cache", "object-store-cache"] resolver = "2" [workspace.package] @@ -47,6 +47,8 @@ datafusion-comet-common = { path = "common" } datafusion-comet-jni-bridge = { path = "jni-bridge" } datafusion-comet-proto = { path = "proto" } datafusion-comet-shuffle = { path = "shuffle" } +datafusion-comet-block-cache = { path = "block-cache" } +datafusion-comet-object-store-cache = { path = "object-store-cache" } chrono = { version = "0.4", default-features = false, features = ["clock"] } chrono-tz = { version = "0.10" } futures = "0.3.32" diff --git a/native/block-cache/Cargo.toml b/native/block-cache/Cargo.toml new file mode 100644 index 0000000000..b35dd49697 --- /dev/null +++ b/native/block-cache/Cargo.toml @@ -0,0 +1,46 @@ +# 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. + +[package] +name = "datafusion-comet-block-cache" +description = "Apache DataFusion Comet: storage-API-neutral block cache core (memory + SSD tiers)" +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +readme = { workspace = true } +license = { workspace = true } +edition = { workspace = true } + +publish = false + +# NOTE: This crate is storage-API-neutral by construction. It MUST NOT depend on +# `object_store`, `opendal`, or any Comet crate. The caller supplies a `RangeFetcher`; +# the core never talks to storage itself. A dependency-tree check in CI guards this +# boundary (see the `dependency_tree_is_storage_neutral` test). +[dependencies] +async-trait = { workspace = true } +bytes = { workspace = true } +futures = { workspace = true } +tokio = { version = "1", features = ["sync"] } + +[dev-dependencies] +tokio = { version = "1", features = ["sync", "rt-multi-thread", "macros", "time"] } + +[lib] +name = "datafusion_comet_block_cache" +path = "src/lib.rs" diff --git a/native/block-cache/src/cache.rs b/native/block-cache/src/cache.rs new file mode 100644 index 0000000000..c3ff675103 --- /dev/null +++ b/native/block-cache/src/cache.rs @@ -0,0 +1,576 @@ +// 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. + +use std::collections::{BTreeSet, HashMap}; +use std::ops::Range; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::{Bytes, BytesMut}; +use futures::future::FutureExt; + +use crate::error::{CacheError, Result}; +use crate::metrics::{Metrics, MetricsSnapshot}; +use crate::sieve::{Block, BlockKey, FetchResult, InFlightFut, Shard}; +use crate::version::{FileKey, FileVersion}; + +/// Minimum / maximum / default block size (the read quantum). Powers of two only. +pub const MIN_BLOCK_SIZE: u64 = 1 << 20; // 1 MiB +pub const MAX_BLOCK_SIZE: u64 = 16 << 20; // 16 MiB +pub const DEFAULT_BLOCK_SIZE: u64 = 4 << 20; // 4 MiB +/// Default number of memory-tier shards. +pub const DEFAULT_NUM_SHARDS: usize = 16; +/// Default cap on a single coalesced upstream fetch (4 default blocks). +pub const DEFAULT_MAX_COALESCE_BYTES: u64 = 16 << 20; // 16 MiB + +/// Configuration for a [`BlockCache`]. +#[derive(Clone, Debug)] +pub struct BlockCacheConfig { + /// Block quantum in bytes. Clamped to a power of two in `[MIN_BLOCK_SIZE, MAX_BLOCK_SIZE]`. + pub block_size: u64, + /// Memory-tier budget in bytes, process-wide. + pub memory_budget: u64, + /// Number of memory-tier shards. + pub num_shards: usize, + /// Cap on bytes fetched in a single coalesced upstream request. + pub max_coalesce_bytes: u64, +} + +impl Default for BlockCacheConfig { + fn default() -> Self { + BlockCacheConfig { + block_size: DEFAULT_BLOCK_SIZE, + memory_budget: 512 << 20, + num_shards: DEFAULT_NUM_SHARDS, + max_coalesce_bytes: DEFAULT_MAX_COALESCE_BYTES, + } + } +} + +/// Round `v` down to the largest power of two `<= v`. +fn floor_pow2(v: u64) -> u64 { + if v == 0 { + return 0; + } + 1u64 << (63 - v.leading_zeros() as u64) +} + +impl BlockCacheConfig { + /// Normalize into a valid config: block size becomes a power of two within bounds, + /// shard count is at least 1, and the coalesce cap is at least one block. + fn normalized(mut self) -> Self { + let clamped = self.block_size.clamp(MIN_BLOCK_SIZE, MAX_BLOCK_SIZE); + self.block_size = floor_pow2(clamped).max(MIN_BLOCK_SIZE); + self.num_shards = self.num_shards.max(1); + self.max_coalesce_bytes = self.max_coalesce_bytes.max(self.block_size); + self + } +} + +/// Fetches absolute byte ranges from the underlying storage on a cache miss. +/// +/// The cache calls this exactly once per block per version regardless of how many tasks +/// concurrently miss it (single-flight). Implementations return the bytes for the +/// requested ranges plus the object version observed by the fetch, which the cache uses +/// to detect in-place overwrites. +#[async_trait] +pub trait RangeFetcher: Send + Sync { + async fn fetch(&self, ranges: &[Range]) -> Result<(Vec, FileVersion)>; +} + +/// Interned file identities and their captured versions. +struct FileTable { + ids: HashMap, + versions: HashMap, + next_id: u64, +} + +/// The decision made after comparing a fetched version against the stored one. +enum VersionDecision { + Unchanged, + FirstSeen, + Overwritten, +} + +/// A block-aligned local data cache (memory tier) sitting behind a caller-supplied +/// [`RangeFetcher`]. Storage-API-neutral: it knows nothing about `object_store`. +pub struct BlockCache { + block_size: u64, + num_shards: usize, + max_coalesce_blocks: u32, + shards: Vec>, + files: Mutex, + memory_budget: AtomicU64, + metrics: Arc, +} + +impl BlockCache { + /// Build a cache from `config` (normalized to valid values). + pub fn new(config: BlockCacheConfig) -> Arc { + let config = config.normalized(); + let per_shard_budget = config.memory_budget / config.num_shards as u64; + let shards = (0..config.num_shards) + .map(|_| Mutex::new(Shard::new(per_shard_budget))) + .collect(); + let max_coalesce_blocks = (config.max_coalesce_bytes / config.block_size).max(1) as u32; + Arc::new(BlockCache { + block_size: config.block_size, + num_shards: config.num_shards, + max_coalesce_blocks, + shards, + files: Mutex::new(FileTable { + ids: HashMap::new(), + versions: HashMap::new(), + next_id: 0, + }), + memory_budget: AtomicU64::new(config.memory_budget), + metrics: Arc::new(Metrics::default()), + }) + } + + /// The block quantum in bytes. + pub fn block_size(&self) -> u64 { + self.block_size + } + + /// A snapshot of the cache counters. + pub fn stats(&self) -> MetricsSnapshot { + self.metrics.snapshot() + } + + /// Serve `ranges` of `file`. Reads are quantized to blocks internally; misses go + /// through `fetcher` exactly once per block regardless of concurrent callers. Returns + /// one `Bytes` per input range, byte-for-byte identical to reading the store directly. + pub async fn get_ranges( + &self, + file: &FileKey, + ranges: &[Range], + fetcher: &dyn RangeFetcher, + ) -> Result> { + if ranges.is_empty() { + return Ok(Vec::new()); + } + let file_id = self.intern(file); + + // Union of blocks touched by any requested range. + let mut needed: BTreeSet = BTreeSet::new(); + for r in ranges { + if r.start >= r.end { + continue; + } + let first = (r.start / self.block_size) as u32; + let last = ((r.end - 1) / self.block_size) as u32; + for b in first..=last { + needed.insert(b); + } + } + + // Probe the memory tier. + let mut have: HashMap = HashMap::with_capacity(needed.len()); + let mut missing: Vec = Vec::new(); + for &b in &needed { + let key = (file_id, b); + let hit = self.shard(key).lock().unwrap().get(&key); + match hit { + Some(block) => { + self.metrics.record_hit(); + have.insert(b, block.data.clone()); + } + None => { + self.metrics.record_miss(); + missing.push(b); + } + } + } + + if !missing.is_empty() { + self.fill_missing(file_id, &missing, fetcher, &mut have) + .await?; + } + + // Assemble each requested range from the blocks now in `have`. + let mut out = Vec::with_capacity(ranges.len()); + for r in ranges { + out.push(self.assemble_range(r, &have)?); + } + Ok(out) + } + + /// Fetch and cache every missing block, deduplicating concurrent fetches (single-flight) + /// and coalescing runs of adjacent missing blocks into one upstream request. + async fn fill_missing( + &self, + file_id: u64, + missing: &[u32], + fetcher: &dyn RangeFetcher, + have: &mut HashMap, + ) -> Result<()> { + // --- claim phase: decide which blocks we own vs. wait on --- + let mut owned: Vec = Vec::new(); + let mut senders: HashMap> = HashMap::new(); + let mut waiters: Vec<(u32, InFlightFut)> = Vec::new(); + + let mut blocks = missing.to_vec(); + blocks.sort_unstable(); + blocks.dedup(); + + for &b in &blocks { + let key = (file_id, b); + let mut shard = self.shard(key).lock().unwrap(); + // Re-check: a concurrent task may have filled the block since our probe. + if let Some(block) = shard.get(&key) { + have.insert(b, block.data.clone()); + continue; + } + if let Some(fut) = shard.in_flight_get(&key) { + waiters.push((b, fut)); + continue; + } + // Claim: install a shared future others can await, and own the fetch. + let (tx, rx) = tokio::sync::oneshot::channel::(); + let fut: InFlightFut = async move { + match rx.await { + Ok(res) => res, + Err(_) => Err(CacheError::Internal( + "fetch owner dropped before delivering block".to_string(), + )), + } + } + .boxed() + .shared(); + shard.in_flight_insert(key, fut); + drop(shard); + owned.push(b); + senders.insert(b, tx); + } + + // --- fetch phase: our owned blocks, coalesced. Must happen BEFORE awaiting other + // owners' futures so that two callers cross-owning each other's blocks cannot + // deadlock (each fetches what it owns first, then waits). --- + let mut first_error: Option = None; + for run in coalesce_runs(&owned, self.max_coalesce_blocks) { + let start_block = run[0]; + let end_block = *run.last().unwrap(); + let abs_start = start_block as u64 * self.block_size; + // Over-read past EOF is fine: the store truncates a partially-out-of-bounds + // range, yielding the short final block. + let abs_end = (end_block as u64 + 1) * self.block_size; + + let fetch_range = abs_start..abs_end; + match fetcher.fetch(std::slice::from_ref(&fetch_range)).await { + Ok((bytes_vec, version)) => { + let full = concat_bytes(bytes_vec); + self.metrics.record_fetch(full.len() as u64); + self.reconcile_version(file_id, &version); + for &b in &run { + let off = ((b - start_block) as u64 * self.block_size) as usize; + let data = if off >= full.len() { + Bytes::new() + } else { + let end = (off + self.block_size as usize).min(full.len()); + full.slice(off..end) + }; + let block = Block::new(data); + self.publish_block(file_id, b, block, have, &mut senders); + } + } + Err(e) => { + // Fail every owned block in this run; remove the in-flight entries so + // the next caller retries. Errors are never cached. + for &b in &run { + let key = (file_id, b); + if let Some(tx) = senders.remove(&b) { + let _ = tx.send(Err(e.clone())); + } + self.shard(key).lock().unwrap().in_flight_remove(&key); + } + if first_error.is_none() { + first_error = Some(e); + } + } + } + } + + // --- wait phase: blocks other tasks own --- + for (b, fut) in waiters { + match fut.await { + Ok(block) => { + block.visited.store(true, Ordering::Relaxed); + have.insert(b, block.data.clone()); + } + Err(e) => { + if first_error.is_none() { + first_error = Some(e); + } + } + } + } + + match first_error { + Some(e) => Err(e), + None => Ok(()), + } + } + + /// Insert a fetched block into the memory tier, hand it to `have`, wake any waiters, + /// and clear the in-flight entry. + fn publish_block( + &self, + file_id: u64, + block_index: u32, + block: Arc, + have: &mut HashMap, + senders: &mut HashMap>, + ) { + let key = (file_id, block_index); + { + let mut shard = self.shard(key).lock().unwrap(); + shard.insert(key, Arc::clone(&block), &self.metrics); + shard.in_flight_remove(&key); + } + have.insert(block_index, block.data.clone()); + if let Some(tx) = senders.remove(&block_index) { + // A closed receiver just means no other task waited; ignore. + let _ = tx.send(Ok(block)); + } + } + + /// Build one requested range's bytes from the (now cached) blocks in `have`. + fn assemble_range(&self, r: &Range, have: &HashMap) -> Result { + if r.start >= r.end { + return Ok(Bytes::new()); + } + let bs = self.block_size; + let first = (r.start / bs) as u32; + let last = ((r.end - 1) / bs) as u32; + + if first == last { + // Single block: zero-copy slice. + let block = have + .get(&first) + .ok_or_else(|| CacheError::Internal("block missing during assembly".to_string()))?; + let block_start = first as u64 * bs; + let lo = (r.start - block_start) as usize; + let hi = ((r.end - block_start).min(block.len() as u64)) as usize; + if lo > block.len() { + return Err(CacheError::Internal( + "range start beyond block during assembly".to_string(), + )); + } + return Ok(block.slice(lo..hi)); + } + + let mut buf = BytesMut::with_capacity((r.end - r.start) as usize); + for b in first..=last { + let block = have + .get(&b) + .ok_or_else(|| CacheError::Internal("block missing during assembly".to_string()))?; + let block_start = b as u64 * bs; + let lo = r.start.max(block_start); + let hi = r.end.min(block_start + block.len() as u64); + if hi <= lo { + continue; + } + buf.extend_from_slice(&block[(lo - block_start) as usize..(hi - block_start) as usize]); + } + Ok(buf.freeze()) + } + + /// Compare a freshly fetched version with the stored one; on the first overwrite + /// detection, drop every cached block of the file and record the new version. + fn reconcile_version(&self, file_id: u64, version: &FileVersion) { + let decision = { + let mut files = self.files.lock().unwrap(); + match files.versions.get(&file_id) { + None => { + files.versions.insert(file_id, version.clone()); + VersionDecision::FirstSeen + } + Some(existing) if existing.matches(version) => VersionDecision::Unchanged, + Some(_) => { + files.versions.insert(file_id, version.clone()); + VersionDecision::Overwritten + } + } + }; + if let VersionDecision::Overwritten = decision { + self.invalidate_blocks(file_id); + self.metrics.record_invalidation(); + } + } + + /// Drop every cached block of `file_id` from all shards. + fn invalidate_blocks(&self, file_id: u64) { + for shard in &self.shards { + shard.lock().unwrap().remove_file_blocks(file_id); + } + } + + /// Drop all cached state for a file (used on `put`/`delete` at the wrapper). + pub fn invalidate_file(&self, file: &FileKey) { + let file_id = { + let files = self.files.lock().unwrap(); + files.ids.get(file).copied() + }; + if let Some(id) = file_id { + self.invalidate_blocks(id); + self.files.lock().unwrap().versions.remove(&id); + } + } + + /// Change the memory-tier budget at runtime. A reduction evicts (SIEVE order) until + /// under the new per-shard budget before returning. Phase-1 callers never invoke this; + /// it is the phase-2 unified-memory contract (grant/shrink) and is exercised by tests. + pub fn set_memory_budget(&self, bytes: u64) { + self.memory_budget.store(bytes, Ordering::Relaxed); + let per_shard = bytes / self.num_shards as u64; + for shard in &self.shards { + shard.lock().unwrap().set_budget(per_shard, &self.metrics); + } + } + + /// The current memory-tier budget in bytes. + pub fn memory_budget(&self) -> u64 { + self.memory_budget.load(Ordering::Relaxed) + } + + /// The version most recently captured for `file`, if the cache has ever fetched it. + /// + /// Used by the `object_store` wrapper to synthesize an `ObjectMeta` for a cache-served + /// `get_opts` without issuing a `head` request. Returns `None` before the first fetch. + pub fn file_version(&self, file: &FileKey) -> Option { + let files = self.files.lock().unwrap(); + let id = files.ids.get(file)?; + files.versions.get(id).cloned() + } + + /// Intern a file key into a compact id, assigning a new one on first sight. + fn intern(&self, file: &FileKey) -> u64 { + let mut files = self.files.lock().unwrap(); + if let Some(id) = files.ids.get(file) { + return *id; + } + let id = files.next_id; + files.next_id += 1; + files.ids.insert(file.clone(), id); + id + } + + /// The shard owning a block key. + fn shard(&self, key: BlockKey) -> &Mutex { + &self.shards[self.shard_index(key)] + } + + fn shard_index(&self, (file_id, block_index): BlockKey) -> usize { + // Cheap integer mix over the 12-byte key. + let mut h = file_id.wrapping_mul(0x9E37_79B9_7F4A_7C15); + h ^= (block_index as u64).wrapping_mul(0xD6E8_FEB8_6659_FD93); + h ^= h >> 29; + (h % self.num_shards as u64) as usize + } +} + +/// Concatenate the pieces returned by a fetch into one `Bytes` (zero-copy for the common +/// single-range fetch this cache issues). +fn concat_bytes(mut pieces: Vec) -> Bytes { + match pieces.len() { + 0 => Bytes::new(), + 1 => pieces.pop().unwrap(), + _ => { + let total: usize = pieces.iter().map(|p| p.len()).sum(); + let mut buf = BytesMut::with_capacity(total); + for p in pieces { + buf.extend_from_slice(&p); + } + buf.freeze() + } + } +} + +/// Group sorted, deduped block indices into runs of consecutive integers, splitting any +/// run longer than `max_blocks` so no single upstream fetch exceeds the coalesce cap. +fn coalesce_runs(sorted: &[u32], max_blocks: u32) -> Vec> { + let mut runs: Vec> = Vec::new(); + let mut cur: Vec = Vec::new(); + for &b in sorted { + let extends = match cur.last() { + Some(&last) => b == last + 1 && (cur.len() as u32) < max_blocks, + None => false, + }; + if !extends && !cur.is_empty() { + runs.push(std::mem::take(&mut cur)); + } + cur.push(b); + } + if !cur.is_empty() { + runs.push(cur); + } + runs +} + +#[cfg(test)] +mod unit_tests { + use super::*; + + #[test] + fn floor_pow2_rounds_down() { + assert_eq!(floor_pow2(1), 1); + assert_eq!(floor_pow2(3), 2); + assert_eq!(floor_pow2(5 << 20), 4 << 20); + assert_eq!(floor_pow2(16 << 20), 16 << 20); + } + + #[test] + fn config_normalizes_block_size_to_pow2_in_range() { + let c = BlockCacheConfig { + block_size: 5 << 20, + ..Default::default() + } + .normalized(); + assert_eq!(c.block_size, 4 << 20); + + let c = BlockCacheConfig { + block_size: 100 << 20, + ..Default::default() + } + .normalized(); + assert_eq!(c.block_size, MAX_BLOCK_SIZE); + + let c = BlockCacheConfig { + block_size: 0, + ..Default::default() + } + .normalized(); + assert_eq!(c.block_size, MIN_BLOCK_SIZE); + } + + #[test] + fn coalesce_runs_splits_gaps_and_caps() { + // adjacent run [1,2,3] and isolated [5] + assert_eq!( + coalesce_runs(&[1, 2, 3, 5], 8), + vec![vec![1, 2, 3], vec![5]] + ); + // cap at 2 blocks per run + assert_eq!( + coalesce_runs(&[1, 2, 3, 4], 2), + vec![vec![1, 2], vec![3, 4]] + ); + assert!(coalesce_runs(&[], 4).is_empty()); + } +} diff --git a/native/block-cache/src/error.rs b/native/block-cache/src/error.rs new file mode 100644 index 0000000000..dc60da4863 --- /dev/null +++ b/native/block-cache/src/error.rs @@ -0,0 +1,52 @@ +// 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. + +use std::fmt; + +/// Error type for the block cache. +/// +/// `Clone` is required because a single fetch drives a shared future that is awaited +/// by every task that missed the same block (single-flight, see `cache.rs`); the +/// result — including the error — is cloned to each waiter. Fetch failures are never +/// cached: the in-flight entry is removed so the next caller retries. +#[derive(Clone, Debug)] +pub enum CacheError { + /// The caller-supplied `RangeFetcher` failed. Carries the upstream message. + Fetch(String), + /// An invariant inside the cache was violated (bug-class, not expected at runtime). + Internal(String), +} + +impl fmt::Display for CacheError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + CacheError::Fetch(msg) => write!(f, "block cache fetch error: {msg}"), + CacheError::Internal(msg) => write!(f, "block cache internal error: {msg}"), + } + } +} + +impl std::error::Error for CacheError {} + +impl CacheError { + /// Build a `Fetch` error from any displayable upstream error. + pub fn fetch(err: impl fmt::Display) -> Self { + CacheError::Fetch(err.to_string()) + } +} + +pub type Result = std::result::Result; diff --git a/native/block-cache/src/lib.rs b/native/block-cache/src/lib.rs new file mode 100644 index 0000000000..5d1f51c4e2 --- /dev/null +++ b/native/block-cache/src/lib.rs @@ -0,0 +1,67 @@ +// 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. + +//! Storage-API-neutral block cache core for DataFusion Comet. +//! +//! This crate implements a block-aligned local data cache (phase 1: memory tier only) with +//! a SIEVE eviction policy, single-flight miss deduplication, coalesced miss fetches, and +//! object-version validation for immutable-file workloads. +//! +//! # Storage neutrality +//! +//! By construction this crate depends on `tokio`/`bytes`/`futures` primitives only. It must +//! never depend on `object_store`, `opendal`, or any Comet crate: the caller supplies a +//! [`RangeFetcher`] and the core never talks to storage itself. This keeps the tier +//! internals swappable and the crate extractable. The boundary is CI-enforced by the +//! `dependency_tree_is_storage_neutral` integration test. +//! +//! # Example +//! +//! ``` +//! # use datafusion_comet_block_cache::*; +//! # use bytes::Bytes; +//! # use std::ops::Range; +//! # use async_trait::async_trait; +//! struct MyFetcher; +//! #[async_trait] +//! impl RangeFetcher for MyFetcher { +//! async fn fetch(&self, ranges: &[Range]) -> Result<(Vec, FileVersion)> { +//! let bytes = ranges.iter().map(|r| Bytes::from(vec![0u8; (r.end - r.start) as usize])).collect(); +//! Ok((bytes, FileVersion { size: 1024, ..Default::default() })) +//! } +//! } +//! # async fn run() -> Result<()> { +//! let cache = BlockCache::new(BlockCacheConfig::default()); +//! let file = FileKey::new(0, "s3://bucket/data.parquet"); +//! let out = cache.get_ranges(&file, &[0..128], &MyFetcher).await?; +//! assert_eq!(out[0].len(), 128); +//! # Ok(()) } +//! ``` + +mod cache; +mod error; +mod metrics; +mod sieve; +mod version; + +pub use cache::{ + BlockCache, BlockCacheConfig, RangeFetcher, DEFAULT_BLOCK_SIZE, DEFAULT_MAX_COALESCE_BYTES, + DEFAULT_NUM_SHARDS, MAX_BLOCK_SIZE, MIN_BLOCK_SIZE, +}; +pub use error::{CacheError, Result}; +pub use metrics::{Metrics, MetricsSnapshot}; +pub use version::{FileKey, FileVersion}; diff --git a/native/block-cache/src/metrics.rs b/native/block-cache/src/metrics.rs new file mode 100644 index 0000000000..b3cec1aeae --- /dev/null +++ b/native/block-cache/src/metrics.rs @@ -0,0 +1,100 @@ +// 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. + +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Process-wide cache counters. Cheap relaxed atomics; read via [`Metrics::snapshot`]. +/// +/// Phase 1 exposes these through a periodic native log line (the wrapper drives it); +/// Spark-visible SQL metrics are a tracked follow-up. +#[derive(Debug, Default)] +pub struct Metrics { + /// Block reads served from the memory tier. + pub hits: AtomicU64, + /// Block reads that had to go upstream. + pub misses: AtomicU64, + /// Upstream fetch calls issued (a coalesced run counts once). + pub fetches: AtomicU64, + /// Bytes returned by upstream fetches. + pub bytes_fetched: AtomicU64, + /// Blocks evicted from the memory tier. + pub evictions: AtomicU64, + /// Files invalidated because their version changed under us. + pub invalidations: AtomicU64, +} + +/// A point-in-time copy of [`Metrics`], safe to format/log. +#[derive(Debug, Clone, Copy, Default)] +pub struct MetricsSnapshot { + pub hits: u64, + pub misses: u64, + pub fetches: u64, + pub bytes_fetched: u64, + pub evictions: u64, + pub invalidations: u64, +} + +impl MetricsSnapshot { + /// Hit ratio over block reads in `[0.0, 1.0]`; `0.0` when nothing has been read yet. + pub fn hit_ratio(&self) -> f64 { + let total = self.hits + self.misses; + if total == 0 { + 0.0 + } else { + self.hits as f64 / total as f64 + } + } +} + +impl Metrics { + #[inline] + pub(crate) fn record_hit(&self) { + self.hits.fetch_add(1, Ordering::Relaxed); + } + + #[inline] + pub(crate) fn record_miss(&self) { + self.misses.fetch_add(1, Ordering::Relaxed); + } + + #[inline] + pub(crate) fn record_fetch(&self, bytes: u64) { + self.fetches.fetch_add(1, Ordering::Relaxed); + self.bytes_fetched.fetch_add(bytes, Ordering::Relaxed); + } + + #[inline] + pub(crate) fn record_eviction(&self) { + self.evictions.fetch_add(1, Ordering::Relaxed); + } + + #[inline] + pub(crate) fn record_invalidation(&self) { + self.invalidations.fetch_add(1, Ordering::Relaxed); + } + + pub fn snapshot(&self) -> MetricsSnapshot { + MetricsSnapshot { + hits: self.hits.load(Ordering::Relaxed), + misses: self.misses.load(Ordering::Relaxed), + fetches: self.fetches.load(Ordering::Relaxed), + bytes_fetched: self.bytes_fetched.load(Ordering::Relaxed), + evictions: self.evictions.load(Ordering::Relaxed), + invalidations: self.invalidations.load(Ordering::Relaxed), + } + } +} diff --git a/native/block-cache/src/sieve.rs b/native/block-cache/src/sieve.rs new file mode 100644 index 0000000000..f24c9757cf --- /dev/null +++ b/native/block-cache/src/sieve.rs @@ -0,0 +1,229 @@ +// 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. + +//! Memory-tier shard: a block map with a SIEVE eviction queue, byte accounting, and the +//! single-flight in-flight map. Each shard is guarded by one `Mutex` at the `BlockCache` +//! level, so all operations here run under that lock. + +use std::collections::{HashMap, VecDeque}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +use bytes::Bytes; +use futures::future::{BoxFuture, Shared}; + +use crate::error::CacheError; +use crate::metrics::Metrics; + +/// A cached block. `data` is the block's bytes (short for the final block of a file); +/// `visited` is SIEVE's per-entry reference bit, set on every hit. +#[derive(Debug)] +pub(crate) struct Block { + pub data: Bytes, + pub visited: AtomicBool, +} + +impl Block { + pub(crate) fn new(data: Bytes) -> Arc { + Arc::new(Block { + data, + visited: AtomicBool::new(false), + }) + } +} + +/// `(file_id, block_index)` — 12 bytes, cheap to hash and shard. +pub(crate) type BlockKey = (u64, u32); + +/// Result delivered to every task waiting on a single-flight fetch. Cloned per waiter, +/// hence the `Clone` bound satisfied by `Arc` and `CacheError`. +pub(crate) type FetchResult = std::result::Result, CacheError>; + +/// A shared in-flight fetch future. The first misser drives the fetch; concurrent +/// missers clone and await this. +pub(crate) type InFlightFut = Shared>; + +/// Fixed per-block bookkeeping charged on top of the block's byte length: the `BlockKey`, +/// the `Arc`, the hash-map slot, and the SIEVE queue slot. Approximate; keeps the +/// budget honest for workloads with many small final blocks. +pub(crate) const BLOCK_OVERHEAD_BYTES: u64 = 64; + +fn block_cost(data_len: usize) -> u64 { + data_len as u64 + BLOCK_OVERHEAD_BYTES +} + +/// One memory-tier shard. +pub(crate) struct Shard { + /// Live blocks, keyed within this shard. + map: HashMap>, + /// SIEVE order. `front` is the head (newest inserts); `back` is the tail (oldest). + queue: VecDeque, + /// SIEVE hand: an index into `queue`. Maintained best-effort across structural + /// changes and clamped on use; only correctness invariant is that it lands on a + /// live entry when eviction runs. + hand: usize, + /// Current accounted bytes (block data + per-block overhead). + current_bytes: u64, + /// Per-shard byte budget (total budget divided across shards). + budget: u64, + /// Single-flight map: block key -> shared fetch future. + in_flight: HashMap, +} + +impl Shard { + pub(crate) fn new(budget: u64) -> Self { + Shard { + map: HashMap::new(), + queue: VecDeque::new(), + hand: 0, + current_bytes: 0, + budget, + in_flight: HashMap::new(), + } + } + + /// Hit path: return a clone of the block if present, setting its `visited` bit. + pub(crate) fn get(&self, key: &BlockKey) -> Option> { + let block = self.map.get(key)?; + block.visited.store(true, Ordering::Relaxed); + Some(Arc::clone(block)) + } + + /// Insert a freshly fetched block, evicting in SIEVE order until back under budget. + /// A block already present (a concurrent fill won the race) is left as-is. + pub(crate) fn insert(&mut self, key: BlockKey, block: Arc, metrics: &Metrics) { + if self.map.contains_key(&key) { + return; + } + let cost = block_cost(block.data.len()); + self.map.insert(key, block); + self.queue.push_front(key); + // push_front shifted every existing index up by one; keep the hand on the same + // logical entry. + if self.queue.len() > 1 { + self.hand += 1; + } + self.current_bytes += cost; + self.evict_to(self.budget, metrics); + } + + /// Change this shard's budget and evict down to it immediately. + pub(crate) fn set_budget(&mut self, budget: u64, metrics: &Metrics) { + self.budget = budget; + self.evict_to(budget, metrics); + } + + /// Evict in SIEVE order until `current_bytes <= target` (or the shard is empty). + fn evict_to(&mut self, target: u64, metrics: &Metrics) { + while self.current_bytes > target && !self.map.is_empty() { + if let Some(cost) = self.evict_one() { + self.current_bytes -= cost; + metrics.record_eviction(); + } else { + break; + } + } + } + + /// Evict exactly one block per SIEVE: walk from the hand giving visited entries a + /// second chance (clearing their bit) until an unvisited entry is found and removed. + /// Returns the evicted block's accounted cost. + fn evict_one(&mut self) -> Option { + if self.queue.is_empty() { + return None; + } + // At most two full sweeps: the first clears every set bit, the second finds an + // unvisited victim. Guard the iteration count regardless. + let max_steps = self.queue.len() * 2 + 1; + for _ in 0..max_steps { + if self.hand >= self.queue.len() { + self.hand = self.queue.len() - 1; // wrap to the tail (oldest) + } + let key = self.queue[self.hand]; + let visited = self + .map + .get(&key) + .map(|b| b.visited.load(Ordering::Relaxed)) + .unwrap_or(false); + if visited { + // Survive: clear the bit and advance toward the head (lower index). + if let Some(b) = self.map.get(&key) { + b.visited.store(false, Ordering::Relaxed); + } + self.hand = if self.hand == 0 { + self.queue.len() - 1 + } else { + self.hand - 1 + }; + } else { + // Evict this entry. + let cost = self + .map + .remove(&key) + .map(|b| block_cost(b.data.len())) + .unwrap_or(0); + self.queue.remove(self.hand); + // The hand now refers to the element that followed the victim toward the + // tail; leave it there (clamped on next entry). Nothing to adjust. + if self.hand >= self.queue.len() && !self.queue.is_empty() { + self.hand = self.queue.len() - 1; + } + return Some(cost); + } + } + None + } + + /// Remove every block belonging to `file_id` (version-mismatch invalidation). + /// Returns bytes reclaimed. In-flight fetches are left untouched — they may be + /// fetching the new version. + pub(crate) fn remove_file_blocks(&mut self, file_id: u64) -> u64 { + let mut reclaimed = 0; + let to_remove: Vec = self + .map + .keys() + .filter(|(fid, _)| *fid == file_id) + .copied() + .collect(); + if to_remove.is_empty() { + return 0; + } + for key in &to_remove { + if let Some(b) = self.map.remove(key) { + reclaimed += block_cost(b.data.len()); + } + } + self.queue.retain(|k| k.0 != file_id); + self.hand = 0; + self.current_bytes = self.current_bytes.saturating_sub(reclaimed); + reclaimed + } + + // --- single-flight in-flight map --- + + pub(crate) fn in_flight_get(&self, key: &BlockKey) -> Option { + self.in_flight.get(key).cloned() + } + + pub(crate) fn in_flight_insert(&mut self, key: BlockKey, fut: InFlightFut) { + self.in_flight.insert(key, fut); + } + + pub(crate) fn in_flight_remove(&mut self, key: &BlockKey) { + self.in_flight.remove(key); + } +} diff --git a/native/block-cache/src/version.rs b/native/block-cache/src/version.rs new file mode 100644 index 0000000000..9cd90b3540 --- /dev/null +++ b/native/block-cache/src/version.rs @@ -0,0 +1,68 @@ +// 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. + +/// Identity of a cached file, as seen by the cache core. +/// +/// `namespace` identifies the logical store (the caller derives it — e.g. Comet hashes +/// its `(url_key, config_hash)` store key). `path` is the object path within that store. +/// The pair is interned to a compact `u64` file id internally (see `cache.rs`). +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct FileKey { + pub namespace: u64, + pub path: String, +} + +impl FileKey { + pub fn new(namespace: u64, path: impl Into) -> Self { + Self { + namespace, + path: path.into(), + } + } +} + +/// The observed version of an object, captured at fetch time from the store's metadata. +/// +/// Used to detect in-place overwrites: a miss fetch piggybacks the current version and, +/// if it disagrees with what the cache stored for the file, all cached blocks of that +/// file are invalidated (see `BlockCache`). The target workloads read immutable objects +/// (Iceberg/Delta/Hive data files), so in steady state versions never change and +/// validation costs nothing. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct FileVersion { + /// Object ETag, if the store reports one. + pub e_tag: Option, + /// Last-modified time in epoch milliseconds, if known. + pub last_modified: Option, + /// Total object size in bytes (full object, not the fetched range). + pub size: u64, +} + +impl FileVersion { + /// Whether `self` and `other` describe the same object version. + /// + /// Prefers the ETag when both sides report one (the strongest signal and what S3/GCS/ + /// Azure expose), and falls back to `(size, last_modified)` otherwise. Two versions + /// that share an ETag but differ in size are treated as different — a defensive guard + /// against weak/duplicated ETags. + pub fn matches(&self, other: &FileVersion) -> bool { + match (&self.e_tag, &other.e_tag) { + (Some(a), Some(b)) => a == b && self.size == other.size, + _ => self.size == other.size && self.last_modified == other.last_modified, + } + } +} diff --git a/native/block-cache/tests/cache_tests.rs b/native/block-cache/tests/cache_tests.rs new file mode 100644 index 0000000000..08809dbbd8 --- /dev/null +++ b/native/block-cache/tests/cache_tests.rs @@ -0,0 +1,432 @@ +// 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. + +// The tests intentionally pass single-range slices (`&[a..b]`) to `get_ranges`; that is exactly +// the call shape the cache serves, not a mistaken range literal. +#![allow(clippy::single_range_in_vec_init)] + +use std::ops::Range; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use bytes::Bytes; +use datafusion_comet_block_cache::*; + +/// A `RangeFetcher` backed by an in-memory byte vector. Counts fetch calls and total bytes +/// fetched, supports swapping its contents (to model an in-place overwrite), and can be +/// told to fail. +struct MockFetcher { + data: Mutex>, + version: Mutex, + fetch_calls: AtomicU64, + bytes_fetched: AtomicU64, + fail: Mutex, +} + +impl MockFetcher { + fn new(bytes: Vec, e_tag: &str) -> Arc { + let size = bytes.len() as u64; + Arc::new(MockFetcher { + data: Mutex::new(bytes), + version: Mutex::new(FileVersion { + e_tag: Some(e_tag.to_string()), + last_modified: Some(1), + size, + }), + fetch_calls: AtomicU64::new(0), + bytes_fetched: AtomicU64::new(0), + fail: Mutex::new(false), + }) + } + + fn calls(&self) -> u64 { + self.fetch_calls.load(Ordering::SeqCst) + } + + fn overwrite(&self, bytes: Vec, e_tag: &str) { + let size = bytes.len() as u64; + *self.data.lock().unwrap() = bytes; + *self.version.lock().unwrap() = FileVersion { + e_tag: Some(e_tag.to_string()), + last_modified: Some(2), + size, + }; + } + + fn set_fail(&self, fail: bool) { + *self.fail.lock().unwrap() = fail; + } +} + +#[async_trait] +impl RangeFetcher for MockFetcher { + async fn fetch(&self, ranges: &[Range]) -> Result<(Vec, FileVersion)> { + if *self.fail.lock().unwrap() { + return Err(CacheError::Fetch("injected failure".to_string())); + } + self.fetch_calls.fetch_add(1, Ordering::SeqCst); + let data = self.data.lock().unwrap(); + let mut out = Vec::with_capacity(ranges.len()); + for r in ranges { + // Emulate object_store: a partially-out-of-bounds range is truncated to EOF. + let start = (r.start as usize).min(data.len()); + let end = (r.end as usize).min(data.len()); + self.bytes_fetched + .fetch_add((end - start) as u64, Ordering::SeqCst); + out.push(Bytes::copy_from_slice(&data[start..end])); + } + Ok((out, self.version.lock().unwrap().clone())) + } +} + +fn seq(n: usize) -> Vec { + (0..n).map(|i| (i % 251) as u8).collect() +} + +fn small_cache() -> Arc { + // 1 MiB blocks, generous budget, few shards for deterministic-ish coverage. + BlockCache::new(BlockCacheConfig { + block_size: MIN_BLOCK_SIZE, + memory_budget: 64 << 20, + num_shards: 4, + max_coalesce_bytes: DEFAULT_MAX_COALESCE_BYTES, + }) +} + +#[tokio::test] +async fn serves_bytes_and_caches_across_calls() { + let bs = MIN_BLOCK_SIZE as usize; + let data = seq(bs * 3 + 12345); // 3 full blocks + a short final block + let fetcher = MockFetcher::new(data.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(7, "s3://b/f.parquet"); + + // A cross-block range with unaligned start and end. + let r: Range = 500..(bs as u64 * 2 + 900); + let out = cache + .get_ranges(&file, std::slice::from_ref(&r), &*fetcher) + .await + .unwrap(); + assert_eq!(&out[0][..], &data[r.start as usize..r.end as usize]); + let calls_after_first = fetcher.calls(); + assert!(calls_after_first >= 1); + + // Second identical read: zero new upstream fetches. + let out2 = cache + .get_ranges(&file, std::slice::from_ref(&r), &*fetcher) + .await + .unwrap(); + assert_eq!(&out2[0][..], &data[r.start as usize..r.end as usize]); + assert_eq!( + fetcher.calls(), + calls_after_first, + "second read must be a full cache hit" + ); +} + +#[tokio::test] +async fn short_final_block_and_multi_range() { + let bs = MIN_BLOCK_SIZE as usize; + let data = seq(bs + 777); // one full block + a short second block + let fetcher = MockFetcher::new(data.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(0, "f"); + + // Read the very end (short block), the footer bytes, and a mid-file range at once. + let ranges = vec![ + (data.len() as u64 - 100)..data.len() as u64, + 10..50, + (bs as u64 - 20)..(bs as u64 + 20), // straddles the block boundary + ]; + let out = cache.get_ranges(&file, &ranges, &*fetcher).await.unwrap(); + for (o, r) in out.iter().zip(&ranges) { + assert_eq!(&o[..], &data[r.start as usize..r.end as usize]); + } +} + +#[tokio::test] +async fn coalesces_adjacent_missing_blocks() { + let bs = MIN_BLOCK_SIZE as usize; + let data = seq(bs * 5); + let fetcher = MockFetcher::new(data.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(0, "f"); + + // Blocks 0..=3 all missing and adjacent -> should coalesce into one fetch (<=16 MiB). + let out = cache + .get_ranges(&file, &[0..(bs as u64 * 4)], &*fetcher) + .await + .unwrap(); + assert_eq!(&out[0][..], &data[0..bs * 4]); + assert_eq!( + fetcher.calls(), + 1, + "four adjacent blocks should coalesce to one fetch" + ); +} + +#[tokio::test] +async fn does_not_refetch_cached_blocks_to_bridge_gaps() { + let bs = MIN_BLOCK_SIZE as usize; + let data = seq(bs * 4); + let fetcher = MockFetcher::new(data.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(0, "f"); + + // Warm block 1 only. + cache + .get_ranges(&file, &[(bs as u64)..(bs as u64 + 10)], &*fetcher) + .await + .unwrap(); + let after_warm = fetcher.calls(); + + // Now read blocks 0,1,2,3. Block 1 is cached, so 0 and 2,3 must fetch as two separate + // runs (the cached block 1 splits them) — not one big run re-fetching block 1. + let out = cache + .get_ranges(&file, &[0..(bs as u64 * 4)], &*fetcher) + .await + .unwrap(); + assert_eq!(&out[0][..], &data[0..bs * 4]); + assert_eq!( + fetcher.calls() - after_warm, + 2, + "cached block must split the fetch into two runs, not be re-fetched" + ); +} + +#[tokio::test] +async fn single_flight_dedups_concurrent_missers() { + let bs = MIN_BLOCK_SIZE as usize; + let data = seq(bs * 2); + let fetcher = MockFetcher::new(data.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(0, "f"); + + // Launch many concurrent readers of the same block. + let mut handles = Vec::new(); + for _ in 0..32 { + let c = Arc::clone(&cache); + let f = Arc::clone(&fetcher); + let file = file.clone(); + handles.push(tokio::spawn(async move { + let out = c.get_ranges(&file, &[0..(bs as u64)], &*f).await.unwrap(); + assert_eq!(out[0].len(), bs); + })); + } + for h in handles { + h.await.unwrap(); + } + assert_eq!( + fetcher.calls(), + 1, + "concurrent missers must trigger exactly one fetch" + ); +} + +#[tokio::test] +async fn fetch_error_propagates_and_is_not_cached() { + let bs = MIN_BLOCK_SIZE as usize; + let data = seq(bs); + let fetcher = MockFetcher::new(data.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(0, "f"); + + fetcher.set_fail(true); + let err = cache.get_ranges(&file, &[0..10], &*fetcher).await; + assert!(err.is_err(), "fetch failure must surface"); + + // Error must not be cached: after recovery, the next read succeeds and fetches. + fetcher.set_fail(false); + let out = cache.get_ranges(&file, &[0..10], &*fetcher).await.unwrap(); + assert_eq!(&out[0][..], &data[0..10]); +} + +#[tokio::test] +async fn version_mismatch_invalidates_and_serves_fresh() { + let bs = MIN_BLOCK_SIZE as usize; + let v1 = seq(bs * 2); + let fetcher = MockFetcher::new(v1.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(0, "f"); + + // Warm block 0. + let out = cache.get_ranges(&file, &[0..100], &*fetcher).await.unwrap(); + assert_eq!(&out[0][..], &v1[0..100]); + + // Overwrite the object in place with different content + new ETag. + let v2: Vec = v1.iter().map(|b| b.wrapping_add(17)).collect(); + fetcher.overwrite(v2.clone(), "v2"); + + // A miss on a DIFFERENT block piggybacks the new version and invalidates block 0. + cache + .get_ranges(&file, &[(bs as u64)..(bs as u64 + 10)], &*fetcher) + .await + .unwrap(); + + // Reading block 0 again now returns the fresh (v2) bytes, not the stale cached ones. + let out = cache.get_ranges(&file, &[0..100], &*fetcher).await.unwrap(); + assert_eq!(&out[0][..], &v2[0..100], "must serve post-overwrite bytes"); + assert!(cache.stats().invalidations >= 1); +} + +#[tokio::test] +async fn sieve_keeps_reused_block_evicts_one_hit_wonder() { + // Budget for exactly ~2 blocks in a single shard, so inserting a 3rd forces eviction. + let bs = MIN_BLOCK_SIZE; + let cache = BlockCache::new(BlockCacheConfig { + block_size: bs, + memory_budget: bs * 2 + 4096, // room for 2 blocks + overhead + num_shards: 1, + max_coalesce_bytes: bs, // force one fetch per block for determinism + }); + let data = seq(bs as usize * 10); + let fetcher = MockFetcher::new(data.clone(), "v1"); + + let read_block = |c: Arc, f: Arc, idx: u64| async move { + let start = idx * bs; + c.get_ranges(&file_clone(), &[start..start + 8], &*f) + .await + .unwrap(); + }; + fn file_clone() -> FileKey { + FileKey::new(0, "f") + } + + // Insert block 0, then re-read it so its `visited` bit is set (it is "reused"). + read_block(Arc::clone(&cache), Arc::clone(&fetcher), 0).await; + read_block(Arc::clone(&cache), Arc::clone(&fetcher), 0).await; // hit -> visited + // Insert block 1 (one-hit-wonder, never revisited). + read_block(Arc::clone(&cache), Arc::clone(&fetcher), 1).await; + // Insert block 2 -> forces one eviction. SIEVE should evict the unvisited block 1, + // keeping the reused block 0. + read_block(Arc::clone(&cache), Arc::clone(&fetcher), 2).await; + + let calls_before = fetcher.calls(); + // Block 0 should still be cached (no fetch); block 1 should have been evicted (fetch). + read_block(Arc::clone(&cache), Arc::clone(&fetcher), 0).await; + assert_eq!( + fetcher.calls(), + calls_before, + "reused block 0 must survive eviction" + ); + read_block(Arc::clone(&cache), Arc::clone(&fetcher), 1).await; + assert_eq!( + fetcher.calls(), + calls_before + 1, + "one-hit block 1 must have been evicted" + ); +} + +#[tokio::test] +async fn set_memory_budget_shrink_evicts_then_growth_readmits() { + let bs = MIN_BLOCK_SIZE; + let cache = BlockCache::new(BlockCacheConfig { + block_size: bs, + memory_budget: bs * 8, + num_shards: 1, + max_coalesce_bytes: bs, + }); + let data = seq(bs as usize * 8); + let fetcher = MockFetcher::new(data.clone(), "v1"); + let file = FileKey::new(0, "f"); + + // Fill several blocks. + for i in 0..6u64 { + cache + .get_ranges(&file, &[i * bs..i * bs + 8], &*fetcher) + .await + .unwrap(); + } + let filled_calls = fetcher.calls(); + + // Shrink to hold ~1 block: must evict down immediately. + cache.set_memory_budget(bs + 4096); + assert!(cache.stats().evictions >= 1); + + // Reading an old block now misses (was evicted) -> a new fetch. + cache.get_ranges(&file, &[0..8], &*fetcher).await.unwrap(); + assert!( + fetcher.calls() > filled_calls, + "shrunk cache must have evicted and re-fetched" + ); + + // Grow back: subsequent fills are re-admitted (cache holds more again). + cache.set_memory_budget(bs * 8); + let before = fetcher.calls(); + cache.get_ranges(&file, &[0..8], &*fetcher).await.unwrap(); // just fetched above -> hit + assert_eq!( + fetcher.calls(), + before, + "recently filled block stays after growth" + ); +} + +#[tokio::test] +async fn invalidate_file_drops_blocks() { + let bs = MIN_BLOCK_SIZE as usize; + let data = seq(bs); + let fetcher = MockFetcher::new(data.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(3, "f"); + + cache.get_ranges(&file, &[0..10], &*fetcher).await.unwrap(); + let before = fetcher.calls(); + cache.invalidate_file(&file); + cache.get_ranges(&file, &[0..10], &*fetcher).await.unwrap(); + assert_eq!( + fetcher.calls(), + before + 1, + "invalidated file must be re-fetched" + ); +} + +#[tokio::test] +async fn byte_exact_over_randomized_ranges() { + let bs = MIN_BLOCK_SIZE as usize; + let data = seq(bs * 4 + 4242); + let fetcher = MockFetcher::new(data.clone(), "v1"); + let cache = small_cache(); + let file = FileKey::new(0, "f"); + + // Deterministic pseudo-random ranges (no rand dep): LCG. + let mut state: u64 = 0x1234_5678; + let mut next = || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + state >> 16 + }; + for _ in 0..200 { + let a = (next() as usize) % data.len(); + let b = (next() as usize) % data.len(); + let (lo, hi) = if a <= b { (a, b) } else { (b, a) }; + let hi = (hi + 1).min(data.len()); + if lo == hi { + continue; + } + let out = cache + .get_ranges(&file, &[lo as u64..hi as u64], &*fetcher) + .await + .unwrap(); + assert_eq!( + &out[0][..], + &data[lo..hi], + "byte mismatch for range {lo}..{hi}" + ); + } +} diff --git a/native/block-cache/tests/neutrality.rs b/native/block-cache/tests/neutrality.rs new file mode 100644 index 0000000000..e229eca444 --- /dev/null +++ b/native/block-cache/tests/neutrality.rs @@ -0,0 +1,69 @@ +// 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. + +//! CI guard for the crate's storage neutrality: `datafusion-comet-block-cache` must never +//! pull in `object_store`, `opendal`, or a Comet crate. This keeps the tier internals +//! swappable and the core extractable, as required by OBJECT_STORE_CACHE_DESIGN.md §2.1. + +use std::process::Command; + +const FORBIDDEN: &[&str] = &["object_store", "opendal", "datafusion-comet-core"]; + +#[test] +fn dependency_tree_is_storage_neutral() { + // Resolve the crate manifest so the test works regardless of the CWD chosen by the + // test harness. + let manifest = concat!(env!("CARGO_MANIFEST_DIR"), "/Cargo.toml"); + let output = Command::new(env!("CARGO")) + .args([ + "tree", + "--manifest-path", + manifest, + "-p", + "datafusion-comet-block-cache", + "--edges", + "normal,build", + "--prefix", + "none", + ]) + .output(); + + let output = match output { + Ok(o) => o, + // If cargo cannot be spawned (unusual outside CI), do not fail spuriously. + Err(e) => { + eprintln!("skipping neutrality check: could not run `cargo tree`: {e}"); + return; + } + }; + assert!( + output.status.success(), + "`cargo tree` failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + let tree = String::from_utf8_lossy(&output.stdout); + for forbidden in FORBIDDEN { + // Match crate names at a line start (name is the first token per `--prefix none`). + let hit = tree + .lines() + .any(|line| line.split_whitespace().next() == Some(*forbidden)); + assert!( + !hit, + "block-cache must stay storage-API-neutral but its dependency tree contains `{forbidden}`:\n{tree}" + ); + } +} diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index e657879d33..f7ecda91ac 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -62,6 +62,8 @@ datafusion-comet-spark-expr = { workspace = true } datafusion-comet-jni-bridge = { workspace = true } datafusion-comet-proto = { workspace = true } datafusion-comet-shuffle = { workspace = true } +datafusion-comet-block-cache = { workspace = true } +datafusion-comet-object-store-cache = { workspace = true } object_store = { workspace = true } url = { workspace = true } aws-config = { workspace = true } diff --git a/native/core/src/execution/data_cache.rs b/native/core/src/execution/data_cache.rs new file mode 100644 index 0000000000..be09269121 --- /dev/null +++ b/native/core/src/execution/data_cache.rs @@ -0,0 +1,129 @@ +// 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. + +//! Process-global object-store data cache (OBJECT_STORE_CACHE_DESIGN.md §2.7). +//! +//! A single [`BlockCache`] lives for the executor's lifetime, initialized from the Spark +//! configs passed on the first `createPlan`. It is `None` when the feature is disabled, so +//! the wrapper is never constructed and there is zero overhead. First plan wins: changing +//! cache configs requires an executor restart (mirrors the existing process-global +//! object-store instance cache). + +use std::collections::HashMap; +use std::sync::{Arc, OnceLock}; + +use datafusion_comet_block_cache::{ + BlockCache, BlockCacheConfig, DEFAULT_BLOCK_SIZE, DEFAULT_MAX_COALESCE_BYTES, + DEFAULT_NUM_SHARDS, +}; +use log::info; + +use crate::execution::spark_config::{ + SparkConfig, COMET_DATA_CACHE_BLOCK_SIZE, COMET_DATA_CACHE_ENABLED, + COMET_DATA_CACHE_MEMORY_LIMIT, COMET_DATA_CACHE_SSD_LIMIT, +}; + +/// Default memory-tier budget when the config is absent (512 MiB, matches `CometConf`). +const DEFAULT_MEMORY_LIMIT: u64 = 512 << 20; + +static DATA_CACHE: OnceLock>> = OnceLock::new(); + +/// Initialize the global data cache from Spark configs, once. Idempotent and cheap on +/// subsequent calls (a `OnceLock` read). `local_dirs` are reserved for the phase-2 SSD tier. +pub(crate) fn init_once(spark_config: &HashMap, local_dirs: &[String]) { + DATA_CACHE.get_or_init(|| build(spark_config, local_dirs)); +} + +/// The process-global cache, or `None` when disabled / not yet initialized. +pub(crate) fn global() -> Option> { + DATA_CACHE.get().and_then(|opt| opt.clone()) +} + +fn build( + spark_config: &HashMap, + _local_dirs: &[String], +) -> Option> { + if !spark_config.get_bool(COMET_DATA_CACHE_ENABLED) { + return None; + } + let memory_budget = spark_config.get_u64(COMET_DATA_CACHE_MEMORY_LIMIT, DEFAULT_MEMORY_LIMIT); + let block_size = spark_config.get_u64(COMET_DATA_CACHE_BLOCK_SIZE, DEFAULT_BLOCK_SIZE); + let ssd_limit = spark_config.get_u64(COMET_DATA_CACHE_SSD_LIMIT, 0); + if ssd_limit > 0 { + // The SSD tier is phase 2; honor memory-tier config now and note the deferral. + info!( + "Comet data cache: SSD tier requested ({} MiB) but not yet implemented (phase 2); \ + using memory tier only", + ssd_limit >> 20 + ); + } + + let cache = BlockCache::new(BlockCacheConfig { + block_size, + memory_budget, + num_shards: DEFAULT_NUM_SHARDS, + max_coalesce_bytes: DEFAULT_MAX_COALESCE_BYTES, + }); + info!( + "Comet object-store data cache enabled: memory budget {} MiB, block size {} MiB, {} shards", + memory_budget >> 20, + cache.block_size() >> 20, + DEFAULT_NUM_SHARDS, + ); + Some(cache) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn conf(pairs: &[(&str, &str)]) -> HashMap { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + #[test] + fn disabled_by_default_builds_nothing() { + assert!(build(&conf(&[]), &[]).is_none()); + assert!(build(&conf(&[(COMET_DATA_CACHE_ENABLED, "false")]), &[]).is_none()); + } + + #[test] + fn enabled_uses_defaults() { + let cache = build(&conf(&[(COMET_DATA_CACHE_ENABLED, "true")]), &[]).unwrap(); + assert_eq!(cache.block_size(), DEFAULT_BLOCK_SIZE); + assert_eq!(cache.memory_budget(), DEFAULT_MEMORY_LIMIT); + } + + #[test] + fn enabled_honors_explicit_byte_configs() { + // Values arrive as byte counts (Scala `bytesConf` serializes bytes over JNI). + let cache = build( + &conf(&[ + (COMET_DATA_CACHE_ENABLED, "true"), + (COMET_DATA_CACHE_MEMORY_LIMIT, &(1024u64 << 20).to_string()), + (COMET_DATA_CACHE_BLOCK_SIZE, &(8u64 << 20).to_string()), + ]), + &[], + ) + .unwrap(); + assert_eq!(cache.block_size(), 8 << 20); + assert_eq!(cache.memory_budget(), 1024 << 20); + } +} diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 6dc00e9cf6..17ae446d40 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -420,6 +420,11 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( local_dirs_vec.push(local_dir); } + // Initialize the process-global object-store data cache on the first plan + // (no-op if disabled or already initialized). Must happen before the planner + // creates any object store so remote stores get wrapped. + crate::execution::data_cache::init_once(&spark_config, &local_dirs_vec); + // We need to keep the session context alive. Some session state like temporary // dictionaries are stored in session context. If it is dropped, the temporary // dictionaries will be dropped as well. diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index ec247f72b7..1dddff4a5e 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -17,6 +17,7 @@ //! PoC of vectorization execution through JNI to Rust. pub mod columnar_to_row; +pub(crate) mod data_cache; pub mod expressions; pub mod jni_api; pub(crate) mod merge_as_partial; diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 4c2811cb5d..b2d7aca0fc 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -26,6 +26,16 @@ pub(crate) const COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED: &str = "spark.comet.parquet.rowFilterPushdown.enabled"; pub(crate) const SPARK_EXECUTOR_CORES: &str = "spark.executor.cores"; +// Object-store data cache (see OBJECT_STORE_CACHE_DESIGN.md §2.8). All read on the native +// side at first `createPlan` to build the process-global block cache. +pub(crate) const COMET_DATA_CACHE_ENABLED: &str = "spark.comet.scan.dataCache.enabled"; +pub(crate) const COMET_DATA_CACHE_MEMORY_LIMIT: &str = "spark.comet.scan.dataCache.memoryLimit"; +pub(crate) const COMET_DATA_CACHE_BLOCK_SIZE: &str = "spark.comet.scan.dataCache.blockSize"; +pub(crate) const COMET_DATA_CACHE_SSD_LIMIT: &str = "spark.comet.scan.dataCache.ssd.limit"; +// Read by the phase-2 SSD tier; declared now so the key set is stable. +#[allow(dead_code)] +pub(crate) const COMET_DATA_CACHE_SSD_PATH: &str = "spark.comet.scan.dataCache.ssd.path"; + pub(crate) trait SparkConfig { fn get_bool(&self, name: &str) -> bool; fn get_u64(&self, name: &str, default_value: u64) -> u64; diff --git a/native/core/src/parquet/parquet_support.rs b/native/core/src/parquet/parquet_support.rs index c75b2ea4f9..2e9e37b18e 100644 --- a/native/core/src/parquet/parquet_support.rs +++ b/native/core/src/parquet/parquet_support.rs @@ -34,6 +34,7 @@ use datafusion::error::DataFusionError; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_object_store_cache::CachingObjectStore; use datafusion_comet_spark_expr::EvalMode; use log::debug; use object_store::path::Path; @@ -552,6 +553,38 @@ fn object_store_cache() -> &'static ObjectStoreCache { CACHE.get_or_init(|| RwLock::new(HashMap::new())) } +/// Wrap `store` with the process-global data cache when the cache is enabled and the store +/// is remote (any scheme but local `file`). Returns `store` unchanged otherwise, so a +/// disabled cache constructs no wrapper. +fn maybe_wrap_with_data_cache( + store: Arc, + url_key: &str, + config_hash: u64, + scheme: &str, +) -> Arc { + // Local filesystem reads are already served from the OS page cache; don't wrap them. + if scheme == "file" { + return store; + } + match crate::execution::data_cache::global() { + Some(cache) => { + let namespace = data_cache_namespace(url_key, config_hash); + debug!("Wrapping object store {url_key} with Comet data cache (namespace {namespace})"); + Arc::new(CachingObjectStore::new(store, cache, namespace)) + } + None => store, + } +} + +/// Derive the cache namespace that isolates one logical store from another. Uses the same +/// distinguishing inputs as the object-store instance cache key `(url_key, config_hash)`. +fn data_cache_namespace(url_key: &str, config_hash: u64) -> u64 { + let mut hasher = DefaultHasher::new(); + url_key.hash(&mut hasher); + config_hash.hash(&mut hasher); + hasher.finish() +} + /// Compute a hash of the object store configuration for cache keying. fn hash_object_store_configs(configs: &HashMap) -> u64 { let mut hasher = DefaultHasher::new(); @@ -619,6 +652,9 @@ pub(crate) fn prepare_object_store_with_configs( .map_err(|e| ExecutionError::GeneralError(e.to_string()))?; let store: Arc = Arc::from(store); + // Wrap remote stores with the process-global data cache when enabled. Wrapping + // before the instance-cache insert means later reuses pick up the cached wrapper. + let store = maybe_wrap_with_data_cache(store, &url_key, config_hash, scheme); // Insert into cache if let Ok(mut cache) = object_store_cache().write() { cache.insert(cache_key, Arc::clone(&store)); diff --git a/native/object-store-cache/Cargo.toml b/native/object-store-cache/Cargo.toml new file mode 100644 index 0000000000..45302007cb --- /dev/null +++ b/native/object-store-cache/Cargo.toml @@ -0,0 +1,46 @@ +# 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. + +[package] +name = "datafusion-comet-object-store-cache" +description = "Apache DataFusion Comet: object_store::ObjectStore wrapper over the block cache core" +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +readme = { workspace = true } +license = { workspace = true } +edition = { workspace = true } + +publish = false + +# This is the ONLY crate that touches the `object_store` trait. It compiles against the +# workspace `object_store` version on every DataFusion bump. +[dependencies] +datafusion-comet-block-cache = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } +chrono = { workspace = true } +futures = { workspace = true } +object_store = { workspace = true } + +[dev-dependencies] +tokio = { version = "1", features = ["rt-multi-thread", "macros"] } + +[lib] +name = "datafusion_comet_object_store_cache" +path = "src/lib.rs" diff --git a/native/object-store-cache/src/lib.rs b/native/object-store-cache/src/lib.rs new file mode 100644 index 0000000000..affaeb8987 --- /dev/null +++ b/native/object-store-cache/src/lib.rs @@ -0,0 +1,477 @@ +// 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. + +//! `object_store::ObjectStore` adapter over the Comet block cache core. +//! +//! [`CachingObjectStore`] wraps any `Arc` and serves block-aligned range +//! reads out of a shared [`BlockCache`]. It is the ONLY crate that touches the +//! `object_store` trait, so a DataFusion / `object_store` version bump compiles the adapter +//! in the same PR (see OBJECT_STORE_CACHE_DESIGN.md §2.1, §2.6). +//! +//! ## What is cached +//! +//! - `get_range` / `get_ranges` — the whole Parquet I/O path in Comet — are served from +//! the cache, block-aligned, with single-flight miss dedup. +//! - `get_opts` is served from the cache only for a plain `GetRange::Bounded` with no +//! conditions / version / head; anything else (conditional gets, suffix/offset ranges, +//! full-object streams, head requests) passes straight through to the inner store. +//! - `put*` / `delete` / `copy*` pass through and invalidate the affected path. +//! - `get` (full-object stream), `list*`, and `head` pass through unchanged. + +use std::fmt::{Debug, Display, Formatter}; +use std::ops::Range; +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::Bytes; +use chrono::{DateTime, Utc}; +use datafusion_comet_block_cache::{ + BlockCache, CacheError, FileKey, FileVersion, RangeFetcher, Result as CacheResult, +}; +use futures::stream::BoxStream; +use futures::StreamExt; +use object_store::path::Path; +use object_store::{ + Attributes, CopyOptions, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, + MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, + PutResult, Result as OsResult, +}; + +/// An [`ObjectStore`] that serves block-aligned range reads from a shared [`BlockCache`]. +/// +/// `namespace` distinguishes logically distinct stores that share the cache (Comet passes a +/// hash of its `(url_key, config_hash)` store key), so identical object paths in different +/// stores never collide in the cache. +pub struct CachingObjectStore { + inner: Arc, + cache: Arc, + namespace: u64, +} + +impl CachingObjectStore { + pub fn new(inner: Arc, cache: Arc, namespace: u64) -> Self { + CachingObjectStore { + inner, + cache, + namespace, + } + } + + fn file_key(&self, location: &Path) -> FileKey { + FileKey::new(self.namespace, location.as_ref()) + } + + fn range_fetcher(&self, location: &Path) -> StoreRangeFetcher { + StoreRangeFetcher { + store: Arc::clone(&self.inner), + path: location.clone(), + } + } + + /// Serve a plain bounded-range `get_opts` from the cache, synthesizing an `ObjectMeta` + /// from the cache's captured version (no extra `head` request). Returns `None` when the + /// version has not been captured yet, so the caller falls back to the inner store. + async fn cached_get_opts( + &self, + location: &Path, + range: Range, + ) -> OsResult> { + let file = self.file_key(location); + let fetcher = self.range_fetcher(location); + let bytes = self + .cache + .get_ranges(&file, std::slice::from_ref(&range), &fetcher) + .await + .map_err(to_os_error)? + .pop() + .unwrap_or_default(); + + let Some(version) = self.cache.file_version(&file) else { + return Ok(None); + }; + let returned = range.start..range.start + bytes.len() as u64; + let meta = version_to_meta(location, &version); + let payload = + GetResultPayload::Stream(futures::stream::once(async move { Ok(bytes) }).boxed()); + Ok(Some(GetResult { + payload, + meta, + range: returned, + attributes: Attributes::default(), + })) + } +} + +impl Display for CachingObjectStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "CachingObjectStore({})", self.inner) + } +} + +impl Debug for CachingObjectStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CachingObjectStore") + .field("inner", &self.inner) + .field("namespace", &self.namespace) + .finish() + } +} + +#[async_trait] +impl ObjectStore for CachingObjectStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> OsResult { + let result = self.inner.put_opts(location, payload, opts).await; + // The object at this path just changed; drop any cached blocks for it. + self.cache.invalidate_file(&self.file_key(location)); + result + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OsResult> { + // A multipart write targets this path; invalidate up front so stale blocks cannot + // be served while/after the upload lands. + self.cache.invalidate_file(&self.file_key(location)); + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> OsResult { + if let Some(range) = plain_bounded_range(&options) { + if let Some(result) = self.cached_get_opts(location, range).await? { + return Ok(result); + } + } + self.inner.get_opts(location, options).await + } + + // `get_range` (the single-range convenience) lives in `ObjectStoreExt` and is not + // overridable; it delegates to `get_opts` with a bounded range, which we cache above. + // `get_ranges` is a core trait method, so we route it straight through the cache. + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> OsResult> { + let fetcher = self.range_fetcher(location); + self.cache + .get_ranges(&self.file_key(location), ranges, &fetcher) + .await + .map_err(to_os_error) + } + + // The single-path `delete` convenience delegates here; invalidate each path as it flows + // to the inner store so a deleted object cannot be served stale from the cache. + fn delete_stream( + &self, + locations: BoxStream<'static, OsResult>, + ) -> BoxStream<'static, OsResult> { + let cache = Arc::clone(&self.cache); + let namespace = self.namespace; + let mapped = locations + .inspect(move |res| { + if let Ok(path) = res { + cache.invalidate_file(&FileKey::new(namespace, path.as_ref())); + } + }) + .boxed(); + self.inner.delete_stream(mapped) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OsResult> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, OsResult> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OsResult { + self.inner.list_with_delimiter(prefix).await + } + + // `copy` and `copy_if_not_exists` are `ObjectStoreExt` conveniences that both delegate + // to `copy_opts`, so invalidating the destination here covers them all. + async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> OsResult<()> { + let result = self.inner.copy_opts(from, to, options).await; + self.cache.invalidate_file(&self.file_key(to)); + result + } +} + +/// A [`RangeFetcher`] that pulls absolute byte ranges from an inner `ObjectStore` and +/// harvests the object version from the response metadata. +struct StoreRangeFetcher { + store: Arc, + path: Path, +} + +#[async_trait] +impl RangeFetcher for StoreRangeFetcher { + async fn fetch(&self, ranges: &[Range]) -> CacheResult<(Vec, FileVersion)> { + let mut out = Vec::with_capacity(ranges.len()); + let mut version = FileVersion::default(); + for r in ranges { + let opts = GetOptions { + range: Some(GetRange::Bounded(r.clone())), + ..Default::default() + }; + let result = self + .store + .get_opts(&self.path, opts) + .await + .map_err(CacheError::fetch)?; + // Capture the version before consuming the payload. + version = FileVersion { + e_tag: result.meta.e_tag.clone(), + last_modified: Some(result.meta.last_modified.timestamp_millis()), + size: result.meta.size, + }; + let bytes = result.bytes().await.map_err(CacheError::fetch)?; + out.push(bytes); + } + Ok((out, version)) + } +} + +/// Extract a plain bounded range from `get_opts` options, or `None` if any condition, +/// version pin, head flag, or non-bounded range is present (those bypass the cache). +fn plain_bounded_range(options: &GetOptions) -> Option> { + if options.if_match.is_some() + || options.if_none_match.is_some() + || options.if_modified_since.is_some() + || options.if_unmodified_since.is_some() + || options.version.is_some() + || options.head + { + return None; + } + match &options.range { + Some(GetRange::Bounded(r)) => Some(r.clone()), + _ => None, + } +} + +/// Reconstruct an `ObjectMeta` from a captured cache [`FileVersion`]. +fn version_to_meta(location: &Path, version: &FileVersion) -> ObjectMeta { + let last_modified = version + .last_modified + .and_then(DateTime::::from_timestamp_millis) + .unwrap_or_else(|| DateTime::::from_timestamp_millis(0).unwrap()); + ObjectMeta { + location: location.clone(), + last_modified, + size: version.size, + e_tag: version.e_tag.clone(), + version: None, + } +} + +/// Map a cache error onto an `object_store::Error`. +fn to_os_error(err: CacheError) -> object_store::Error { + object_store::Error::Generic { + store: "CachingObjectStore", + source: Box::new(err), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_comet_block_cache::BlockCacheConfig; + use object_store::memory::InMemory; + use object_store::{ObjectStoreExt, PutPayload}; + use std::sync::atomic::{AtomicU64, Ordering}; + + /// An `ObjectStore` that counts `get_opts` calls, wrapping an inner store. + #[derive(Debug)] + struct CountingStore { + inner: Arc, + get_opts_calls: AtomicU64, + } + + impl CountingStore { + fn new(inner: Arc) -> Arc { + Arc::new(CountingStore { + inner, + get_opts_calls: AtomicU64::new(0), + }) + } + fn calls(&self) -> u64 { + self.get_opts_calls.load(Ordering::SeqCst) + } + } + + impl Display for CountingStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "CountingStore({})", self.inner) + } + } + + #[async_trait] + impl ObjectStore for CountingStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> OsResult { + self.inner.put_opts(location, payload, opts).await + } + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> OsResult> { + self.inner.put_multipart_opts(location, opts).await + } + async fn get_opts(&self, location: &Path, options: GetOptions) -> OsResult { + self.get_opts_calls.fetch_add(1, Ordering::SeqCst); + self.inner.get_opts(location, options).await + } + fn delete_stream( + &self, + locations: BoxStream<'static, OsResult>, + ) -> BoxStream<'static, OsResult> { + self.inner.delete_stream(locations) + } + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OsResult> { + self.inner.list(prefix) + } + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OsResult { + self.inner.list_with_delimiter(prefix).await + } + async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> OsResult<()> { + self.inner.copy_opts(from, to, options).await + } + } + + fn setup(data: &[u8]) -> (Arc, CachingObjectStore, Path) { + let mem = Arc::new(InMemory::new()); + let path = Path::from("data.parquet"); + futures::executor::block_on(mem.put(&path, PutPayload::from(data.to_vec()))).unwrap(); + let counting = CountingStore::new(mem); + let cache = BlockCache::new(BlockCacheConfig { + block_size: 1 << 20, + memory_budget: 64 << 20, + num_shards: 4, + ..Default::default() + }); + let store = CachingObjectStore::new(counting.clone(), cache, 42); + (counting, store, path) + } + + fn seq(n: usize) -> Vec { + (0..n).map(|i| (i % 251) as u8).collect() + } + + #[tokio::test] + async fn second_get_ranges_issues_zero_upstream_requests() { + let data = seq((1 << 20) * 2 + 500); + let (counting, store, path) = setup(&data); + let ranges = vec![10..2000u64, (1 << 20) - 5..(1 << 20) + 300]; + + let a = store.get_ranges(&path, &ranges).await.unwrap(); + let calls_after_first = counting.calls(); + assert!(calls_after_first >= 1); + + let b = store.get_ranges(&path, &ranges).await.unwrap(); + assert_eq!( + counting.calls(), + calls_after_first, + "second read must be all hits" + ); + assert_eq!(a, b); + for (bytes, r) in b.iter().zip(&ranges) { + assert_eq!(&bytes[..], &data[r.start as usize..r.end as usize]); + } + } + + #[tokio::test] + async fn byte_exact_vs_unwrapped_store() { + let data = seq((1 << 20) * 3 + 4242); + let (_counting, store, path) = setup(&data); + let mem = Arc::new(InMemory::new()); + futures::executor::block_on(mem.put(&path, PutPayload::from(data.clone()))).unwrap(); + + let mut state: u64 = 99; + let mut next = || { + state = state.wrapping_mul(6364136223846793005).wrapping_add(1); + state >> 20 + }; + for _ in 0..64 { + let a = (next() as usize) % data.len(); + let b = (next() as usize) % data.len(); + let (lo, hi) = if a <= b { (a, b + 1) } else { (b, a + 1) }; + let hi = hi.min(data.len()); + if lo == hi { + continue; + } + let cached = store.get_range(&path, lo as u64..hi as u64).await.unwrap(); + let direct = mem.get_range(&path, lo as u64..hi as u64).await.unwrap(); + assert_eq!(cached, direct, "mismatch for {lo}..{hi}"); + } + } + + #[tokio::test] + async fn get_opts_conditional_passes_through() { + let data = seq(1 << 20); + let (counting, store, path) = setup(&data); + // Warm the cache. + store.get_range(&path, 0..100).await.unwrap(); + let before = counting.calls(); + + // A conditional get must bypass the cache and hit the inner store. + let opts = GetOptions { + range: Some(GetRange::Bounded(0..100)), + if_match: Some("\"nope\"".to_string()), + ..Default::default() + }; + let _ = store.get_opts(&path, opts).await; // precondition may fail; that's fine + assert!( + counting.calls() > before, + "conditional get must reach inner store" + ); + } + + #[tokio::test] + async fn put_invalidates_cached_blocks() { + let data = seq(1 << 20); + let (counting, store, path) = setup(&data); + store.get_range(&path, 0..100).await.unwrap(); + let before = counting.calls(); + + // Overwrite through the wrapper: must invalidate, so the next read re-fetches. + let new = seq(1 << 20) + .iter() + .map(|b| b.wrapping_add(1)) + .collect::>(); + store + .put(&path, PutPayload::from(new.clone())) + .await + .unwrap(); + let out = store.get_range(&path, 0..100).await.unwrap(); + assert!(counting.calls() > before, "read after put must re-fetch"); + assert_eq!(&out[..], &new[0..100]); + } +} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 8e47151358..1d0ace7f12 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -132,6 +132,77 @@ object CometConf extends ShimCometConf { .checkValue(v => v > 0, "Data file concurrency limit must be positive") .createWithDefault(1) + val COMET_DATA_CACHE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.scan.dataCache.enabled") + .category(CATEGORY_SCAN) + .doc( + "Whether to cache object-store data bytes locally (memory tier) behind the native " + + "scan's object store, so repeated reads of the same remote Parquet files avoid " + + "re-fetching from S3/HDFS/ABFS. Experimental; default false. The memory-tier " + + "budget is charged outside Spark's task memory pools, so size it from " + + "spark.executor.memoryOverhead headroom (see dataCache.memoryLimit).") + .booleanConf + .createWithDefault(false) + + val COMET_DATA_CACHE_MEMORY_LIMIT: ConfigEntry[Long] = + conf("spark.comet.scan.dataCache.memoryLimit") + .category(CATEGORY_SCAN) + .doc( + "Memory-tier budget in bytes for the object-store data cache, process-wide per " + + "executor (accepts size suffixes, e.g. 512m, 1g). This memory sits outside " + + "Comet's task memory pools and must come out of spark.executor.memoryOverhead " + + "headroom. Only used when spark.comet.scan.dataCache.enabled is true.") + .bytesConf(ByteUnit.BYTE) + .checkValue(v => v > 0, "Data cache memory limit must be positive") + .createWithDefault(512L * 1024 * 1024) + + val COMET_DATA_CACHE_BLOCK_SIZE: ConfigEntry[Long] = + conf("spark.comet.scan.dataCache.blockSize") + .category(CATEGORY_SCAN) + .doc( + "Block quantum in bytes for the object-store data cache (accepts size suffixes, " + + "e.g. 4m). Reads are aligned to this size before caching. Rounded down to a power " + + "of two and clamped to [1 MiB, 16 MiB] natively. Only used when " + + "spark.comet.scan.dataCache.enabled is true.") + .bytesConf(ByteUnit.BYTE) + .checkValue( + v => v >= (1L << 20) && v <= (16L << 20), + "Block size must be between 1 and 16 MiB") + .createWithDefault(4L * 1024 * 1024) + + val COMET_DATA_CACHE_SSD_LIMIT: ConfigEntry[Long] = + conf("spark.comet.scan.dataCache.ssd.limit") + .category(CATEGORY_SCAN) + .doc( + "SSD-tier budget in bytes for the object-store data cache; 0 disables the tier " + + "(accepts size suffixes, e.g. 10g). The SSD tier is not yet implemented (phase 2) " + + "and this value is currently ignored beyond a log note. Only used when " + + "spark.comet.scan.dataCache.enabled is true.") + .bytesConf(ByteUnit.BYTE) + .checkValue(v => v >= 0, "SSD limit must be non-negative") + .createWithDefault(0L) + + val COMET_DATA_CACHE_SSD_PATH: OptionalConfigEntry[String] = + conf("spark.comet.scan.dataCache.ssd.path") + .category(CATEGORY_SCAN) + .doc( + "Directory for the object-store data cache SSD tier (phase 2). Defaults to the first " + + "Spark block-manager local directory when unset. Only used when the SSD tier is " + + "enabled.") + .stringConf + .createOptional + + val COMET_DATA_CACHE_LOCALITY_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.scan.dataCache.locality.enabled") + .category(CATEGORY_SCAN) + .doc( + "Whether to declare preferred locations for native scan partitions when the data " + + "cache is enabled, routing repeat reads of a Parquet file back to the host that " + + "cached it. Consumed only on the driver (never crosses JNI). No effect when " + + "spark.comet.scan.dataCache.enabled is false.") + .booleanConf + .createWithDefault(true) + val COMET_CSV_V2_NATIVE_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.scan.csv.v2.enabled") .category(CATEGORY_TESTING) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index e59dec7ebf..dfefc68236 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -278,6 +278,24 @@ object CometExecIterator extends Logging { CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.key, CometConf.COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED.get(SQLConf.get).toString) + // Object-store data cache configs read by the native side at first `createPlan`. The + // native cache is only built when `enabled` is true; the rest supply the memory-tier + // budget/block size (and the phase-2 SSD budget note) so native and Scala defaults never + // drift. `dataCache.ssd.path` is carried only if explicitly set (phase 2), and + // `dataCache.locality.enabled` is driver-only and deliberately never crosses JNI. + builder.putEntries( + CometConf.COMET_DATA_CACHE_ENABLED.key, + CometConf.COMET_DATA_CACHE_ENABLED.get(SQLConf.get).toString) + builder.putEntries( + CometConf.COMET_DATA_CACHE_MEMORY_LIMIT.key, + CometConf.COMET_DATA_CACHE_MEMORY_LIMIT.get(SQLConf.get).toString) + builder.putEntries( + CometConf.COMET_DATA_CACHE_BLOCK_SIZE.key, + CometConf.COMET_DATA_CACHE_BLOCK_SIZE.get(SQLConf.get).toString) + builder.putEntries( + CometConf.COMET_DATA_CACHE_SSD_LIMIT.key, + CometConf.COMET_DATA_CACHE_SSD_LIMIT.get(SQLConf.get).toString) + builder.build().toByteArray } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala index e8d2ef0006..150bc79dc7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometExecRDD.scala @@ -144,6 +144,17 @@ private[spark] class CometExecRDD( // Duplicates logic from Spark's ZippedPartitionsBaseRDD.getPreferredLocations override def getPreferredLocations(split: Partition): Seq[String] = { + // Cache-affinity: for a file-backed scan partition, prefer the host(s) that own its + // files in the data cache. Pure lookup (assignment happens at scan planning), so this + // stays cheap - it is called per partition by the DAGScheduler. When the data cache / + // locality is disabled no files are ever assigned, so this returns Nil and we fall back. + split match { + case p: CometExecPartition if p.filePaths.nonEmpty => + val hosts = CometFileLocalityManager.preferredHostsForPartition(p.filePaths) + if (hosts.nonEmpty) return hosts + case _ => + } + if (inputRDDs == null || inputRDDs.isEmpty) return Nil val idx = split.index diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometFileLocalityManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometFileLocalityManager.scala new file mode 100644 index 0000000000..a86aa351e7 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometFileLocalityManager.scala @@ -0,0 +1,212 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import java.util.concurrent.ConcurrentHashMap + +import scala.collection.mutable + +import org.apache.spark.SparkContext +import org.apache.spark.internal.Logging + +/** + * Driver-side cache-affinity scheduler for the object-store data cache + * (OBJECT_STORE_CACHE_DESIGN.md section 2.10). + * + * Without scheduler affinity a private per-executor cache is ~`1/K` effective at K executors, + * because Spark schedules S3-backed scan tasks with no locality. This manager assigns each + * Parquet file a sticky owner host and exposes it through `RDD.getPreferredLocations`, so repeat + * reads of a file route back to the executor that already cached it - turning K private caches + * into one coherent cluster cache. + * + * The driver *decides* placement rather than *observing* it (assignment-follows-scheduling): + * there is no executor->driver reporting channel and no broadcast; executor caches simply fill + * where tasks run. Modeled on `indextables_spark`'s `DriverSplitLocalityManager`. + * + * Preferred locations are hints: if the owner host is busy past `spark.locality.wait` the task + * runs elsewhere, reads through that host's cache, and is still correct (worst case a cache + * miss). Locality can never fail or stall a query. + */ +object CometFileLocalityManager extends Logging { + + /** Driver-lifetime sticky assignment: file path -> owner host. */ + private val assignments = new ConcurrentHashMap[String, String]() + + /** Hosts observed on a previous query, to detect newly-added executors for rebalance. */ + @volatile private var knownHosts: Set[String] = Set.empty + + /** Guards the multi-step assign/rebalance so concurrent queries don't interleave it. */ + private val lock = new Object + + /** + * Assign a sticky owner host to each file referenced by a query, balancing this query's work + * across `availableHosts`. + * + * 1. Files whose assigned host is still available keep it (sticky locality). 2. Unassigned + * files, and files orphaned by a lost host, go to the host with the fewest files *in this + * query* (least-loaded, balancing current work not history). 3. When new executors have + * appeared since the last query, a fair-share rebalance moves just enough files from the + * most-overloaded hosts to bring newcomers up to `ceil(filesInQuery / hosts)`. + */ + def assignFilesForQuery(paths: Seq[String], availableHosts: Seq[String]): Unit = { + if (paths.isEmpty || availableHosts.isEmpty) { + return + } + lock.synchronized { + val hosts = availableHosts.distinct + val hostSet = hosts.toSet + val hadNewHosts = hostSet.exists(h => !knownHosts.contains(h)) + knownHosts = hostSet + + // Per-query load: how many of THIS query's files each host owns. + val queryLoad = mutable.HashMap[String, Int]() + hosts.foreach(h => queryLoad(h) = 0) + + // Step 1: keep valid existing assignments; collect files needing (re)assignment. + val needsAssignment = mutable.ArrayBuffer[String]() + for (p <- paths.distinct) { + val cur = assignments.get(p) + if (cur != null && hostSet.contains(cur)) { + queryLoad(cur) += 1 + } else { + needsAssignment += p + } + } + + // Step 2: assign the rest to the least-loaded host (by current query load). + for (p <- needsAssignment) { + val host = leastLoaded(hosts, queryLoad) + assignments.put(p, host) + queryLoad(host) += 1 + } + + // Step 3: only rebalance to newcomers when the host set actually grew. + if (hadNewHosts) { + rebalanceForNewHosts(paths.distinct, hosts, queryLoad) + } + } + } + + /** + * Preferred hosts for a partition, most-preferred first. A partition whose files map to several + * hosts returns them ordered by how many of the partition's files each host owns (majority + * preference). Spark treats the result as an ordered preference list. + */ + def preferredHostsForPartition(filePaths: Seq[String]): Seq[String] = { + if (filePaths.isEmpty) { + return Nil + } + val counts = mutable.HashMap[String, Int]() + for (p <- filePaths) { + val h = assignments.get(p) + if (h != null) { + counts(h) = counts.getOrElse(h, 0) + 1 + } + } + counts.toSeq.sortBy { case (_, c) => -c }.map { case (host, _) => host } + } + + /** + * Available executor hosts: `SparkContext.getExecutorMemoryStatus` minus the driver, with a + * localhost fallback for local mode / before any executor has registered. + */ + def getAvailableHosts(sc: SparkContext): Seq[String] = { + val driverHost = sc.getConf.get("spark.driver.host", "localhost") + val hosts = sc.getExecutorMemoryStatus.keys.iterator + .map(blockManagerId => blockManagerId.split(":").head) + .filter(_ != driverHost) + .toSeq + .distinct + if (hosts.nonEmpty) hosts else Seq("localhost") + } + + /** + * Drop assignments for files no longer referenced. Pure hygiene to bound the map (`O(distinct + * files seen)`); not required for correctness. + */ + def pruneStaleAssignments(referencedPaths: Set[String]): Unit = { + val it = assignments.keySet.iterator + while (it.hasNext) { + val path = it.next() + if (!referencedPaths.contains(path)) { + it.remove() + } + } + } + + /** Current assignment for a file, if any. Exposed for tests. */ + private[comet] def assignedHost(path: String): Option[String] = + Option(assignments.get(path)) + + /** Number of live assignments. Exposed for tests. */ + private[comet] def size: Int = assignments.size + + /** Reset all state. Exposed for tests. */ + private[comet] def clear(): Unit = lock.synchronized { + assignments.clear() + knownHosts = Set.empty + } + + private def leastLoaded(hosts: Seq[String], load: mutable.HashMap[String, Int]): String = + hosts.minBy(h => load.getOrElse(h, 0)) + + /** + * Move files from over-fair-share hosts to under-fair-share hosts (the newcomers) until every + * under-loaded host reaches fair share or no overloaded host has spare files. Only this query's + * files move; untouched assignments keep their locality. + */ + private def rebalanceForNewHosts( + paths: Seq[String], + hosts: Seq[String], + queryLoad: mutable.HashMap[String, Int]): Unit = { + val fairShare = math.ceil(paths.size.toDouble / hosts.size).toInt.max(1) + + // Files of this query grouped by their (already-assigned) owner host. + val byHost = mutable.HashMap[String, mutable.ArrayBuffer[String]]() + hosts.foreach(h => byHost(h) = mutable.ArrayBuffer[String]()) + for (p <- paths) { + val h = assignments.get(p) + if (h != null && byHost.contains(h)) { + byHost(h) += p + } + } + + val underloaded = mutable.Queue[String]() + hosts.foreach(h => if (byHost(h).size < fairShare) underloaded.enqueue(h)) + + for (host <- hosts) { + while (byHost(host).size > fairShare && underloaded.nonEmpty) { + val target = underloaded.head + if (byHost(target).size >= fairShare) { + underloaded.dequeue() + } else { + val moved = byHost(host).remove(byHost(host).size - 1) + assignments.put(moved, target) + byHost(target) += moved + queryLoad(host) = (queryLoad.getOrElse(host, 1) - 1).max(0) + queryLoad(target) = queryLoad.getOrElse(target, 0) + 1 + if (byHost(target).size >= fairShare) { + underloaded.dequeue() + } + } + } + } + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index ff52c6a778..4222518a86 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -37,6 +37,7 @@ import org.apache.spark.util.collection._ import com.google.common.base.Objects +import org.apache.comet.CometConf import org.apache.comet.parquet.CometParquetUtils import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.QueryPlanSerde.exprToProto @@ -256,6 +257,18 @@ case class CometNativeScanExec( (None, Seq.empty) } + // Cache-affinity: when the data cache and locality are enabled, assign each scanned file + // a sticky owner host so repeat reads route back to the executor that cached it. This is + // a driver-side hint consumed by CometExecRDD.getPreferredLocations (section 2.10). + if (CometConf.COMET_DATA_CACHE_ENABLED.get() && + CometConf.COMET_DATA_CACHE_LOCALITY_ENABLED.get()) { + val allPaths = perPartitionFilePaths.iterator.flatten.toSeq + if (allPaths.nonEmpty) { + val hosts = CometFileLocalityManager.getAvailableHosts(sparkContext) + CometFileLocalityManager.assignFilesForQuery(allPaths, hosts) + } + } + new CometExecRDD( sparkContext, Seq.empty, diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometFileLocalityManagerSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometFileLocalityManagerSuite.scala new file mode 100644 index 0000000000..d20baaaac0 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometFileLocalityManagerSuite.scala @@ -0,0 +1,132 @@ +/* + * 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. + */ + +package org.apache.spark.sql.comet + +import org.scalatest.BeforeAndAfterEach +import org.scalatest.funsuite.AnyFunSuite + +/** + * Unit tests for [[CometFileLocalityManager]] cache-affinity semantics (section 2.10). Pure + * driver logic - no SparkContext required. + */ +class CometFileLocalityManagerSuite extends AnyFunSuite with BeforeAndAfterEach { + + override def beforeEach(): Unit = { + super.beforeEach() + CometFileLocalityManager.clear() + } + + private def files(n: Int): Seq[String] = (0 until n).map(i => s"s3://bucket/f$i.parquet") + + test("assignments are sticky across calls") { + val fs = files(6) + CometFileLocalityManager.assignFilesForQuery(fs, Seq("h1", "h2", "h3")) + val first = fs.map(f => CometFileLocalityManager.assignedHost(f)) + // Same hosts still available -> every file keeps its owner. + CometFileLocalityManager.assignFilesForQuery(fs, Seq("h1", "h2", "h3")) + val second = fs.map(f => CometFileLocalityManager.assignedHost(f)) + assert(first == second) + assert(first.forall(_.isDefined)) + } + + test("new files are assigned least-loaded per query") { + val hosts = Seq("h1", "h2", "h3") + CometFileLocalityManager.assignFilesForQuery(files(9), hosts) + // 9 files over 3 hosts -> perfectly balanced (3 each). + val counts = files(9) + .flatMap(f => CometFileLocalityManager.assignedHost(f)) + .groupBy(identity) + .map { case (host, occurrences) => (host, occurrences.size) } + assert(counts.values.forall(_ == 3), s"expected 3 files per host, got $counts") + } + + test("lost host reassigns its files, survivors untouched") { + val fs = files(6) + CometFileLocalityManager.assignFilesForQuery(fs, Seq("h1", "h2", "h3")) + val before = fs.map(f => CometFileLocalityManager.assignedHost(f).get) + + // h3 disappears. Files owned by h3 must move; files on h1/h2 must stay. + CometFileLocalityManager.assignFilesForQuery(fs, Seq("h1", "h2")) + val after = fs.map(f => CometFileLocalityManager.assignedHost(f).get) + + for ((f, i) <- fs.zipWithIndex) { + if (before(i) == "h3") { + assert(after(i) != "h3", s"$f should have been reassigned off the lost host") + assert(Set("h1", "h2").contains(after(i))) + } else { + assert(after(i) == before(i), s"$f on a surviving host should not move") + } + } + assert(!after.contains("h3")) + } + + test("new host triggers fair-share rebalance moving only the excess") { + val fs = files(6) + // Start with two hosts: 3 files each. + CometFileLocalityManager.assignFilesForQuery(fs, Seq("h1", "h2")) + assert(hostCounts(fs).values.forall(_ == 3)) + + // A third host appears. fairShare = ceil(6/3) = 2. Each old host should shed 1 file to + // the newcomer; the newcomer ends with 2. + CometFileLocalityManager.assignFilesForQuery(fs, Seq("h1", "h2", "h3")) + val counts = hostCounts(fs) + assert(counts.getOrElse("h3", 0) == 2, s"newcomer should reach fair share: $counts") + assert(counts.getOrElse("h1", 0) == 2, s"$counts") + assert(counts.getOrElse("h2", 0) == 2, s"$counts") + } + + test("preferredHostsForPartition orders hosts by majority ownership") { + val fs = files(3) + // Force a known layout: f0,f1 -> hA ; f2 -> hB by assigning on a single host first then + // reassigning is awkward; instead assign across two hosts and read back the mapping. + CometFileLocalityManager.assignFilesForQuery(fs, Seq("hA", "hB")) + val prefs = CometFileLocalityManager.preferredHostsForPartition(fs) + // The most-owning host for this partition must come first. + val counts = hostCounts(fs) + val expectedFirst = counts.maxBy(_._2)._1 + assert(prefs.headOption.contains(expectedFirst)) + assert(prefs.toSet == counts.keySet) + } + + test("preferredHostsForPartition returns Nil for unassigned files") { + assert(CometFileLocalityManager.preferredHostsForPartition(files(2)).isEmpty) + } + + test("prune drops only unreferenced paths") { + val fs = files(4) + CometFileLocalityManager.assignFilesForQuery(fs, Seq("h1", "h2")) + assert(CometFileLocalityManager.size == 4) + CometFileLocalityManager.pruneStaleAssignments(Set(fs(0), fs(1))) + assert(CometFileLocalityManager.size == 2) + assert(CometFileLocalityManager.assignedHost(fs(0)).isDefined) + assert(CometFileLocalityManager.assignedHost(fs(2)).isEmpty) + } + + test("empty inputs are no-ops") { + CometFileLocalityManager.assignFilesForQuery(Seq.empty, Seq("h1")) + CometFileLocalityManager.assignFilesForQuery(files(1), Seq.empty) + assert(CometFileLocalityManager.size == 0) + } + + private def hostCounts(fs: Seq[String]): Map[String, Int] = + fs.flatMap(f => CometFileLocalityManager.assignedHost(f)) + .groupBy(identity) + .map { case (host, occurrences) => (host, occurrences.size) } +}