HIVE-29368: limiting the NDV fix to the CASE clause handler only

ql/src/java/org/apache/hadoop/hive/ql/udf/generic/GenericUDFWhen.java

@@ -18,8 +18,10 @@
 
 package org.apache.hadoop.hive.ql.udf.generic;
 
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 
 import org.apache.hadoop.hive.ql.exec.Description;
 import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
@@ -29,6 +31,7 @@
 import org.apache.hadoop.hive.ql.stats.estimator.StatEstimator;
 import org.apache.hadoop.hive.ql.stats.estimator.StatEstimatorProvider;
 import org.apache.hadoop.hive.serde.serdeConstants;
+import org.apache.hadoop.hive.serde2.objectinspector.ConstantObjectInspector;
 import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
 import org.apache.hadoop.hive.serde2.objectinspector.primitive.BooleanObjectInspector;
 
@@ -59,13 +62,16 @@
 public class GenericUDFWhen extends GenericUDF implements StatEstimatorProvider {
   private transient ObjectInspector[] argumentOIs;
   private transient GenericUDFUtils.ReturnObjectInspectorResolver returnOIResolver;
+  private transient int numberOfDistinctConstants;
 
   @Override
   public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentTypeException {
 
     argumentOIs = arguments;
     returnOIResolver = new GenericUDFUtils.ReturnObjectInspectorResolver(true);
 
+    Set<Object> distinctConstants = new HashSet<>();
+
     for (int i = 0; i + 1 < arguments.length; i += 2) {
       if (!arguments[i].getTypeName().equals(serdeConstants.BOOLEAN_TYPE_NAME)) {
         throw new UDFArgumentTypeException(i, "\""
@@ -79,6 +85,7 @@ public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumen
             + "\" is expected but \"" + arguments[i + 1].getTypeName()
             + "\" is found");
       }
+      collectNonNullConstant(arguments[i + 1], distinctConstants);
     }
     if (arguments.length % 2 == 1) {
       int i = arguments.length - 2;
@@ -89,11 +96,23 @@ public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumen
             + "\" is expected but \"" + arguments[i + 1].getTypeName()
             + "\" is found");
       }
+      collectNonNullConstant(arguments[i + 1], distinctConstants);
     }
 
+    numberOfDistinctConstants = distinctConstants.size();
+
     return returnOIResolver.get();
   }
 
