fix(promql): stop fabricating histogram fold hash partitioning (#9509)

* fix(promql): stop fabricating histogram fold hash partitioning Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): cover histogram fraction binary join missing rows Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): normalize dictionary labels in histogram regression Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: cover histogram joins on logical metric tables Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: specify insert columns for logical metric fixture Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: use supported histogram child replacement API Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test(promql): keep histogram fold fixtures truthful Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
2026-10-16 00:42:29 +00:00 · 2026-10-13 07:58:31 +00:00
parent d86741d0ad
commit 2039f508fc
2 changed files with 192 additions and 13 deletions
@@ -385,10 +385,7 @@ impl HistogramFold {
        let output_schema: SchemaRef = self.output_schema.inner().clone();
        let properties = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(output_schema.clone()),
-            Partitioning::Hash(
-                partition_exprs.clone(),
-                exec_input.output_partitioning().partition_count(),
-            ),
+            Partitioning::UnknownPartitioning(exec_input.output_partitioning().partition_count()),
            EmissionType::Incremental,
            Boundedness::Bounded,
        ));
@@ -604,10 +601,7 @@ impl ExecutionPlan for HistogramFoldExec {
        let new_input = children[0].clone();
        let properties = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(self.output_schema.clone()),
-            Partitioning::Hash(
-                self.partition_exprs.clone(),
-                new_input.output_partitioning().partition_count(),
-            ),
+            Partitioning::UnknownPartitioning(new_input.output_partitioning().partition_count()),
            EmissionType::Incremental,
            Boundedness::Bounded,
        ));
