feat: Add indexed hash join optimization

datafusion/physical-expr/src/partitioning.rs

@@ -213,10 +213,13 @@ impl Partitioning {
                 PartitioningSatisfaction::Exact
             }
             // When partition count is 1, hash requirement is satisfied.
-            Distribution::HashPartitioned(_) if self.partition_count() == 1 => {
+            Distribution::HashPartitioned(_) | Distribution::IndexedHashPartitioned(_)
+                if self.partition_count() == 1 =>
+            {
                 PartitioningSatisfaction::Exact
             }
-            Distribution::HashPartitioned(required_exprs) => match self {
+            Distribution::HashPartitioned(required_exprs)
+            | Distribution::IndexedHashPartitioned(required_exprs) => match self {
                 // Here we do not check the partition count for hash partitioning and assumes the partition count
                 // and hash functions in the system are the same. In future if we plan to support storage partition-wise joins,
                 // then we need to have the partition count and hash functions validation.
@@ -322,6 +325,9 @@ pub enum Distribution {
     /// Requires children to be distributed in such a way that the same
     /// values of the keys end up in the same partition
     HashPartitioned(Vec<Arc<dyn PhysicalExpr>>),
+    /// A stream of `IndexedBatch` is distributed amongst partitions based on the hash of the expressions.
+    /// This is a special case of `HashPartitioned` that is used to avoid copying data.
+    IndexedHashPartitioned(Vec<Arc<dyn PhysicalExpr>>),
 }
 
 impl Distribution {
@@ -335,6 +341,10 @@ impl Distribution {
             Distribution::HashPartitioned(expr) => {
                 Partitioning::Hash(expr, partition_count)
             }
+            Distribution::IndexedHashPartitioned(expr) => {
+                // This is a special case that RepartitionExec knows how to handle
+                Partitioning::Hash(expr, partition_count)
+            }
         }
     }
 }
@@ -347,6 +357,13 @@ impl Display for Distribution {
             Distribution::HashPartitioned(exprs) => {
                 write!(f, "HashPartitioned[{}])", format_physical_expr_list(exprs))
             }
+            Distribution::IndexedHashPartitioned(exprs) => {
+                write!(
+                    f,
+                    "IndexedHashPartitioned[{}])",
+                    format_physical_expr_list(exprs)
+                )
+            }
         }
     }
 }

datafusion/physical-plan/src/repartition/indexed_batch.rs

@@ -0,0 +1,32 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A batch of data that is indexed by a set of indices.
+
+use std::sync::Arc;
+use arrow::array::{PrimitiveArray, UInt32Type};
+use arrow::record_batch::RecordBatch;
+
+/// A batch of data that is indexed by a set of indices.
+/// This is used to avoid copying data when repartitioning.
+#[derive(Debug, Clone)]
+pub struct IndexedBatch {
+    /// The batch of data.
+    pub batch: Arc<RecordBatch>,
+    /// The indices into the batch.
+    pub indices: PrimitiveArray<UInt32Type>,
+}