sunchao commented on code in PR #4828: URL: https://github.com/apache/datafusion-comet/pull/4828#discussion_r4104765743
########## 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<u64>]) -> Result<(Vec<Bytes>, FileVersion)>; +} + +/// Interned file identities and their captured versions. +struct FileTable { + ids: HashMap<FileKey, u64>, + versions: HashMap<u64, FileVersion>, + 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<Mutex<Shard>>, + files: Mutex<FileTable>, + memory_budget: AtomicU64, + metrics: Arc<Metrics>, +} + +impl BlockCache { + /// Build a cache from `config` (normalized to valid values). + pub fn new(config: BlockCacheConfig) -> Arc<Self> { + 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<u64>], + fetcher: &dyn RangeFetcher, + ) -> Result<Vec<Bytes>> { + 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<u32> = 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<u32, Bytes> = HashMap::with_capacity(needed.len()); + let mut missing: Vec<u32> = 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<u32, Bytes>, + ) -> Result<()> { + // --- claim phase: decide which blocks we own vs. wait on --- + let mut owned: Vec<u32> = Vec::new(); + let mut senders: HashMap<u32, tokio::sync::oneshot::Sender<FetchResult>> = 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::<FetchResult>(); + 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<CacheError> = 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; Review Comment: [P2] Handle EOF before sending rounded ranges to the legacy HDFS backend. In a build using `--no-default-features --features hdfs`, `HadoopFileSystem::get_opts` forwards bounded ranges unchanged and `read_range` rejects short reads. Enabling this cache therefore makes a valid footer read fail whenever the file's final block is partial. The assumption that every wrapped store truncates at EOF does not hold for this supported backend. Clamp requests using the file size or make that backend support truncated final ranges. Evidence: `native/hdfs/src/object_store/hdfs.rs:170` passes the requested range through, and lines 107–117 return an error when actual bytes differ from requested bytes. A source-matched strict-read mock successfully read footer range `1016..1024` directly. Through the cache it failed with `expected size 4194304 and actual size 1024`. No live HDFS cluster was used. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala: ########## @@ -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() && Review Comment: [P2] Wire affinity into fused native-plan construction as well. For a native projection, filter, or aggregate above this scan, the parent `CometNativeExec` executes the fused subtree without calling this scan's `doExecuteColumnar`. Consequently, these assignments never run. The parent's `CometExecRDD` also receives no `perPartitionFilePaths`, so it returns no cache-affinity hints even if assignments already exist. Carry the paths and assignment step through `NativeExecContext` so ordinary fused queries receive the advertised locality behavior. Evidence: Exact-head source trace for a native `SELECT id + 1 FROM parquet...`: `operators.scala:buildNativeContext` treats `CometNativeScanExec` as a `CometNativeExec` and skips executing it. `findAllPlanData` reads only its common/per-partition bytes. `executeColumnarWithContext` constructs the RDD without file paths, leaving `CometExecPartition.filePaths` empty. This was verified statically, not by a local JVM query. ########## 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<u64>]) -> Result<(Vec<Bytes>, FileVersion)>; +} + +/// Interned file identities and their captured versions. +struct FileTable { + ids: HashMap<FileKey, u64>, + versions: HashMap<u64, FileVersion>, + 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<Mutex<Shard>>, + files: Mutex<FileTable>, + memory_budget: AtomicU64, + metrics: Arc<Metrics>, +} + +impl BlockCache { + /// Build a cache from `config` (normalized to valid values). + pub fn new(config: BlockCacheConfig) -> Arc<Self> { + 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<u64>], + fetcher: &dyn RangeFetcher, + ) -> Result<Vec<Bytes>> { + 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<u32> = 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<u32, Bytes> = HashMap::with_capacity(needed.len()); + let mut missing: Vec<u32> = 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<u32, Bytes>, + ) -> Result<()> { + // --- claim phase: decide which blocks we own vs. wait on --- + let mut owned: Vec<u32> = Vec::new(); + let mut senders: HashMap<u32, tokio::sync::oneshot::Sender<FetchResult>> = 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::<FetchResult>(); + 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<CacheError> = 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) Review Comment: [P2] Account for the backing allocation retained by each cached slice. Blocks produced by `full.slice` share the entire coalesced fetch allocation, but `Shard::block_cost` charges only each slice's length. Evicting sibling blocks therefore subtracts memory that remains allocated. This breaks the configured memory limit even after callers release their results, consuming executor overhead outside Spark's pools. Give blocks independent allocations or account for and evict shared backing allocations together. Evidence: A tracked `Bytes` owner reproduced this with the normal 16 shards, default 4 MiB blocks and 16 MiB coalescing, and a 128 MiB budget. After reading 32 distinct 16 MiB files and dropping every result, 256 MiB of fetched allocations remained live despite 112 evictions. Reducing the budget to zero released them. ########## 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<u64>]) -> Result<(Vec<Bytes>, FileVersion)>; +} + +/// Interned file identities and their captured versions. +struct FileTable { + ids: HashMap<FileKey, u64>, + versions: HashMap<u64, FileVersion>, + 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<Mutex<Shard>>, + files: Mutex<FileTable>, + memory_budget: AtomicU64, + metrics: Arc<Metrics>, +} + +impl BlockCache { + /// Build a cache from `config` (normalized to valid values). + pub fn new(config: BlockCacheConfig) -> Arc<Self> { + 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<u64>], + fetcher: &dyn RangeFetcher, + ) -> Result<Vec<Bytes>> { + 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<u32> = 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<u32, Bytes> = HashMap::with_capacity(needed.len()); + let mut missing: Vec<u32> = 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<u32, Bytes>, + ) -> Result<()> { + // --- claim phase: decide which blocks we own vs. wait on --- + let mut owned: Vec<u32> = Vec::new(); + let mut senders: HashMap<u32, tokio::sync::oneshot::Sender<FetchResult>> = 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::<FetchResult>(); + 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); Review Comment: [P1] Clean up in-flight claims when their owner is cancelled. If a read is dropped while awaiting `fetcher.fetch`, its senders disappear but these entries remain in `in_flight`. Every subsequent read of those blocks then receives `fetch owner dropped before delivering block`, without retrying storage. Because the cache is process-wide, one cancelled read can break later queries until executor restart. Add cancellation-safe ownership guards or remove cancelled claims so another caller can fetch. Evidence: At this head, a disposable Tokio test started a blocked fetch, aborted its owner, then retried three times with a healthy fetcher. All three retries returned the same Internal error and the healthy fetcher recorded zero calls. `invalidate_file` also leaves in-flight entries untouched. ########## 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<u64>]) -> Result<(Vec<Bytes>, FileVersion)>; +} + +/// Interned file identities and their captured versions. +struct FileTable { + ids: HashMap<FileKey, u64>, + versions: HashMap<u64, FileVersion>, + 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<Mutex<Shard>>, + files: Mutex<FileTable>, + memory_budget: AtomicU64, + metrics: Arc<Metrics>, +} + +impl BlockCache { + /// Build a cache from `config` (normalized to valid values). + pub fn new(config: BlockCacheConfig) -> Arc<Self> { + 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<u64>], + fetcher: &dyn RangeFetcher, + ) -> Result<Vec<Bytes>> { + 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<u32> = 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<u32, Bytes> = HashMap::with_capacity(needed.len()); + let mut missing: Vec<u32> = 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<u32, Bytes>, + ) -> Result<()> { + // --- claim phase: decide which blocks we own vs. wait on --- + let mut owned: Vec<u32> = Vec::new(); + let mut senders: HashMap<u32, tokio::sync::oneshot::Sender<FetchResult>> = 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::<FetchResult>(); + 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<CacheError> = 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 { Review Comment: [P2] Preserve concurrency between independent missing runs. This await completes each upstream request before starting the next. Previously, `ObjectStore::get_ranges` fetched disjoint ranges concurrently, and Parquet's reader uses that path for column reads. With the cache enabled, cold scans and scans exceeding cache capacity accumulate serial request latency. Fetch owned runs with bounded concurrency while retaining the existing rule that owned work completes before awaiting other owners. Evidence: An adapter reproduction requested eight disjoint ranges from the same delayed in-memory store, with a controlled 25 ms delay per GET. Direct `get_ranges` reached concurrency 8 and completed in 26.4 ms. The cold cached call reached concurrency 1 and took 210.6 ms, returning identical bytes. `object_store` 0.13.2 uses `buffered(OBJECT_STORE_COALESCE_PARALLEL)` for the direct path. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
