Java bindings for cudf::hash_join

java/src/main/java/ai/rapids/cudf/HashJoin.java

@@ -0,0 +1,127 @@
+/*
+ * Copyright (c) 2021, NVIDIA CORPORATION.
+ *
+ * Licensed 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.
+ */
+
+package ai.rapids.cudf;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * This class represents a hash table built from the join keys of the right-side table for a
+ * join operation. This hash table can then be reused across a series of left probe tables
+ * to compute gather maps for joins more efficiently when the right-side table is not changing.
+ * It can also be used to query the output row count of a join and then pass that result to the
+ * operation that generates the join gather maps to avoid redundant computation when the output
+ * row count must be checked before manifesting the join gather maps.
+ */
+public class HashJoin implements AutoCloseable {
+  static {
+    NativeDepsLoader.loadNativeDeps();
+  }
+
+  private static final Logger log = LoggerFactory.getLogger(HashJoin.class);
+
+  private static class HashJoinCleaner extends MemoryCleaner.Cleaner {
+    private Table buildKeys;
+    private long nativeHandle;
+
+    HashJoinCleaner(Table buildKeys, long nativeHandle) {
+      this.buildKeys = buildKeys;
+      this.nativeHandle = nativeHandle;
+      addRef();
+    }
+
+    @Override
+    protected synchronized boolean cleanImpl(boolean logErrorIfNotClean) {
+      long origAddress = nativeHandle;
+      boolean neededCleanup = nativeHandle != 0;
+      if (neededCleanup) {
+        try {
+          destroy(nativeHandle);
+          buildKeys.close();
+          buildKeys = null;
+        } finally {
+          nativeHandle = 0;
+        }
+        if (logErrorIfNotClean) {
+          log.error("A HASH TABLE WAS LEAKED (ID: " + id + " " + Long.toHexString(origAddress));
+        }
+      }
+      return neededCleanup;
+    }
+
+    @Override
+    public boolean isClean() {
+      return nativeHandle == 0;
+    }
+  }
+
+  private final HashJoinCleaner cleaner;
+  private final boolean compareNulls;
+  private boolean isClosed = false;
+
+  /**
+   * Construct a hash table for a join from a table representing the join key columns from the
+   * right-side table in the join. The resulting instance must be closed to release the
+   * GPU resources associated with the instance.
+   * @param buildKeys table view containing the join keys for the right-side join table
+   * @param compareNulls true if null key values should match otherwise false
+   */
+  public HashJoin(Table buildKeys, boolean compareNulls) {
+    this.compareNulls = compareNulls;
+    Table buildTable = new Table(buildKeys.getColumns());
+    try {
+      long handle = create(buildTable.getNativeView(), compareNulls);
+      this.cleaner = new HashJoinCleaner(buildTable, handle);
+      MemoryCleaner.register(this, cleaner);
+    } catch (Throwable t) {
+      try {
+        buildTable.close();
+      } catch (Throwable t2) {
+        t.addSuppressed(t2);
+      }
+      throw t;
+    }
+  }
+
+  @Override
+  public synchronized void close() {
+    cleaner.delRef();
+    if (isClosed) {
+      cleaner.logRefCountDebug("double free " + this);
+      throw new IllegalStateException("Close called too many times " + this);
+    }
+    cleaner.clean(false);
+    isClosed = true;
+  }
+
+  long getNativeView() {
+    return cleaner.nativeHandle;
+  }
+
+  /** Get the number of join key columns for the table that was used to generate the has table. */
+  public long getNumberOfColumns() {
+    return cleaner.buildKeys.getNumberOfColumns();
+  }
+
+  /** Returns true if the hash table was built to match on nulls otherwise false. */
+  public boolean getCompareNulls() {
+    return compareNulls;
+  }
+
+  private static native long create(long tableView, boolean nullEqual);
+  private static native void destroy(long handle);
+}

java/src/main/java/ai/rapids/cudf/MemoryCleaner.java

@@ -277,6 +277,10 @@ public static void register(CompiledExpression expr, Cleaner cleaner) {
     all.add(new CleanerWeakReference(expr, cleaner, collected, false));
   }
 
