query_planning/
execution_contracts.rs

1//! Data-only memory-control and source-chunk planning contracts.
2//!
3//! Planning sizes source chunks from these policies; the runtime enforces the
4//! device budget they carry.
5
6use crate::error::{QueryError, QueryErrorCode, Result};
7
8/// The device budget and source chunking one execution runs under. An admitted
9/// attempt's budget is its grant.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub struct MemoryControlPolicy {
12    /// Allocation cap of the attempt. The runtime's device ledger enforces it;
13    /// planning only sizes chunks against it.
14    pub device_budget_bytes: u64,
15    /// Source-byte grouping target, compared with each chunk's decoded size
16    /// estimate (`crate::chunk::decoded`). `u64::MAX` leaves the configured
17    /// target unset; source tasks derive a finite target from the admitted grant.
18    pub max_source_chunk_bytes: u64,
19    /// Row groups coalesced into one source chunk. The throughput planning
20    /// mode may raise it to the row-group cap `crate::chunk` owns.
21    pub max_row_groups_per_chunk: usize,
22    /// Attempts the engine's retry may spend splitting one chunk after an
23    /// allocation failure.
24    pub max_retry_attempts: usize,
25    /// Smallest chunk, in row groups, a retry may split down to.
26    pub min_retry_chunk_row_groups: usize,
27}
28
29impl MemoryControlPolicy {
30    /// No configured source-byte ceiling; the immutable grant sizes batches.
31    pub const DEFAULT_MAX_SOURCE_CHUNK_BYTES: u64 = u64::MAX;
32    /// One row group per chunk. The default throughput planning mode raises
33    /// this cap for local Parquet sources; the conservative mode keeps it.
34    pub const DEFAULT_MAX_ROW_GROUPS_PER_CHUNK: usize = 1;
35    /// Default retry limit for splitting a source chunk after allocation failure.
36    pub const DEFAULT_MAX_RETRY_ATTEMPTS: usize = 3;
37    /// Default minimum retry chunk of one row group.
38    pub const DEFAULT_MIN_RETRY_CHUNK_ROW_GROUPS: usize = 1;
39
40    /// Rejects a zero in any field; the error names the field in its `field`
41    /// fact. Zero is not a "disabled" value here, since each field bounds a
42    /// loop or a division in chunk planning.
43    pub fn validate(self) -> Result<()> {
44        let Self {
45            device_budget_bytes,
46            max_source_chunk_bytes,
47            max_row_groups_per_chunk,
48            max_retry_attempts,
49            min_retry_chunk_row_groups,
50        } = self;
51        validate_nonzero_policy_field("device_budget_bytes", device_budget_bytes == 0)?;
52        validate_nonzero_policy_field("max_source_chunk_bytes", max_source_chunk_bytes == 0)?;
53        validate_nonzero_policy_field("max_row_groups_per_chunk", max_row_groups_per_chunk == 0)?;
54        validate_nonzero_policy_field("max_retry_attempts", max_retry_attempts == 0)?;
55        validate_nonzero_policy_field("min_retry_chunk_row_groups", min_retry_chunk_row_groups == 0)
56    }
57}
58
59/// How source chunk planning sizes local Parquet chunks. The chosen mode and
60/// the reason for it are recorded in each chunk plan.
61///
62/// In both modes the source-byte target bounds each chunk's decoded size, and
63/// a chunk whose allocation the query cap rejects is split by row group and
64/// retried within the policy's retry limits.
65#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
66pub enum SourceChunkPlanningMode {
67    /// Chunks follow `max_row_groups_per_chunk`, one row group by default.
68    /// Each chunk is then the smallest retry unit, but per-chunk work such as
69    /// partial aggregation grows with the table's row-group count.
70    Conservative,
71    /// Raises a one-row-group cap so a chunk of a local Parquet source can hold
72    /// many row groups up to the source-byte target; `crate::chunk` owns the
73    /// row-group cap and the source conditions. A source that does not qualify
74    /// keeps conservative chunks.
75    #[default]
76    LocalParquetThroughput,
77}
78
79impl SourceChunkPlanningMode {
80    /// Stable source-planning mode label used in reports and configuration.
81    pub const fn as_str(self) -> &'static str {
82        match self {
83            Self::Conservative => "conservative",
84            Self::LocalParquetThroughput => "local_parquet_throughput",
85        }
86    }
87}
88
89impl std::fmt::Display for SourceChunkPlanningMode {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        f.write_str(self.as_str())
92    }
93}
94
95fn validate_nonzero_policy_field(field: &'static str, is_zero: bool) -> Result<()> {
96    if is_zero {
97        return Err(QueryError::new(
98            QueryErrorCode::InvalidConfiguration,
99            "query_engine.memory_policy.validate",
100            "bounded memory-control policy fields must be greater than zero",
101        )
102        .with_fact("policy", "bounded")
103        .with_fact("field", field)
104        .with_fact("reason", "zero"));
105    }
106    Ok(())
107}
query_planning/execution_contracts.rs

query_planning/
execution_contracts.rs
