Fix Unit Tests

Author: cgivre

exec/java-exec/src/main/java/org/apache/drill/exec/planner/RuleInstance.java

@@ -29,7 +29,7 @@
 import org.apache.calcite.rel.rules.FilterSetOpTransposeRule;
 import org.apache.calcite.rel.rules.JoinPushExpressionsRule;
 import org.apache.calcite.rel.rules.JoinPushThroughJoinRule;
-import org.apache.calcite.rel.rules.JoinPushTransitivePredicatesRule;
+import org.apache.drill.exec.planner.logical.DrillJoinPushTransitivePredicatesRule;
 import org.apache.calcite.rel.rules.ProjectRemoveRule;
 import org.apache.calcite.rel.rules.ProjectSetOpTransposeRule;
 import org.apache.calcite.rel.rules.ProjectToWindowRule;
@@ -171,12 +171,9 @@ public boolean matches(RelOptRuleCall call) {
    * {@link org.apache.calcite.rel.core.Join} and creates
    * {@link org.apache.calcite.rel.core.Filter}s if those predicates can be pushed
    * to its inputs.
+   * Uses Drill-specific version that excludes SEMI joins to preserve partition pruning behavior.
    */
-  RelOptRule DRILL_JOIN_PUSH_TRANSITIVE_PREDICATES_RULE =
-    JoinPushTransitivePredicatesRule.Config.DEFAULT
-      .withRelBuilderFactory(DrillRelBuilder.proto(
-        DrillRelFactories.DRILL_LOGICAL_JOIN_FACTORY, DrillRelFactories.DRILL_LOGICAL_FILTER_FACTORY))
-      .toRule();
+  RelOptRule DRILL_JOIN_PUSH_TRANSITIVE_PREDICATES_RULE = DrillJoinPushTransitivePredicatesRule.INSTANCE;
 
   RelOptRule REMOVE_IS_NOT_DISTINCT_FROM_RULE =
     FilterRemoveIsNotDistinctFromRule.Config.DEFAULT

exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillJoinPushTransitivePredicatesRule.java

@@ -0,0 +1,54 @@
+/*
+ * 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.drill.exec.planner.logical;
+
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.rel.core.Join;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.rel.rules.JoinPushTransitivePredicatesRule;
+import org.apache.drill.exec.planner.DrillRelBuilder;
+
+/**
+ * Drill-specific version of JoinPushTransitivePredicatesRule that excludes SEMI joins.
+ * SEMI joins are used for IN clause conversion and pushing predicates through them
+ * can interfere with partition pruning logic.
+ */
+public class DrillJoinPushTransitivePredicatesRule extends JoinPushTransitivePredicatesRule {
+
+  protected DrillJoinPushTransitivePredicatesRule(Config config) {
+    super(config);
+  }
+
+  @Override
+  public boolean matches(RelOptRuleCall call) {
+    Join join = call.rel(0);
+    // Don't apply this rule to SEMI joins to preserve partition pruning behavior
+    if (join.getJoinType() == JoinRelType.SEMI) {
+      return false;
+    }
+    return super.matches(call);
+  }
+
+  public static final DrillJoinPushTransitivePredicatesRule INSTANCE =
+      new DrillJoinPushTransitivePredicatesRule(
+          JoinPushTransitivePredicatesRule.Config.DEFAULT
+              .withRelBuilderFactory(DrillRelBuilder.proto(
+                  DrillRelFactories.DRILL_LOGICAL_JOIN_FACTORY,
+                  DrillRelFactories.DRILL_LOGICAL_FILTER_FACTORY))
+              .as(JoinPushTransitivePredicatesRule.Config.class));
+}

exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillRelFactories.java

@@ -27,7 +27,9 @@
 import org.apache.calcite.rel.core.RelFactories;
 import org.apache.calcite.rel.hint.RelHint;
 import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexShuttle;
 import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.tools.RelBuilderFactory;
@@ -136,7 +138,25 @@ public RelNode createProject(RelNode input, List<RelHint> hints, List<? extends
   private static class DrillFilterFactoryImpl implements RelFactories.FilterFactory {
     @Override
     public RelNode createFilter(RelNode child, RexNode condition, Set<CorrelationId> variablesSet) {
-      return DrillFilterRel.create(child, condition);
+      // Normalize nullability of RexInputRef nodes to match the input's row type
+      // This is necessary for Calcite 1.37+ which has stricter type checking
+      RexNode normalizedCondition = condition.accept(new RexShuttle() {
+        @Override
+        public RexNode visitInputRef(RexInputRef inputRef) {
+          int index = inputRef.getIndex();
+          if (index >= child.getRowType().getFieldCount()) {
+            return inputRef;
+          }
+          RelDataType actualType = child.getRowType().getFieldList().get(index).getType();
+          // If nullability differs, create a new RexInputRef with correct nullability
+          if (inputRef.getType().isNullable() != actualType.isNullable() ||
+              !inputRef.getType().equals(actualType)) {
+            return new RexInputRef(index, actualType);
+          }
+          return inputRef;
+        }
+      });
+      return DrillFilterRel.create(child, normalizedCondition);
     }
   }