fix(promql): align aggregate row adapters with current native APIs

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

src/promql/src/extension_plan/aggregate_finalize/exec_tests.rs

@@ -129,6 +129,7 @@ fn window_on(grid: &Arc<EvalGrid>) -> NativeNode {
                     range_ms: 300_000,
                     offset_ms: 0,
                     anchor_ms: None,
+                    modifier: None,
                 },
                 input_grid: None,
                 output_grid: Arc::clone(grid),

src/promql/src/extension_plan/aggregate_finalize/tests.rs

@@ -71,6 +71,7 @@ fn window() -> NativeNode {
                     range_ms: 300_000,
                     offset_ms: 0,
                     anchor_ms: None,
+                    modifier: None,
                 },
                 input_grid: None,
                 output_grid: grid(),
@@ -410,6 +411,7 @@ fn a_subquery_above_an_aggregate_is_refused() {
                     range_ms: 300_000,
                     offset_ms: 0,
                     anchor_ms: None,
+                    modifier: None,
                 },
                 input_grid: Some(grid()),
                 output_grid: grid(),

src/promql/src/extension_plan/fused_series.rs

@@ -64,7 +64,7 @@ use datafusion::physical_plan::{
 };
 use datatypes::series::SeriesOrdering;
 use datatypes::series::raw::RawSeriesChunk;
-use datatypes::series::samples::{DenseSamplesRef, SampleRef, Samples, SeriesSchema};
+use datatypes::series::samples::{DenseSamplesRef, FieldSpec, SampleRef, Samples, SeriesSchema};
 use datatypes::series::time::EvalGrid;
 use futures::Stream;
 use store_api::region_engine::{RawLabelSource, SendableRawSeriesStream};
@@ -590,7 +590,7 @@ fn validate_physical_tree(tree: &NativeTree) -> DataFusionResult<Vec<Arc<SeriesS
 
 /// Where one output column of [`FusedSeriesExec`] comes from.
 #[derive(Clone, Copy, Debug)]
-enum OutputColumn {
+pub(crate) enum OutputColumn {
     /// The evaluation grid timestamps of the current chunk.
     Time,
     /// Series field by index in the bound series schema.
@@ -600,9 +600,9 @@ enum OutputColumn {
 }
 
 /// Classifies every output column, rejecting anything the conversion cannot produce.
-fn classify_output(
+pub(crate) fn classify_output(
     output_schema: &SchemaRef,
-    series_schema: &SeriesSchema,
+    series_fields: &[FieldSpec],
     tags: &[String],
     time_column: &str,
 ) -> DataFusionResult<Vec<OutputColumn>> {
@@ -617,11 +617,7 @@ fn classify_output(
                 "tag column {tag:?} of FusedSeriesExec cannot be the time column"
             )));
         }
-        if series_schema
-            .fields
-            .iter()
-            .any(|field| field.name.as_ref() == tag)
-        {
+        if series_fields.iter().any(|field| field.name.as_ref() == tag) {
             return Err(DataFusionError::Plan(format!(
                 "tag column {tag:?} of FusedSeriesExec clashes with a series field"
             )));
@@ -629,7 +625,7 @@ fn classify_output(
     }
 
     let mut time_seen = false;
-    let mut field_seen = vec![false; series_schema.fields.len()];
+    let mut field_seen = vec![false; series_fields.len()];
     let mut tag_seen = vec![false; tags.len()];
     let mut classification = Vec::with_capacity(output_schema.fields().len());
     for field in output_schema.fields() {
@@ -657,8 +653,7 @@ fn classify_output(
             classification.push(OutputColumn::Time);
             continue;
         }
-        if let Some(index) = series_schema
-            .fields
+        if let Some(index) = series_fields
             .iter()
             .position(|field| field.name.as_ref() == name)
         {
@@ -719,7 +714,7 @@ fn classify_output(
         if !seen {
             return Err(DataFusionError::Plan(format!(
                 "missing field column {:?} in FusedSeriesExec output schema",
-                series_schema.fields[index].name
+                series_fields[index].name
             )));
         }
     }
@@ -833,7 +828,12 @@ impl FusedSeriesExec {
         let classification = if partial_fold.is_some() {
             Vec::new()
         } else {
-            classify_output(&output_schema, output_series, &output_tags, &time_column)?
+            classify_output(
+                &output_schema,
+                &output_series.fields,
+                &output_tags,
+                &time_column,
+            )?
         };
         let partition_count = inputs[0].0.properties().partitioning.partition_count();
         for (index, (input, label_source, ordering)) in inputs.iter().enumerate() {
@@ -1070,7 +1070,7 @@ impl ExecutionPlan for FusedSeriesExec {
 ///
 /// The time column is the absolute grid positions the chunk carries, missing samples are nulls and
 /// ordinary NaN stays a value.
-fn series_to_record_batch(
+pub(crate) fn series_to_record_batch(
     classification: &[OutputColumn],
     tags: &[String],
     output_schema: &SchemaRef,

src/promql/src/extension_plan/fused_series/tests.rs

@@ -3039,8 +3039,8 @@ async fn repeated_metrics_reads_register_nothing_and_outlive_the_execution() {
 
 /// One chunk of the series keyed `key`, carrying `labels` and closing when `last`.
 ///
-/// The key and the labels are separate so a test can hold two raw series under one label set: the
-/// collision check reads the labels, and two different keys are what makes such a pair two series.
+/// The key and labels are separate so a test can hold two raw series under one label set; different
+/// keys make such a pair two series.
 fn labelled_chunk(
     key: &str,
     labels: &[(&str, &str)],
@@ -3069,7 +3069,11 @@ fn partial_op(kind: AggregateKind, grouped_by: &[&str]) -> BoundAggregate {
 
 /// The binding of the prefix tree, declaring `tags` as the node's labels.
 fn partial_binding_of(tree: &NativeTree, tags: &[&str]) -> FusedSeriesBinding {
-    binding_of(tree, tags.iter().map(|tag| (*tag).to_string()).collect())
+    binding_of(
+        tree,
+        raw_schema(),
+        tags.iter().map(|tag| (*tag).to_string()).collect(),
+    )
 }
 
 /// The logical node that runs the prefix and exports `kind`'s state, grouped by `host`.