Fix column picks for LeftMark and RightMark

Author: tglanz

datafusion/physical-plan/src/joins/sort_merge_join/stream.rs

@@ -1559,29 +1559,32 @@ impl SortMergeJoinStream {
                 null_joined_batch.num_rows(),
             );
 
-            let columns = if !matches!(self.join_type, JoinType::Right) {
-                // For left joins, the first columns are the left columns.
-                // Critical: There is a bug here still because the match directions.
-                let mut left_columns = null_joined_batch
-                    .columns()
-                    .iter()
-                    .take(left_columns_length)
-                    .cloned()
-                    .collect::<Vec<_>>();
-
-                left_columns.extend(right_columns);
-                left_columns
-            } else {
-                // For right joins, the first columns are the right columns.
-                let left_columns = null_joined_batch
-                    .columns()
-                    .iter()
-                    .skip(right_columns_length)
-                    .cloned()
-                    .collect::<Vec<_>>();
-
-                right_columns.extend(left_columns);
-                right_columns
+            let columns = match self.join_type {
+                JoinType::Right => {
+                    // The first columns are the right columns.
+                    let left_columns = null_joined_batch
+                        .columns()
+                        .iter()
+                        .skip(right_columns_length)
+                        .cloned()
+                        .collect::<Vec<_>>();
+
+                    right_columns.extend(left_columns);
+                    right_columns
+                }
+                JoinType::Left | JoinType::LeftMark | JoinType::RightMark => {
+                    // The first columns are the left columns.
+                    let mut left_columns = null_joined_batch
+                        .columns()
+                        .iter()
+                        .take(left_columns_length)
+                        .cloned()
+                        .collect::<Vec<_>>();
+
+                    left_columns.extend(right_columns);
+                    left_columns
+                }
+                _ => exec_err!("Did not expect join type {}", self.join_type)?,
             };
 
             // Push the streamed/buffered batch joined nulls to the output

datafusion/physical-plan/src/joins/sort_merge_join/tests.rs

@@ -644,10 +644,8 @@ async fn join_right_one() -> Result<()> {
     Ok(())
 }
 
-
 #[tokio::test]
 async fn join_right_different_columns_count_with_filter() -> Result<()> {
-
     // select *
     // from t1
     // right join t2 on t1.b1 = t2.b1 and t1.a1 > t2.a2
@@ -707,7 +705,6 @@ async fn join_right_different_columns_count_with_filter() -> Result<()> {
 
 #[tokio::test]
 async fn join_left_different_columns_count_with_filter() -> Result<()> {
-
     // select *
     // from t2
     // left join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1
@@ -765,6 +762,128 @@ async fn join_left_different_columns_count_with_filter() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn join_left_mark_different_columns_count_with_filter() -> Result<()> {
+    // select *
+    // from t2
+    // left mark join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1
+
+    let left = build_table_two_cols(
+        ("a2", &vec![10, 20, 30]),
+        ("b1", &vec![4, 5, 6]), // 6 does not exist on the right
+    );
+
+    let right = build_table(
+        ("a1", &vec![1, 21, 3]), // 20(t2.a2) > 1(t1.a1)
+        ("b1", &vec![4, 5, 7]),
+        ("c1", &vec![7, 8, 9]),
+    );
+
+    let on = vec![(
+        Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("b1", &right.schema())?) as _,
+    )];
+
+    let filter = JoinFilter::new(
+        Arc::new(BinaryExpr::new(
+            Arc::new(Column::new("a2", 0)),
+            Operator::Gt,
+            Arc::new(Column::new("a1", 1)),
+        )),
+        vec![
+            ColumnIndex {
+                index: 0,
+                side: JoinSide::Left,
+            },
+            ColumnIndex {
+                index: 0,
+                side: JoinSide::Right,
+            },
+        ],
+        Arc::new(Schema::new(vec![
+            Field::new("a2", DataType::Int32, true),
+            Field::new("a1", DataType::Int32, true),
+        ])),
+    );
+
+    let (_, batches) =
+        join_collect_with_filter(left, right, on, filter, LeftMark).await?;
+
+    // The output order is important as SMJ preserves sortedness
+    // LeftMark returns all left rows with a boolean mark column
+    assert_snapshot!(batches_to_string(&batches), @r#"
+            +----+----+-------+
+            | a2 | b1 | mark  |
+            +----+----+-------+
+            | 10 | 4  | true  |
+            | 20 | 5  | false |
+            | 30 | 6  | false |
+            +----+----+-------+
+            "#);
+    Ok(())
+}
+
+#[tokio::test]
+async fn join_right_mark_different_columns_count_with_filter() -> Result<()> {
+    // select *
+    // from t1
+    // right mark join t2 on t1.b1 = t2.b1 and t1.a1 > t2.a2
+
+    let left = build_table(
+        ("a1", &vec![1, 21, 3]), // 21(t1.a1) > 20(t2.a2)
+        ("b1", &vec![4, 5, 7]),
+        ("c1", &vec![7, 8, 9]),
+    );
+
+    let right = build_table_two_cols(
+        ("a2", &vec![10, 20, 30]),
+        ("b1", &vec![4, 5, 6]), // 6 does not exist on the left
+    );
+
+    let on = vec![(
+        Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("b1", &right.schema())?) as _,
+    )];
+
+    let filter = JoinFilter::new(
+        Arc::new(BinaryExpr::new(
+            Arc::new(Column::new("a1", 0)),
+            Operator::Gt,
+            Arc::new(Column::new("a2", 1)),
+        )),
+        vec![
+            ColumnIndex {
+                index: 0,
+                side: JoinSide::Left,
+            },
+            ColumnIndex {
+                index: 0,
+                side: JoinSide::Right,
+            },
+        ],
+        Arc::new(Schema::new(vec![
+            Field::new("a1", DataType::Int32, true),
+            Field::new("a2", DataType::Int32, true),
+        ])),
+    );
+
+    let (_, batches) =
+        join_collect_with_filter(left, right, on, filter, RightMark).await?;
+
+    // The output order is important as SMJ preserves sortedness
+    // RightMark returns all right rows with a boolean mark column
+    assert_snapshot!(batches_to_string(&batches), @r#"
+            +----+----+-------+
+            | a2 | b1 | mark  |
+            +----+----+-------+
+            | 10 | 4  | false |
+            | 20 | 5  | true  |
+            | 30 | 6  | false |
+            +----+----+-------+
+            "#);
+    Ok(())
+}
+
 #[tokio::test]
 async fn join_full_one() -> Result<()> {
     let left = build_table(