@@ -1442,6 +1436,7 @@ mod test {
    use datafusion::datasource::memory::MemorySourceConfig;
    use datafusion::datasource::source::DataSourceExec;
    use datafusion::logical_expr::EmptyRelation;
+    use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
    use datafusion::prelude::SessionContext;
    use datatypes::arrow_array::StringArray;
    use futures::FutureExt;
@@ -1461,7 +1456,9 @@ mod test {
        RecordBatch::try_new(schema, columns).unwrap()
    }

-    fn prepare_test_data() -> DataSourceExec {
+    /// The batches and schema shared by [`prepare_test_data`] and
+    /// [`prepare_test_data_with_two_partitions`].
+    fn prepare_test_data_batches() -> (Vec<RecordBatch>, SchemaRef) {
        let schema = Arc::new(Schema::new(vec![
            Field::new("host", DataType::Utf8, true),
            Field::new("le", DataType::Utf8, true),
@@ -1514,8 +1511,24 @@ mod test {
        )
        .unwrap();

+        (vec![data_1, data_2, data_3], schema)
+    }
+
+    fn prepare_test_data() -> DataSourceExec {
+        let (batches, schema) = prepare_test_data_batches();
        DataSourceExec::new(Arc::new(
-            MemorySourceConfig::try_new(&[vec![data_1, data_2, data_3]], schema, None).unwrap(),
+            MemorySourceConfig::try_new(&[batches], schema, None).unwrap(),
+        ))
+    }
+
+    /// The batches from [`prepare_test_data_batches`], laid out over two
+    /// partitions, so that partitioning metadata is checked against a real
+    /// multi-partition input.
+    fn prepare_test_data_with_two_partitions() -> DataSourceExec {
+        let (mut batches, schema) = prepare_test_data_batches();
+        let first_partition = vec![batches.remove(0)];
+        DataSourceExec::new(Arc::new(
+            MemorySourceConfig::try_new(&[first_partition, batches], schema, None).unwrap(),
        ))
    }

@@ -1603,9 +1616,8 @@ mod test {

        let properties = PlanProperties::new(
            EquivalenceProperties::new(output_schema.clone()),
-            Partitioning::Hash(
-                partition_exprs.clone(),
-                input.output_partitioning().partition_count(),
+            Partitioning::UnknownPartitioning(
+                input.properties().output_partitioning().partition_count(),
            ),
            EmissionType::Incremental,
            Boundedness::Bounded,
@@ -1808,6 +1820,39 @@ mod test {
        assert!(broken_result.is_err());
    }

+    #[test]
+    fn fold_does_not_claim_hash_partitioning() {
+        let input: Arc<dyn ExecutionPlan> = Arc::new(prepare_test_data_with_two_partitions());
+        let logical_input = LogicalPlan::EmptyRelation(EmptyRelation {
+            produce_one_row: false,
+            schema: Arc::new(input.schema().to_dfschema().unwrap()),
+        });
+        let logical_fold = HistogramFold::new(
+            "le".to_string(),
+            "val".to_string(),
+            "host".to_string(),
+            0.5,
+            logical_input,
+        )
+        .unwrap();
+        let constructed = logical_fold.to_execution_plan(input.clone());
+        assert!(matches!(
+            constructed.output_partitioning(),
+            Partitioning::UnknownPartitioning(2)
+        ));
+
+        let rebuilt = constructed
+            .replace_children(
+                vec![input],
+                ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+            )
+            .unwrap();
+        assert!(matches!(
+            rebuilt.output_partitioning(),
+            Partitioning::UnknownPartitioning(2)
+        ));
+    }
+
    #[test]
    fn confirm_schema() {
        let input_schema = Schema::new(vec![
@@ -26,6 +26,7 @@ use datatypes::arrow::array::{Float64Array, Int64Array, StringArray, TimestampMi
 use datatypes::arrow::compute::cast;
 use datatypes::arrow::datatypes::DataType;
 use frontend::instance::Instance;
+use query::datafusion::QUERY_PARALLELISM_HINT;
 use query::parser::{PromQuery, QueryLanguageParser, QueryStatement};
 use rstest::rstest;
 use rstest_reuse::apply;
@@ -193,6 +194,139 @@ async fn promql_query_as_batches(
    }
 }

+#[tokio::test(flavor = "multi_thread")]
+async fn histogram_fraction_binary_join_preserves_partitioned_rows() {
+    common_telemetry::init_default_ut_logging();
+
+    let distributed = crate::tests::create_distributed_instance(
+        "histogram_fraction_binary_join_preserves_partitioned_rows",
+    )
+    .await;
+    let instance = distributed.frontend();
+
+    for (table, create) in [
+        (
+            "pa_promql_m",
+            r#"CREATE TABLE pa_promql_m (
+                ts TIMESTAMP TIME INDEX,
+                host STRING,
+                shard STRING,
+                val DOUBLE,
+                PRIMARY KEY (host, shard)
+            ) PARTITION ON COLUMNS (shard) (
+                shard < 'b',
+                shard >= 'b' AND shard < 'c',
+                shard >= 'c' AND shard < 'd',
+                shard >= 'd'
+            );"#,
+        ),
+        (
+            "lg_hist_m_2",
+            r#"CREATE TABLE phy_hist_m_2 (
+                ts TIMESTAMP TIME INDEX,
+                host STRING,
+                shard STRING,
+                val DOUBLE,
+                PRIMARY KEY (host, shard)
+            ) PARTITION ON COLUMNS (shard) (
+                shard < 'b',
+                shard >= 'b' AND shard < 'c',
+                shard >= 'c' AND shard < 'd',
+                shard >= 'd'
+            ) ENGINE=metric WITH ('physical_metric_table' = 'true');
+            CREATE TABLE lg_hist_m_2 (
+                ts TIMESTAMP TIME INDEX,
+                host STRING,
+                shard STRING,
+                val DOUBLE,
+                PRIMARY KEY (host, shard)
+            ) PARTITION ON COLUMNS (shard) (
+                shard < 'b',
+                shard >= 'b' AND shard < 'c',
+                shard >= 'c' AND shard < 'd',
+                shard >= 'd'
+            ) ENGINE=metric WITH ('on_physical_table' = 'phy_hist_m_2');"#,
+        ),
+    ] {
+        execute_all(
+            &instance,
+            &format!(
+                "{create}\nINSERT INTO {table} (ts, host, shard, val) VALUES\n                    (0, 'h', 'a', 1),\n                    (0, 'h', 'b', 1),\n                    (0, 'h', 'c', 1),\n                    (0, 'h', 'd', 1);"
+            ),
+            QueryContext::arc(),
+        )
+        .await;
+
+        let expr = format!(
+            r#"histogram_fraction(-Inf,+Inf,abs(label_replace(({table} + on(host) group_left() (sum by(host)({table}))),"le","+Inf","missing_label",".*")))"#
+        );
+        for parallelism in [1, 4] {
+            let mut query_ctx = QueryContext::with_db_name(None);
+            query_ctx.set_extension(QUERY_PARALLELISM_HINT, parallelism.to_string());
+            let query_ctx = Arc::new(query_ctx);
+            for (query, expected_value) in [
+                (format!("TQL EVAL (0, 0, '1s') {expr}"), 1.0),
+                (
+                    format!("TQL EVAL (0, 0, '1s') {expr} + on(host, shard) {table}"),
+                    2.0,
+                ),
+            ] {
+                let output = instance
+                    .do_query(&query, query_ctx.clone())
+                    .await
+                    .remove(0)
+                    .unwrap();
+                let batches = match output.data {
+                    OutputData::Stream(stream) => collect_batches(stream).await.unwrap(),
+                    OutputData::RecordBatches(batches) => batches,
+                    other => panic!("expected result batches, got {other:?}"),
+                };
+                let mut rows = batches
+                    .iter()
+                    .flat_map(|batch| {
+                        let ts = batch
+                            .column_by_name("ts")
+                            .unwrap()
+                            .as_any()
+                            .downcast_ref::<TimestampMillisecondArray>()
+                            .unwrap();
+                        let host_array =
+                            cast(batch.column_by_name("host").unwrap(), &DataType::Utf8).unwrap();
+                        let host = host_array.as_any().downcast_ref::<StringArray>().unwrap();
+                        let shard_array =
+                            cast(batch.column_by_name("shard").unwrap(), &DataType::Utf8).unwrap();
+                        let shard = shard_array.as_any().downcast_ref::<StringArray>().unwrap();
+                        let value = batch
+                            .columns()
+                            .iter()
+                            .find_map(|column| column.as_any().downcast_ref::<Float64Array>())
+                            .unwrap();
+                        (0..batch.num_rows())
+                            .map(move |row| {
+                                (
+                                    ts.value(row),
+                                    host.value(row).to_owned(),
+                                    shard.value(row).to_owned(),
+                                    value.value(row),
+                                )
+                            })
+                            .collect::<Vec<_>>()
+                    })
+                    .collect::<Vec<_>>();
+                rows.sort_by(|left, right| left.2.cmp(&right.2));
+                assert_eq!(
+                    rows,
+                    ["a", "b", "c", "d"]
+                        .into_iter()
+                        .map(|shard| (0, "h".to_owned(), shard.to_owned(), expected_value))
+                        .collect::<Vec<_>>(),
+                    "query_parallelism={parallelism}, query={query}"
+                );
+            }
+        }
+    }
+}
+
 const ANON_PROMQL_RATIO_REPRO_DB: &str = "repro_db";

 const ANON_PROMQL_RATIO_REPRO_CREATE: &str = r#"