alopex_sql/distributed_read/
assembler.rs

1//! Atomic coordinator-side assembly for distributed SQL reads.
2//!
3//! Range workers never expose a client-visible stream directly. This module
4//! accepts only already fenced range payloads, verifies that every planned
5//! range has acknowledged cleanup, and then produces an immutable result.
6
7use std::cmp::Ordering;
8use std::collections::{BTreeMap, BTreeSet, HashSet};
9use std::sync::Arc;
10
11use alopex_core::sql::stream::ByteSized;
12
13use crate::catalog::ColumnMetadata;
14use crate::executor::ColumnInfo;
15use crate::executor::query::{
16    aggregate::{create_accumulator, encode_group_key, merge_exact_aggregate_states},
17    projected_columns,
18};
19use crate::planner::aggregate_expr::AggregateExpr;
20use crate::planner::typed_expr::Projection;
21use crate::storage::SqlValue;
22
23const ASSEMBLER_ENTRY_OVERHEAD_BYTES: u64 = 32;
24
25/// Bounds coordinator-owned materialization. Spill is intentionally not
26/// implemented yet; a non-zero spill budget remains a declared upper bound,
27/// never permission for an unbounded temporary file.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct DistributedReadBudget {
30    pub max_assembler_bytes: u64,
31    pub max_spill_bytes: u64,
32}
33
34impl Default for DistributedReadBudget {
35    fn default() -> Self {
36        Self {
37            max_assembler_bytes: 64 * 1024 * 1024,
38            max_spill_bytes: 0,
39        }
40    }
41}
42
43/// One normalized global ORDER BY key. Values are produced by a worker's
44/// closed descriptor; this type deliberately carries no executable SQL.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct GlobalOrder {
47    pub ascending: bool,
48    pub nulls_first: bool,
49}
50
51/// Final shaping shared by row and aggregate assembly.
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct ResultPresentation {
54    pub columns: Vec<ColumnInfo>,
55    pub distinct: bool,
56    pub order: Vec<GlobalOrder>,
57    /// Final-row column positions used by aggregate modes. Row mode instead
58    /// receives precomputed keys so ORDER BY expressions need not be decoded.
59    pub final_order_key_indexes: Vec<usize>,
60    pub offset: u64,
61    pub limit: Option<u64>,
62}
63
64impl ResultPresentation {
65    /// Construct output metadata with the existing local projection kernel.
66    pub fn from_projection(
67        projection: &Projection,
68        input_schema: &[ColumnMetadata],
69        distinct: bool,
70        order: Vec<GlobalOrder>,
71        final_order_key_indexes: Vec<usize>,
72        offset: u64,
73        limit: Option<u64>,
74    ) -> Result<Self, DistributedReadAssemblyError> {
75        Ok(Self {
76            columns: projected_columns(projection, input_schema)
77                .map_err(|error| DistributedReadAssemblyError::Projection(error.to_string()))?,
78            distinct,
79            order,
80            final_order_key_indexes,
81            offset,
82            limit,
83        })
84    }
85}
86
87/// Row-mode assembly for already projected range rows.
88#[derive(Debug, Clone, PartialEq, Eq)]
89pub struct RowMergePlan {
90    pub presentation: ResultPresentation,
91}
92
93/// Aggregate assembly whose worker states can be merged exactly.
94#[derive(Debug, Clone)]
95pub struct ExactAggregatePlan {
96    pub presentation: ResultPresentation,
97    pub group_column_count: usize,
98    pub aggregates: Vec<AggregateExpr>,
99}
100
101/// Aggregate assembly that must replay globally ordered inputs to preserve
102/// local floating-point, DISTINCT, and string-aggregate semantics.
103#[derive(Debug, Clone)]
104pub struct OrderedAggregatePlan {
105    pub presentation: ResultPresentation,
106    pub group_column_count: usize,
107    pub aggregates: Vec<AggregateExpr>,
108    pub logical_input_order: Vec<GlobalOrder>,
109}
110
111/// Closed assembly mode selected by the SQL catalog/coordinator.
112#[derive(Debug, Clone)]
113pub enum AssemblyPlan {
114    Rows(RowMergePlan),
115    ExactAggregates(ExactAggregatePlan),
116    OrderedAggregates(OrderedAggregatePlan),
117}
118
119/// One final-projection row emitted by a range worker during preparation.
120#[derive(Debug, Clone, PartialEq)]
121pub struct AssemblerRow {
122    pub values: Vec<SqlValue>,
123    pub order_keys: Vec<SqlValue>,
124    /// Stable physical row key used only to break otherwise equal global keys.
125    pub row_key: u64,
126}
127
128/// One group-local vector of exact accumulator states.
129#[derive(Debug, Clone, PartialEq)]
130pub struct ExactAggregatePartial {
131    pub group_key: Vec<SqlValue>,
132    /// One serializable state vector per aggregate expression.
133    pub states: Vec<Vec<SqlValue>>,
134}
135
136/// One raw aggregate input. The worker has already evaluated the closed
137/// descriptor's aggregate arguments; the coordinator only replays values.
138#[derive(Debug, Clone, PartialEq)]
139pub struct OrderedAggregateInput {
140    pub group_key: Vec<SqlValue>,
141    pub aggregate_arguments: Vec<Option<SqlValue>>,
142    pub logical_order_keys: Vec<SqlValue>,
143    pub row_key: u64,
144}
145
146/// Typed payloads accepted from a completed range session.
147#[derive(Debug, Clone, PartialEq)]
148pub enum RangeAssemblerPayload {
149    Rows(Vec<AssemblerRow>),
150    ExactAggregatePartial(ExactAggregatePartial),
151    OrderedAggregateInput(OrderedAggregateInput),
152}
153
154/// Worker terminal state supplied by the transport coordinator. A completed
155/// range is usable only when it explicitly acknowledges cleanup.
156#[derive(Debug, Clone, PartialEq, Eq)]
157pub enum RangeTerminal {
158    Pending,
159    Completed { cleanup_acknowledged: bool },
160    Failed { reason: String },
161}
162
163/// All preparation payloads and terminal evidence for one planned range.
164#[derive(Debug, Clone, PartialEq)]
165pub struct RangeAssemblerInput {
166    pub range_id: String,
167    /// Output schema for row payloads. Aggregate payload schemas are defined
168    /// by the assembly plan because partial states are not public rows.
169    pub columns: Vec<ColumnInfo>,
170    pub payloads: Vec<RangeAssemblerPayload>,
171    pub terminal: RangeTerminal,
172}
173
174/// Observable lifecycle state. A failed preparation never retains a prepared
175/// result and therefore cannot be opened as a row stream.
176#[derive(Debug, Clone, Copy, PartialEq, Eq)]
177pub enum AssemblerTerminalStatus {
178    Pending,
179    Prepared,
180    Failed,
181    Closed,
182}
183
184/// Classified preparation failure. No variant carries partial SQL rows.
185#[derive(Debug, thiserror::Error, PartialEq, Eq)]
186pub enum DistributedReadAssemblyError {
187    #[error("distributed assembler budget must be non-zero")]
188    ZeroAssemblerBudget,
189    #[error("distributed assembler requires at least one planned range")]
190    EmptyRangeSet,
191    #[error("range '{range_id}' is outside the immutable read fence")]
192    UnexpectedRange { range_id: String },
193    #[error("range '{range_id}' was supplied more than once")]
194    DuplicateRange { range_id: String },
195    #[error("range '{range_id}' did not reach a terminal End")]
196    MissingRangeEnd { range_id: String },
197    #[error("range '{range_id}' reached End without cleanup acknowledgement")]
198    CleanupNotAcknowledged { range_id: String },
199    #[error("range '{range_id}' failed during preparation: {reason}")]
200    RangeFailed { range_id: String, reason: String },
201    #[error("range '{range_id}' output schema differs from the planned result schema")]
202    SchemaMismatch { range_id: String },
203    #[error("range '{range_id}' supplied an incompatible {payload} payload")]
204    PayloadMismatch {
205        range_id: String,
206        payload: &'static str,
207    },
208    #[error("range '{range_id}' supplied {actual} values where {expected} were required")]
209    ValueArity {
210        range_id: String,
211        expected: usize,
212        actual: usize,
213    },
214    #[error("range '{range_id}' supplied {actual} order keys where {expected} were required")]
215    OrderKeyArity {
216        range_id: String,
217        expected: usize,
218        actual: usize,
219    },
220    #[error(
221        "distributed assembler requires {required_bytes} bytes, exceeding the {max_assembler_bytes}-byte memory budget (spill limit {max_spill_bytes} bytes)"
222    )]
223    ResourceLimitExceeded {
224        required_bytes: u64,
225        max_assembler_bytes: u64,
226        max_spill_bytes: u64,
227    },
228    #[error("distributed projection metadata could not be built: {0}")]
229    Projection(String),
230    #[error("distributed aggregate assembly failed: {0}")]
231    Aggregate(String),
232    #[error("the assembler is not prepared")]
233    NotPrepared,
234    #[error("the assembler has already been closed")]
235    Closed,
236    #[error("the assembler has already terminated with a failure")]
237    Failed,
238}
239
240/// Immutable result made visible only after every range has cleaned up.
241#[derive(Debug, Clone, PartialEq)]
242pub struct PreparedResult {
243    columns: Arc<[ColumnInfo]>,
244    rows: Arc<[Vec<SqlValue>]>,
245}
246
247impl PreparedResult {
248    fn new(columns: Vec<ColumnInfo>, rows: Vec<Vec<SqlValue>>) -> Self {
249        Self {
250            columns: columns.into(),
251            rows: rows.into(),
252        }
253    }
254
255    pub fn columns(&self) -> &[ColumnInfo] {
256        &self.columns
257    }
258
259    pub fn row_count(&self) -> usize {
260        self.rows.len()
261    }
262
263    /// Open an independent read-only stream over the immutable result.
264    pub fn open_prepared_stream(&self) -> PreparedResultStream {
265        PreparedResultStream {
266            result: self.clone(),
267            index: 0,
268        }
269    }
270
271    pub fn query_result(&self) -> crate::executor::QueryResult {
272        crate::executor::QueryResult::new(self.columns.to_vec(), self.rows.to_vec())
273    }
274}
275
276/// Reader for a [`PreparedResult`]. It has no transport/session dependency.
277pub struct PreparedResultStream {
278    result: PreparedResult,
279    index: usize,
280}
281
282impl PreparedResultStream {
283    pub fn next_row(&mut self) -> Option<Vec<SqlValue>> {
284        let row = self.result.rows.get(self.index)?.clone();
285        self.index += 1;
286        Some(row)
287    }
288}
289
290/// Coordinator-owned atomic result assembler.
291pub struct GlobalResultAssembler {
292    plan: AssemblyPlan,
293    budget: DistributedReadBudget,
294    expected_ranges: BTreeSet<String>,
295    ranges: BTreeMap<String, RangeAssemblerInput>,
296    prepared: Option<PreparedResult>,
297    status: AssemblerTerminalStatus,
298}
299
300impl GlobalResultAssembler {
301    pub fn new(
302        expected_ranges: impl IntoIterator<Item = String>,
303        plan: AssemblyPlan,
304        budget: DistributedReadBudget,
305    ) -> Result<Self, DistributedReadAssemblyError> {
306        if budget.max_assembler_bytes == 0 {
307            return Err(DistributedReadAssemblyError::ZeroAssemblerBudget);
308        }
309        let expected_ranges = expected_ranges.into_iter().collect::<BTreeSet<_>>();
310        if expected_ranges.is_empty() {
311            return Err(DistributedReadAssemblyError::EmptyRangeSet);
312        }
313        Ok(Self {
314            plan,
315            budget,
316            expected_ranges,
317            ranges: BTreeMap::new(),
318            prepared: None,
319            status: AssemblerTerminalStatus::Pending,
320        })
321    }
322
323    /// Register all buffered payloads for one range. Callers must not open a
324    /// client stream before [`Self::prepare`] returns successfully.
325    pub fn push_range(
326        &mut self,
327        input: RangeAssemblerInput,
328    ) -> Result<(), DistributedReadAssemblyError> {
329        self.ensure_pending()?;
330        if !self.expected_ranges.contains(&input.range_id) {
331            return Err(DistributedReadAssemblyError::UnexpectedRange {
332                range_id: input.range_id,
333            });
334        }
335        if self.ranges.contains_key(&input.range_id) {
336            return Err(DistributedReadAssemblyError::DuplicateRange {
337                range_id: input.range_id,
338            });
339        }
340        self.ranges.insert(input.range_id.clone(), input);
341        Ok(())
342    }
343
344    /// Validate all terminal acknowledgements and atomically materialize the
345    /// final result. On every error, coordinator-owned payloads are discarded.
346    pub fn prepare(&mut self) -> Result<PreparedResult, DistributedReadAssemblyError> {
347        self.ensure_pending()?;
348        let ranges = std::mem::take(&mut self.ranges);
349        let outcome = self.prepare_from_ranges(ranges);
350        match outcome {
351            Ok(result) => {
352                self.status = AssemblerTerminalStatus::Prepared;
353                self.prepared = Some(result.clone());
354                Ok(result)
355            }
356            Err(error) => {
357                self.status = AssemblerTerminalStatus::Failed;
358                self.prepared = None;
359                Err(error)
360            }
361        }
362    }
363
364    pub fn open_prepared_stream(
365        &self,
366    ) -> Result<PreparedResultStream, DistributedReadAssemblyError> {
367        self.prepared
368            .as_ref()
369            .map(PreparedResult::open_prepared_stream)
370            .ok_or(DistributedReadAssemblyError::NotPrepared)
371    }
372
373    pub fn terminal_status(&self) -> AssemblerTerminalStatus {
374        self.status
375    }
376
377    /// Idempotently discard preparation and prepared-result state.
378    pub fn close(&mut self) {
379        self.ranges.clear();
380        self.prepared = None;
381        self.status = AssemblerTerminalStatus::Closed;
382    }
383
384    fn ensure_pending(&self) -> Result<(), DistributedReadAssemblyError> {
385        match self.status {
386            AssemblerTerminalStatus::Pending => Ok(()),
387            AssemblerTerminalStatus::Closed => Err(DistributedReadAssemblyError::Closed),
388            AssemblerTerminalStatus::Failed => Err(DistributedReadAssemblyError::Failed),
389            AssemblerTerminalStatus::Prepared => Err(DistributedReadAssemblyError::NotPrepared),
390        }
391    }
392
393    fn prepare_from_ranges(
394        &self,
395        ranges: BTreeMap<String, RangeAssemblerInput>,
396    ) -> Result<PreparedResult, DistributedReadAssemblyError> {
397        self.validate_range_terminals(&ranges)?;
398        let mut budget = BudgetAccount::new(self.budget);
399        match &self.plan {
400            AssemblyPlan::Rows(plan) => prepare_rows(ranges, plan, &mut budget),
401            AssemblyPlan::ExactAggregates(plan) => {
402                prepare_exact_aggregates(ranges, plan, &mut budget)
403            }
404            AssemblyPlan::OrderedAggregates(plan) => {
405                prepare_ordered_aggregates(ranges, plan, &mut budget)
406            }
407        }
408    }
409
410    fn validate_range_terminals(
411        &self,
412        ranges: &BTreeMap<String, RangeAssemblerInput>,
413    ) -> Result<(), DistributedReadAssemblyError> {
414        for range_id in &self.expected_ranges {
415            let input = ranges.get(range_id).ok_or_else(|| {
416                DistributedReadAssemblyError::MissingRangeEnd {
417                    range_id: range_id.clone(),
418                }
419            })?;
420            match &input.terminal {
421                RangeTerminal::Pending => {
422                    return Err(DistributedReadAssemblyError::MissingRangeEnd {
423                        range_id: range_id.clone(),
424                    });
425                }
426                RangeTerminal::Completed {
427                    cleanup_acknowledged: false,
428                } => {
429                    return Err(DistributedReadAssemblyError::CleanupNotAcknowledged {
430                        range_id: range_id.clone(),
431                    });
432                }
433                RangeTerminal::Completed {
434                    cleanup_acknowledged: true,
435                } => {}
436                RangeTerminal::Failed { reason } => {
437                    return Err(DistributedReadAssemblyError::RangeFailed {
438                        range_id: range_id.clone(),
439                        reason: reason.clone(),
440                    });
441                }
442            }
443        }
444        Ok(())
445    }
446}
447
448struct BudgetAccount {
449    budget: DistributedReadBudget,
450    used_bytes: u64,
451}
452
453impl BudgetAccount {
454    fn new(budget: DistributedReadBudget) -> Self {
455        Self {
456            budget,
457            used_bytes: 0,
458        }
459    }
460
461    fn reserve_values(
462        &mut self,
463        values: impl IntoIterator<Item = SqlValue>,
464    ) -> Result<(), DistributedReadAssemblyError> {
465        let bytes = values
466            .into_iter()
467            .map(|value| value.estimated_bytes())
468            .sum::<u64>();
469        self.reserve(bytes)
470    }
471
472    fn reserve(&mut self, bytes: u64) -> Result<(), DistributedReadAssemblyError> {
473        let required_bytes = self.used_bytes.saturating_add(bytes);
474        if required_bytes > self.budget.max_assembler_bytes {
475            return Err(DistributedReadAssemblyError::ResourceLimitExceeded {
476                required_bytes,
477                max_assembler_bytes: self.budget.max_assembler_bytes,
478                max_spill_bytes: self.budget.max_spill_bytes,
479            });
480        }
481        self.used_bytes = required_bytes;
482        Ok(())
483    }
484}
485
486#[derive(Debug)]
487struct RankedRow {
488    values: Vec<SqlValue>,
489    order_keys: Vec<SqlValue>,
490    range_id: String,
491    row_key: u64,
492}
493
494fn prepare_rows(
495    ranges: BTreeMap<String, RangeAssemblerInput>,
496    plan: &RowMergePlan,
497    budget: &mut BudgetAccount,
498) -> Result<PreparedResult, DistributedReadAssemblyError> {
499    let mut rows = Vec::new();
500    for (range_id, input) in ranges {
501        if input.columns != plan.presentation.columns {
502            return Err(DistributedReadAssemblyError::SchemaMismatch { range_id });
503        }
504        for payload in input.payloads {
505            let RangeAssemblerPayload::Rows(payload_rows) = payload else {
506                return Err(DistributedReadAssemblyError::PayloadMismatch {
507                    range_id,
508                    payload: "non-row",
509                });
510            };
511            for row in payload_rows {
512                validate_row(&range_id, &row, &plan.presentation)?;
513                budget.reserve_values(row.values.iter().chain(row.order_keys.iter()).cloned())?;
514                budget.reserve(ASSEMBLER_ENTRY_OVERHEAD_BYTES)?;
515                rows.push(RankedRow {
516                    values: row.values,
517                    order_keys: row.order_keys,
518                    range_id: range_id.clone(),
519                    row_key: row.row_key,
520                });
521            }
522        }
523    }
524    materialize_presentation(plan.presentation.clone(), rows, budget)
525}
526
527fn prepare_exact_aggregates(
528    ranges: BTreeMap<String, RangeAssemblerInput>,
529    plan: &ExactAggregatePlan,
530    budget: &mut BudgetAccount,
531) -> Result<PreparedResult, DistributedReadAssemblyError> {
532    let mut groups = BTreeMap::<Vec<u8>, ExactGroup>::new();
533    for (range_id, input) in ranges {
534        for payload in input.payloads {
535            let RangeAssemblerPayload::ExactAggregatePartial(partial) = payload else {
536                return Err(DistributedReadAssemblyError::PayloadMismatch {
537                    range_id,
538                    payload: "non-exact-aggregate",
539                });
540            };
541            validate_exact_partial(&range_id, &partial, plan)?;
542            budget.reserve_values(
543                partial
544                    .group_key
545                    .iter()
546                    .chain(partial.states.iter().flatten())
547                    .cloned(),
548            )?;
549            budget.reserve(ASSEMBLER_ENTRY_OVERHEAD_BYTES)?;
550            let encoded_group = encode_group_key(&partial.group_key)
551                .map_err(|error| DistributedReadAssemblyError::Aggregate(error.to_string()))?;
552            let group = groups.entry(encoded_group).or_insert_with(|| ExactGroup {
553                values: partial.group_key.clone(),
554                state_rows: Vec::new(),
555            });
556            group.state_rows.push(partial.states);
557        }
558    }
559    if groups.is_empty() && plan.group_column_count == 0 {
560        groups.insert(
561            Vec::new(),
562            ExactGroup {
563                values: Vec::new(),
564                state_rows: Vec::new(),
565            },
566        );
567    }
568
569    let mut rows = Vec::with_capacity(groups.len());
570    for (_, group) in groups {
571        let aggregate_values = merge_exact_aggregate_states(&plan.aggregates, group.state_rows)
572            .map_err(|error| DistributedReadAssemblyError::Aggregate(error.to_string()))?;
573        let mut values = group.values;
574        values.extend(aggregate_values);
575        validate_final_width(&values, &plan.presentation)?;
576        let order_keys = final_order_keys(&values, &plan.presentation)?;
577        budget.reserve_values(values.iter().chain(order_keys.iter()).cloned())?;
578        budget.reserve(ASSEMBLER_ENTRY_OVERHEAD_BYTES)?;
579        rows.push(RankedRow {
580            values,
581            order_keys,
582            range_id: String::new(),
583            row_key: 0,
584        });
585    }
586    materialize_presentation(plan.presentation.clone(), rows, budget)
587}
588
589fn prepare_ordered_aggregates(
590    ranges: BTreeMap<String, RangeAssemblerInput>,
591    plan: &OrderedAggregatePlan,
592    budget: &mut BudgetAccount,
593) -> Result<PreparedResult, DistributedReadAssemblyError> {
594    let mut inputs = Vec::new();
595    for (range_id, input) in ranges {
596        for payload in input.payloads {
597            let RangeAssemblerPayload::OrderedAggregateInput(payload) = payload else {
598                return Err(DistributedReadAssemblyError::PayloadMismatch {
599                    range_id,
600                    payload: "non-ordered-aggregate",
601                });
602            };
603            validate_ordered_input(&range_id, &payload, plan)?;
604            budget.reserve_values(
605                payload
606                    .group_key
607                    .iter()
608                    .chain(payload.aggregate_arguments.iter().flatten())
609                    .chain(payload.logical_order_keys.iter())
610                    .cloned(),
611            )?;
612            budget.reserve(ASSEMBLER_ENTRY_OVERHEAD_BYTES)?;
613            let encoded_group = encode_group_key(&payload.group_key)
614                .map_err(|error| DistributedReadAssemblyError::Aggregate(error.to_string()))?;
615            inputs.push(RankedOrderedInput {
616                input: payload,
617                range_id: range_id.clone(),
618                encoded_group,
619            });
620        }
621    }
622    inputs.sort_by(|left, right| compare_ordered_input(left, right, plan));
623
624    let mut groups = BTreeMap::<Vec<u8>, OrderedGroup>::new();
625    for ranked in inputs {
626        let group = groups
627            .entry(ranked.encoded_group)
628            .or_insert_with(|| OrderedGroup {
629                values: ranked.input.group_key.clone(),
630                accumulators: plan
631                    .aggregates
632                    .iter()
633                    .map(|aggregate| create_accumulator(&aggregate.function, aggregate.distinct))
634                    .collect(),
635            });
636        for (accumulator, value) in group
637            .accumulators
638            .iter_mut()
639            .zip(ranked.input.aggregate_arguments)
640        {
641            accumulator
642                .update(value)
643                .map_err(|error| DistributedReadAssemblyError::Aggregate(error.to_string()))?;
644        }
645    }
646    if groups.is_empty() && plan.group_column_count == 0 {
647        groups.insert(
648            Vec::new(),
649            OrderedGroup {
650                values: Vec::new(),
651                accumulators: plan
652                    .aggregates
653                    .iter()
654                    .map(|aggregate| create_accumulator(&aggregate.function, aggregate.distinct))
655                    .collect(),
656            },
657        );
658    }
659
660    let mut rows = Vec::with_capacity(groups.len());
661    for (_, group) in groups {
662        let mut values = group.values;
663        let aggregate_values = group
664            .accumulators
665            .iter()
666            .map(|accumulator| accumulator.finalize())
667            .collect::<std::result::Result<Vec<_>, _>>()
668            .map_err(|error| DistributedReadAssemblyError::Aggregate(error.to_string()))?;
669        values.extend(aggregate_values);
670        validate_final_width(&values, &plan.presentation)?;
671        let order_keys = final_order_keys(&values, &plan.presentation)?;
672        budget.reserve_values(values.iter().chain(order_keys.iter()).cloned())?;
673        budget.reserve(ASSEMBLER_ENTRY_OVERHEAD_BYTES)?;
674        rows.push(RankedRow {
675            values,
676            order_keys,
677            range_id: String::new(),
678            row_key: 0,
679        });
680    }
681    materialize_presentation(plan.presentation.clone(), rows, budget)
682}
683
684fn validate_row(
685    range_id: &str,
686    row: &AssemblerRow,
687    presentation: &ResultPresentation,
688) -> Result<(), DistributedReadAssemblyError> {
689    validate_final_width(&row.values, presentation).map_err(|_| {
690        DistributedReadAssemblyError::ValueArity {
691            range_id: range_id.to_string(),
692            expected: presentation.columns.len(),
693            actual: row.values.len(),
694        }
695    })?;
696    if row.order_keys.len() != presentation.order.len() {
697        return Err(DistributedReadAssemblyError::OrderKeyArity {
698            range_id: range_id.to_string(),
699            expected: presentation.order.len(),
700            actual: row.order_keys.len(),
701        });
702    }
703    Ok(())
704}
705
706fn validate_exact_partial(
707    range_id: &str,
708    partial: &ExactAggregatePartial,
709    plan: &ExactAggregatePlan,
710) -> Result<(), DistributedReadAssemblyError> {
711    if partial.group_key.len() != plan.group_column_count {
712        return Err(DistributedReadAssemblyError::ValueArity {
713            range_id: range_id.to_string(),
714            expected: plan.group_column_count,
715            actual: partial.group_key.len(),
716        });
717    }
718    if partial.states.len() != plan.aggregates.len() {
719        return Err(DistributedReadAssemblyError::ValueArity {
720            range_id: range_id.to_string(),
721            expected: plan.aggregates.len(),
722            actual: partial.states.len(),
723        });
724    }
725    Ok(())
726}
727
728fn validate_ordered_input(
729    range_id: &str,
730    input: &OrderedAggregateInput,
731    plan: &OrderedAggregatePlan,
732) -> Result<(), DistributedReadAssemblyError> {
733    if input.group_key.len() != plan.group_column_count {
734        return Err(DistributedReadAssemblyError::ValueArity {
735            range_id: range_id.to_string(),
736            expected: plan.group_column_count,
737            actual: input.group_key.len(),
738        });
739    }
740    if input.aggregate_arguments.len() != plan.aggregates.len() {
741        return Err(DistributedReadAssemblyError::ValueArity {
742            range_id: range_id.to_string(),
743            expected: plan.aggregates.len(),
744            actual: input.aggregate_arguments.len(),
745        });
746    }
747    if input.logical_order_keys.len() != plan.logical_input_order.len() {
748        return Err(DistributedReadAssemblyError::OrderKeyArity {
749            range_id: range_id.to_string(),
750            expected: plan.logical_input_order.len(),
751            actual: input.logical_order_keys.len(),
752        });
753    }
754    Ok(())
755}
756
757fn validate_final_width(
758    values: &[SqlValue],
759    presentation: &ResultPresentation,
760) -> Result<(), DistributedReadAssemblyError> {
761    if values.len() == presentation.columns.len() {
762        Ok(())
763    } else {
764        Err(DistributedReadAssemblyError::Aggregate(format!(
765            "final row has {} value(s), expected {}",
766            values.len(),
767            presentation.columns.len()
768        )))
769    }
770}
771
772fn materialize_presentation(
773    presentation: ResultPresentation,
774    mut rows: Vec<RankedRow>,
775    budget: &mut BudgetAccount,
776) -> Result<PreparedResult, DistributedReadAssemblyError> {
777    if presentation.distinct {
778        let mut seen = HashSet::new();
779        let mut distinct_rows = Vec::with_capacity(rows.len());
780        for row in rows {
781            let key = encode_group_key(&row.values)
782                .map_err(|error| DistributedReadAssemblyError::Aggregate(error.to_string()))?;
783            budget.reserve(key.len() as u64 + ASSEMBLER_ENTRY_OVERHEAD_BYTES)?;
784            if seen.insert(key) {
785                distinct_rows.push(row);
786            }
787        }
788        rows = distinct_rows;
789    }
790    rows.sort_by(|left, right| compare_ranked_rows(left, right, &presentation.order));
791    let start = usize::try_from(presentation.offset).unwrap_or(usize::MAX);
792    let limit = presentation
793        .limit
794        .map(|limit| usize::try_from(limit).unwrap_or(usize::MAX));
795    let rows = rows
796        .into_iter()
797        .skip(start)
798        .take(limit.unwrap_or(usize::MAX))
799        .map(|row| row.values)
800        .collect();
801    Ok(PreparedResult::new(presentation.columns, rows))
802}
803
804fn final_order_keys(
805    values: &[SqlValue],
806    presentation: &ResultPresentation,
807) -> Result<Vec<SqlValue>, DistributedReadAssemblyError> {
808    if presentation.order.is_empty() {
809        return Ok(Vec::new());
810    }
811    if presentation.final_order_key_indexes.len() != presentation.order.len() {
812        return Err(DistributedReadAssemblyError::Aggregate(
813            "aggregate final ordering requires one output-column index per ORDER BY key".into(),
814        ));
815    }
816    presentation
817        .final_order_key_indexes
818        .iter()
819        .map(|index| {
820            values.get(*index).cloned().ok_or_else(|| {
821                DistributedReadAssemblyError::Aggregate(format!(
822                    "aggregate final ORDER BY index {index} is outside the result row"
823                ))
824            })
825        })
826        .collect()
827}
828
829fn compare_ranked_rows(left: &RankedRow, right: &RankedRow, order: &[GlobalOrder]) -> Ordering {
830    compare_order_keys(&left.order_keys, &right.order_keys, order)
831        .then_with(|| left.range_id.cmp(&right.range_id))
832        .then_with(|| left.row_key.cmp(&right.row_key))
833}
834
835struct ExactGroup {
836    values: Vec<SqlValue>,
837    state_rows: Vec<Vec<Vec<SqlValue>>>,
838}
839
840struct RankedOrderedInput {
841    input: OrderedAggregateInput,
842    range_id: String,
843    encoded_group: Vec<u8>,
844}
845
846struct OrderedGroup {
847    values: Vec<SqlValue>,
848    accumulators: Vec<Box<dyn crate::executor::query::aggregate::Accumulator>>,
849}
850
851fn compare_ordered_input(
852    left: &RankedOrderedInput,
853    right: &RankedOrderedInput,
854    plan: &OrderedAggregatePlan,
855) -> Ordering {
856    left.encoded_group
857        .cmp(&right.encoded_group)
858        .then_with(|| {
859            compare_order_keys(
860                &left.input.logical_order_keys,
861                &right.input.logical_order_keys,
862                &plan.logical_input_order,
863            )
864        })
865        .then_with(|| left.range_id.cmp(&right.range_id))
866        .then_with(|| left.input.row_key.cmp(&right.input.row_key))
867}
868
869fn compare_order_keys(left: &[SqlValue], right: &[SqlValue], order: &[GlobalOrder]) -> Ordering {
870    for ((left, right), order) in left.iter().zip(right).zip(order) {
871        let comparison = compare_sql_value(left, right, *order);
872        if comparison != Ordering::Equal {
873            return comparison;
874        }
875    }
876    Ordering::Equal
877}
878
879fn compare_sql_value(left: &SqlValue, right: &SqlValue, order: GlobalOrder) -> Ordering {
880    match (left, right) {
881        (SqlValue::Null, SqlValue::Null) => Ordering::Equal,
882        (SqlValue::Null, _) => {
883            if order.nulls_first {
884                Ordering::Less
885            } else {
886                Ordering::Greater
887            }
888        }
889        (_, SqlValue::Null) => {
890            if order.nulls_first {
891                Ordering::Greater
892            } else {
893                Ordering::Less
894            }
895        }
896        _ => match left.partial_cmp(right).unwrap_or(Ordering::Equal) {
897            Ordering::Equal => Ordering::Equal,
898            comparison if order.ascending => comparison,
899            comparison => comparison.reverse(),
900        },
901    }
902}
903
904#[cfg(test)]
905mod tests {
906    use super::*;
907    use crate::executor::query::aggregate::create_accumulator;
908    use crate::planner::ResolvedType;
909    use crate::planner::aggregate_expr::AggregateFunction;
910
911    fn columns(names: &[(&str, ResolvedType)]) -> Vec<ColumnInfo> {
912        names
913            .iter()
914            .map(|(name, data_type)| ColumnInfo::new(*name, data_type.clone()))
915            .collect()
916    }
917
918    fn presentation(
919        columns: Vec<ColumnInfo>,
920        distinct: bool,
921        order: Vec<GlobalOrder>,
922        final_order_key_indexes: Vec<usize>,
923        offset: u64,
924        limit: Option<u64>,
925    ) -> ResultPresentation {
926        ResultPresentation {
927            columns,
928            distinct,
929            order,
930            final_order_key_indexes,
931            offset,
932            limit,
933        }
934    }
935
936    fn completed_range(
937        range_id: &str,
938        columns: Vec<ColumnInfo>,
939        payloads: Vec<RangeAssemblerPayload>,
940    ) -> RangeAssemblerInput {
941        RangeAssemblerInput {
942            range_id: range_id.into(),
943            columns,
944            payloads,
945            terminal: RangeTerminal::Completed {
946                cleanup_acknowledged: true,
947            },
948        }
949    }
950
951    fn row(value: &str, key: i32, row_key: u64) -> AssemblerRow {
952        AssemblerRow {
953            values: vec![SqlValue::Text(value.into())],
954            order_keys: vec![SqlValue::Integer(key)],
955            row_key,
956        }
957    }
958
959    #[test]
960    fn rows_remain_private_until_every_range_cleanup_acknowledges() {
961        let output_columns = columns(&[("name", ResolvedType::Text)]);
962        let plan = AssemblyPlan::Rows(RowMergePlan {
963            presentation: presentation(output_columns.clone(), false, vec![], vec![], 0, None),
964        });
965        let mut assembler = GlobalResultAssembler::new(
966            vec!["range-a".into(), "range-b".into()],
967            plan,
968            DistributedReadBudget::default(),
969        )
970        .unwrap();
971        assembler
972            .push_range(completed_range(
973                "range-a",
974                output_columns.clone(),
975                vec![RangeAssemblerPayload::Rows(vec![AssemblerRow {
976                    values: vec![SqlValue::Text("private".into())],
977                    order_keys: vec![],
978                    row_key: 1,
979                }])],
980            ))
981            .unwrap();
982        assembler
983            .push_range(RangeAssemblerInput {
984                range_id: "range-b".into(),
985                columns: output_columns,
986                payloads: vec![],
987                terminal: RangeTerminal::Completed {
988                    cleanup_acknowledged: false,
989                },
990            })
991            .unwrap();
992
993        assert!(matches!(
994            assembler.prepare(),
995            Err(DistributedReadAssemblyError::CleanupNotAcknowledged { range_id }) if range_id == "range-b"
996        ));
997        assert_eq!(assembler.terminal_status(), AssemblerTerminalStatus::Failed);
998        assert!(matches!(
999            assembler.open_prepared_stream(),
1000            Err(DistributedReadAssemblyError::NotPrepared)
1001        ));
1002    }
1003
1004    #[test]
1005    fn row_merge_applies_global_distinct_order_offset_and_limit_before_opening_stream() {
1006        let output_columns = columns(&[("name", ResolvedType::Text)]);
1007        let plan = AssemblyPlan::Rows(RowMergePlan {
1008            presentation: presentation(
1009                output_columns.clone(),
1010                true,
1011                vec![GlobalOrder {
1012                    ascending: true,
1013                    nulls_first: false,
1014                }],
1015                vec![],
1016                1,
1017                Some(1),
1018            ),
1019        });
1020        let mut assembler = GlobalResultAssembler::new(
1021            vec!["range-a".into(), "range-b".into()],
1022            plan,
1023            DistributedReadBudget::default(),
1024        )
1025        .unwrap();
1026        assembler
1027            .push_range(completed_range(
1028                "range-a",
1029                output_columns.clone(),
1030                vec![RangeAssemblerPayload::Rows(vec![
1031                    row("charlie", 3, 1),
1032                    row("alpha", 1, 2),
1033                ])],
1034            ))
1035            .unwrap();
1036        assembler
1037            .push_range(completed_range(
1038                "range-b",
1039                output_columns,
1040                vec![RangeAssemblerPayload::Rows(vec![
1041                    row("alpha", 1, 3),
1042                    row("bravo", 2, 4),
1043                ])],
1044            ))
1045            .unwrap();
1046
1047        let prepared = assembler.prepare().unwrap();
1048        assert_eq!(prepared.row_count(), 1);
1049        let mut stream = prepared.open_prepared_stream();
1050        assert_eq!(
1051            stream.next_row(),
1052            Some(vec![SqlValue::Text("bravo".into())])
1053        );
1054        assert_eq!(stream.next_row(), None);
1055    }
1056
1057    #[test]
1058    fn exact_partial_merge_is_limited_to_proven_accumulators() {
1059        let aggregate = AggregateExpr::count_star();
1060        let output_columns = columns(&[("count", ResolvedType::BigInt)]);
1061        let plan = AssemblyPlan::ExactAggregates(ExactAggregatePlan {
1062            presentation: presentation(output_columns.clone(), false, vec![], vec![], 0, None),
1063            group_column_count: 0,
1064            aggregates: vec![aggregate],
1065        });
1066        let mut assembler = GlobalResultAssembler::new(
1067            vec!["range-a".into(), "range-b".into()],
1068            plan,
1069            DistributedReadBudget::default(),
1070        )
1071        .unwrap();
1072        for (range_id, count) in [("range-a", 2), ("range-b", 3)] {
1073            let mut accumulator = create_accumulator(&AggregateFunction::Count, false);
1074            for _ in 0..count {
1075                accumulator.update(None).unwrap();
1076            }
1077            assembler
1078                .push_range(completed_range(
1079                    range_id,
1080                    vec![],
1081                    vec![RangeAssemblerPayload::ExactAggregatePartial(
1082                        ExactAggregatePartial {
1083                            group_key: vec![],
1084                            states: vec![accumulator.state().unwrap()],
1085                        },
1086                    )],
1087                ))
1088                .unwrap();
1089        }
1090
1091        assert_eq!(
1092            assembler.prepare().unwrap().query_result().rows,
1093            vec![vec![SqlValue::BigInt(5)]]
1094        );
1095    }
1096
1097    #[test]
1098    fn aggregate_results_apply_global_order_and_limit_after_all_range_merges() {
1099        let aggregate = AggregateExpr::count_star();
1100        let output_columns = columns(&[
1101            ("team", ResolvedType::Text),
1102            ("count", ResolvedType::BigInt),
1103        ]);
1104        let plan = AssemblyPlan::ExactAggregates(ExactAggregatePlan {
1105            presentation: presentation(
1106                output_columns.clone(),
1107                false,
1108                vec![GlobalOrder {
1109                    ascending: false,
1110                    nulls_first: false,
1111                }],
1112                vec![1],
1113                0,
1114                Some(1),
1115            ),
1116            group_column_count: 1,
1117            aggregates: vec![aggregate],
1118        });
1119        let mut assembler = GlobalResultAssembler::new(
1120            vec!["range-a".into(), "range-b".into()],
1121            plan,
1122            DistributedReadBudget::default(),
1123        )
1124        .unwrap();
1125        for (range_id, team, count) in [("range-a", "alpha", 1), ("range-b", "beta", 3)] {
1126            let mut accumulator = create_accumulator(&AggregateFunction::Count, false);
1127            for _ in 0..count {
1128                accumulator.update(None).unwrap();
1129            }
1130            assembler
1131                .push_range(completed_range(
1132                    range_id,
1133                    output_columns.clone(),
1134                    vec![RangeAssemblerPayload::ExactAggregatePartial(
1135                        ExactAggregatePartial {
1136                            group_key: vec![SqlValue::Text(team.into())],
1137                            states: vec![accumulator.state().unwrap()],
1138                        },
1139                    )],
1140                ))
1141                .unwrap();
1142        }
1143
1144        assert_eq!(
1145            assembler.prepare().unwrap().query_result().rows,
1146            vec![vec![SqlValue::Text("beta".into()), SqlValue::BigInt(3)]]
1147        );
1148    }
1149
1150    #[test]
1151    fn ordered_inputs_replay_distinct_string_aggregates_in_global_order() {
1152        let aggregate = AggregateExpr {
1153            function: AggregateFunction::GroupConcat {
1154                separator: Some("|".into()),
1155            },
1156            arg: None,
1157            distinct: true,
1158            result_type: ResolvedType::Text,
1159        };
1160        let output_columns = columns(&[("names", ResolvedType::Text)]);
1161        let plan = AssemblyPlan::OrderedAggregates(OrderedAggregatePlan {
1162            presentation: presentation(output_columns.clone(), false, vec![], vec![], 0, None),
1163            group_column_count: 0,
1164            aggregates: vec![aggregate],
1165            logical_input_order: vec![GlobalOrder {
1166                ascending: true,
1167                nulls_first: false,
1168            }],
1169        });
1170        let mut assembler = GlobalResultAssembler::new(
1171            vec!["range-a".into(), "range-b".into()],
1172            plan,
1173            DistributedReadBudget::default(),
1174        )
1175        .unwrap();
1176        assembler
1177            .push_range(completed_range(
1178                "range-a",
1179                output_columns.clone(),
1180                vec![
1181                    RangeAssemblerPayload::OrderedAggregateInput(OrderedAggregateInput {
1182                        group_key: vec![],
1183                        aggregate_arguments: vec![Some(SqlValue::Text("b".into()))],
1184                        logical_order_keys: vec![SqlValue::Integer(2)],
1185                        row_key: 1,
1186                    }),
1187                    RangeAssemblerPayload::OrderedAggregateInput(OrderedAggregateInput {
1188                        group_key: vec![],
1189                        aggregate_arguments: vec![Some(SqlValue::Text("a".into()))],
1190                        logical_order_keys: vec![SqlValue::Integer(1)],
1191                        row_key: 2,
1192                    }),
1193                ],
1194            ))
1195            .unwrap();
1196        assembler
1197            .push_range(completed_range(
1198                "range-b",
1199                output_columns,
1200                vec![RangeAssemblerPayload::OrderedAggregateInput(
1201                    OrderedAggregateInput {
1202                        group_key: vec![],
1203                        aggregate_arguments: vec![Some(SqlValue::Text("b".into()))],
1204                        logical_order_keys: vec![SqlValue::Integer(3)],
1205                        row_key: 3,
1206                    },
1207                )],
1208            ))
1209            .unwrap();
1210
1211        assert_eq!(
1212            assembler.prepare().unwrap().query_result().rows,
1213            vec![vec![SqlValue::Text("a|b".into())]]
1214        );
1215    }
1216
1217    #[test]
1218    fn unproven_exact_partial_form_is_rejected_without_a_prepared_result() {
1219        let aggregate = AggregateExpr {
1220            function: AggregateFunction::Sum,
1221            arg: None,
1222            distinct: false,
1223            result_type: ResolvedType::Double,
1224        };
1225        let output_columns = columns(&[("sum", ResolvedType::Double)]);
1226        let plan = AssemblyPlan::ExactAggregates(ExactAggregatePlan {
1227            presentation: presentation(output_columns, false, vec![], vec![], 0, None),
1228            group_column_count: 0,
1229            aggregates: vec![aggregate],
1230        });
1231        let mut assembler = GlobalResultAssembler::new(
1232            vec!["range-a".into()],
1233            plan,
1234            DistributedReadBudget::default(),
1235        )
1236        .unwrap();
1237        assembler
1238            .push_range(completed_range("range-a", vec![], vec![]))
1239            .unwrap();
1240
1241        assert!(matches!(
1242            assembler.prepare(),
1243            Err(DistributedReadAssemblyError::Aggregate(reason)) if reason.contains("requires ordered input replay")
1244        ));
1245        assert!(matches!(
1246            assembler.open_prepared_stream(),
1247            Err(DistributedReadAssemblyError::NotPrepared)
1248        ));
1249    }
1250
1251    #[test]
1252    fn budget_exhaustion_exposes_no_prepared_rows() {
1253        let output_columns = columns(&[("name", ResolvedType::Text)]);
1254        let plan = AssemblyPlan::Rows(RowMergePlan {
1255            presentation: presentation(output_columns.clone(), false, vec![], vec![], 0, None),
1256        });
1257        let mut assembler = GlobalResultAssembler::new(
1258            vec!["range-a".into()],
1259            plan,
1260            DistributedReadBudget {
1261                max_assembler_bytes: 3,
1262                max_spill_bytes: 0,
1263            },
1264        )
1265        .unwrap();
1266        assembler
1267            .push_range(completed_range(
1268                "range-a",
1269                output_columns,
1270                vec![RangeAssemblerPayload::Rows(vec![AssemblerRow {
1271                    values: vec![SqlValue::Text("four".into())],
1272                    order_keys: vec![],
1273                    row_key: 1,
1274                }])],
1275            ))
1276            .unwrap();
1277
1278        assert!(matches!(
1279            assembler.prepare(),
1280            Err(DistributedReadAssemblyError::ResourceLimitExceeded { .. })
1281        ));
1282        assert!(matches!(
1283            assembler.open_prepared_stream(),
1284            Err(DistributedReadAssemblyError::NotPrepared)
1285        ));
1286    }
1287}
alopex_sql/distributed_read/assembler.rs

alopex_sql/distributed_read/
assembler.rs
