diff --git a/src/promql/src/extension_plan/histogram_fold.rs b/src/promql/src/extension_plan/histogram_fold.rs index 709117df294..fb691cbc00d 100644 --- a/src/promql/src/extension_plan/histogram_fold.rs +++ b/src/promql/src/extension_plan/histogram_fold.rs @@ -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, 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 = 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![ diff --git a/tests-integration/src/tests/promql_test.rs b/tests-integration/src/tests/promql_test.rs index 74114bbd6c4..8cf27d165e9 100644 --- a/tests-integration/src/tests/promql_test.rs +++ b/tests-integration/src/tests/promql_test.rs @@ -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::() + .unwrap(); + let host_array = + cast(batch.column_by_name("host").unwrap(), &DataType::Utf8).unwrap(); + let host = host_array.as_any().downcast_ref::().unwrap(); + let shard_array = + cast(batch.column_by_name("shard").unwrap(), &DataType::Utf8).unwrap(); + let shard = shard_array.as_any().downcast_ref::().unwrap(); + let value = batch + .columns() + .iter() + .find_map(|column| column.as_any().downcast_ref::()) + .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::>() + }) + .collect::>(); + 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::>(), + "query_parallelism={parallelism}, query={query}" + ); + } + } + } +} + const ANON_PROMQL_RATIO_REPRO_DB: &str = "repro_db"; const ANON_PROMQL_RATIO_REPRO_CREATE: &str = r#"