+  private static void collectNonNullConstant(ObjectInspector oi, Set<Object> distinctConstants) {
+    if (oi instanceof ConstantObjectInspector) {
+      Object value = ((ConstantObjectInspector) oi).getWritableConstantValue();
+      if (value != null) {
+        distinctConstants.add(value);
+      }
+    }
+  }
+
   @Override
   public Object evaluate(DeferredObject[] arguments) throws HiveException {
     for (int i = 0; i + 1 < arguments.length; i += 2) {
@@ -137,10 +156,15 @@ public String getDisplayString(String[] children) {
 
   @Override
   public StatEstimator getStatEstimator() {
-    return new WhenStatEstimator();
+    return new WhenStatEstimator(numberOfDistinctConstants);
   }
 
   static class WhenStatEstimator implements StatEstimator {
+    private final int numberOfDistinctConstants;
+
+    WhenStatEstimator(int numberOfDistinctConstants) {
+      this.numberOfDistinctConstants = numberOfDistinctConstants;
+    }
 
     @Override
     public Optional<ColStatistics> estimate(List<ColStatistics> argStats) {
@@ -151,7 +175,14 @@ public Optional<ColStatistics> estimate(List<ColStatistics> argStats) {
       if (argStats.size() % 2 == 1) {
         combiner.add(argStats.get(argStats.size() - 1));
       }
-      return combiner.getResult();
+      Optional<ColStatistics> result = combiner.getResult();
+      if (result.isPresent()) {
+        ColStatistics stat = result.get();
+        if (numberOfDistinctConstants > stat.getCountDistint()) {
+          stat.setCountDistint(numberOfDistinctConstants);
+        }
+      }
+      return result;
     }
   }
 }

ql/src/test/org/apache/hadoop/hive/ql/udf/generic/TestGenericUDFWhenStatEstimator.java

@@ -0,0 +1,223 @@
+/*
+ * 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.udf.generic;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Optional;
+
+import org.apache.hadoop.hive.ql.plan.ColStatistics;
+import org.apache.hadoop.hive.ql.stats.estimator.StatEstimator;
+import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
+import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
+import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory;
+import org.apache.hadoop.io.Text;
+import org.junit.Test;
+
+public class TestGenericUDFWhenStatEstimator {
+
+  private static ObjectInspector booleanOI =
+      PrimitiveObjectInspectorFactory.writableBooleanObjectInspector;
+
+  private static ObjectInspector stringConstant(String value) {
+    return PrimitiveObjectInspectorFactory.getPrimitiveWritableConstantObjectInspector(
+        TypeInfoFactory.stringTypeInfo, value == null ? null : new Text(value));
+  }
+
+  private static ObjectInspector stringColumn() {
+    return PrimitiveObjectInspectorFactory.writableStringObjectInspector;
+  }
+
+  private static ColStatistics colStats(long ndv) {
+    ColStatistics cs = new ColStatistics("col", "string");
+    cs.setCountDistint(ndv);
+    cs.setNumNulls(0);
+    return cs;
+  }
+
+  @Test
+  public void testAllDistinctConstants() throws Exception {
+    // CASE WHEN cond THEN 'A' WHEN cond THEN 'B' ELSE 'C' END
+    GenericUDFWhen udf = new GenericUDFWhen();
+    ObjectInspector[] args = {
+        booleanOI, stringConstant("A"),
+        booleanOI, stringConstant("B"),
+        stringConstant("C")
+    };
+    udf.initialize(args);
+
+    StatEstimator estimator = udf.getStatEstimator();
+    List<ColStatistics> argStats = Arrays.asList(
+        colStats(2), colStats(1),
+        colStats(2), colStats(1),
+        colStats(1)
+    );
+
+    Optional<ColStatistics> result = estimator.estimate(argStats);
+    assertTrue(result.isPresent());
+    assertEquals(3, result.get().getCountDistint());
+  }
+
+  @Test
+  public void testDuplicateConstants() throws Exception {
+    // CASE WHEN cond THEN 'A' WHEN cond THEN 'A' ELSE 'B' END
+    GenericUDFWhen udf = new GenericUDFWhen();
+    ObjectInspector[] args = {
+        booleanOI, stringConstant("A"),
+        booleanOI, stringConstant("A"),
+        stringConstant("B")
+    };
+    udf.initialize(args);
+
+    StatEstimator estimator = udf.getStatEstimator();
+    List<ColStatistics> argStats = Arrays.asList(
+        colStats(2), colStats(1),
+        colStats(2), colStats(1),
+        colStats(1)
+    );
+
+    Optional<ColStatistics> result = estimator.estimate(argStats);
+    assertTrue(result.isPresent());
+    assertEquals(2, result.get().getCountDistint());
+  }
+
+  @Test
+  public void testConstantsWithNullElse() throws Exception {
+    // CASE WHEN cond THEN 'A' WHEN cond THEN 'B' ELSE NULL END
+    // NULL constant has NDV=0 from StatsUtils, but we have 2 non-null constants
+    GenericUDFWhen udf = new GenericUDFWhen();
+    ObjectInspector[] args = {
+        booleanOI, stringConstant("A"),
+        booleanOI, stringConstant("B"),
+        stringConstant(null)
+    };
+    udf.initialize(args);
+
+    StatEstimator estimator = udf.getStatEstimator();
+    List<ColStatistics> argStats = Arrays.asList(
+        colStats(2), colStats(1),
+        colStats(2), colStats(1),
+        colStats(0)  // NULL constant gets NDV=0
+    );
+
+    Optional<ColStatistics> result = estimator.estimate(argStats);
+    assertTrue(result.isPresent());
+    assertEquals(2, result.get().getCountDistint());
+  }
+
+  @Test
+  public void testConstantsFloorOverColumnNdv() throws Exception {
+    // CASE WHEN cond THEN 'A' WHEN cond THEN 'B' WHEN cond THEN 'C' ELSE column END
+    // 3 constants should floor NDV even if column has lower NDV
+    GenericUDFWhen udf = new GenericUDFWhen();
+    ObjectInspector[] args = {
+        booleanOI, stringConstant("A"),
+        booleanOI, stringConstant("B"),
+        booleanOI, stringConstant("C"),
+        stringColumn()
+    };
+    udf.initialize(args);
+
+    StatEstimator estimator = udf.getStatEstimator();
+    List<ColStatistics> argStats = Arrays.asList(
+        colStats(2), colStats(1),
+        colStats(2), colStats(1),
+        colStats(2), colStats(1),
+        colStats(2)  // column with NDV=2
+    );
+
+    Optional<ColStatistics> result = estimator.estimate(argStats);
+    assertTrue(result.isPresent());
+    assertEquals(3, result.get().getCountDistint());
+  }
+
+  @Test
+  public void testManyConstantsWithNullElse() throws Exception {
+    // CASE with 20 distinct constants + NULL ELSE
+    // Should return NDV=20, not NDV=1 (the bug we're fixing)
+    GenericUDFWhen udf = new GenericUDFWhen();
+    ObjectInspector[] args = new ObjectInspector[41]; // 20 pairs + 1 else
+    ColStatistics[] stats = new ColStatistics[41];
+
+    for (int i = 0; i < 20; i++) {
+      args[i * 2] = booleanOI;
+      args[i * 2 + 1] = stringConstant(String.valueOf((char) ('A' + i)));
+      stats[i * 2] = colStats(2);
+      stats[i * 2 + 1] = colStats(1);
+    }
+    args[40] = stringConstant(null);
+    stats[40] = colStats(0);  // NULL gets NDV=0
+
+    udf.initialize(args);
+
+    StatEstimator estimator = udf.getStatEstimator();
+    Optional<ColStatistics> result = estimator.estimate(Arrays.asList(stats));
+
+    assertTrue(result.isPresent());
+    assertEquals(20, result.get().getCountDistint());
+  }
+
+  @Test
+  public void testNoElseClause() throws Exception {
+    // CASE WHEN cond THEN 'A' WHEN cond THEN 'B' END (implicit NULL else)
+    GenericUDFWhen udf = new GenericUDFWhen();
+    ObjectInspector[] args = {
+        booleanOI, stringConstant("A"),
+        booleanOI, stringConstant("B")
+    };
+    udf.initialize(args);
+
+    StatEstimator estimator = udf.getStatEstimator();
+    List<ColStatistics> argStats = Arrays.asList(
+        colStats(2), colStats(1),
+        colStats(2), colStats(1)
+    );
+
+    Optional<ColStatistics> result = estimator.estimate(argStats);
+    assertTrue(result.isPresent());
+    assertEquals(2, result.get().getCountDistint());
+  }
+
+  @Test
+  public void testNullInThenBranch() throws Exception {
+    // CASE WHEN cond THEN NULL WHEN cond THEN 'A' ELSE 'B' END
+    GenericUDFWhen udf = new GenericUDFWhen();
+    ObjectInspector[] args = {
+        booleanOI, stringConstant(null),
+        booleanOI, stringConstant("A"),
+        stringConstant("B")
+    };
+    udf.initialize(args);
+
+    StatEstimator estimator = udf.getStatEstimator();
+    List<ColStatistics> argStats = Arrays.asList(
+        colStats(2), colStats(0),
+        colStats(2), colStats(1),
+        colStats(1)
+    );
+
+    Optional<ColStatistics> result = estimator.estimate(argStats);
+    assertTrue(result.isPresent());
+    assertEquals(2, result.get().getCountDistint());
+  }
+
+}

ql/src/test/queries/clientpositive/ndv_case_const.q

@@ -0,0 +1,27 @@
+CREATE TABLE t (cond INT, c2 STRING, c100 STRING);
+ALTER TABLE t UPDATE STATISTICS SET('numRows'='10000','rawDataSize'='1000000');
+ALTER TABLE t UPDATE STATISTICS FOR COLUMN cond SET('numDVs'='10','numNulls'='0');
+ALTER TABLE t UPDATE STATISTICS FOR COLUMN c2 SET('numDVs'='2','numNulls'='0','avgColLen'='5','maxColLen'='10');
+ALTER TABLE t UPDATE STATISTICS FOR COLUMN c100 SET('numDVs'='100','numNulls'='0','avgColLen'='5','maxColLen'='10');
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN 'A' WHEN cond=2 THEN 'B' ELSE 'C' END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN 'A' WHEN cond=2 THEN 'B' WHEN cond=3 THEN 'A' ELSE 'B' END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN 'A' WHEN cond=2 THEN 'B' ELSE NULL END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN NULL WHEN cond=2 THEN 'A' ELSE 'B' END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN NULL WHEN cond=2 THEN NULL WHEN cond=3 THEN 'A' ELSE 'B' END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN NULL WHEN cond=2 THEN 'A' WHEN cond=3 THEN NULL ELSE 'B' END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN NULL WHEN cond=2 THEN NULL WHEN cond=3 THEN c100 ELSE 'A' END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN c2 WHEN cond=2 THEN c100 ELSE 'A' END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN c2 WHEN cond=2 THEN c100 ELSE c2 END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN 'A' WHEN cond=2 THEN 'B' WHEN cond=3 THEN 'C' ELSE c2 END x FROM t) sub GROUP BY x;
+
+EXPLAIN SELECT x FROM (SELECT CASE WHEN cond=1 THEN 'A' WHEN cond=2 THEN 'B' ELSE c100 END x FROM t) sub GROUP BY x;