Skip to content

Commit c21c291

Browse files
authored
feat(admission): count the on-disk liquid tier toward the gate budget (#9)
* feat(admission): count the on-disk liquid tier toward the gate budget The footprint admission gate sized its budget against RAM only, so with the disk tier enabled a scan whose liquid footprint exceeded memory was bypassed to the fallback mount even though it would fit on the on-disk liquid tier without thrashing. Size the gate against both tiers, weighted differently: memory keeps the compaction overcommit (tolerance, the measured RAM crossover) while the disk tier counts at face value — threshold = memory * tolerance + disk. With the disk tier off (max_disk_bytes == 0) the threshold is exactly memory * tolerance, so gate behaviour is unchanged there. Extract the threshold into a pure admission_threshold() and add unit tests pinning the formula (these run everywhere; the cache-based gate tests remain Linux-only via direct_io). * docs(admission): don't link public docs to private admission_threshold cargo doc --document-private-items with -D warnings denies public items linking to a private fn (rustdoc::private_intra_doc_links). State the formula inline in the public docs; keep the link only from private should_bypass.
1 parent d04d2b3 commit c21c291

1 file changed

Lines changed: 159 additions & 33 deletions

File tree

  • src/datafusion/src/optimizers

‎src/datafusion/src/optimizers/mod.rs‎

Lines changed: 159 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -32,29 +32,35 @@ use crate::{LiquidCacheParquetRef, LiquidParquetSource, cache::ColumnSqueezeHint
3232

3333
/// Parameters for the footprint-based admission gate.
3434
///
35-
/// A parquet scan is routed through LiquidCache only if its estimated in-memory
36-
/// liquid footprint stays within `budget × tolerance`; larger scans are left as
35+
/// A parquet scan is routed through LiquidCache only if its estimated liquid
36+
/// footprint stays within the admission threshold; larger scans are left as
3737
/// vanilla parquet reads (which, if the object store is a cached mount, read
3838
/// from it).
3939
///
40+
/// The threshold spans both cache tiers, weighted differently:
41+
/// `memory × tolerance + disk`. A scan that overflows
42+
/// RAM spills to the on-disk liquid tier rather than thrashing, so disk capacity
43+
/// counts toward what fits — at face value, since only the RAM tier has the
44+
/// measured compaction overcommit. With the disk tier off it is just
45+
/// `memory × tolerance`.
46+
///
4047
/// The estimate multiplies the raw required parquet bytes by `expansion`
4148
/// (parquet -> liquid in-memory blow-up) and `safety` (extra margin); both are
42-
/// `>= 1.0` so the estimate is conservative (over-counts). `tolerance` (`>= 1.0`)
43-
/// is the one knob applied to the *budget*: it encodes that LiquidCache compacts
44-
/// in RAM and still beats the fallback mount until footprint reaches ~5x the
45-
/// budget, so the gate tolerates that overcommit before bypassing.
49+
/// `>= 1.0` so the estimate is conservative (over-counts).
4650
#[derive(Debug, Clone, Copy)]
4751
pub struct AdmissionGate {
4852
/// Parquet-bytes -> liquid-in-memory-bytes multiplier (>= 1.0). Inflates the
4953
/// estimate (conservative direction).
5054
pub expansion: f64,
5155
/// Extra safety margin on the estimate (>= 1.0). Conservative direction.
5256
pub safety: f64,
53-
/// Overcommit tolerance: how far the estimated liquid footprint may exceed
54-
/// the budget before the scan is denied, in multiples of the budget.
55-
/// LiquidCache compacts in RAM up to ~5x over budget before it starts
56-
/// thrashing, so caching still wins in that band. This is the one *relaxing*
57-
/// knob, so it is clamped to `[1.0, 5.0]` (5.0 = the measured crossover).
57+
/// Overcommit tolerance on the *memory* tier: how far the estimated liquid
58+
/// footprint may exceed the RAM budget before it counts against the scan, in
59+
/// multiples of the RAM budget. LiquidCache compacts in RAM up to ~5x over
60+
/// budget before it starts thrashing, so caching still wins in that band.
61+
/// This is the one *relaxing* knob, so it is clamped to `[1.0, 5.0]` (5.0 =
62+
/// the measured crossover). It applies only to memory; the disk tier is
63+
/// counted at face value (`memory × tolerance + disk`), never scaled by it.
5864
pub tolerance: f64,
5965
/// Fail-loud mode. When `true`, a panic during footprint estimation is *not*
6066
/// caught — it aborts the query — so estimation bugs surface immediately.
@@ -96,8 +102,9 @@ impl LocalModeOptimizer {
96102

97103
/// Enable the footprint-based admission gate. A parquet scan is cached only
98104
/// when its estimated liquid footprint (raw required bytes x `expansion` x
99-
/// `safety`) stays within `budget × tolerance`; otherwise it is read directly
100-
/// from the parquet source, bypassing the cache. See [`AdmissionGate`].
105+
/// `safety`) stays within the admission threshold (`memory × tolerance +
106+
/// disk`); otherwise it is read directly from the parquet source, bypassing
107+
/// the cache. See [`AdmissionGate`].
101108
///
102109
/// Inputs are sanitized so a misconfigured value can never make the gate
103110
/// unsound: `expansion`/`safety` are forced finite and `>= 1.0` (their only
@@ -138,13 +145,21 @@ impl PhysicalOptimizerRule for LocalModeOptimizer {
138145
let analysis = HintAnalyzer::analyze(&plan);
139146
let cache = self.cache.clone();
140147
let admission = self.admission;
141-
let budget = self.cache.max_memory_bytes() as u64;
148+
// The gate sizes against both cache tiers, not just RAM: when a scan
149+
// overflows memory its entries spill to the on-disk liquid tier (NVMe)
150+
// instead of thrashing. They are weighted differently (see
151+
// `admission_threshold`) — RAM carries the compaction overcommit, disk
152+
// counts at face value — so they are passed through separately. With the
153+
// disk tier off (`max_disk_bytes == 0`) the threshold is exactly the RAM
154+
// budget, so gate behaviour is unchanged there.
155+
let memory_budget = self.cache.max_memory_bytes() as u64;
156+
let disk_budget = self.cache.max_disk_bytes() as u64;
142157
let mut convert = |node: &Arc<dyn ExecutionPlan>, hints: ColumnSqueezeHints| {
143158
// Leave scans whose estimated liquid footprint exceeds the budget as
144159
// vanilla parquet reads, so oversized scans don't thrash the cache.
145160
if let Some(gate) = admission
146161
&& let Some((cfg, src)) = parquet_scan_parts(node)
147-
&& should_bypass_guarded(cfg, src, gate, budget)
162+
&& should_bypass_guarded(cfg, src, gate, memory_budget, disk_budget)
148163
{
149164
return None;
150165
}
@@ -324,37 +339,48 @@ fn estimate_required_bytes(cfg: &FileScanConfig, src: &ParquetSource) -> Footpri
324339
}
325340
}
326341

342+
/// The admission threshold in bytes: the largest estimated liquid footprint the
343+
/// cache admits before a scan is bypassed.
344+
///
345+
/// The two tiers are weighted differently. The memory tier carries the
346+
/// compaction overcommit — `tolerance`, the measured RAM crossover where
347+
/// LiquidCache still beats the fallback mount despite spilling. The on-disk
348+
/// liquid tier is counted at **face value (1×)**: a scan that overflows RAM
349+
/// spills to NVMe without thrashing, so its capacity needs no overcommit — and
350+
/// extending the RAM crossover factor to disk has no evidence behind it. Hence
351+
/// `memory × tolerance + disk` rather than `(memory + disk) × tolerance`.
352+
///
353+
/// Computed in `f64` so the `memory × tolerance` product cannot overflow `u64`.
354+
/// A zero threshold (both tiers unsized / cache disabled) bypasses any non-empty
355+
/// scan and admits only zero-footprint ones.
356+
fn admission_threshold(memory_budget: u64, disk_budget: u64, tolerance: f64) -> f64 {
357+
(memory_budget as f64) * tolerance + (disk_budget as f64)
358+
}
359+
327360
/// Decide whether a parquet scan should bypass the cache (read as vanilla
328361
/// parquet) rather than be transcoded into LiquidCache.
329362
///
330363
/// The scan's estimated liquid footprint is `raw × expansion × safety`, where
331364
/// `raw` is [`estimate_required_bytes`]. It bypasses when that footprint exceeds
332-
/// `budget × tolerance` — equivalently, when the pressure ratio
333-
/// `footprint / budget` exceeds `tolerance`. The comparison is done in finite
334-
/// `f64` space to avoid the overflow/truncation of integer `budget × tolerance`.
335-
///
336-
/// `budget == 0` (cache disabled / unsized) bypasses any non-empty scan and
337-
/// admits only zero-footprint ones.
365+
/// the [`admission_threshold`] (`memory × tolerance + disk`). The comparison is
366+
/// done in finite `f64` space to avoid integer overflow.
338367
fn should_bypass(
339368
cfg: &FileScanConfig,
340369
src: &ParquetSource,
341370
gate: AdmissionGate,
342-
budget: u64,
371+
memory_budget: u64,
372+
disk_budget: u64,
343373
) -> bool {
344374
let est = estimate_required_bytes(cfg, src);
345375
// Multipliers are already sanitized to finite, >= 1.0 in `with_admission_gate`.
346376
let footprint = (est.raw_bytes as f64) * gate.expansion * gate.safety;
347-
let bypass = if budget == 0 {
348-
footprint > 0.0
349-
} else {
350-
footprint / (budget as f64) > gate.tolerance
351-
};
377+
let threshold = admission_threshold(memory_budget, disk_budget, gate.tolerance);
378+
let bypass = footprint > threshold;
352379

353380
// One line per admission decision. Without this the gate is a black box and
354381
// every tuning cycle costs a benchmark run to infer decisions from side
355382
// effects. `fallback_files > 0` is the flag that the catalog has no
356383
// per-column sizes for this table (estimate is coarse / over-counts).
357-
let threshold = (budget as f64) * gate.tolerance;
358384
let path = cfg
359385
.file_groups
360386
.iter()
@@ -366,15 +392,18 @@ fn should_bypass(
366392
target: "liquid_cache::admission",
367393
"admission {verdict}: file={path} projected_cols={cols} \
368394
charged_files={charged} partitioned_files={total} \
369-
fallback_files={fb} raw_bytes={raw} footprint_bytes={fp} budget_bytes={budget} \
370-
threshold_bytes={thr} expansion={exp} safety={saf} tolerance={tol}",
395+
fallback_files={fb} raw_bytes={raw} footprint_bytes={fp} \
396+
memory_bytes={mem} disk_bytes={disk} threshold_bytes={thr} \
397+
expansion={exp} safety={saf} tolerance={tol}",
371398
verdict = if bypass { "BYPASS" } else { "ADMIT" },
372399
cols = est.required_cols,
373400
charged = est.charged_files,
374401
total = est.partitioned_files,
375402
fb = est.fallback_files,
376403
raw = est.raw_bytes,
377404
fp = footprint as u64,
405+
mem = memory_budget,
406+
disk = disk_budget,
378407
thr = threshold as u64,
379408
exp = gate.expansion,
380409
saf = gate.safety,
@@ -405,10 +434,11 @@ fn should_bypass_guarded(
405434
cfg: &FileScanConfig,
406435
src: &ParquetSource,
407436
gate: AdmissionGate,
408-
budget: u64,
437+
memory_budget: u64,
438+
disk_budget: u64,
409439
) -> bool {
410440
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
411-
should_bypass(cfg, src, gate, budget)
441+
should_bypass(cfg, src, gate, memory_budget, disk_budget)
412442
}));
413443
match result {
414444
Ok(bypass) => bypass,
@@ -689,6 +719,13 @@ mod tests {
689719
}
690720

691721
async fn build_cache_with_budget(max_memory_bytes: usize) -> LiquidCacheParquetRef {
722+
build_cache_with_mem_disk(max_memory_bytes, 0).await
723+
}
724+
725+
async fn build_cache_with_mem_disk(
726+
max_memory_bytes: usize,
727+
max_disk_bytes: usize,
728+
) -> LiquidCacheParquetRef {
692729
let tmp_dir = tempfile::tempdir().unwrap();
693730
let store = t4::mount(tmp_dir.path().join("liquid_cache.t4"))
694731
.await
@@ -697,7 +734,7 @@ mod tests {
697734
LiquidCacheParquet::new(
698735
8192,
699736
max_memory_bytes,
700-
usize::MAX,
737+
max_disk_bytes,
701738
store,
702739
Box::new(LiquidPolicy::new()),
703740
Box::new(TranscodeSqueezeEvict),
@@ -768,6 +805,52 @@ mod tests {
768805
"fitting scan should be wrapped in LiquidParquetSource"
769806
);
770807
}
808+
809+
/// The budget counts the on-disk liquid tier, not just memory: the same scan
810+
/// that a memory-only budget bypasses is admitted once the disk tier has room
811+
/// for it (evicted entries spill to disk instead of thrashing).
812+
#[tokio::test]
813+
async fn test_admission_gate_counts_disk_tier() {
814+
let ctx = SessionContext::new();
815+
ctx.register_parquet(
816+
"nano_hits",
817+
"../../examples/nano_hits.parquet",
818+
Default::default(),
819+
)
820+
.await
821+
.unwrap();
822+
let plan = ctx
823+
.sql("SELECT * FROM nano_hits WHERE \"URL\" like 'https://%' limit 10")
824+
.await
825+
.unwrap()
826+
.create_physical_plan()
827+
.await
828+
.unwrap();
829+
let config = ConfigOptions::default();
830+
831+
// A huge expansion inflates the estimate past the 1 MB memory budget, and
832+
// with no disk tier there is nowhere else for it to fit → bypass.
833+
let mem_only = LocalModeOptimizer::new(build_cache_with_mem_disk(1_000_000, 0).await)
834+
.with_admission_gate(1e9, 1.0, 1.0, false)
835+
.optimize(plan.clone(), &config)
836+
.unwrap();
837+
assert!(
838+
!has_liquid_source(&mem_only),
839+
"with a memory-only budget the oversized estimate should bypass"
840+
);
841+
842+
// Same 1 MB memory and same estimate, but now a large disk tier — the
843+
// budget is memory + disk, so the scan fits and is cached.
844+
let with_disk =
845+
LocalModeOptimizer::new(build_cache_with_mem_disk(1_000_000, usize::MAX).await)
846+
.with_admission_gate(1e9, 1.0, 1.0, false)
847+
.optimize(plan, &config)
848+
.unwrap();
849+
assert!(
850+
has_liquid_source(&with_disk),
851+
"the disk tier's capacity should count toward the budget and admit the scan"
852+
);
853+
}
771854
}
772855

773856
/// Pure unit tests for the footprint byte-math (no cache / no t4 mount, so they
@@ -850,3 +933,46 @@ mod footprint_tests {
850933
assert_eq!(file_required_bytes(Some(&s), 2000, &[0]), (2000, true));
851934
}
852935
}
936+
937+
/// Pure unit tests for the admission threshold arithmetic (no cache / no t4
938+
/// mount, so they run everywhere — including where `direct_io` is unavailable).
939+
#[cfg(test)]
940+
mod threshold_tests {
941+
use super::admission_threshold;
942+
943+
#[test]
944+
fn memory_carries_tolerance_disk_is_face_value() {
945+
// memory × tolerance + disk = 10×3 + 100 = 130. This pins the exact
946+
// formula: it is NOT (memory+disk)×tolerance (=330), max(mem,disk) (=100),
947+
// nor disk alone (=100).
948+
assert_eq!(admission_threshold(10, 100, 3.0), 130.0);
949+
}
950+
951+
#[test]
952+
fn disk_off_is_memory_times_tolerance() {
953+
// With no disk tier the threshold is exactly the RAM budget × tolerance,
954+
// so gate behaviour matches the pre-disk model.
955+
assert_eq!(admission_threshold(1_000, 0, 3.0), 3_000.0);
956+
}
957+
958+
#[test]
959+
fn disk_gets_no_overcommit() {
960+
// A pure-disk budget contributes at 1×, never scaled by tolerance.
961+
assert_eq!(admission_threshold(0, 500, 5.0), 500.0);
962+
}
963+
964+
#[test]
965+
fn both_tiers_unsized_is_zero() {
966+
// Zero threshold => any non-empty scan bypasses (cache disabled/unsized).
967+
assert_eq!(admission_threshold(0, 0, 3.0), 0.0);
968+
}
969+
970+
#[test]
971+
fn large_memory_times_tolerance_does_not_overflow() {
972+
// memory × tolerance is computed in f64, so a budget near u64::MAX cannot
973+
// wrap (the bug an integer `budget * tolerance` would have).
974+
let t = admission_threshold(u64::MAX, 0, 5.0);
975+
assert!(t > u64::MAX as f64);
976+
assert!(t.is_finite());
977+
}
978+
}

0 commit comments

Comments
 (0)