linkrs_query/executor/streaming/operators/blocking/
helpers.rs

1use crate::executor::streaming::chunk::DataChunk;
2use crate::executor::streaming::chunk::TypedColumn;
3use crate::executor::streaming::operators::source_operator::OperatorConfig;
4use crate::executor::streaming::runtime::ExecutionRuntime;
5use crate::executor::streaming::slot::SlotLayout;
6use linkrs_core::columnar::MaterializedBatch;
7use linkrs_core::types::expr::Expression;
8use linkrs_core::types::operators::AggregateFunction;
9use std::sync::Arc;
10
11pub(super) struct BlockingContext<'a> {
12    pub runtime: &'a Option<Arc<ExecutionRuntime>>,
13    pub output_layout: &'a Arc<SlotLayout>,
14    pub config: &'a OperatorConfig,
15}
16
17/// Whether a `Count` over a bare variable can skip argument evaluation for
18/// this chunk: identity typed columns are provably non-null, and `Count`
19/// only observes null-ness, so every visible row counts without evaluating
20/// or materializing the argument.
21pub(super) fn count_identity_arg(
22    func: &AggregateFunction,
23    args: &[Expression],
24    chunk: &DataChunk,
25) -> bool {
26    if !matches!(func, AggregateFunction::Count) {
27        return false;
28    }
29    let Some(Expression::Variable(name)) = args.first() else {
30        return false;
31    };
32    let Some(slot) = chunk.get_layout().slot_id(name) else {
33        return false;
34    };
35    matches!(
36        chunk.typed_column(slot),
37        Some(TypedColumn::VertexIdentity(_)) | Some(TypedColumn::EdgeHeader(_))
38    )
39}
40
41/// Extract the field name from an aggregate function's args, if any.
42/// COUNT(*) has no args; other aggregates have the field expression at index 0.
43pub(crate) fn aggregate_arg_field_name(
44    func: &AggregateFunction,
45    args: &[Expression],
46) -> Option<String> {
47    match func {
48        AggregateFunction::Count => None,
49        _ => {
50            if let Some(Expression::Variable(name)) = args.first() {
51                Some(name.clone())
52            } else {
53                None
54            }
55        }
56    }
57}
58
59/// Shared outlet for blocking operators: replaces the drained
60/// `IntoIter<Vec<Value>>` pattern while keeping chunking (2048 rows),
61/// logical row order, and outlet observability identical.
62pub(super) fn emit_batch_slice(
63    batch: &MaterializedBatch,
64    offset: &mut usize,
65    ctx: &BlockingContext<'_>,
66) -> DataChunk {
67    let total = batch.num_rows();
68    let len = 2048.min(total.saturating_sub(*offset));
69    let chunk = DataChunk::slice_from_batch(batch, *offset, len, Arc::clone(ctx.output_layout));
70    *offset += len;
71    if let Some(rt) = ctx.runtime.as_ref() {
72        rt.columnar_stats().record_batch_outlet();
73    }
74    chunk
75}
linkrs_query/executor/streaming/operators/blocking/helpers.rs

linkrs_query/executor/streaming/operators/blocking/
helpers.rs