+  static void register(HashJoin hashJoin, Cleaner cleaner) {
+    all.add(new CleanerWeakReference(hashJoin, cleaner, collected, true));
+  }
+
   /**
    * This is not 100% perfect and we can still run into situations where RMM buffers were not
    * collected and this returns false because of thread race conditions. This is just a best effort.

java/src/main/java/ai/rapids/cudf/Table.java

@@ -505,18 +505,48 @@ private static native long[] leftJoin(long leftTable, int[] leftJoinCols, long r
   private static native long[] leftJoinGatherMaps(long leftKeys, long rightKeys,
                                                   boolean compareNullsEqual) throws CudfException;
 
+  private static native long leftJoinRowCount(long leftTable, long rightHashJoin,
+                                              boolean nullsEqual) throws CudfException;
+
+  private static native long[] leftHashJoinGatherMaps(long leftTable, long rightHashJoin,
+                                                      boolean nullsEqual) throws CudfException;
+
+  private static native long[] leftHashJoinGatherMapsWithCount(long leftTable, long rightHashJoin,
+                                                               boolean nullsEqual,
+                                                               long outputRowCount) throws CudfException;
+
   private static native long[] innerJoin(long leftTable, int[] leftJoinCols, long rightTable,
                                          int[] rightJoinCols, boolean compareNullsEqual) throws CudfException;
 
   private static native long[] innerJoinGatherMaps(long leftKeys, long rightKeys,
                                                    boolean compareNullsEqual) throws CudfException;
 
+  private static native long innerJoinRowCount(long table, long hashJoin,
+                                               boolean nullsEqual) throws CudfException;
+
+  private static native long[] innerHashJoinGatherMaps(long table, long hashJoin,
+                                                       boolean nullsEqual) throws CudfException;
+
+  private static native long[] innerHashJoinGatherMapsWithCount(long table, long hashJoin,
+                                                                boolean nullsEqual,
+                                                                long outputRowCount) throws CudfException;
+
   private static native long[] fullJoin(long leftTable, int[] leftJoinCols, long rightTable,
                                          int[] rightJoinCols, boolean compareNullsEqual) throws CudfException;
 
   private static native long[] fullJoinGatherMaps(long leftKeys, long rightKeys,
                                                   boolean compareNullsEqual) throws CudfException;
 
+  private static native long fullJoinRowCount(long leftTable, long rightHashJoin,
+                                              boolean nullsEqual) throws CudfException;
+
+  private static native long[] fullHashJoinGatherMaps(long leftTable, long rightHashJoin,
+                                                      boolean nullsEqual) throws CudfException;
+
+  private static native long[] fullHashJoinGatherMapsWithCount(long leftTable, long rightHashJoin,
+                                                               boolean nullsEqual,
+                                                               long outputRowCount) throws CudfException;
+
   private static native long[] leftSemiJoin(long leftTable, int[] leftJoinCols, long rightTable,
       int[] rightJoinCols, boolean compareNullsEqual) throws CudfException;
 
@@ -2040,6 +2070,69 @@ public GatherMap[] leftJoinGatherMaps(Table rightKeys, boolean compareNullsEqual
     return buildJoinGatherMaps(gatherMapData);
   }
 
+  /**
+   * Computes the number of rows resulting from a left equi-join between two tables.
+   * It is assumed this table instance holds the key columns from the left table, and the
+   * {@link HashJoin} argument has been constructed from the key columns from the right table.
+   * @param rightHash hash table built from join key columns from the right table
+   * @return row count of the join result
+   */
+  public long leftJoinRowCount(HashJoin rightHash) {
+    if (getNumberOfColumns() != rightHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "rightKeys: " + rightHash.getNumberOfColumns());
+    }
+    return leftJoinRowCount(getNativeView(), rightHash.getNativeView(),
+        rightHash.getCompareNulls());
+  }
+
+  /**
+   * Computes the gather maps that can be used to manifest the result of a left equi-join between
+   * two tables. It is assumed this table instance holds the key columns from the left table, and
+   * the {@link HashJoin} argument has been constructed from the key columns from the right table.
+   * Two {@link GatherMap} instances will be returned that can be used to gather the left and right
+   * tables, respectively, to produce the result of the left join.
+   * It is the responsibility of the caller to close the resulting gather map instances.
+   * @param rightHash hash table built from join key columns from the right table
+   * @return left and right table gather maps
+   */
+  public GatherMap[] leftJoinGatherMaps(HashJoin rightHash) {
+    if (getNumberOfColumns() != rightHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "rightKeys: " + rightHash.getNumberOfColumns());
+    }
+    long[] gatherMapData =
+        leftHashJoinGatherMaps(getNativeView(), rightHash.getNativeView(),
+            rightHash.getCompareNulls());
+    return buildJoinGatherMaps(gatherMapData);
+  }
+
+  /**
+   * Computes the gather maps that can be used to manifest the result of a left equi-join between
+   * two tables. It is assumed this table instance holds the key columns from the left table, and
+   * the {@link HashJoin} argument has been constructed from the key columns from the right table.
+   * Two {@link GatherMap} instances will be returned that can be used to gather the left and right
+   * tables, respectively, to produce the result of the left join.
+   * It is the responsibility of the caller to close the resulting gather map instances.
+   * This interface allows passing an output row count that was previously computed from
+   * {@link #leftJoinRowCount(HashJoin)}.
+   * WARNING: Passing a row count that is smaller than the actual row count will result
+   * in undefined behavior.
+   * @param rightHash hash table built from join key columns from the right table
+   * @param outputRowCount number of output rows in the join result
+   * @return left and right table gather maps
+   */
+  public GatherMap[] leftJoinGatherMaps(HashJoin rightHash, long outputRowCount) {
+    if (getNumberOfColumns() != rightHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "rightKeys: " + rightHash.getNumberOfColumns());
+    }
+    long[] gatherMapData =
+        leftHashJoinGatherMapsWithCount(getNativeView(), rightHash.getNativeView(),
+            rightHash.getCompareNulls(), outputRowCount);
+    return buildJoinGatherMaps(gatherMapData);
+  }
+
   /**
    * Computes the number of rows from the result of a left join between two tables when a
    * conditional expression is true. It is assumed this table instance holds the columns from
@@ -2124,6 +2217,67 @@ public GatherMap[] innerJoinGatherMaps(Table rightKeys, boolean compareNullsEqua
     return buildJoinGatherMaps(gatherMapData);
   }
 
+  /**
+   * Computes the number of rows resulting from an inner equi-join between two tables.
+   * @param otherHash hash table built from join key columns from the other table
+   * @return row count of the join result
+   */
+  public long innerJoinRowCount(HashJoin otherHash) {
+    if (getNumberOfColumns() != otherHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "otherKeys: " + otherHash.getNumberOfColumns());
+    }
+    return innerJoinRowCount(getNativeView(), otherHash.getNativeView(),
+        otherHash.getCompareNulls());
+  }
+
+  /**
+   * Computes the gather maps that can be used to manifest the result of an inner equi-join between
+   * two tables. It is assumed this table instance holds the key columns from the left table, and
+   * the {@link HashJoin} argument has been constructed from the key columns from the right table.
+   * Two {@link GatherMap} instances will be returned that can be used to gather the left and right
+   * tables, respectively, to produce the result of the inner join.
+   * It is the responsibility of the caller to close the resulting gather map instances.
+   * @param rightHash hash table built from join key columns from the right table
+   * @return left and right table gather maps
+   */
+  public GatherMap[] innerJoinGatherMaps(HashJoin rightHash) {
+    if (getNumberOfColumns() != rightHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "rightKeys: " + rightHash.getNumberOfColumns());
+    }
+    long[] gatherMapData =
+        innerHashJoinGatherMaps(getNativeView(), rightHash.getNativeView(),
+            rightHash.getCompareNulls());
+    return buildJoinGatherMaps(gatherMapData);
+  }
+
+  /**
+   * Computes the gather maps that can be used to manifest the result of an inner equi-join between
+   * two tables. It is assumed this table instance holds the key columns from the left table, and
+   * the {@link HashJoin} argument has been constructed from the key columns from the right table.
+   * Two {@link GatherMap} instances will be returned that can be used to gather the left and right
+   * tables, respectively, to produce the result of the inner join.
+   * It is the responsibility of the caller to close the resulting gather map instances.
+   * This interface allows passing an output row count that was previously computed from
+   * {@link #innerJoinRowCount(HashJoin)}.
+   * WARNING: Passing a row count that is smaller than the actual row count will result
+   * in undefined behavior.
+   * @param rightHash hash table built from join key columns from the right table
+   * @param outputRowCount number of output rows in the join result
+   * @return left and right table gather maps
+   */
+  public GatherMap[] innerJoinGatherMaps(HashJoin rightHash, long outputRowCount) {
+    if (getNumberOfColumns() != rightHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "rightKeys: " + rightHash.getNumberOfColumns());
+    }
+    long[] gatherMapData =
+        innerHashJoinGatherMapsWithCount(getNativeView(), rightHash.getNativeView(),
+            rightHash.getCompareNulls(), outputRowCount);
+    return buildJoinGatherMaps(gatherMapData);
+  }
+
   /**
    * Computes the number of rows from the result of an inner join between two tables when a
    * conditional expression is true. It is assumed this table instance holds the columns from
@@ -2209,6 +2363,72 @@ public GatherMap[] fullJoinGatherMaps(Table rightKeys, boolean compareNullsEqual
     return buildJoinGatherMaps(gatherMapData);
   }
 
+  /**
+   * Computes the number of rows resulting from a full equi-join between two tables.
+   * It is assumed this table instance holds the key columns from the left table, and the
+   * {@link HashJoin} argument has been constructed from the key columns from the right table.
+   * Note that unlike {@link #leftJoinRowCount(HashJoin)} and {@link #innerJoinRowCount(HashJoin),
+   * this will perform some redundant calculations compared to
+   * {@link #fullJoinGatherMaps(HashJoin, long)}.
+   * @param rightHash hash table built from join key columns from the right table
+   * @return row count of the join result
+   */
+  public long fullJoinRowCount(HashJoin rightHash) {
+    if (getNumberOfColumns() != rightHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "rightKeys: " + rightHash.getNumberOfColumns());
+    }
+    return fullJoinRowCount(getNativeView(), rightHash.getNativeView(),
+        rightHash.getCompareNulls());
+  }
+
+  /**
+   * Computes the gather maps that can be used to manifest the result of a full equi-join between
+   * two tables. It is assumed this table instance holds the key columns from the left table, and
+   * the {@link HashJoin} argument has been constructed from the key columns from the right table.
+   * Two {@link GatherMap} instances will be returned that can be used to gather the left and right
+   * tables, respectively, to produce the result of the full join.
+   * It is the responsibility of the caller to close the resulting gather map instances.
+   * @param rightHash hash table built from join key columns from the right table
+   * @return left and right table gather maps
+   */
+  public GatherMap[] fullJoinGatherMaps(HashJoin rightHash) {
+    if (getNumberOfColumns() != rightHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "rightKeys: " + rightHash.getNumberOfColumns());
+    }
+    long[] gatherMapData =
+        fullHashJoinGatherMaps(getNativeView(), rightHash.getNativeView(),
+            rightHash.getCompareNulls());
+    return buildJoinGatherMaps(gatherMapData);
+  }
+
+  /**
+   * Computes the gather maps that can be used to manifest the result of a full equi-join between
+   * two tables. It is assumed this table instance holds the key columns from the left table, and
+   * the {@link HashJoin} argument has been constructed from the key columns from the right table.
+   * Two {@link GatherMap} instances will be returned that can be used to gather the left and right
+   * tables, respectively, to produce the result of the full join.
+   * It is the responsibility of the caller to close the resulting gather map instances.
+   * This interface allows passing an output row count that was previously computed from
+   * {@link #fullJoinRowCount(HashJoin)}.
+   * WARNING: Passing a row count that is smaller than the actual row count will result
+   * in undefined behavior.
+   * @param rightHash hash table built from join key columns from the right table
+   * @param outputRowCount number of output rows in the join result
+   * @return left and right table gather maps
+   */
+  public GatherMap[] fullJoinGatherMaps(HashJoin rightHash, long outputRowCount) {
+    if (getNumberOfColumns() != rightHash.getNumberOfColumns()) {
+      throw new IllegalArgumentException("column count mismatch, this: " + getNumberOfColumns() +
+          "rightKeys: " + rightHash.getNumberOfColumns());
+    }
+    long[] gatherMapData =
+        fullHashJoinGatherMapsWithCount(getNativeView(), rightHash.getNativeView(),
+            rightHash.getCompareNulls(), outputRowCount);
+    return buildJoinGatherMaps(gatherMapData);
+  }
+
   /**
    * Computes the gather maps that can be used to manifest the result of a full join between
    * two tables when a conditional expression is true. It is assumed this table instance holds