verify on for both sides

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

src/query/src/promql/planner.rs

@@ -782,6 +782,8 @@ impl PromPlanner {
                     // under this case we only join on time index
                     left_context.tag_columns.is_empty() || right_context.tag_columns.is_empty(),
                     modifier,
+                    &left_context,
+                    &right_context,
                 )?;
                 let join_plan_schema = join_plan.schema().clone();
 
@@ -3479,13 +3481,6 @@ impl PromPlanner {
         (output_field_columns, field_pairs)
     }
 
-    fn plan_has_tsid_column(plan: &LogicalPlan) -> bool {
-        plan.schema()
-            .fields()
-            .iter()
-            .any(|field| field.name() == DATA_SCHEMA_TSID_COLUMN_NAME)
-    }
-
     fn optional_tsid_projection(
         schema: &DFSchemaRef,
         table_ref: Option<&TableReference>,
@@ -3501,15 +3496,15 @@ impl PromPlanner {
 
     fn binary_join_key_columns(
         &self,
-        left: &LogicalPlan,
-        right: &LogicalPlan,
+        left_context: &PromPlannerContext,
+        right_context: &PromPlannerContext,
         only_join_time_index: bool,
         modifier: &Option<BinModifier>,
     ) -> (BTreeSet<String>, BTreeSet<String>) {
         let use_tsid_join = !only_join_time_index
-            && self.binary_modifier_preserves_tsid_join_key(left, right, modifier)
-            && Self::plan_has_tsid_column(left)
-            && Self::plan_has_tsid_column(right);
+            && self.binary_modifier_preserves_tsid_join_key(left_context, right_context, modifier)
+            && left_context.use_tsid
+            && right_context.use_tsid;
 
         let (mut left_tag_columns, mut right_tag_columns) = if use_tsid_join {
             (
@@ -3554,8 +3549,8 @@ impl PromPlanner {
 
     fn binary_modifier_preserves_tsid_join_key(
         &self,
-        left: &LogicalPlan,
-        right: &LogicalPlan,
+        left_context: &PromPlannerContext,
+        right_context: &PromPlannerContext,
         modifier: &Option<BinModifier>,
     ) -> bool {
         let Some(modifier) = modifier else {
@@ -3569,23 +3564,23 @@ impl PromPlanner {
         match &modifier.matching {
             None => true,
             Some(LabelModifier::Exclude(ignoring)) => ignoring.labels.iter().all(|label| {
-                !left.schema().has_column_with_unqualified_name(label)
-                    && !right.schema().has_column_with_unqualified_name(label)
+                !left_context.tag_columns.contains(label)
+                    && !right_context.tag_columns.contains(label)
             }),
             Some(LabelModifier::Include(on)) => {
                 let on_labels = on.labels.iter().cloned().collect::<BTreeSet<_>>();
-                let full_label_set = self
-                    .ctx
+                let left_labels = left_context
+                    .tag_columns
+                    .iter()
+                    .cloned()
+                    .collect::<BTreeSet<_>>();
+                let right_labels = right_context
                     .tag_columns
                     .iter()
                     .cloned()
                     .collect::<BTreeSet<_>>();
 
-                on_labels == full_label_set
-                    && on_labels.iter().all(|label| {
-                        left.schema().has_column_with_unqualified_name(label)
-                            && right.schema().has_column_with_unqualified_name(label)
-                    })
+                on_labels == left_labels && on_labels == right_labels
             }
         }
     }
@@ -3603,9 +3598,15 @@ impl PromPlanner {
         right_time_index_column: Option<String>,
         only_join_time_index: bool,
         modifier: &Option<BinModifier>,
+        left_context: &PromPlannerContext,
+        right_context: &PromPlannerContext,
     ) -> Result<LogicalPlan> {
-        let (mut left_tag_columns, mut right_tag_columns) =
-            self.binary_join_key_columns(&left, &right, only_join_time_index, modifier);
+        let (mut left_tag_columns, mut right_tag_columns) = self.binary_join_key_columns(
+            left_context,
+            right_context,
+            only_join_time_index,
+            modifier,
+        );
 
         // push time index column if it exists
         if let (Some(left_time_index_column), Some(right_time_index_column)) =
@@ -4401,14 +4402,29 @@ mod test {
     async fn build_test_table_provider_with_tsid_fields(
         table_specs: &[((String, String), usize)],
         num_tag: usize,
+    ) -> DfTableSourceProvider {
+        let table_specs = table_specs
+            .iter()
+            .map(|(table_name_tuple, num_field)| (table_name_tuple.clone(), num_tag, *num_field))
+            .collect::<Vec<_>>();
+        build_test_table_provider_with_tsid_tag_fields(&table_specs).await
+    }
+
+    async fn build_test_table_provider_with_tsid_tag_fields(
+        table_specs: &[((String, String), usize, usize)],
     ) -> DfTableSourceProvider {
         let catalog_list = MemoryCatalogManager::with_default_setup();
 
         let physical_table_name = "phy";
         let physical_table_id = 999u32;
+        let physical_num_tag = table_specs
+            .iter()
+            .map(|(_, num_tag, _)| *num_tag)
+            .max()
+            .unwrap_or(0);
         let physical_num_field = table_specs
             .iter()
-            .map(|(_, num_field)| *num_field)
+            .map(|(_, _, num_field)| *num_field)
             .max()
             .unwrap_or(0);
 
@@ -4426,7 +4442,7 @@ mod test {
                     false,
                 ),
             ];
-            for i in 0..num_tag {
+            for i in 0..physical_num_tag {
                 columns.push(ColumnSchema::new(
                     format!("tag_{i}"),
                     ConcreteDataType::string_datatype(),
@@ -4450,11 +4466,13 @@ mod test {
             }
 
             let schema = Arc::new(Schema::new(columns));
-            let primary_key_indices = (0..(2 + num_tag)).collect::<Vec<_>>();
+            let primary_key_indices = (0..(2 + physical_num_tag)).collect::<Vec<_>>();
             let table_meta = TableMetaBuilder::empty()
                 .schema(schema)
                 .primary_key_indices(primary_key_indices)
-                .value_indices((2 + num_tag..2 + num_tag + 1 + physical_num_field).collect())
+                .value_indices(
+                    (2 + physical_num_tag..2 + physical_num_tag + 1 + physical_num_field).collect(),
+                )
                 .engine(METRIC_ENGINE_NAME.to_string())
                 .next_column_id(1024)
                 .build()
@@ -4481,9 +4499,10 @@ mod test {
         }
 
         // Register metric engine logical tables without `__tsid`, referencing the physical table.
-        for (idx, ((schema_name, table_name), num_field)) in table_specs.iter().enumerate() {
+        for (idx, ((schema_name, table_name), num_tag, num_field)) in table_specs.iter().enumerate()
+        {
             let mut columns = vec![];
-            for i in 0..num_tag {
+            for i in 0..*num_tag {
                 columns.push(ColumnSchema::new(
                     format!("tag_{i}"),
                     ConcreteDataType::string_datatype(),
@@ -4515,8 +4534,8 @@ mod test {
             let table_id = 1024u32 + idx as u32;
             let table_meta = TableMetaBuilder::empty()
                 .schema(schema)
-                .primary_key_indices((0..num_tag).collect())
-                .value_indices((num_tag + 1..num_tag + 1 + *num_field).collect())
+                .primary_key_indices((0..*num_tag).collect())
+                .value_indices((*num_tag + 1..*num_tag + 1 + *num_field).collect())
                 .engine(METRIC_ENGINE_NAME.to_string())
                 .options(options)
                 .next_column_id(1024)
@@ -5283,6 +5302,40 @@ mod test {
         assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
     }
 
+    #[tokio::test]
+    async fn on_label_set_must_cover_both_sides_to_use_tsid_binary_join() {
+        let eval_stmt = build_eval_stmt("some_metric / on(tag_0) some_alt_metric");
+
+        let table_provider = build_test_table_provider_with_tsid_tag_fields(&[
+            (
+                (DEFAULT_SCHEMA_NAME.to_string(), "some_metric".to_string()),
+                2,
+                1,
+            ),
+            (
+                (
+                    DEFAULT_SCHEMA_NAME.to_string(),
+                    "some_alt_metric".to_string(),
+                ),
+                1,
+                1,
+            ),
+        ])
+        .await;
+        let plan =
+            PromPlanner::stmt_to_plan(table_provider, &eval_stmt, &build_query_engine_state())
+                .await
+                .unwrap();
+
+        let plan_str = plan.display_indent_schema().to_string();
+        assert!(!plan_str.contains("__tsid ="), "{plan_str}");
+        assert!(
+            plan_str.contains("some_metric.tag_0 = some_alt_metric.tag_0"),
+            "{plan_str}"
+        );
+        assert!(!plan_str.contains("tag_1 ="), "{plan_str}");
+    }
+
     #[tokio::test]
     async fn comparison_binary_join_uses_tsid_and_keeps_it_in_filtered_result() {
         let eval_stmt = build_eval_stmt("some_metric > some_alt_metric");

tests/cases/standalone/common/promql/tsid_binary_join_regression.result

@@ -38,6 +38,20 @@ WITH(
 
 Affected Rows: 0
 
+CREATE TABLE tsid_binary_join_right_by_job (
+  job STRING NULL,
+  ts TIMESTAMP(3) NOT NULL,
+  greptime_value DOUBLE NULL,
+  TIME INDEX (ts),
+  PRIMARY KEY(job),
+)
+ENGINE = metric
+WITH(
+  on_physical_table = 'tsid_binary_join_physical'
+);
+
+Affected Rows: 0
+
 INSERT INTO tsid_binary_join_left (host, job, ts, greptime_value) VALUES
   ('host1', 'job1', 0, 12),
   ('host2', 'job2', 0, 18),
@@ -54,6 +68,14 @@ INSERT INTO tsid_binary_join_right (host, job, ts, greptime_value) VALUES
 
 Affected Rows: 4
 
+INSERT INTO tsid_binary_join_right_by_job (job, ts, greptime_value) VALUES
+  ('job1', 0, 3),
+  ('job2', 0, 6),
+  ('job1', 5000, 5),
+  ('job2', 5000, 7);
+
+Affected Rows: 4
+
 -- Default vector-vector arithmetic should join on `__tsid` and time index.
 -- SQLNESS REPLACE (metrics.*) REDACTED
 -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED
@@ -132,6 +154,45 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join
 |_|_| Total rows: 4_|
 +-+-+-+
 
+-- `on(job)` must stay label-based when only the left side has extra row-key labels.
+-- SQLNESS REPLACE (metrics.*) REDACTED
+-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED
+-- SQLNESS REPLACE (-+) -
+-- SQLNESS REPLACE (\s\s+) _
+-- SQLNESS REPLACE (peers.*) REDACTED
+-- SQLNESS REPLACE Hash\(\[job@1,\sts@2\],.* Hash([job@1, ts@2],REDACTED
+-- SQLNESS REPLACE Hash\(\[job@2,\sts@4\],.* Hash([job@2, ts@4],REDACTED
+-- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED
+-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED
+-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED
+TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_by_job;
+
++-+-+-+
+| stage | node | plan_|
++-+-+-+
+| 0_| 0_|_ProjectionExec: expr=[job@2 as job, ts@4 as ts, __tsid@3 as __tsid, greptime_value@0 / greptime_value@1 as tsid_binary_join_left.greptime_value / tsid_binary_join_right_by_job.greptime_value] REDACTED
+|_|_|_HashJoinExec: mode=Partitioned, join_type=Inner, on=[(job@1, job@1), (ts@2, ts@3)], projection=[greptime_value@0, greptime_value@3, job@4, __tsid@5, ts@6], NullsEqual: true REDACTED
+|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@2],REDACTED
+|_|_|_ProjectionExec: expr=[greptime_value@0 as greptime_value, job@2 as job, ts@4 as ts] REDACTED
+|_|_|_MergeScanExec: REDACTED
+|_|_|_RepartitionExec: partitioning=Hash([job@1, ts@3], 32), input_partitions=REDACTED REDACTED
+|_|_|_MergeScanExec: REDACTED
+|_|_|_|
+| 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED
+|_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED
+|_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, host@3 as host, job@4 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED
+|_|_|_CooperativeExec REDACTED
+|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED
+|_|_|_|
+| 1_| 0_|_PromInstantManipulateExec: range=[0..5000], lookback=[300000], interval=[5000], time index=[ts] REDACTED
+|_|_|_PromSeriesDivideExec: tags=["__tsid"] REDACTED
+|_|_|_ProjectionExec: expr=[greptime_value@1 as greptime_value, job@3 as job, __tsid@2 as __tsid, ts@0 as ts] REDACTED
+|_|_|_CooperativeExec REDACTED
+|_|_|_SeriesScan: region=REDACTED, "partition_count":REDACTED, "distribution":"PerSeries" REDACTED
+|_|_|_|
+|_|_| Total rows: 4_|
++-+-+-+
+
 -- Comparison filters can join on `__tsid`, but the filtered result must still behave like
 -- a regular derived vector downstream.
 -- SQLNESS REPLACE (metrics.*) REDACTED
@@ -223,6 +284,22 @@ TQL EVAL (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right;
 | host2 | job2 | 1970-01-01T00:00:05 | 3.0                                                                          |
 +-------+------+---------------------+------------------------------------------------------------------------------+
 
+-- SQLNESS SORT_RESULT 3 1
+TQL EVAL (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_by_job;
+
++------+---------------------+-------------------------------------------------------------------------------------+
+| job  | ts                  | tsid_binary_join_left.greptime_value / tsid_binary_join_right_by_job.greptime_value |
++------+---------------------+-------------------------------------------------------------------------------------+
+| job1 | 1970-01-01T00:00:00 | 4.0                                                                                 |
+| job1 | 1970-01-01T00:00:05 | 3.0                                                                                 |
+| job2 | 1970-01-01T00:00:00 | 3.0                                                                                 |
+| job2 | 1970-01-01T00:00:05 | 3.0                                                                                 |
++------+---------------------+-------------------------------------------------------------------------------------+
+
+DROP TABLE tsid_binary_join_right_by_job;
+
+Affected Rows: 0
+
 DROP TABLE tsid_binary_join_right;
 
 Affected Rows: 0

tests/cases/standalone/common/promql/tsid_binary_join_regression.sql

@@ -33,6 +33,18 @@ WITH(
   on_physical_table = 'tsid_binary_join_physical'
 );
 
+CREATE TABLE tsid_binary_join_right_by_job (
+  job STRING NULL,
+  ts TIMESTAMP(3) NOT NULL,
+  greptime_value DOUBLE NULL,
+  TIME INDEX (ts),
+  PRIMARY KEY(job),
+)
+ENGINE = metric
+WITH(
+  on_physical_table = 'tsid_binary_join_physical'
+);
+
 INSERT INTO tsid_binary_join_left (host, job, ts, greptime_value) VALUES
   ('host1', 'job1', 0, 12),
   ('host2', 'job2', 0, 18),
@@ -45,6 +57,12 @@ INSERT INTO tsid_binary_join_right (host, job, ts, greptime_value) VALUES
   ('host1', 'job1', 5000, 5),
   ('host2', 'job2', 5000, 7);
 
+INSERT INTO tsid_binary_join_right_by_job (job, ts, greptime_value) VALUES
+  ('job1', 0, 3),
+  ('job2', 0, 6),
+  ('job1', 5000, 5),
+  ('job2', 5000, 7);
+
 -- Default vector-vector arithmetic should join on `__tsid` and time index.
 -- SQLNESS REPLACE (metrics.*) REDACTED
 -- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED
@@ -71,6 +89,19 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right;
 -- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED
 TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / ignoring(host) tsid_binary_join_right;
 
+-- `on(job)` must stay label-based when only the left side has extra row-key labels.
+-- SQLNESS REPLACE (metrics.*) REDACTED
+-- SQLNESS REPLACE (RoundRobinBatch.*) REDACTED
+-- SQLNESS REPLACE (-+) -
+-- SQLNESS REPLACE (\s\s+) _
+-- SQLNESS REPLACE (peers.*) REDACTED
+-- SQLNESS REPLACE Hash\(\[job@1,\sts@2\],.* Hash([job@1, ts@2],REDACTED
+-- SQLNESS REPLACE Hash\(\[job@2,\sts@4\],.* Hash([job@2, ts@4],REDACTED
+-- SQLNESS REPLACE input_partitions=\d+ input_partitions=REDACTED
+-- SQLNESS REPLACE "partition_count":\{(.*?)\} "partition_count":REDACTED
+-- SQLNESS REPLACE region=\d+\(\d+,\s+\d+\) region=REDACTED
+TQL ANALYZE (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_by_job;
+
 -- Comparison filters can join on `__tsid`, but the filtered result must still behave like
 -- a regular derived vector downstream.
 -- SQLNESS REPLACE (metrics.*) REDACTED
@@ -101,6 +132,10 @@ TQL ANALYZE (0, 5, '5s') tsid_binary_join_left > bool tsid_binary_join_right;
 -- SQLNESS SORT_RESULT 3 1
 TQL EVAL (0, 5, '5s') tsid_binary_join_left / tsid_binary_join_right;
 
+-- SQLNESS SORT_RESULT 3 1
+TQL EVAL (0, 5, '5s') tsid_binary_join_left / on(job) tsid_binary_join_right_by_job;
+
+DROP TABLE tsid_binary_join_right_by_job;
 DROP TABLE tsid_binary_join_right;
 DROP TABLE tsid_binary_join_left;
 DROP TABLE tsid_binary_join_physical;