HIVE-29702: Reduce-side merge join can produce NULL values or fail with IllegalStateException under Tez container reuse

itests/src/test/resources/testconfiguration.properties

@@ -34,6 +34,7 @@ minitez.query.files=\
   orc_vectorization_ppd.q,\
   partition_default_name_change_numeric.q,\
   tez_complextype_with_null.q,\
+  tez_dummystore_reduce_side_reuse.q,\
   tez_tag.q,\
   tez_union_udtf.q,\
   tez_union_with_udf.q,\

ql/src/java/org/apache/hadoop/hive/ql/exec/TezDummyStoreOperator.java

@@ -18,8 +18,11 @@
 
 package org.apache.hadoop.hive.ql.exec;
 
+import java.util.List;
+
 import org.apache.hadoop.hive.ql.CompilationOpContext;
 import org.apache.hadoop.hive.ql.metadata.HiveException;
+import org.apache.hadoop.hive.ql.plan.OperatorDesc;
 
 /**
  * A dummy store operator same as the dummy store operator but for tez. This is required so that we
@@ -59,6 +62,16 @@ public void setFetchDone(boolean fetchDone) {
 
   @Override
   public void closeOp(boolean abort) throws HiveException {
+    List<Operator<? extends OperatorDesc>> children = getChildOperators();
+    if (children != null && !children.isEmpty()) {
+      for (Operator<? extends OperatorDesc> child : children) {
+        List<Operator<? extends OperatorDesc>> parents = child.getParentOperators();
+        if (parents != null) {
+          parents.remove(this);
+        }
+      }
+      children.clear();
+    }
     super.closeOp(abort);
     fetchDone = false;
   }

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/ReduceRecordProcessor.java

@@ -358,13 +358,16 @@ void close() {
       if (abort) {
         setAborted(true);
       }
-      reducer.close(abort);
+      // Close Tree1 (DummyStore chain) before Tree2 (MergeJoin chain).
+      // DummyStore.closeOp() removes itself from MergeJoin.parentOperators, so that when
+      // Tree2 closes, MergeJoin.allInitializedParentsAreClosed() returns true and
+      // checkAndGenObject() fires to emit the last join group.
       if (mergeWorkList != null) {
         for (BaseWork redWork : mergeWorkList) {
           ((ReduceWork) redWork).getReducer().close(abort);
         }
       }
-
+      reducer.close(abort);
       // Need to close the dummyOps as well. The operator pipeline
       // is not considered "closed/done" unless all operators are
       // done. For broadcast joins that includes the dummy parents.

ql/src/java/org/apache/hadoop/hive/ql/parse/GenTezUtils.java

@@ -31,6 +31,7 @@
 import org.apache.hadoop.hive.ql.ddl.DDLUtils;
 import org.apache.hadoop.hive.ql.exec.AbstractFileMergeOperator;
 import org.apache.hadoop.hive.ql.exec.AppMasterEventOperator;
+import org.apache.hadoop.hive.ql.exec.CommonMergeJoinOperator;
 import org.apache.hadoop.hive.ql.exec.FetchTask;
 import org.apache.hadoop.hive.ql.exec.FileSinkOperator;
 import org.apache.hadoop.hive.ql.exec.FilterOperator;
@@ -110,6 +111,14 @@
     assert context.parentOfRoot instanceof ReduceSinkOperator;
     ReduceSinkOperator reduceSink = (ReduceSinkOperator) context.parentOfRoot;
 
+    // If this RS feeds into a CommonMergeJoinOp (possibly via intermediate ops like GroupBy
+    // or DummyStore), normalize all RS inputs to the same numReducers so that every mapper
+    // writes to the same number of partitions and join keys land in the same task. Also
+    // disable Tez's dynamic auto-parallelism for these RS: ShuffleVertexManager can still
+    // shrink the actual task count at runtime independently of our numReducers value, which
+    // would silently break the same-key-same-partition guarantee this normalization relies on.
+    boolean isMergeJoinInput = normalizeMergeJoinReducers(reduceSink);
+
     reduceWork.setNumReduceTasks(reduceSink.getConf().getNumReducers());
     reduceWork.setSlowStart(reduceSink.getConf().isSlowStart());
     float minSrcFraction = context.conf.getFloat(
@@ -127,7 +136,8 @@
       reduceSink.getConf().getReducerTraits().remove(AUTOPARALLEL);
     }
 
-    if (isAutoReduceParallelism && reduceSink.getConf().getReducerTraits().contains(AUTOPARALLEL)) {
+    if (isAutoReduceParallelism && !isMergeJoinInput
+        && reduceSink.getConf().getReducerTraits().contains(AUTOPARALLEL)) {
 
       // configured limit for reducers
       final int maxReducers = context.conf.getIntVar(HiveConf.ConfVars.MAX_REDUCERS);
@@ -187,6 +197,87 @@
     return reduceWork;
   }
 
+  /**
+   * If the given RS leads (through child ops) to a CommonMergeJoinOperator, normalizes all RS
+   * ancestors of that MergeJoinOp's parents to use the same maximum numReducers. This is a
+   * last-resort fix: even if earlier optimizer passes failed to align numReducers, we correct
+   * it at edge-creation time so each mapper writes the same partition count.
+   */
+  private static boolean normalizeMergeJoinReducers(ReduceSinkOperator reduceSink) {
+    CommonMergeJoinOperator mergeJoin = findDownstreamMergeJoin(reduceSink, 8);
+    if (mergeJoin == null) {
+      return false;
+    }
+    List<ReduceSinkOperator> allRS = new ArrayList<>();
+    for (Operator<?> parent : mergeJoin.getParentOperators()) {
+      ReduceSinkOperator rsAnc = findUpstreamRS(parent, 8);
+      if (rsAnc != null) {
+        allRS.add(rsAnc);
+      }
+    }
+    if (allRS.isEmpty()) {
+      return true;
+    }
+    int maxNR = 0;
+    for (ReduceSinkOperator rs : allRS) {
+      if (rs.getConf().getNumReducers() > maxNR) {
+        maxNR = rs.getConf().getNumReducers();
+      }
+    }
+    if (maxNR <= 0) {
+      return true;
+    }
+    for (ReduceSinkOperator rs : allRS) {
+      int orig = rs.getConf().getNumReducers();
+      boolean hasAutoParallel = rs.getConf().getReducerTraits().contains(AUTOPARALLEL);
+      boolean hasUniform = rs.getConf().getReducerTraits().contains(UNIFORM);
+      if (orig != maxNR || hasAutoParallel || hasUniform) {
+        rs.getConf().setNumReducers(maxNR);
+        // ReduceSinkDesc.setReducerTraits() is ADDITIVE (this.reduceTraits.addAll(traits)) in
+        // every branch except when the passed set contains FIXED, in which case it explicitly
+        // removes AUTOPARALLEL/UNIFORM first. Passing a traits set with those already removed
+        // is therefore a no-op against the existing (unremoved) traits - FIXED must be passed
+        // to actually clear them.
+        rs.getConf().setReducerTraits(EnumSet.of(ReduceSinkDesc.ReducerTraits.FIXED));
+      }
+    }
+    return true;
+  }
+
+  private static CommonMergeJoinOperator findDownstreamMergeJoin(
+      ReduceSinkOperator rs, int maxDepth) {
+    Operator<?> curr = rs;
+    for (int d = 0; d < maxDepth; d++) {
+      List<Operator<?>> children = curr.getChildOperators();
+      if (children == null || children.isEmpty()) {
+        break;
+      }
+      curr = children.get(0);
+      if (curr instanceof CommonMergeJoinOperator) {
+        return (CommonMergeJoinOperator) curr;
+      }
+    }
+    return null;
+  }
+
+  private static ReduceSinkOperator findUpstreamRS(Operator<?> op, int maxDepth) {
+    if (op instanceof ReduceSinkOperator) {
+      return (ReduceSinkOperator) op;
+    }
+    Operator<?> curr = op;
+    for (int d = 0; d < maxDepth; d++) {
+      List<Operator<?>> parents = curr.getParentOperators();
+      if (parents == null || parents.isEmpty()) {
+        break;
+      }
+      curr = parents.get(0);
+      if (curr instanceof ReduceSinkOperator) {
+        return (ReduceSinkOperator) curr;
+      }
+    }
+    return null;
+  }
+
   /**
    * Checks if there is a Bucket Map Join (BMJ) following a hierarchy:
    * ReduceSinkOperator (RS) -> ReduceSinkOperator (RS) -> MapJoinOperator (MJ).

ql/src/test/org/apache/hadoop/hive/ql/exec/TestTezDummyStoreOperator.java

@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+package org.apache.hadoop.hive.ql.exec;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.hadoop.hive.ql.metadata.HiveException;
+import org.apache.hadoop.hive.ql.plan.OperatorDesc;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+/**
+ * Unit tests for {@link TezDummyStoreOperator#closeOp(boolean)}.
+ *
+ * When a Tez container is reused for a subsequent task of the same reducer vertex, Hive's
+ * per-query ObjectCache (keyed by vertex name) can hand the new task the same already-
+ * initialized DummyStore/MergeJoin operator instances the previous task used. If closeOp()
+ * leaves stale wiring behind (this operator still listed as a parent of its former child,
+ * and/or that child still referenced from childOperators), the reused instances carry over
+ * incorrect state into the next task's execution, causing
+ * ReduceRecordProcessor.getJoinParentOp() to walk into the wrong operator and throw
+ * IllegalStateException("Was expecting dummy store operator but found: ...").
+ */
+public class TestTezDummyStoreOperator {
+
+  private TezDummyStoreOperator dummyStore;
+
+  @Before
+  public void setUp() {
+    dummyStore = new TezDummyStoreOperator();
+  }
+
+  @Test
+  public void testCloseOpRemovesSelfFromChildParents() throws HiveException {
+    CommonMergeJoinOperator mergeJoin = new CommonMergeJoinOperator();
+
+    wireParentChild(dummyStore, mergeJoin);
+
+    Assert.assertTrue("precondition: mergeJoin's parents must contain dummyStore",
+        mergeJoin.getParentOperators().contains(dummyStore));
+    Assert.assertTrue("precondition: dummyStore's children must contain mergeJoin",
+        dummyStore.getChildOperators().contains(mergeJoin));
+
+    dummyStore.closeOp(false);
+
+    Assert.assertFalse("dummyStore must be removed from its former child's parentOperators",
+        mergeJoin.getParentOperators().contains(dummyStore));
+    Assert.assertTrue("dummyStore's own childOperators must be cleared",
+        dummyStore.getChildOperators().isEmpty());
+  }
+
+  @Test
+  public void testCloseOpRemovesSelfFromAllChildrenWhenMultiple() throws HiveException {
+    CommonMergeJoinOperator child1 = new CommonMergeJoinOperator();
+    CommonMergeJoinOperator child2 = new CommonMergeJoinOperator();
+
+    List<Operator<? extends OperatorDesc>> children = new ArrayList<>();
+    children.add(child1);
+    children.add(child2);
+    dummyStore.setChildOperators(children);
+
+    List<Operator<? extends OperatorDesc>> parentsOfChild1 = new ArrayList<>();
+    parentsOfChild1.add(dummyStore);
+    child1.setParentOperators(parentsOfChild1);
+
+    List<Operator<? extends OperatorDesc>> parentsOfChild2 = new ArrayList<>();
+    parentsOfChild2.add(dummyStore);
+    child2.setParentOperators(parentsOfChild2);
+
+    dummyStore.closeOp(false);
+
+    Assert.assertFalse(child1.getParentOperators().contains(dummyStore));
+    Assert.assertFalse(child2.getParentOperators().contains(dummyStore));
+    Assert.assertTrue(dummyStore.getChildOperators().isEmpty());
+  }
+
+  @Test
+  public void testCloseOpWithNoChildrenDoesNotThrow() throws HiveException {
+    dummyStore.setChildOperators(new ArrayList<>());
+    dummyStore.closeOp(false);
+    Assert.assertTrue(dummyStore.getChildOperators().isEmpty());
+  }
+
+  @Test
+  public void testCloseOpWithNullChildrenDoesNotThrow() throws HiveException {
+    dummyStore.setChildOperators(null);
+    // setChildOperators(null) replaces null with an empty list (see Operator#setChildOperators).
+    dummyStore.closeOp(false);
+    Assert.assertTrue(dummyStore.getChildOperators().isEmpty());
+  }
+
+  @Test
+  public void testCloseOpDoesNotAffectUnrelatedParentReferences() throws HiveException {
+    CommonMergeJoinOperator mergeJoin = new CommonMergeJoinOperator();
+    TezDummyStoreOperator otherDummyStore = new TezDummyStoreOperator();
+
+    List<Operator<? extends OperatorDesc>> parents = new ArrayList<>();
+    parents.add(otherDummyStore);
+    parents.add(dummyStore);
+    mergeJoin.setParentOperators(parents);
+
+    List<Operator<? extends OperatorDesc>> children = new ArrayList<>();
+    children.add(mergeJoin);
+    dummyStore.setChildOperators(children);
+
+    dummyStore.closeOp(false);
+
+    Assert.assertFalse(mergeJoin.getParentOperators().contains(dummyStore));
+    Assert.assertTrue("unrelated parent reference must be left untouched",
+        mergeJoin.getParentOperators().contains(otherDummyStore));
+  }
+
+  private static void wireParentChild(
+      Operator<? extends OperatorDesc> parent, Operator<? extends OperatorDesc> child) {
+    List<Operator<? extends OperatorDesc>> children = new ArrayList<>();
+    children.add(child);
+    parent.setChildOperators(children);
+
+    List<Operator<? extends OperatorDesc>> parents = new ArrayList<>();
+    parents.add(parent);
+    child.setParentOperators(parents);
+  }
+}