[HSTACK] Bugfix apache#10978: Inconsistent behavior in HashJoin Projections

Author: adragomir

datafusion/physical-plan/src/joins/hash_join.rs

@@ -26,7 +26,8 @@ use std::{any::Any, vec};
 
 use super::utils::{
     asymmetric_join_output_partitioning, get_final_indices_from_shared_bitmap,
-    reorder_output_after_swap, swap_join_projection,
+    project_index_to_exprs, remap_join_projections_join_to_output,
+    swap_join_projection, swap_reverting_projection,
 };
 use super::{
     utils::{OnceAsync, OnceFut},
@@ -625,7 +626,29 @@ impl HashJoinExec {
         {
             Ok(Arc::new(new_join))
         } else {
-            reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema())
+            // TODO avoid adding ProjectionExec again and again, only adding Final Projection
+            // ADR: FIXME the projection inside the hash join functionality is not consistent
+            // see https://github.com/apache/datafusion/commit/afddb321e9a98ffc1947005c38b6b50a6ef2a401
+            // Failing to do the below code will create a projection exec with a projection that is
+            // possibly outside the schema.
+            let actual_projection = if new_join.projection.is_some() {
+                let tmp = remap_join_projections_join_to_output(
+                    new_join.left().clone(),
+                    new_join.right().clone(),
+                    new_join.join_type(),
+                    new_join.projection.clone(),
+                )?
+                .unwrap();
+                project_index_to_exprs(&tmp, &new_join.schema())
+            } else {
+                swap_reverting_projection(&left.schema(), &right.schema())
+            };
+            // let swap_proj = swap_reverting_projection(&left.schema(), &right.schema());
+
+            let proj = ProjectionExec::try_new(actual_projection, Arc::new(new_join))?;
+            Ok(Arc::new(proj))
+
+            // reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema())
         }
     }
 }
@@ -1655,7 +1678,7 @@ mod tests {
     use datafusion_execution::runtime_env::RuntimeEnvBuilder;
     use datafusion_expr::Operator;
     use datafusion_physical_expr::expressions::{BinaryExpr, Literal};
-    use datafusion_physical_expr::PhysicalExpr;
+    use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
     use hashbrown::HashTable;
     use rstest::*;
     use rstest_reuse::*;

datafusion/physical-plan/src/joins/utils.rs

@@ -41,14 +41,15 @@ use arrow::compute;
 use arrow::datatypes::{
     ArrowNativeType, Field, Schema, SchemaBuilder, UInt32Type, UInt64Type,
 };
+use arrow_schema::SchemaRef;
 use datafusion_common::cast::as_boolean_array;
 use datafusion_common::stats::Precision;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{
     plan_err, DataFusionError, JoinSide, JoinType, Result, SharedResult,
 };
 use datafusion_expr::interval_arithmetic::Interval;
-use datafusion_physical_expr::equivalence::add_offset_to_expr;
+use datafusion_physical_expr::equivalence::{add_offset_to_expr, ProjectionMapping};
 use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr::utils::{collect_columns, merge_vectors};
 use datafusion_physical_expr::{
@@ -62,6 +63,7 @@ use crate::projection::ProjectionExec;
 use futures::future::{BoxFuture, Shared};
 use futures::{ready, FutureExt};
 use parking_lot::Mutex;
+use crate::common::can_project;
 
 /// Maps a `u64` hash value based on the build side ["on" values] to a list of indices with this key's value.
 ///
@@ -649,6 +651,66 @@ pub fn build_join_schema(
     (fields.finish().with_metadata(metadata), column_indices)
 }
 
+/// This assumes that the projections are relative to the join schema.
+/// We need to redo them to point to the actual hash join output schema
+pub fn remap_join_projections_join_to_output(
+    left: Arc<dyn ExecutionPlan>,
+    right: Arc<dyn ExecutionPlan>,
+    join_type: &JoinType,
+    projection: Option<Vec<usize>>,
+) -> Result<Option<Vec<usize>>> {
+    match projection {
+        Some(ref projection) => {
+            let (join_schema, _) = build_join_schema(
+                left.schema().as_ref(),
+                right.schema().as_ref(),
+                join_type
+            );
+
+            let join_schema = Arc::new(join_schema);
+            can_project(&join_schema, Some(projection.clone()).as_ref())?;
+
+            let projection_exprs = project_index_to_exprs(
+                &projection.clone(),
+                &join_schema
+            );
+            let projection_mapping =
+                ProjectionMapping::try_new(&projection_exprs, &join_schema)?;
+
+            // projection mapping contains from and to, get the second one
+            let dest_physical_exprs = projection_mapping.map.iter().map(|(_, t)| t.clone()).collect::<Vec<_>>();
+            let dest_columns = dest_physical_exprs.iter().map(|pe| pe.as_any().downcast_ref::<Column>()).collect::<Vec<_>>();
+            let output = dest_physical_exprs.iter().enumerate().map(|(idx, _)| {
+                // :Vec<(Arc<dyn PhysicalExpr>, String)>
+                // (pe.clone(), dest_column.name().to_owned())
+                let dest_column = dest_columns.get(idx).unwrap().unwrap();
+                dest_column.index()
+            }).collect::<Vec<_>>();
+            Ok(Some(output))
+        },
+        None => Ok(None)
+    }
+}
+
+pub fn project_index_to_exprs(
+    projection_index: &[usize],
+    schema: &SchemaRef,
+) -> Vec<(Arc<dyn PhysicalExpr>, String)> {
+    projection_index
+        .iter()
+        .map(|index| {
+            let field = schema.field(*index);
+            (
+                Arc::new(Column::new(
+                    field.name(),
+                    *index,
+                )) as Arc<dyn PhysicalExpr>,
+                field.name().to_owned(),
+            )
+        })
+        .collect::<Vec<_>>()
+}
+
 /// A [`OnceAsync`] runs an `async` closure once, where multiple calls to
 /// [`OnceAsync::once`] return a [`OnceFut`] that resolves to the result of the
 /// same computation.
@@ -1755,6 +1817,7 @@ pub(crate) fn reorder_output_after_swap(
     left_schema: &Schema,
     right_schema: &Schema,
 ) -> Result<Arc<dyn ExecutionPlan>> {
+    //////////////////////
     let proj = ProjectionExec::try_new(
         swap_reverting_projection(left_schema, right_schema),
         plan,
@@ -1767,7 +1830,7 @@ pub(crate) fn reorder_output_after_swap(
 ///
 /// Returns the expressions that will allow to swap back the values from the
 /// original left as the first columns and those on the right next.
-fn swap_reverting_projection(
+pub fn swap_reverting_projection(
     left_schema: &Schema,
     right_schema: &Schema,
 ) -> Vec<(Arc<dyn PhysicalExpr>, String)> {