feat(promql): resolve runtime input ids from source bindings

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/promql/src/native/execute.rs

@@ -2054,6 +2054,7 @@ struct NodeState {
 
 /// The caller-owned state of one input leaf.
 struct InputSlot {
+    source_index: usize,
     schema: Arc<SeriesSchema>,
     open_source_series: bool,
     state: NodeState,
@@ -2100,6 +2101,7 @@ impl RuntimeNode {
                 })?;
                 let id = InputId(inputs.len());
                 inputs.push(InputSlot {
+                    source_index: bound.source_index,
                     schema,
                     open_source_series: false,
                     state: NodeState::new("native.input", Arc::clone(&metrics), pool),
@@ -2862,6 +2864,15 @@ impl NativeExecutionState {
         self.root.register_metrics("root", set, partition);
     }
 
+    /// Returns the runtime input slots bound to `source_index`.
+    pub fn input_ids(&self, source_index: usize) -> impl Iterator<Item = InputId> + '_ {
+        self.inputs
+            .iter()
+            .enumerate()
+            .filter(move |(_, slot)| slot.source_index == source_index)
+            .map(|(index, _)| InputId::new(index))
+    }
+
     pub fn reserved_bytes(&self) -> usize {
         self.inputs
             .iter()

src/promql/src/native/execute/tests.rs

@@ -106,11 +106,18 @@ fn input() -> NativeNode {
 
 /// The same raw source, declaring `ordering`.
 fn raw_input_with_ordering(ordering: SeriesOrdering) -> NativeNode {
+    raw_input_with_source_index_and_ordering(0, ordering)
+}
+
+fn raw_input_with_source_index_and_ordering(
+    source_index: usize,
+    ordering: SeriesOrdering,
+) -> NativeNode {
     let output = NativeType::Samples(raw_schema());
     NativeNode {
         expr: NativeExpr::Input(BoundInput {
             output: output.clone(),
-            source_index: 0,
+            source_index,
             ordering,
         }),
         output,
@@ -132,16 +139,28 @@ fn raw_input_of(count: usize) -> NativeNode {
 
 /// A source whose chunks are already evaluated, such as the inner result of a subquery.
 fn evaluated_input(grid: Arc<EvalGrid>) -> NativeNode {
-    evaluated_input_with_ordering(grid, SeriesOrdering::KeyThenTime)
+    evaluated_input_with_source_index(grid, 0)
+}
+
+fn evaluated_input_with_source_index(grid: Arc<EvalGrid>, source_index: usize) -> NativeNode {
+    evaluated_input_with_source_index_and_ordering(grid, source_index, SeriesOrdering::KeyThenTime)
 }
 
 /// A source of evaluated chunks that guarantees `ordering`.
 fn evaluated_input_with_ordering(grid: Arc<EvalGrid>, ordering: SeriesOrdering) -> NativeNode {
+    evaluated_input_with_source_index_and_ordering(grid, 0, ordering)
+}
+
+fn evaluated_input_with_source_index_and_ordering(
+    grid: Arc<EvalGrid>,
+    source_index: usize,
+    ordering: SeriesOrdering,
+) -> NativeNode {
     let output = evaluated(grid);
     NativeNode {
         expr: NativeExpr::Input(BoundInput {
             output: output.clone(),
-            source_index: 0,
+            source_index,
             ordering,
         }),
         output,
@@ -487,6 +506,7 @@ fn raw_shared_histogram_chunk(
 /// rather than off a node: a source reserves for what the caller queued at it and for nothing else.
 fn source_runtime(pool: &Arc<dyn MemoryPool>) -> InputSlot {
     InputSlot {
+        source_index: 0,
         schema: raw_schema(),
         open_source_series: false,
         state: NodeState::new("native.input", Arc::new(NodeMetrics::default()), pool),
@@ -1549,6 +1569,85 @@ fn an_error_releases_the_reservation() {
     assert_eq!(pool.reserved(), 0);
 }
 
+/// Runtime routing keeps a raw input's source index and feeds its actual leaf id.
+#[test]
+fn runtime_input_ids_route_raw_sources() {
+    let out_grid = grid(0, 60_000, 2);
+    let tree = NativeTree::try_new(window_with(
+        raw_input_with_source_index_and_ordering(7, SeriesOrdering::KeyThenTime),
+        RangeFunction::LastOverTime,
+        1,
+        None,
+        Arc::clone(&out_grid),
+    ))
+    .unwrap();
+    let mut execution = executor(&tree, pool());
+    let id = execution.state().input_ids(7).next().unwrap();
+    execution
+        .push_raw(
+            id,
+            raw_chunk(&raw_schema(), &[0, 60_000], &[4.0, 5.0], true),
+        )
+        .unwrap();
+    assert_eq!(
+        addressed(&run(&mut execution)),
+        vec![(0, vec![Some(4.0), Some(5.0)])]
+    );
+}
+
+/// A fanout enqueue failure stops the whole execution and releases earlier recipients' reservations.
+#[test]
+fn a_fanout_enqueue_failure_releases_every_recipient() {
+    let out_grid = grid(0, 60_000, 1);
+    let tree = NativeTree::try_new(binary_of(
+        window_with(
+            raw_input_with_source_index_and_ordering(0, SeriesOrdering::KeyThenTime),
+            RangeFunction::LastOverTime,
+            1,
+            None,
+            Arc::clone(&out_grid),
+        ),
+        window_with(
+            raw_input_with_source_index_and_ordering(0, SeriesOrdering::KeyThenTime),
+            RangeFunction::LastOverTime,
+            1,
+            None,
+            out_grid,
+        ),
+        binary_matching(BinaryOp::Add),
+    ))
+    .unwrap();
+    let schema = raw_schema();
+    let chunk = raw_chunk(&schema, &[0], &[1.0], true);
+    let first_charge = raw_bytes(&chunk);
+    let pool = bounded_pool(first_charge);
+    let mut execution = executor(&tree, Arc::clone(&pool));
+    let ids: Vec<_> = execution.state().input_ids(0).collect();
+    let first = ids[0];
+    let second = ids[1];
+    let error = execution
+        .push_raw(first, chunk.clone())
+        .and_then(|()| execution.push_raw(second, chunk))
+        .expect_err("second fanout recipient exceeded the shared pool");
+    assert!(matches!(error, DataFusionError::ResourcesExhausted(_)));
+    assert_eq!(execution.state().reserved_bytes(), 0);
+    assert_eq!(pool.reserved(), 0);
+    assert!(
+        execution
+            .next_chunk()
+            .unwrap_err()
+            .to_string()
+            .contains("stopped after an earlier error"),
+        "a failed fanout execution continued"
+    );
+    assert!(
+        execution
+            .push_raw(first, raw_chunk(&schema, &[1], &[2.0], true))
+            .is_err(),
+        "a failed fanout execution accepted another input"
+    );
+}
+
 /// A chunk the pool cannot admit fails the execution and gives back everything it already held.
 ///
 /// The window holds the first chunk's samples, so the second one is refused against a reservation
@@ -4844,6 +4943,50 @@ fn binary_tree(op: BinaryOp, points: u64) -> (NativeTree, Arc<EvalGrid>, InputId
     )
 }
 
+/// Runtime source indexes map to all matching leaf slots, including below unary and scalar nodes.
+#[test]
+fn runtime_input_ids_route_repeated_and_distinct_source_indexes() {
+    let out_grid = grid(0, 60_000, 2);
+    let tree = NativeTree::try_new(binary_of(
+        binary_of(
+            map(
+                evaluated_input_with_source_index(Arc::clone(&out_grid), 0),
+                vec![MathOp::Abs],
+            ),
+            evaluated_input_with_source_index(Arc::clone(&out_grid), 7),
+            binary_matching(BinaryOp::Add),
+        ),
+        broadcast_right(
+            evaluated_input_with_source_index(Arc::clone(&out_grid), 0),
+            3.0,
+            BinaryOp::Add,
+            false,
+        ),
+        binary_matching(BinaryOp::Add),
+    ))
+    .unwrap();
+    let mut execution = executor(&tree, pool());
+    let ids_0: Vec<_> = execution.state().input_ids(0).collect();
+    let ids_7: Vec<_> = execution.state().input_ids(7).collect();
+    assert_eq!(ids_0.len(), 2);
+    assert_eq!(ids_7.len(), 1);
+    assert_ne!(ids_0[0], ids_0[1]);
+    assert_ne!(ids_0[0], ids_7[0]);
+    assert_ne!(ids_0[1], ids_7[0]);
+    assert!(execution.state().input_ids(99).next().is_none());
+
+    let schema = evaluated_schema(Arc::clone(&out_grid));
+    for (id, value) in [(ids_0[0], 1.0), (ids_7[0], 2.0), (ids_0[1], 3.0)] {
+        execution
+            .push_evaluated(id, evaluated_chunk(&schema, 0, vec![value; 2], true))
+            .unwrap();
+    }
+    assert_eq!(
+        addressed(&run(&mut execution)),
+        vec![(0, vec![Some(6.0); 2])]
+    );
+}
+
 /// Two evaluated sources under a binary node add the series they match, point by point.
 #[test]
 fn a_binary_of_two_inputs_adds_the_matched_series() {
@@ -9382,6 +9525,8 @@ fn a_chunk_pushed_after_the_input_ended_is_refused() {
     let schema = evaluated_schema(Arc::clone(&grid));
     let pool = bounded_pool(1 << 22);
     let mut execution = executor(&tree, Arc::clone(&pool));
+    let left = execution.state().input_ids(0).next().unwrap();
+    let right = execution.state().input_ids(0).nth(1).unwrap();
     // Two matched series on each side, so the answer is two chunks and pulling one leaves one
     // queued: `{job="a"}` adds two to one, `{job="b"}` adds three to two.
     for (side, metric, value) in [
@@ -9430,7 +9575,7 @@ fn a_chunk_pushed_after_the_input_ended_is_refused() {
 
     let rejection = execution
         .push_evaluated(
-            InputId::new(0),
+            left,
             evaluated_chunk_of(
                 &schema,
                 named_metric("m", "late"),