x

Author: forsaken628

src/query/service/src/pipelines/processors/transforms/hash_join/desc.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use databend_common_column::bitmap::Bitmap;
+use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::arrow::and_validities;
 use databend_common_expression::type_check::check_function;
@@ -64,6 +65,7 @@ pub struct HashJoinDesc {
     pub(crate) probe_projections: ColumnSet,
     pub(crate) probe_to_build: Vec<(usize, (bool, bool))>,
     pub(crate) build_schema: DataSchemaRef,
+    pub(crate) probe_schema: DataSchemaRef,
 }
 
 #[derive(Debug, Clone)]
@@ -139,6 +141,7 @@ impl HashJoinDesc {
             build_projection: join.build_projections.clone(),
             probe_projections: join.probe_projections.clone(),
             build_schema: join.build.output_schema()?,
+            probe_schema: join.probe.output_schema()?,
         })
     }
 
@@ -265,30 +268,50 @@ impl HashJoinDesc {
         &self,
         function_ctx: &FunctionContext,
         block_size: usize,
-    ) -> Result<FilterExecutor> {
-        let offset = self.probe_projections.len();
+    ) -> Result<Option<FilterExecutor>> {
+        let probe_len = self.probe_schema.num_fields();
 
-        let predicate = self
+        let projection = self
+            .probe_projections
+            .iter()
+            .copied()
+            .chain(self.build_projection.iter().map(|idx| probe_len + *idx))
+            .collect::<Vec<_>>();
+
+        let eq_predicates = self
             .probe_keys
             .iter()
             .zip(&self.build_keys)
             .map(|(probe, build)| {
-                let build = build.project_column_ref(|old| Ok(old + offset))?;
+                let build = build.project_column_ref(|old| Ok(old + probe_len))?;
                 check_function(None, "eq", &[], &[probe.clone(), build], &BUILTIN_FUNCTIONS)
+            });
+
+        let other_predicate = self.other_predicate.as_ref().map(|expr| {
+            expr.project_column_ref(|index| {
+                projection.get(*index).copied().ok_or_else(|| {
+                    ErrorCode::Internal(format!(
+                        "Invalid column index {} for nested loop predicate projection",
+                        index
+                    ))
+                })
             })
-            .chain(self.other_predicate.iter().map(|expr| Ok(expr.clone())))
+        });
+
+        let predicate = eq_predicates
+            .chain(other_predicate)
             .reduce(|lhs, rhs| {
                 check_function(None, "and_filters", &[], &[lhs?, rhs?], &BUILTIN_FUNCTIONS)
             })
             .unwrap()?;
 
-        Ok(FilterExecutor::new(
+        Ok(Some(FilterExecutor::new(
             predicate,
             function_ctx.clone(),
             block_size,
-            None,
+            Some(projection),
             &BUILTIN_FUNCTIONS,
             false,
-        ))
+        )))
     }
 }

src/query/service/src/pipelines/processors/transforms/new_hash_join/memory/inner_join.rs

@@ -126,8 +126,7 @@ impl Join for InnerHashJoin {
                 ))
             }
             HashJoinHashTable::NestedLoop(build_blocks) => {
-                let probe_block = data.project(&self.desc.probe_projections);
-                let nested = Box::new(LoopJoinStream::new(probe_block, build_blocks));
+                let nested = Box::new(LoopJoinStream::new(data, build_blocks));
                 return match self.nested_loop_filter.as_mut() {
                     Some(filter_executor) => {
                         Ok(InnerHashJoinFilterStream::create(nested, filter_executor))