uqa_execution/batch/
materialization.rs1use 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 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 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 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 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}