linkrs_query/executor/streaming/chunk/columnar_batch/
operations.rs

1use std::cmp::Ordering;
2use std::sync::Arc;
3
4use crate::executor::streaming::chunk::columnar_common::{
5    bitmap_is_valid, gather_bitmap, gather_column,
6};
7use crate::executor::streaming::chunk::typed::{TypedColumn, TypedKind};
8use crate::executor::streaming::helpers::compare_values;
9use linkrs_core::value::date_time::{DateTimeValue, DateValue};
10use linkrs_core::value::NullType;
11use linkrs_core::Value;
12
13use super::BatchColumn;
14
15/// Compare two rows of a nullable column with NULL-last ordering (mirrors
16/// [`compare_values`]: NULL equals NULL and sorts last).
17fn nullable_cmp_at<T>(
18    bitmap: &[u64],
19    a: &T,
20    b: &T,
21    a_idx: usize,
22    b_idx: usize,
23    cmp: fn(&T, &T) -> Ordering,
24) -> Ordering {
25    let a_valid = bitmap_is_valid(bitmap, a_idx);
26    let b_valid = bitmap_is_valid(bitmap, b_idx);
27    match (a_valid, b_valid) {
28        (true, true) => cmp(a, b),
29        (false, false) => Ordering::Equal,
30        (false, true) => Ordering::Greater,
31        (true, false) => Ordering::Less,
32    }
33}
34
35/// Append validity bits (packed, one bit per row) to a bitmap starting at
36/// row `rows_before`.
37fn extend_bitmap(bm: &mut Vec<u64>, rows_before: usize, valid: impl Iterator<Item = bool>) {
38    for (row, is_valid) in (rows_before..).zip(valid) {
39        let word = row / 64;
40        if word >= bm.len() {
41            bm.resize(word + 1, 0u64);
42        }
43        if is_valid {
44            bm[word] |= 1u64 << (row % 64);
45        }
46    }
47}
48
49impl BatchColumn {
50    /// Whether this column uses a typed (non-fallback) representation.
51    pub fn is_typed(&self) -> bool {
52        !matches!(self, BatchColumn::Empty | BatchColumn::Fallback(_))
53    }
54
55    /// The raw kind of a typed column (None for Empty/Fallback and for the
56    /// identity columns, which have no scalar `TypedKind`).
57    pub fn kind(&self) -> Option<TypedKind> {
58        match self {
59            BatchColumn::Empty
60            | BatchColumn::Fallback(_)
61            | BatchColumn::VertexIdentity(_)
62            | BatchColumn::EdgeHeader(_) => None,
63            BatchColumn::I64(_) | BatchColumn::NullableI64(..) => Some(TypedKind::I64),
64            BatchColumn::F64(_) | BatchColumn::NullableF64(..) => Some(TypedKind::F64),
65            BatchColumn::I32(_) | BatchColumn::NullableI32(..) => Some(TypedKind::I32),
66            BatchColumn::Bool(_) | BatchColumn::NullableBool(..) => Some(TypedKind::Bool),
67            BatchColumn::Date(_) | BatchColumn::NullableDate(..) => Some(TypedKind::Date),
68            BatchColumn::DateTime(_) | BatchColumn::NullableDateTime(..) => {
69                Some(TypedKind::DateTime)
70            }
71            BatchColumn::Utf8(_) | BatchColumn::NullableUtf8(..) => Some(TypedKind::Utf8),
72            BatchColumn::Decimal(_) | BatchColumn::NullableDecimal(..) => Some(TypedKind::Decimal),
73        }
74    }
75
76    /// Materialize the value at `idx` (O(1) for typed variants; NULL for
77    /// invalid rows of the `Nullable*` variants).
78    pub fn value_at(&self, idx: usize) -> Value {
79        match self {
80            BatchColumn::Empty => Value::Null(NullType::Null),
81            BatchColumn::I64(v) => Value::BigInt(v[idx]),
82            BatchColumn::F64(v) => Value::Double(v[idx]),
83            BatchColumn::I32(v) => Value::Int(v[idx]),
84            BatchColumn::Bool(v) => Value::Bool(v[idx]),
85            BatchColumn::Date(v) => Value::Date(DateValue::from_days(v[idx])),
86            BatchColumn::DateTime(v) => Value::DateTime(DateTimeValue::from_micros(v[idx])),
87            BatchColumn::Utf8(v) => Value::String(v[idx].as_ref().into()),
88            BatchColumn::Decimal(v) => Value::Decimal128(v[idx].clone()),
89            BatchColumn::NullableI64(v, b) => {
90                if bitmap_is_valid(b, idx) {
91                    Value::BigInt(v[idx])
92                } else {
93                    Value::Null(NullType::Null)
94                }
95            }
96            BatchColumn::NullableF64(v, b) => {
97                if bitmap_is_valid(b, idx) {
98                    Value::Double(v[idx])
99                } else {
100                    Value::Null(NullType::Null)
101                }
102            }
103            BatchColumn::NullableI32(v, b) => {
104                if bitmap_is_valid(b, idx) {
105                    Value::Int(v[idx])
106                } else {
107                    Value::Null(NullType::Null)
108                }
109            }
110            BatchColumn::NullableBool(v, b) => {
111                if bitmap_is_valid(b, idx) {
112                    Value::Bool(v[idx])
113                } else {
114                    Value::Null(NullType::Null)
115                }
116            }
117            BatchColumn::NullableDate(v, b) => {
118                if bitmap_is_valid(b, idx) {
119                    Value::Date(DateValue::from_days(v[idx]))
120                } else {
121                    Value::Null(NullType::Null)
122                }
123            }
124            BatchColumn::NullableDateTime(v, b) => {
125                if bitmap_is_valid(b, idx) {
126                    Value::DateTime(DateTimeValue::from_micros(v[idx]))
127                } else {
128                    Value::Null(NullType::Null)
129                }
130            }
131            BatchColumn::NullableUtf8(v, b) => {
132                if bitmap_is_valid(b, idx) {
133                    Value::String(v[idx].as_ref().into())
134                } else {
135                    Value::Null(NullType::Null)
136                }
137            }
138            BatchColumn::NullableDecimal(v, b) => {
139                if bitmap_is_valid(b, idx) {
140                    Value::Decimal128(v[idx].clone())
141                } else {
142                    Value::Null(NullType::Null)
143                }
144            }
145            BatchColumn::VertexIdentity(v) => v
146                .get(idx)
147                .map(|&id| Value::VertexId(id))
148                .unwrap_or(Value::Null(NullType::Null)),
149            BatchColumn::EdgeHeader(v) => v
150                .header_at(idx)
151                .map(Value::edge_header)
152                .unwrap_or(Value::Null(NullType::Null)),
153            BatchColumn::Fallback(v) => v[idx].clone(),
154        }
155    }
156
157    /// Compare two rows on this column.
158    ///
159    /// Typed columns compare on raw scalars (identical ordering to
160    /// [`compare_values`] for same-kind values: i64/i32/bool use the
161    /// primitive order, f64 mirrors `Value` float ordering, strings are
162    /// lexicographic). NULL ordering matches [`compare_values`] (NULLs last).
163    /// Date columns and fallback columns delegate to [`compare_values`] (the
164    /// row path falls back to the string representation there, which
165    /// diverges from the day order for pre-epoch dates).
166    pub fn compare_at(&self, a: usize, b: usize) -> Ordering {
167        match self {
168            BatchColumn::Empty => Ordering::Equal,
169            BatchColumn::I64(v) => v[a].cmp(&v[b]),
170            BatchColumn::F64(v) => {
171                let x = v[a];
172                let y = v[b];
173                if x < y {
174                    Ordering::Less
175                } else if x > y {
176                    Ordering::Greater
177                } else {
178                    Ordering::Equal
179                }
180            }
181            BatchColumn::I32(v) => v[a].cmp(&v[b]),
182            BatchColumn::Bool(v) => v[a].cmp(&v[b]),
183            BatchColumn::Date(v) => compare_values(
184                &Value::Date(DateValue::from_days(v[a])),
185                &Value::Date(DateValue::from_days(v[b])),
186            ),
187            BatchColumn::DateTime(v) => compare_values(
188                &Value::DateTime(DateTimeValue::from_micros(v[a])),
189                &Value::DateTime(DateTimeValue::from_micros(v[b])),
190            ),
191            BatchColumn::Utf8(v) => v[a].cmp(&v[b]),
192            BatchColumn::Decimal(v) => v[a].cmp(&v[b]),
193            BatchColumn::NullableI64(v, bm) => {
194                nullable_cmp_at(bm, &v[a], &v[b], a, b, |x, y| x.cmp(y))
195            }
196            BatchColumn::NullableF64(v, bm) => nullable_cmp_at(bm, &v[a], &v[b], a, b, |x, y| {
197                if x < y {
198                    Ordering::Less
199                } else if x > y {
200                    Ordering::Greater
201                } else {
202                    Ordering::Equal
203                }
204            }),
205            BatchColumn::NullableI32(v, bm) => {
206                nullable_cmp_at(bm, &v[a], &v[b], a, b, |x, y| x.cmp(y))
207            }
208            BatchColumn::NullableBool(v, bm) => {
209                nullable_cmp_at(bm, &v[a], &v[b], a, b, |x, y| x.cmp(y))
210            }
211            BatchColumn::NullableDate(v, bm) => {
212                nullable_cmp_at(bm, &v[a], &v[b], a, b, |x, y| x.cmp(y))
213            }
214            BatchColumn::NullableDateTime(v, bm) => {
215                nullable_cmp_at(bm, &v[a], &v[b], a, b, |x, y| x.cmp(y))
216            }
217            BatchColumn::NullableUtf8(v, bm) => {
218                nullable_cmp_at(bm, &v[a], &v[b], a, b, |x, y| x.cmp(y))
219            }
220            BatchColumn::NullableDecimal(v, bm) => {
221                nullable_cmp_at(bm, &v[a], &v[b], a, b, |x, y| x.cmp(y))
222            }
223            BatchColumn::VertexIdentity(v) => v[a].cmp(&v[b]),
224            // Edge headers have no raw `Ord`: compare the materialized
225            // values exactly as the `Fallback` path does.
226            BatchColumn::EdgeHeader(v) => compare_values(
227                &v.header_at(a)
228                    .map(Value::edge_header)
229                    .unwrap_or(Value::Null(NullType::Null)),
230                &v.header_at(b)
231                    .map(Value::edge_header)
232                    .unwrap_or(Value::Null(NullType::Null)),
233            ),
234            BatchColumn::Fallback(v) => compare_values(&v[a], &v[b]),
235        }
236    }
237
238    /// Compare the value `v` against row `idx` of this column.
239    ///
240    /// Uses the raw fast path when `v` matches the typed column kind and the
241    /// row is valid; otherwise delegates to [`compare_values`] exactly.
242    pub fn compare_value_at(&self, v: &Value, idx: usize) -> Ordering {
243        let raw = match (self, v) {
244            (BatchColumn::I64(col), Value::BigInt(x)) => Some(x.cmp(&col[idx])),
245            (BatchColumn::F64(col), Value::Double(x)) => {
246                let y = col[idx];
247                Some(if *x < y {
248                    Ordering::Less
249                } else if *x > y {
250                    Ordering::Greater
251                } else {
252                    Ordering::Equal
253                })
254            }
255            (BatchColumn::I32(col), Value::Int(x)) => Some(x.cmp(&col[idx])),
256            (BatchColumn::Bool(col), Value::Bool(x)) => Some(x.cmp(&col[idx])),
257            (BatchColumn::Utf8(col), Value::String(x)) => Some(x.as_str().cmp(&col[idx])),
258            (BatchColumn::NullableI64(col, b), Value::BigInt(x)) => {
259                bitmap_is_valid(b, idx).then(|| x.cmp(&col[idx]))
260            }
261            (BatchColumn::NullableF64(col, b), Value::Double(x)) => {
262                bitmap_is_valid(b, idx).then(|| {
263                    let y = col[idx];
264                    if *x < y {
265                        Ordering::Less
266                    } else if *x > y {
267                        Ordering::Greater
268                    } else {
269                        Ordering::Equal
270                    }
271                })
272            }
273            (BatchColumn::NullableI32(col, b), Value::Int(x)) => {
274                bitmap_is_valid(b, idx).then(|| x.cmp(&col[idx]))
275            }
276            (BatchColumn::NullableBool(col, b), Value::Bool(x)) => {
277                bitmap_is_valid(b, idx).then(|| x.cmp(&col[idx]))
278            }
279            (BatchColumn::NullableUtf8(col, b), Value::String(x)) => {
280                bitmap_is_valid(b, idx).then(|| x.as_str().cmp(&col[idx]))
281            }
282            _ => None,
283        };
284        match raw {
285            Some(ordering) => ordering,
286            None => compare_values(v, &self.value_at(idx)),
287        }
288    }
289
290    /// Append the entries of a chunk column at `indices`.
291    ///
292    /// When `self` is Empty the kind is taken from the chunk column (a
293    /// fallback chunk column starts a fallback batch column). A kind
294    /// mismatch degrades the accumulated column to [`BatchColumn::Fallback`].
295    /// A NULL introduced by a `Nullable*` chunk column upgrades a plain
296    /// typed column to its `Nullable*` form (past rows become valid).
297    pub(super) fn append_typed(&mut self, col: &TypedColumn, indices: &[usize]) {
298        match self {
299            BatchColumn::Empty => {
300                *self = Self::gather(col, indices);
301            }
302            BatchColumn::I64(buf) => match col {
303                TypedColumn::I64(src) => buf.extend(indices.iter().map(|&i| src[i])),
304                TypedColumn::NullableI64(src, bm) => {
305                    let rows_before = buf.len();
306                    *self = Self::to_nullable(self);
307                    if let BatchColumn::NullableI64(buf, bm_out) = self {
308                        buf.extend(indices.iter().map(|&i| src[i]));
309                        extend_bitmap(
310                            bm_out,
311                            rows_before,
312                            indices.iter().map(|&i| bitmap_is_valid(bm, i)),
313                        );
314                    }
315                }
316                _ => *self = Self::degraded_append(self, col, indices),
317            },
318            BatchColumn::F64(buf) => match col {
319                TypedColumn::F64(src) => buf.extend(indices.iter().map(|&i| src[i])),
320                TypedColumn::NullableF64(src, bm) => {
321                    let rows_before = buf.len();
322                    *self = Self::to_nullable(self);
323                    if let BatchColumn::NullableF64(buf, bm_out) = self {
324                        buf.extend(indices.iter().map(|&i| src[i]));
325                        extend_bitmap(
326                            bm_out,
327                            rows_before,
328                            indices.iter().map(|&i| bitmap_is_valid(bm, i)),
329                        );
330                    }
331                }
332                _ => *self = Self::degraded_append(self, col, indices),
333            },
334            BatchColumn::I32(buf) => match col {
335                TypedColumn::I32(src) => buf.extend(indices.iter().map(|&i| src[i])),
336                TypedColumn::NullableI32(src, bm) => {
337                    let rows_before = buf.len();
338                    *self = Self::to_nullable(self);
339                    if let BatchColumn::NullableI32(buf, bm_out) = self {
340                        buf.extend(indices.iter().map(|&i| src[i]));
341                        extend_bitmap(
342                            bm_out,
343                            rows_before,
344                            indices.iter().map(|&i| bitmap_is_valid(bm, i)),
345                        );
346                    }
347                }
348                _ => *self = Self::degraded_append(self, col, indices),
349            },
350            BatchColumn::Bool(buf) => match col {
351                TypedColumn::Bool(src) => buf.extend(indices.iter().map(|&i| src[i])),
352                TypedColumn::NullableBool(src, bm) => {
353                    let rows_before = buf.len();
354                    *self = Self::to_nullable(self);
355                    if let BatchColumn::NullableBool(buf, bm_out) = self {
356                        buf.extend(indices.iter().map(|&i| src[i]));
357                        extend_bitmap(
358                            bm_out,
359                            rows_before,
360                            indices.iter().map(|&i| bitmap_is_valid(bm, i)),
361                        );
362                    }
363                }
364                _ => *self = Self::degraded_append(self, col, indices),
365            },
366            BatchColumn::Date(buf) => match col {
367                TypedColumn::Date(src) => buf.extend(indices.iter().map(|&i| src[i])),
368                TypedColumn::NullableDate(src, bm) => {
369                    let rows_before = buf.len();
370                    *self = Self::to_nullable(self);
371                    if let BatchColumn::NullableDate(buf, bm_out) = self {
372                        buf.extend(indices.iter().map(|&i| src[i]));
373                        extend_bitmap(
374                            bm_out,
375                            rows_before,
376                            indices.iter().map(|&i| bitmap_is_valid(bm, i)),
377                        );
378                    }
379                }
380                _ => *self = Self::degraded_append(self, col, indices),
381            },
382            BatchColumn::DateTime(buf) => match col {
383                TypedColumn::DateTime(src) => buf.extend(indices.iter().map(|&i| src[i])),
384                TypedColumn::NullableDateTime(src, bm) => {
385                    let rows_before = buf.len();
386                    *self = Self::to_nullable(self);
387                    if let BatchColumn::NullableDateTime(buf, bm_out) = self {
388                        buf.extend(indices.iter().map(|&i| src[i]));
389                        extend_bitmap(
390                            bm_out,
391                            rows_before,
392                            indices.iter().map(|&i| bitmap_is_valid(bm, i)),
393                        );
394                    }
395                }
396                _ => *self = Self::degraded_append(self, col, indices),
397            },
398            BatchColumn::Utf8(buf) => match col {
399                TypedColumn::Utf8(src) => buf.extend(indices.iter().map(|&i| src[i].clone())),
400                TypedColumn::NullableUtf8(src, bm) => {
401                    let rows_before = buf.len();
402                    *self = Self::to_nullable(self);
403                    if let BatchColumn::NullableUtf8(buf, bm_out) = self {
404                        buf.extend(indices.iter().map(|&i| src[i].clone()));
405                        extend_bitmap(
406                            bm_out,
407                            rows_before,
408                            indices.iter().map(|&i| bitmap_is_valid(bm, i)),
409                        );
410                    }
411                }
412                _ => *self = Self::degraded_append(self, col, indices),
413            },
414            BatchColumn::Decimal(buf) => match col {
415                TypedColumn::Decimal(src) => buf.extend(indices.iter().map(|&i| src[i].clone())),
416                TypedColumn::NullableDecimal(src, bm) => {
417                    let rows_before = buf.len();
418                    *self = Self::to_nullable(self);
419                    if let BatchColumn::NullableDecimal(buf, bm_out) = self {
420                        buf.extend(indices.iter().map(|&i| src[i].clone()));
421                        extend_bitmap(
422                            bm_out,
423                            rows_before,
424                            indices.iter().map(|&i| bitmap_is_valid(bm, i)),
425                        );
426                    }
427                }
428                _ => *self = Self::degraded_append(self, col, indices),
429            },
430            BatchColumn::NullableI64(buf, bm_out) => match col {
431                TypedColumn::NullableI64(src, bm) => {
432                    let rows_before = buf.len();
433                    buf.extend(indices.iter().map(|&i| src[i]));
434                    extend_bitmap(
435                        bm_out,
436                        rows_before,
437                        indices.iter().map(|&i| bitmap_is_valid(bm, i)),
438                    );
439                }
440                TypedColumn::I64(src) => {
441                    let rows_before = buf.len();
442                    buf.extend(indices.iter().map(|&i| src[i]));
443                    extend_bitmap(bm_out, rows_before, indices.iter().map(|_| true));
444                }
445                _ => *self = Self::degraded_append(self, col, indices),
446            },
447            BatchColumn::NullableF64(buf, bm_out) => match col {
448                TypedColumn::NullableF64(src, bm) => {
449                    let rows_before = buf.len();
450                    buf.extend(indices.iter().map(|&i| src[i]));
451                    extend_bitmap(
452                        bm_out,
453                        rows_before,
454                        indices.iter().map(|&i| bitmap_is_valid(bm, i)),
455                    );
456                }
457                TypedColumn::F64(src) => {
458                    let rows_before = buf.len();
459                    buf.extend(indices.iter().map(|&i| src[i]));
460                    extend_bitmap(bm_out, rows_before, indices.iter().map(|_| true));
461                }
462                _ => *self = Self::degraded_append(self, col, indices),
463            },
464            BatchColumn::NullableI32(buf, bm_out) => match col {
465                TypedColumn::NullableI32(src, bm) => {
466                    let rows_before = buf.len();
467                    buf.extend(indices.iter().map(|&i| src[i]));
468                    extend_bitmap(
469                        bm_out,
470                        rows_before,
471                        indices.iter().map(|&i| bitmap_is_valid(bm, i)),
472                    );
473                }
474                TypedColumn::I32(src) => {
475                    let rows_before = buf.len();
476                    buf.extend(indices.iter().map(|&i| src[i]));
477                    extend_bitmap(bm_out, rows_before, indices.iter().map(|_| true));
478                }
479                _ => *self = Self::degraded_append(self, col, indices),
480            },
481            BatchColumn::NullableBool(buf, bm_out) => match col {
482                TypedColumn::NullableBool(src, bm) => {
483                    let rows_before = buf.len();
484                    buf.extend(indices.iter().map(|&i| src[i]));
485                    extend_bitmap(
486                        bm_out,
487                        rows_before,
488                        indices.iter().map(|&i| bitmap_is_valid(bm, i)),
489                    );
490                }
491                TypedColumn::Bool(src) => {
492                    let rows_before = buf.len();
493                    buf.extend(indices.iter().map(|&i| src[i]));
494                    extend_bitmap(bm_out, rows_before, indices.iter().map(|_| true));
495                }
496                _ => *self = Self::degraded_append(self, col, indices),
497            },
498            BatchColumn::NullableDate(buf, bm_out) => match col {
499                TypedColumn::NullableDate(src, bm) => {
500                    let rows_before = buf.len();
501                    buf.extend(indices.iter().map(|&i| src[i]));
502                    extend_bitmap(
503                        bm_out,
504                        rows_before,
505                        indices.iter().map(|&i| bitmap_is_valid(bm, i)),
506                    );
507                }
508                TypedColumn::Date(src) => {
509                    let rows_before = buf.len();
510                    buf.extend(indices.iter().map(|&i| src[i]));
511                    extend_bitmap(bm_out, rows_before, indices.iter().map(|_| true));
512                }
513                _ => *self = Self::degraded_append(self, col, indices),
514            },
515            BatchColumn::NullableDateTime(buf, bm_out) => match col {
516                TypedColumn::NullableDateTime(src, bm) => {
517                    let rows_before = buf.len();
518                    buf.extend(indices.iter().map(|&i| src[i]));
519                    extend_bitmap(
520                        bm_out,
521                        rows_before,
522                        indices.iter().map(|&i| bitmap_is_valid(bm, i)),
523                    );
524                }
525                TypedColumn::DateTime(src) => {
526                    let rows_before = buf.len();
527                    buf.extend(indices.iter().map(|&i| src[i]));
528                    extend_bitmap(bm_out, rows_before, indices.iter().map(|_| true));
529                }
530                _ => *self = Self::degraded_append(self, col, indices),
531            },
532            BatchColumn::NullableUtf8(buf, bm_out) => match col {
533                TypedColumn::NullableUtf8(src, bm) => {
534                    let rows_before = buf.len();
535                    buf.extend(indices.iter().map(|&i| src[i].clone()));
536                    extend_bitmap(
537                        bm_out,
538                        rows_before,
539                        indices.iter().map(|&i| bitmap_is_valid(bm, i)),
540                    );
541                }
542                TypedColumn::Utf8(src) => {
543                    let rows_before = buf.len();
544                    buf.extend(indices.iter().map(|&i| src[i].clone()));
545                    extend_bitmap(bm_out, rows_before, indices.iter().map(|_| true));
546                }
547                _ => *self = Self::degraded_append(self, col, indices),
548            },
549            BatchColumn::NullableDecimal(buf, bm_out) => match col {
550                TypedColumn::NullableDecimal(src, bm) => {
551                    let rows_before = buf.len();
552                    buf.extend(indices.iter().map(|&i| src[i].clone()));
553                    extend_bitmap(
554                        bm_out,
555                        rows_before,
556                        indices.iter().map(|&i| bitmap_is_valid(bm, i)),
557                    );
558                }
559                TypedColumn::Decimal(src) => {
560                    let rows_before = buf.len();
561                    buf.extend(indices.iter().map(|&i| src[i].clone()));
562                    extend_bitmap(bm_out, rows_before, indices.iter().map(|_| true));
563                }
564                _ => *self = Self::degraded_append(self, col, indices),
565            },
566            BatchColumn::VertexIdentity(buf) => match col {
567                TypedColumn::VertexIdentity(src) => buf.extend(indices.iter().map(|&i| src[i])),
568                _ => *self = Self::degraded_append(self, col, indices),
569            },
570            BatchColumn::EdgeHeader(buf) => match col {
571                TypedColumn::EdgeHeader(src) => buf.extend_indexed(src, indices),
572                _ => *self = Self::degraded_append(self, col, indices),
573            },
574            BatchColumn::Fallback(buf) => {
575                buf.extend(indices.iter().map(|&i| {
576                    col.value_at(i)
577                        .unwrap_or_else(|| Value::Null(NullType::Null))
578                }));
579            }
580        }
581    }
582
583    /// Build a batch column from a chunk column at `indices` (kind taken
584    /// from the chunk column).
585    fn gather(col: &TypedColumn, indices: &[usize]) -> Self {
586        gather_column!(BatchColumn, col, indices)
587    }
588
589    /// Upgrade a plain typed column to its `Nullable*` form (past rows all
590    /// valid), used when a later chunk introduces NULLs into the column.
591    fn to_nullable(current: &Self) -> Self {
592        match current {
593            BatchColumn::I64(v) => {
594                BatchColumn::NullableI64(v.clone(), vec![!0u64; v.len().div_ceil(64)])
595            }
596            BatchColumn::F64(v) => {
597                BatchColumn::NullableF64(v.clone(), vec![!0u64; v.len().div_ceil(64)])
598            }
599            BatchColumn::I32(v) => {
600                BatchColumn::NullableI32(v.clone(), vec![!0u64; v.len().div_ceil(64)])
601            }
602            BatchColumn::Bool(v) => {
603                BatchColumn::NullableBool(v.clone(), vec![!0u64; v.len().div_ceil(64)])
604            }
605            BatchColumn::Date(v) => {
606                BatchColumn::NullableDate(v.clone(), vec![!0u64; v.len().div_ceil(64)])
607            }
608            BatchColumn::DateTime(v) => {
609                BatchColumn::NullableDateTime(v.clone(), vec![!0u64; v.len().div_ceil(64)])
610            }
611            BatchColumn::Utf8(v) => {
612                BatchColumn::NullableUtf8(v.clone(), vec![!0u64; v.len().div_ceil(64)])
613            }
614            BatchColumn::Decimal(v) => {
615                BatchColumn::NullableDecimal(v.clone(), vec![!0u64; v.len().div_ceil(64)])
616            }
617            _ => current.clone(),
618        }
619    }
620
621    /// Degrade an existing typed column to `Fallback` (keeping accumulated
622    /// rows) and append the chunk column values.
623    fn degraded_append(current: &Self, col: &TypedColumn, indices: &[usize]) -> Self {
624        let mut values: Vec<Value> = (0..current.len()).map(|i| current.value_at(i)).collect();
625        values.extend(indices.iter().map(|&i| {
626            col.value_at(i)
627                .unwrap_or_else(|| Value::Null(NullType::Null))
628        }));
629        BatchColumn::Fallback(values)
630    }
631
632    pub(super) fn append_row_value(&mut self, value: &Value) {
633        match self {
634            BatchColumn::Empty => {
635                // No kind established yet: start as a single-value fallback
636                // (rows do not carry typed raw data on the append path).
637                *self = BatchColumn::Fallback(vec![value.clone()]);
638            }
639            BatchColumn::I64(buf) => {
640                if let Value::BigInt(x) = value {
641                    buf.push(*x);
642                } else {
643                    *self = BatchColumn::degraded_push(self, value);
644                }
645            }
646            BatchColumn::F64(buf) => {
647                if let Value::Double(x) = value {
648                    buf.push(*x);
649                } else {
650                    *self = BatchColumn::degraded_push(self, value);
651                }
652            }
653            BatchColumn::I32(buf) => {
654                if let Value::Int(x) = value {
655                    buf.push(*x);
656                } else {
657                    *self = BatchColumn::degraded_push(self, value);
658                }
659            }
660            BatchColumn::Bool(buf) => {
661                if let Value::Bool(x) = value {
662                    buf.push(*x);
663                } else {
664                    *self = BatchColumn::degraded_push(self, value);
665                }
666            }
667            BatchColumn::Date(buf) => {
668                if let Value::Date(x) = value {
669                    buf.push(x.to_days());
670                } else {
671                    *self = BatchColumn::degraded_push(self, value);
672                }
673            }
674            BatchColumn::DateTime(buf) => {
675                if let Value::DateTime(x) = value {
676                    buf.push(x.to_micros());
677                } else {
678                    *self = BatchColumn::degraded_push(self, value);
679                }
680            }
681            BatchColumn::Utf8(buf) => {
682                if let Value::String(x) = value {
683                    buf.push(Arc::from(x.as_str()));
684                } else {
685                    *self = BatchColumn::degraded_push(self, value);
686                }
687            }
688            BatchColumn::Decimal(buf) => {
689                if let Value::Decimal128(x) = value {
690                    buf.push(x.clone());
691                } else {
692                    *self = BatchColumn::degraded_push(self, value);
693                }
694            }
695            BatchColumn::NullableI64(buf, bm) => {
696                if let Value::BigInt(x) = value {
697                    buf.push(*x);
698                    extend_bitmap(bm, buf.len() - 1, std::iter::once(true));
699                } else {
700                    *self = BatchColumn::degraded_push(self, value);
701                }
702            }
703            BatchColumn::NullableF64(buf, bm) => {
704                if let Value::Double(x) = value {
705                    buf.push(*x);
706                    extend_bitmap(bm, buf.len() - 1, std::iter::once(true));
707                } else {
708                    *self = BatchColumn::degraded_push(self, value);
709                }
710            }
711            BatchColumn::NullableI32(buf, bm) => {
712                if let Value::Int(x) = value {
713                    buf.push(*x);
714                    extend_bitmap(bm, buf.len() - 1, std::iter::once(true));
715                } else {
716                    *self = BatchColumn::degraded_push(self, value);
717                }
718            }
719            BatchColumn::NullableBool(buf, bm) => {
720                if let Value::Bool(x) = value {
721                    buf.push(*x);
722                    extend_bitmap(bm, buf.len() - 1, std::iter::once(true));
723                } else {
724                    *self = BatchColumn::degraded_push(self, value);
725                }
726            }
727            BatchColumn::NullableDate(buf, bm) => {
728                if let Value::Date(x) = value {
729                    buf.push(x.to_days());
730                    extend_bitmap(bm, buf.len() - 1, std::iter::once(true));
731                } else {
732                    *self = BatchColumn::degraded_push(self, value);
733                }
734            }
735            BatchColumn::NullableDateTime(buf, bm) => {
736                if let Value::DateTime(x) = value {
737                    buf.push(x.to_micros());
738                    extend_bitmap(bm, buf.len() - 1, std::iter::once(true));
739                } else {
740                    *self = BatchColumn::degraded_push(self, value);
741                }
742            }
743            BatchColumn::NullableUtf8(buf, bm) => {
744                if let Value::String(x) = value {
745                    buf.push(Arc::from(x.as_str()));
746                    extend_bitmap(bm, buf.len() - 1, std::iter::once(true));
747                } else {
748                    *self = BatchColumn::degraded_push(self, value);
749                }
750            }
751            BatchColumn::NullableDecimal(buf, bm) => {
752                if let Value::Decimal128(x) = value {
753                    buf.push(x.clone());
754                    extend_bitmap(bm, buf.len() - 1, std::iter::once(true));
755                } else {
756                    *self = BatchColumn::degraded_push(self, value);
757                }
758            }
759            BatchColumn::VertexIdentity(buf) => {
760                if let Value::VertexId(id) = value {
761                    buf.push(*id);
762                } else {
763                    *self = BatchColumn::degraded_push(self, value);
764                }
765            }
766            BatchColumn::EdgeHeader(buf) => {
767                if let Value::EdgeHeader(header) = value {
768                    buf.push(header.src, header.dst, &header.edge_type, header.ranking);
769                } else {
770                    *self = BatchColumn::degraded_push(self, value);
771                }
772            }
773            BatchColumn::Fallback(buf) => buf.push(value.clone()),
774        }
775    }
776
777    fn degraded_push(current: &Self, value: &Value) -> Self {
778        let mut values: Vec<Value> = (0..current.len()).map(|i| current.value_at(i)).collect();
779        values.push(value.clone());
780        BatchColumn::Fallback(values)
781    }
782
783    pub(super) fn truncate(&mut self, len: usize) {
784        match self {
785            BatchColumn::Empty => {}
786            BatchColumn::I64(v) => v.truncate(len),
787            BatchColumn::F64(v) => v.truncate(len),
788            BatchColumn::I32(v) => v.truncate(len),
789            BatchColumn::Bool(v) => v.truncate(len),
790            BatchColumn::Date(v) => v.truncate(len),
791            BatchColumn::DateTime(v) => v.truncate(len),
792            BatchColumn::Utf8(v) => v.truncate(len),
793            BatchColumn::Decimal(v) => v.truncate(len),
794            BatchColumn::NullableI64(v, bm) => {
795                v.truncate(len);
796                bm.truncate(len.div_ceil(64));
797            }
798            BatchColumn::NullableF64(v, bm) => {
799                v.truncate(len);
800                bm.truncate(len.div_ceil(64));
801            }
802            BatchColumn::NullableI32(v, bm) => {
803                v.truncate(len);
804                bm.truncate(len.div_ceil(64));
805            }
806            BatchColumn::NullableBool(v, bm) => {
807                v.truncate(len);
808                bm.truncate(len.div_ceil(64));
809            }
810            BatchColumn::NullableDate(v, bm) => {
811                v.truncate(len);
812                bm.truncate(len.div_ceil(64));
813            }
814            BatchColumn::NullableDateTime(v, bm) => {
815                v.truncate(len);
816                bm.truncate(len.div_ceil(64));
817            }
818            BatchColumn::NullableUtf8(v, bm) => {
819                v.truncate(len);
820                bm.truncate(len.div_ceil(64));
821            }
822            BatchColumn::NullableDecimal(v, bm) => {
823                v.truncate(len);
824                bm.truncate(len.div_ceil(64));
825            }
826            BatchColumn::VertexIdentity(v) => v.truncate(len),
827            BatchColumn::EdgeHeader(v) => v.truncate(len),
828            BatchColumn::Fallback(v) => v.truncate(len),
829        }
830    }
831
832    pub(super) fn permute(&mut self, perm: &[usize]) {
833        match self {
834            BatchColumn::Empty => {}
835            BatchColumn::I64(v) => {
836                let old = std::mem::take(v);
837                *v = perm.iter().map(|&i| old[i]).collect();
838            }
839            BatchColumn::F64(v) => {
840                let old = std::mem::take(v);
841                *v = perm.iter().map(|&i| old[i]).collect();
842            }
843            BatchColumn::I32(v) => {
844                let old = std::mem::take(v);
845                *v = perm.iter().map(|&i| old[i]).collect();
846            }
847            BatchColumn::Bool(v) => {
848                let old = std::mem::take(v);
849                *v = perm.iter().map(|&i| old[i]).collect();
850            }
851            BatchColumn::Date(v) => {
852                let old = std::mem::take(v);
853                *v = perm.iter().map(|&i| old[i]).collect();
854            }
855            BatchColumn::DateTime(v) => {
856                let old = std::mem::take(v);
857                *v = perm.iter().map(|&i| old[i]).collect();
858            }
859            BatchColumn::Utf8(v) => {
860                let old = std::mem::take(v);
861                *v = perm.iter().map(|&i| old[i].clone()).collect();
862            }
863            BatchColumn::Decimal(v) => {
864                let old = std::mem::take(v);
865                *v = perm.iter().map(|&i| old[i].clone()).collect();
866            }
867            BatchColumn::NullableI64(v, bm) => {
868                let (old_v, old_bm) = (std::mem::take(v), std::mem::take(bm));
869                *v = perm.iter().map(|&i| old_v[i]).collect();
870                *bm = gather_bitmap(&old_bm, perm);
871            }
872            BatchColumn::NullableF64(v, bm) => {
873                let (old_v, old_bm) = (std::mem::take(v), std::mem::take(bm));
874                *v = perm.iter().map(|&i| old_v[i]).collect();
875                *bm = gather_bitmap(&old_bm, perm);
876            }
877            BatchColumn::NullableI32(v, bm) => {
878                let (old_v, old_bm) = (std::mem::take(v), std::mem::take(bm));
879                *v = perm.iter().map(|&i| old_v[i]).collect();
880                *bm = gather_bitmap(&old_bm, perm);
881            }
882            BatchColumn::NullableBool(v, bm) => {
883                let (old_v, old_bm) = (std::mem::take(v), std::mem::take(bm));
884                *v = perm.iter().map(|&i| old_v[i]).collect();
885                *bm = gather_bitmap(&old_bm, perm);
886            }
887            BatchColumn::NullableDate(v, bm) => {
888                let (old_v, old_bm) = (std::mem::take(v), std::mem::take(bm));
889                *v = perm.iter().map(|&i| old_v[i]).collect();
890                *bm = gather_bitmap(&old_bm, perm);
891            }
892            BatchColumn::NullableDateTime(v, bm) => {
893                let (old_v, old_bm) = (std::mem::take(v), std::mem::take(bm));
894                *v = perm.iter().map(|&i| old_v[i]).collect();
895                *bm = gather_bitmap(&old_bm, perm);
896            }
897            BatchColumn::NullableUtf8(v, bm) => {
898                let (old_v, old_bm) = (std::mem::take(v), std::mem::take(bm));
899                *v = perm.iter().map(|&i| old_v[i].clone()).collect();
900                *bm = gather_bitmap(&old_bm, perm);
901            }
902            BatchColumn::NullableDecimal(v, bm) => {
903                let (old_v, old_bm) = (std::mem::take(v), std::mem::take(bm));
904                *v = perm.iter().map(|&i| old_v[i].clone()).collect();
905                *bm = gather_bitmap(&old_bm, perm);
906            }
907            BatchColumn::VertexIdentity(v) => {
908                let old = std::mem::take(v);
909                *v = perm.iter().map(|&i| old[i]).collect();
910            }
911            BatchColumn::EdgeHeader(v) => v.permute(perm),
912            BatchColumn::Fallback(v) => {
913                let old = std::mem::take(v);
914                *v = perm.iter().map(|&i| old[i].clone()).collect();
915            }
916        }
917    }
918}
linkrs_query/executor/streaming/chunk/columnar_batch/operations.rs

linkrs_query/executor/streaming/chunk/columnar_batch/
operations.rs
