query_planning/chunk/
policy.rs

1use super::memory::SourceTaskMemoryModel;
2use crate::{
3    error::{QueryError, QueryErrorCode, Result as QueryResult},
4    execution_contracts::{MemoryControlPolicy, SourceChunkPlanningMode},
5    source::{LocalParquetFile, LocalParquetSource, ParquetReadStrategy},
6};
7
8/// Row-group cap used by local-throughput source planning when the configured
9/// policy is the conservative one-row-group shape. It bounds how much source
10/// one chunk retry reads again; the source-byte target sizes the chunk.
11pub const LOCAL_PARQUET_THROUGHPUT_MAX_ROW_GROUPS_PER_CHUNK: usize =
12    SourceChunkPlanningOptions::DEFAULT_LOCAL_PARQUET_THROUGHPUT_MAX_ROW_GROUPS_PER_CHUNK;
13
14/// Expert policy for local Parquet source chunk planning heuristics.
15///
16/// This controls when [`SourceChunkPlanningMode::LocalParquetThroughput`] may
17/// widen conservative one-row-group chunks. It is intentionally a Rust-only
18/// policy until production evidence shows a stable tuning model.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20#[non_exhaustive]
21pub struct SourceChunkPlanningOptions {
22    local_parquet_throughput_max_row_groups_per_chunk: usize,
23}
24
25impl SourceChunkPlanningOptions {
26    pub(crate) const DEFAULT_LOCAL_PARQUET_THROUGHPUT_MAX_ROW_GROUPS_PER_CHUNK: usize = 1024;
27
28    pub(crate) fn local_parquet_throughput_max_row_groups_per_chunk(&self) -> usize {
29        self.local_parquet_throughput_max_row_groups_per_chunk
30    }
31
32    /// Sets the row-group cap available when local-throughput planning widens a
33    /// conservative chunk. The value must be nonzero; call [`Self::validate`]
34    /// after changing options.
35    pub fn with_local_parquet_throughput_max_row_groups_per_chunk(mut self, row_groups: usize) -> Self {
36        self.local_parquet_throughput_max_row_groups_per_chunk = row_groups;
37        self
38    }
39
40    /// Rejects a zero row-group limit with [`QueryErrorCode::InvalidConfiguration`].
41    pub fn validate(&self) -> QueryResult<()> {
42        validate_source_chunk_planning_nonzero_usize(
43            "local_parquet_throughput_max_row_groups_per_chunk",
44            self.local_parquet_throughput_max_row_groups_per_chunk,
45        )
46    }
47}
48
49impl Default for SourceChunkPlanningOptions {
50    fn default() -> Self {
51        Self {
52            local_parquet_throughput_max_row_groups_per_chunk:
53                Self::DEFAULT_LOCAL_PARQUET_THROUGHPUT_MAX_ROW_GROUPS_PER_CHUNK,
54        }
55    }
56}
57
58fn validate_source_chunk_planning_nonzero_usize(field: &'static str, value: usize) -> QueryResult<()> {
59    if value == 0 {
60        return Err(source_chunk_planning_zero_error(field, value));
61    }
62    Ok(())
63}
64
65fn source_chunk_planning_zero_error(field: &'static str, value: impl std::fmt::Display) -> QueryError {
66    QueryError::new(
67        QueryErrorCode::InvalidConfiguration,
68        "query_engine.source_chunk_planning_options.validate",
69        "source chunk planning option must be greater than zero",
70    )
71    .with_fact("field", field)
72    .with_fact("value", value)
73    .with_fact("reason", "zero")
74}
75
76/// Policy that controls how a local Parquet source is split into chunks.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub struct SourceChunkPlanningPolicy {
79    /// Target upper bound, in decoded bytes ([`crate::chunk::decoded`]), for
80    /// one chunk's read columns. A row group without a decoded estimate counts
81    /// its encoded source-read bytes.
82    ///
83    /// `None` means row-group grouping is constrained only by
84    /// [`max_row_groups_per_chunk`](Self::max_row_groups_per_chunk).
85    pub max_source_chunk_bytes: Option<u64>,
86    /// Maximum number of row groups assigned to one chunk.
87    ///
88    /// Callers must pass a value greater than zero; unbounded policies use
89    /// `usize::MAX`.
90    pub max_row_groups_per_chunk: usize,
91    source_chunk_planning_mode: SourceChunkPlanningMode,
92    source_chunk_planning_reason: &'static str,
93}
94
95impl SourceChunkPlanningPolicy {
96    pub const fn new(max_source_chunk_bytes: Option<u64>, max_row_groups_per_chunk: usize) -> Self {
97        Self {
98            max_source_chunk_bytes,
99            max_row_groups_per_chunk,
100            source_chunk_planning_mode: SourceChunkPlanningMode::Conservative,
101            source_chunk_planning_reason: "explicit_policy",
102        }
103    }
104
105    pub(crate) fn source_chunk_planning_mode(self) -> SourceChunkPlanningMode {
106        self.source_chunk_planning_mode
107    }
108
109    pub fn source_chunk_planning_reason(self) -> &'static str {
110        self.source_chunk_planning_reason
111    }
112
113    pub fn mark_downstream_cap(&mut self, reason: &'static str) {
114        self.source_chunk_planning_reason = reason;
115    }
116
117    /// Derives a source-task policy from the bounded memory policy and slot model.
118    ///
119    /// The default source-byte setting uses
120    /// `device_budget_bytes / model.live_slot_count()`. Explicit targets at or
121    /// below the grant are bounded by the same live-slot target. An explicit
122    /// target above the grant remains available for allocation retry. Both
123    /// targets bound the same decoded estimate.
124    pub fn for_source_task(memory_policy: MemoryControlPolicy, model: SourceTaskMemoryModel) -> Self {
125        let base = Self::from(memory_policy);
126        let MemoryControlPolicy { device_budget_bytes, max_source_chunk_bytes, max_row_groups_per_chunk, .. } =
127            memory_policy;
128        if max_source_chunk_bytes > device_budget_bytes
129            && max_source_chunk_bytes != MemoryControlPolicy::DEFAULT_MAX_SOURCE_CHUNK_BYTES
130        {
131            return base;
132        }
133        let budget_target = device_budget_bytes / model.live_slot_count().max(1);
134        let task_target = max_source_chunk_bytes.min(budget_target.max(1));
135        let mut policy = Self::new(Some(task_target), max_row_groups_per_chunk);
136        policy.source_chunk_planning_reason =
137            if task_target < max_source_chunk_bytes { "memory_cap" } else { "configured_policy" };
138        policy
139    }
140
141    /// Derives a source-task policy with explicit local-throughput planning options.
142    pub fn for_source_task_with_mode_and_options(
143        source: &LocalParquetSource,
144        memory_policy: MemoryControlPolicy,
145        model: SourceTaskMemoryModel,
146        mode: SourceChunkPlanningMode,
147        options: SourceChunkPlanningOptions,
148    ) -> Self {
149        Self::for_source_task(memory_policy, model).with_source_planning_mode(source, mode, options)
150    }
151
152    /// Derives a source policy from memory policy and source planning mode.
153    pub fn from_memory_policy_with_mode(
154        source: &LocalParquetSource,
155        memory_policy: MemoryControlPolicy,
156        mode: SourceChunkPlanningMode,
157    ) -> Self {
158        Self::from_memory_policy_with_mode_and_options(
159            source,
160            memory_policy,
161            mode,
162            SourceChunkPlanningOptions::default(),
163        )
164    }
165
166    /// Derives a source policy with explicit local-throughput planning options.
167    pub fn from_memory_policy_with_mode_and_options(
168        source: &LocalParquetSource,
169        memory_policy: MemoryControlPolicy,
170        mode: SourceChunkPlanningMode,
171        options: SourceChunkPlanningOptions,
172    ) -> Self {
173        Self::from(memory_policy).with_source_planning_mode(source, mode, options)
174    }
175
176    pub(super) fn effective_max_row_groups_per_chunk(self) -> usize {
177        self.max_row_groups_per_chunk.max(1)
178    }
179
180    fn with_source_planning_mode(
181        mut self,
182        source: &LocalParquetSource,
183        mode: SourceChunkPlanningMode,
184        options: SourceChunkPlanningOptions,
185    ) -> Self {
186        if mode != SourceChunkPlanningMode::LocalParquetThroughput {
187            return self;
188        }
189        self.source_chunk_planning_mode = mode;
190        if !local_parquet_throughput_planning_admitted(source, self) {
191            self.source_chunk_planning_reason = "local_parquet_throughput_rejected";
192            return self;
193        }
194        self.max_row_groups_per_chunk =
195            self.max_row_groups_per_chunk.max(options.local_parquet_throughput_max_row_groups_per_chunk());
196        self.source_chunk_planning_reason = "local_parquet_throughput_admitted";
197        self
198    }
199}
200
201impl Default for SourceChunkPlanningPolicy {
202    fn default() -> Self {
203        Self {
204            max_source_chunk_bytes: None,
205            max_row_groups_per_chunk: usize::MAX,
206            source_chunk_planning_mode: SourceChunkPlanningMode::Conservative,
207            source_chunk_planning_reason: "unbounded_policy",
208        }
209    }
210}
211
212impl From<MemoryControlPolicy> for SourceChunkPlanningPolicy {
213    fn from(memory_policy: MemoryControlPolicy) -> Self {
214        let target = if memory_policy.max_source_chunk_bytes == MemoryControlPolicy::DEFAULT_MAX_SOURCE_CHUNK_BYTES {
215            (memory_policy.device_budget_bytes / SourceTaskMemoryModel::aggregate_pipeline(false).live_slot_count())
216                .max(1)
217        } else {
218            memory_policy.max_source_chunk_bytes
219        };
220        let mut policy = Self::new(Some(target), memory_policy.max_row_groups_per_chunk);
221        policy.source_chunk_planning_reason =
222            if memory_policy.max_source_chunk_bytes == MemoryControlPolicy::DEFAULT_MAX_SOURCE_CHUNK_BYTES {
223                "memory_cap"
224            } else {
225                "configured_policy"
226            };
227        policy
228    }
229}
230
231/// Whether throughput planning may raise a one-row-group cap for `source`.
232/// The grant does not enter here: the source-byte target already scales the
233/// widened chunks with it.
234fn local_parquet_throughput_planning_admitted(source: &LocalParquetSource, policy: SourceChunkPlanningPolicy) -> bool {
235    if policy.max_row_groups_per_chunk != 1 || policy.max_source_chunk_bytes.is_none() {
236        return false;
237    }
238    if source.files.is_empty()
239        || source.is_zero_column_count_source()
240        || source.scan_limit.is_some()
241        || source.read_columns.is_empty()
242        || !source.partition_columns.is_empty()
243    {
244        return false;
245    }
246    source.files.iter().all(|file| {
247        local_parquet_throughput_file_uses_supported_reader_path(file)
248            && file.partition_values.is_empty()
249            && matches!(
250                file.resolved_source.read_strategy(),
251                ParquetReadStrategy::LocalKvikio | ParquetReadStrategy::LocalCacheKvikio
252            )
253            && local_parquet_throughput_file_has_source_read_bytes(source, file)
254    })
255}
256
257fn local_parquet_throughput_file_uses_supported_reader_path(file: &LocalParquetFile) -> bool {
258    match file.resolved_source.read_strategy() {
259        ParquetReadStrategy::LocalKvikio => file.location.is_local_path(),
260        ParquetReadStrategy::LocalCacheKvikio => file.path.is_file(),
261        ParquetReadStrategy::RemoteDirect => false,
262    }
263}
264
265fn local_parquet_throughput_file_has_source_read_bytes(source: &LocalParquetSource, file: &LocalParquetFile) -> bool {
266    let selected_row_groups = match &file.row_groups {
267        Some(row_groups) if !row_groups.is_empty() => row_groups.clone(),
268        Some(_) => return false,
269        None => {
270            let Some(count) = file.projected_row_group_uncompressed_bytes.as_ref().map(Vec::len) else {
271                return false;
272            };
273            if count == 0 {
274                return false;
275            }
276            (0..count).collect()
277        },
278    };
279    let Some(max_row_group) = selected_row_groups.iter().copied().max() else {
280        return false;
281    };
282    let projected_covers_selection =
283        file.projected_row_group_uncompressed_bytes.as_ref().is_some_and(|bytes| max_row_group < bytes.len());
284    let pure_filter_covers_selection = source.pure_filter_columns.is_empty()
285        || file.pure_filter_row_group_uncompressed_bytes.as_ref().is_some_and(|bytes| max_row_group < bytes.len());
286    projected_covers_selection && pure_filter_covers_selection
287}
query_planning/chunk/policy.rs

query_planning/chunk/
policy.rs
