micromegas_analytics/lakehouse/
partitioned_execution_plan.rs

1use super::{partition::Partition, reader_factory::ReaderFactory, view::ScanSortColumn};
2use crate::{dfext::predicate::filters_to_predicate, time::datetime_to_scalar};
3use chrono::{DateTime, Utc};
4use datafusion::{
5    arrow::{compute::SortOptions, datatypes::SchemaRef},
6    catalog::{Session, memory::DataSourceExec},
7    common::stats::Precision,
8    datasource::{
9        listing::PartitionedFile,
10        physical_plan::{FileScanConfigBuilder, ParquetSource},
11    },
12    execution::object_store::ObjectStoreUrl,
13    physical_expr::{LexOrdering, PhysicalSortExpr},
14    physical_plan::{ColumnStatistics, ExecutionPlan, Statistics},
15    prelude::*,
16};
17use micromegas_tracing::prelude::*;
18use std::sync::Arc;
19
20/// Which pair of bounds on `Partition` a declared ordering's leading column is checked against.
21#[derive(Clone, Copy, Debug)]
22pub enum OrderingBounds {
23    /// `min_event_time()` / `max_event_time()` -- `Option`, absent for empty partitions.
24    EventTime,
25    /// `begin_insert_time()` / `end_insert_time()` -- always present.
26    InsertTime,
27}
28
29/// Reads the pair of bounds a declared ordering's leading column is checked against, per
30/// `OrderingBounds`. `InsertTime` bounds are always present; `EventTime` bounds are `None` for
31/// empty partitions (callers are expected to have already filtered those out).
32fn partition_bounds(
33    p: &Partition,
34    bounds: OrderingBounds,
35) -> Option<(DateTime<Utc>, DateTime<Utc>)> {
36    match bounds {
37        OrderingBounds::EventTime => p.min_event_time().zip(p.max_event_time()),
38        OrderingBounds::InsertTime => Some((p.begin_insert_time(), p.end_insert_time())),
39    }
40}
41
42/// Sorts the non-empty partitions by their leading-column bound ascending (tiebreak `file_path`)
43/// and verifies that adjacent partitions' ranges do not overlap. This makes the declared scan
44/// ordering self-contained: the file group is guaranteed to concatenate in globally-sorted order,
45/// independent of the order the partition cache returned.
46///
47/// Returns an error if any adjacent pair overlaps: the declared ordering cannot be honored, so we
48/// fail loudly instead of silently emitting a mis-ordered scan. For `OrderingBounds::EventTime`
49/// the most likely cause is TSC-frequency estimation drift across materialization epochs (for
50/// `tsc_frequency == 0` processes whose blocks were materialized under different clock
51/// estimates); the fix is to retire the affected stream's partitions so they rebuild with a
52/// single, consistent converter. For `OrderingBounds::InsertTime` an overlap indicates a genuine
53/// partitioning bug -- input partitions are expected to be non-overlapping in insert_time by
54/// construction.
55fn sort_and_check_non_overlapping(
56    mut partitions: Vec<&Partition>,
57    bounds: OrderingBounds,
58) -> datafusion::error::Result<Vec<&Partition>> {
59    partitions.sort_by(|a, b| {
60        partition_bounds(a, bounds)
61            .map(|(begin, _)| begin)
62            .cmp(&partition_bounds(b, bounds).map(|(begin, _)| begin))
63            .then_with(|| a.file_path.cmp(&b.file_path))
64    });
65    for pair in partitions.windows(2) {
66        let prev = pair[0];
67        let next = pair[1];
68        if let (Some((_, prev_max)), Some((next_min, _))) = (
69            partition_bounds(prev, bounds),
70            partition_bounds(next, bounds),
71        ) && prev_max > next_min
72        {
73            return Err(datafusion::error::DataFusionError::Execution(format!(
74                "declared scan ordering violated: partition {:?} (range ending {prev_max}) overlaps partition {:?} (range starting {next_min}). \
75                 For event-time ordering this can happen when a stream's blocks were registered out of event-time order, or -- for tsc_frequency == 0 processes -- when TSC-frequency \
76                 re-estimation drifted across materialization epochs spanning a clock adjustment (see the ordering-invariant notes on View::get_scan_output_ordering in view.rs). \
77                 Retire the affected stream's partitions so they rebuild with a single, consistent time converter.",
78                prev.file_path, next.file_path
79            )));
80        }
81    }
82    Ok(partitions)
83}
84
85/// Attaches the leading `output_ordering` column's min/max statistics to a `PartitionedFile`,
86/// using `Precision::Inexact` since the bounds read from `Partition` (per `OrderingBounds`) are
87/// not necessarily the column's exact min/max. DataFusion's multi-file-group ordering validation
88/// (`is_ordering_valid_for_file_groups`) requires these statistics to be present -- without them
89/// the declared ordering is silently dropped for any file group with more than one file.
90fn attach_ordering_statistics(
91    mut file: PartitionedFile,
92    schema: &SchemaRef,
93    leading_column: &ScanSortColumn,
94    partition: &Partition,
95    bounds: OrderingBounds,
96) -> datafusion::error::Result<PartitionedFile> {
97    let mut stats = Statistics::new_unknown(schema);
98    if let Some((min_time, max_time)) = partition_bounds(partition, bounds) {
99        let idx = schema.index_of(&leading_column.column)?;
100        stats.column_statistics[idx] = ColumnStatistics::new_unknown()
101            .with_min_value(Precision::Inexact(datetime_to_scalar(min_time)))
102            .with_max_value(Precision::Inexact(datetime_to_scalar(max_time)));
103    }
104    file = file.with_statistics(Arc::new(stats));
105    Ok(file)
106}
107
108/// Builds the `LexOrdering` declaring the already-satisfied output ordering of the scan, matching
109/// DataFusion's default `ORDER BY` semantics (ASC NULLS LAST unless `descending`).
110fn make_lex_ordering(
111    schema: &SchemaRef,
112    output_ordering: &[ScanSortColumn],
113) -> datafusion::error::Result<Option<LexOrdering>> {
114    let sort_exprs = output_ordering
115        .iter()
116        .map(|c| {
117            let col =
118                datafusion::physical_expr::expressions::Column::new_with_schema(&c.column, schema)?;
119            Ok(PhysicalSortExpr::new(
120                Arc::new(col),
121                SortOptions {
122                    descending: c.descending,
123                    // Match DataFusion's default ORDER BY semantics: ASC NULLS LAST, DESC NULLS
124                    // FIRST. Hardcoding `false` here would declare `DESC NULLS LAST`, which fails
125                    // to satisfy a descending query's `DESC NULLS FIRST` requirement and silently
126                    // keeps a redundant Sort.
127                    nulls_first: c.descending,
128                },
129            ))
130        })
131        .collect::<datafusion::error::Result<Vec<_>>>()?;
132    Ok(LexOrdering::new(sort_exprs))
133}
134
135/// Creates a partitioned execution plan for scanning Parquet files.
136///
137/// `output_ordering` declares an ordering the scan's rows already satisfy (see
138/// `View::get_scan_output_ordering`). When non-empty, the file group is sorted by the leading
139/// column's bound (read per `ordering_bounds`) and checked for non-overlap (erroring if
140/// violated), per-file min/max statistics are attached so DataFusion accepts the declared
141/// ordering, and the ordering is attached to the resulting `FileScanConfig` so `EnforceSorting`
142/// can elide a redundant `Sort` node. When empty, behavior is unchanged from before this
143/// parameter existed, and `ordering_bounds` is unused.
144#[span_fn]
145#[expect(clippy::too_many_arguments)]
146pub fn make_partitioned_execution_plan(
147    schema: SchemaRef,
148    reader_factory: Arc<ReaderFactory>,
149    state: &dyn Session,
150    projection: Option<&Vec<usize>>,
151    filters: &[Expr],
152    limit: Option<usize>,
153    partitions: Arc<Vec<Partition>>,
154    output_ordering: &[ScanSortColumn],
155    ordering_bounds: OrderingBounds,
156) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
157    let predicate = filters_to_predicate(schema.clone(), state, filters)?;
158
159    let non_empty_partitions: Vec<&Partition> =
160        partitions.iter().filter(|p| !p.is_empty()).collect();
161    let non_empty_partitions = if output_ordering.is_empty() {
162        non_empty_partitions
163    } else {
164        sort_and_check_non_overlapping(non_empty_partitions, ordering_bounds)?
165    };
166
167    let mut file_group = vec![];
168    for part in &non_empty_partitions {
169        let file_path = part.file_path.as_ref().ok_or_else(|| {
170            datafusion::error::DataFusionError::Internal(format!(
171                "non-empty partition has no file_path: num_rows={}",
172                part.num_rows
173            ))
174        })?;
175        let mut pf = PartitionedFile::new(file_path, part.file_size as u64);
176        if let Some(leading_column) = output_ordering.first() {
177            pf = attach_ordering_statistics(pf, &schema, leading_column, part, ordering_bounds)?;
178        }
179        file_group.push(pf);
180    }
181
182    // If all partitions are empty, return EmptyExec with projected schema
183    if file_group.is_empty() {
184        use datafusion::physical_plan::empty::EmptyExec;
185        let projected_schema = if let Some(projection) = projection {
186            Arc::new(schema.project(projection)?)
187        } else {
188            schema
189        };
190        return Ok(Arc::new(EmptyExec::new(projected_schema)));
191    }
192
193    let object_store_url = ObjectStoreUrl::parse("obj://lakehouse/").unwrap();
194    let source = Arc::new(
195        ParquetSource::new(schema.clone())
196            .with_predicate(predicate)
197            .with_parquet_file_reader_factory(reader_factory),
198    );
199    let mut builder = FileScanConfigBuilder::new(object_store_url, source)
200        .with_limit(limit)
201        .with_projection_indices(projection.cloned())?
202        .with_file_groups(vec![file_group.into()]);
203
204    if let Some(lex) = make_lex_ordering(&schema, output_ordering)? {
205        builder = builder.with_output_ordering(vec![lex]);
206    }
207    let file_scan_config = builder.build();
208    Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
209}
micromegas_analytics/lakehouse/partitioned_execution_plan.rs

micromegas_analytics/lakehouse/
partitioned_execution_plan.rs
