Skip to main content

uqa_execution/batch/
materialization.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Final conversion from positional physical rows to named result rows.
8
9use std::collections::HashMap;
10
11use super::{
12    Arc, PhysicalRow, ResultRow, RowFragment, RowSchema, SmallVec, Value, INLINE_ROW_FRAGMENTS,
13    NULL_SLOT,
14};
15
16struct ResultMaterializationPlan {
17    entries: Box<[ResultMaterializationEntry]>,
18}
19
20struct ResultMaterializationEntry {
21    column: String,
22    logical: usize,
23    take: bool,
24}
25
26impl RowSchema {
27    pub(super) fn materialize_result_row(&self, row: PhysicalRow) -> ResultRow {
28        if self.index.cold.identity_layout {
29            return self.materialize_identity_result_row(row);
30        }
31        self.materialize_remapped_result_row(row)
32    }
33
34    pub(super) fn materialize_remapped_result_row(&self, row: PhysicalRow) -> ResultRow {
35        let plan = self.result_materialization_plan();
36        self.materialize_remapped_result_row_with_plan(row, &plan)
37    }
38
39    pub(super) fn materialize_remapped_result_rows(
40        &self,
41        rows: Vec<PhysicalRow>,
42    ) -> Vec<ResultRow> {
43        let plan = self.result_materialization_plan();
44        rows.into_iter()
45            .map(|row| self.materialize_remapped_result_row_with_plan(row, &plan))
46            .collect()
47    }
48
49    fn result_materialization_plan(&self) -> ResultMaterializationPlan {
50        let columns = self.columns();
51        let mut last_logical_by_label = HashMap::<&str, usize>::with_capacity(columns.len());
52        for (logical, column) in columns.iter().enumerate() {
53            last_logical_by_label.insert(column, logical);
54        }
55        let mut logical_order = last_logical_by_label.into_values().collect::<Vec<_>>();
56        logical_order.sort_unstable_by(|left, right| columns[*left].cmp(&columns[*right]));
57        let mut remaining_reads = vec![0usize; self.physical_width()];
58        for logical in &logical_order {
59            let slot = self.index.slots[*logical];
60            if slot != NULL_SLOT {
61                remaining_reads[slot] += 1;
62            }
63        }
64        let take = logical_order
65            .iter()
66            .map(|logical| {
67                let slot = self.index.slots[*logical];
68                if slot == NULL_SLOT {
69                    return false;
70                }
71                remaining_reads[slot] -= 1;
72                remaining_reads[slot] == 0
73            })
74            .collect::<Vec<_>>();
75        let entries = logical_order
76            .into_iter()
77            .zip(take)
78            .map(|(logical, take)| ResultMaterializationEntry {
79                column: columns[logical].clone(),
80                logical,
81                take,
82            })
83            .collect::<Vec<_>>();
84        ResultMaterializationPlan {
85            entries: entries.into_boxed_slice(),
86        }
87    }
88
89    fn materialize_remapped_result_row_with_plan(
90        &self,
91        row: PhysicalRow,
92        plan: &ResultMaterializationPlan,
93    ) -> ResultRow {
94        // A shared scan fragment cannot donate its values, so read only requested output slots instead of cloning an intermediate positional vector that is immediately dismantled.
95        if row.fragments.len() == 1 && Arc::strong_count(&row.fragments[0].values) > 1 {
96            return self.materialize_remapped_shared_fragment_row(&row.fragments[0], plan);
97        }
98        let mut fragments = row.into_value_fragments();
99        debug_assert_eq!(
100            self.physical_width(),
101            fragments.iter().map(Vec::len).sum::<usize>()
102        );
103        let mut result = ResultRow::new();
104        for entry in &plan.entries {
105            let slot = self.index.slots[entry.logical];
106            let value = if slot == NULL_SLOT {
107                Value::Null
108            } else {
109                materialize_fragment_slot(&mut fragments, slot, entry.take)
110            };
111            result.insert(entry.column.clone(), value);
112        }
113        result
114    }
115
116    fn materialize_remapped_shared_fragment_row(
117        &self,
118        fragment: &RowFragment,
119        plan: &ResultMaterializationPlan,
120    ) -> ResultRow {
121        debug_assert_eq!(self.physical_width(), fragment.len());
122        let mut result = ResultRow::new();
123        for entry in &plan.entries {
124            let slot = self.index.slots[entry.logical];
125            let value = if slot == NULL_SLOT {
126                Value::Null
127            } else {
128                fragment.get(slot).cloned().unwrap_or(Value::Null)
129            };
130            result.insert(entry.column.clone(), value);
131        }
132        result
133    }
134
135    pub(super) fn materialize_identity_result_row(&self, row: PhysicalRow) -> ResultRow {
136        debug_assert_eq!(
137            self.len(),
138            row.fragments.iter().map(RowFragment::len).sum::<usize>()
139        );
140        let mut columns = self.columns().iter();
141        let mut result = ResultRow::new();
142        for fragment in row.fragments {
143            fragment.materialize_into(&mut columns, &mut result);
144        }
145        result
146    }
147}
148
149fn materialize_fragment_slot(fragments: &mut [Vec<Value>], mut slot: usize, take: bool) -> Value {
150    for fragment in fragments {
151        if slot < fragment.len() {
152            return if take {
153                std::mem::replace(&mut fragment[slot], Value::Null)
154            } else {
155                fragment[slot].clone()
156            };
157        }
158        slot -= fragment.len();
159    }
160    Value::Null
161}
162
163impl RowFragment {
164    /// Insert this fragment directly into a named result row. Shared scan projections clone only the selected values; they do not allocate an intermediate positional row at the final materialization boundary.
165    fn materialize_into<'a>(
166        self,
167        columns: &mut impl Iterator<Item = &'a String>,
168        result: &mut ResultRow,
169    ) {
170        let Self { values, projection } = self;
171        let Some(projection) = projection else {
172            match Arc::try_unwrap(values) {
173                Ok(values) => insert_values(columns, result, values),
174                Err(values) => insert_values(columns, result, values.iter().cloned()),
175            }
176            return;
177        };
178        match Arc::try_unwrap(values) {
179            Ok(mut values) => {
180                if projection
181                    .iter()
182                    .enumerate()
183                    .all(|(position, slot)| position == *slot)
184                {
185                    values.truncate(projection.len());
186                    insert_values(columns, result, values);
187                    return;
188                }
189                let mut remaining = vec![0usize; values.len()];
190                for slot in projection.iter().copied().filter(|slot| *slot != NULL_SLOT) {
191                    if let Some(count) = remaining.get_mut(slot) {
192                        *count += 1;
193                    }
194                }
195                for slot in projection.iter().copied() {
196                    let value = if slot == NULL_SLOT {
197                        Value::Null
198                    } else {
199                        let Some(count) = remaining.get_mut(slot) else {
200                            insert_value(columns, result, Value::Null);
201                            continue;
202                        };
203                        *count -= 1;
204                        if *count == 0 {
205                            values
206                                .get_mut(slot)
207                                .map(|value| std::mem::replace(value, Value::Null))
208                                .unwrap_or(Value::Null)
209                        } else {
210                            values.get(slot).cloned().unwrap_or(Value::Null)
211                        }
212                    };
213                    insert_value(columns, result, value);
214                }
215            }
216            Err(values) => {
217                for slot in projection.iter().copied() {
218                    let value = if slot == NULL_SLOT {
219                        Value::Null
220                    } else {
221                        values.get(slot).cloned().unwrap_or(Value::Null)
222                    };
223                    insert_value(columns, result, value);
224                }
225            }
226        }
227    }
228
229    /// Consume this fragment at an explicit row-materialization boundary. Unshared contiguous values, and the common prefix projection emitted by blocking operators, retain their existing allocations instead of being cloned one value at a time.
230    fn into_values(self) -> Vec<Value> {
231        let Self { values, projection } = self;
232        let Some(projection) = projection else {
233            return Arc::try_unwrap(values).unwrap_or_else(|values| values.as_ref().clone());
234        };
235        let mut values = match Arc::try_unwrap(values) {
236            Ok(values) => values,
237            Err(values) => {
238                return projection
239                    .iter()
240                    .map(|slot| {
241                        if *slot == NULL_SLOT {
242                            Value::Null
243                        } else {
244                            values.get(*slot).cloned().unwrap_or(Value::Null)
245                        }
246                    })
247                    .collect();
248            }
249        };
250        if projection
251            .iter()
252            .enumerate()
253            .all(|(position, slot)| position == *slot)
254        {
255            values.truncate(projection.len());
256            return values;
257        }
258
259        let mut remaining = vec![0usize; values.len()];
260        for slot in projection.iter().copied().filter(|slot| *slot != NULL_SLOT) {
261            if let Some(count) = remaining.get_mut(slot) {
262                *count += 1;
263            }
264        }
265        let mut values = values.into_iter().map(Some).collect::<Vec<_>>();
266        projection
267            .iter()
268            .map(|slot| {
269                if *slot == NULL_SLOT {
270                    return Value::Null;
271                }
272                let Some(count) = remaining.get_mut(*slot) else {
273                    return Value::Null;
274                };
275                *count -= 1;
276                if *count == 0 {
277                    values[*slot].take().unwrap_or(Value::Null)
278                } else {
279                    values[*slot].clone().unwrap_or(Value::Null)
280                }
281            })
282            .collect()
283    }
284}
285
286fn insert_values<'a>(
287    columns: &mut impl Iterator<Item = &'a String>,
288    result: &mut ResultRow,
289    values: impl IntoIterator<Item = Value>,
290) {
291    for value in values {
292        insert_value(columns, result, value);
293    }
294}
295
296fn insert_value<'a>(
297    columns: &mut impl Iterator<Item = &'a String>,
298    result: &mut ResultRow,
299    value: Value,
300) {
301    if let Some(column) = columns.next() {
302        result.insert(column.clone(), value);
303    }
304}
305
306impl PhysicalRow {
307    /// Consume an owned row at a positional state boundary, moving uniquely owned fragments and cloning only shared or multiply referenced values.
308    pub fn into_physical_values(self) -> Vec<Value> {
309        let mut fragments = self.into_value_fragments();
310        if fragments.len() == 1 {
311            return fragments.pop().unwrap_or_default();
312        }
313        let capacity = fragments.iter().map(Vec::len).sum();
314        let mut values = Vec::with_capacity(capacity);
315        for fragment in fragments {
316            values.extend(fragment);
317        }
318        values
319    }
320
321    fn into_value_fragments(self) -> SmallVec<[Vec<Value>; INLINE_ROW_FRAGMENTS]> {
322        self.fragments
323            .into_iter()
324            .map(RowFragment::into_values)
325            .collect()
326    }
327}