From de810286cd7f64ebe3e4b18555ccdb8672c7e87a Mon Sep 17 00:00:00 2001
From: Clint Wylie <cjwylie@gmail.com>
Date: Wed, 30 Jan 2019 07:13:07 -0800
Subject: [PATCH] fix bug with expression virtual column selectors backed by a
 single long column (#6957)

* fix issue with SingleLongInputCachingExpressionColumnValueSelector when sql compatible null handling enabled

* add test with doubles to show same behavior for floats/doubles that lack the optimization of longs

* simplify

* fix import
---
 ...tCachingExpressionColumnValueSelector.java |  4 +
 .../virtual/ExpressionVirtualColumnTest.java  | 75 +++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/processing/src/main/java/org/apache/druid/segment/virtual/SingleLongInputCachingExpressionColumnValueSelector.java b/processing/src/main/java/org/apache/druid/segment/virtual/SingleLongInputCachingExpressionColumnValueSelector.java
index fa002c02203..f05329bd367 100644
--- a/processing/src/main/java/org/apache/druid/segment/virtual/SingleLongInputCachingExpressionColumnValueSelector.java
+++ b/processing/src/main/java/org/apache/druid/segment/virtual/SingleLongInputCachingExpressionColumnValueSelector.java
@@ -98,6 +98,10 @@ public class SingleLongInputCachingExpressionColumnValueSelector implements Colu
   @Override
   public ExprEval getObject()
   {
+    // things can still call this even when underlying selector is null (e.g. ExpressionColumnValueSelector#isNull)
+    if (selector.isNull()) {
+      return ExprEval.ofLong(null);
+    }
     // No assert for null handling, as the delegate selector already has it.
     final long input = selector.getLong();
     final boolean cached = input == lastInput && lastOutput != null;
diff --git a/processing/src/test/java/org/apache/druid/segment/virtual/ExpressionVirtualColumnTest.java b/processing/src/test/java/org/apache/druid/segment/virtual/ExpressionVirtualColumnTest.java
index 382de0bd4c8..b28fd5b8af9 100644
--- a/processing/src/test/java/org/apache/druid/segment/virtual/ExpressionVirtualColumnTest.java
+++ b/processing/src/test/java/org/apache/druid/segment/virtual/ExpressionVirtualColumnTest.java
@@ -112,6 +112,18 @@ public class ExpressionVirtualColumnTest
       ValueType.LONG,
       TestExprMacroTable.INSTANCE
   );
+  private static final ExpressionVirtualColumn SCALE_LONG = new ExpressionVirtualColumn(
+      "expr",
+      "x * 2",
+      ValueType.LONG,
+      TestExprMacroTable.INSTANCE
+  );
+  private static final ExpressionVirtualColumn SCALE_FLOAT = new ExpressionVirtualColumn(
+      "expr",
+      "x * 2",
+      ValueType.FLOAT,
+      TestExprMacroTable.INSTANCE
+  );
 
   private static final ThreadLocal<Row> CURRENT_ROW = new ThreadLocal<>();
   private static final ColumnSelectorFactory COLUMN_SELECTOR_FACTORY = RowBasedColumnSelectorFactory.create(
@@ -491,4 +503,67 @@ public class ExpressionVirtualColumnTest
     Assert.assertEquals(ImmutableList.of("z"), Z_LIKE.requiredColumns());
     Assert.assertEquals(ImmutableList.of("z", "x"), Z_CONCAT_X.requiredColumns());
   }
+
+  @Test
+  public void testExprEvalSelectorWithLongsAndNulls()
+  {
+    final ColumnValueSelector<ExprEval> selector = ExpressionSelectors.makeExprEvalSelector(
+        RowBasedColumnSelectorFactory.create(
+            CURRENT_ROW,
+            ImmutableMap.of("x", ValueType.LONG)
+        ),
+        Parser.parse(SCALE_LONG.getExpression(), TestExprMacroTable.INSTANCE)
+    );
+
+    CURRENT_ROW.set(ROW0);
+    if (NullHandling.replaceWithDefault()) {
+      Assert.assertEquals(0, selector.getLong(), 0.0f);
+      Assert.assertFalse(selector.isNull());
+    } else {
+      Assert.assertTrue(selector.isNull());
+      Assert.assertTrue(selector.getObject().isNumericNull());
+    }
+  }
+
+  @Test
+  public void testExprEvalSelectorWithDoublesAndNulls()
+  {
+    final ColumnValueSelector<ExprEval> selector = ExpressionSelectors.makeExprEvalSelector(
+        RowBasedColumnSelectorFactory.create(
+            CURRENT_ROW,
+            ImmutableMap.of("x", ValueType.DOUBLE)
+        ),
+        Parser.parse(SCALE_FLOAT.getExpression(), TestExprMacroTable.INSTANCE)
+    );
+
+    CURRENT_ROW.set(ROW0);
+    if (NullHandling.replaceWithDefault()) {
+      Assert.assertEquals(0, selector.getDouble(), 0.0f);
+      Assert.assertFalse(selector.isNull());
+    } else {
+      Assert.assertTrue(selector.isNull());
+      Assert.assertTrue(selector.getObject().isNumericNull());
+    }
+  }
+
+  @Test
+  public void testExprEvalSelectorWithFloatAndNulls()
+  {
+    final ColumnValueSelector<ExprEval> selector = ExpressionSelectors.makeExprEvalSelector(
+        RowBasedColumnSelectorFactory.create(
+            CURRENT_ROW,
+            ImmutableMap.of("x", ValueType.FLOAT)
+        ),
+        Parser.parse(SCALE_FLOAT.getExpression(), TestExprMacroTable.INSTANCE)
+    );
+
+    CURRENT_ROW.set(ROW0);
+    if (NullHandling.replaceWithDefault()) {
+      Assert.assertEquals(0, selector.getFloat(), 0.0f);
+      Assert.assertFalse(selector.isNull());
+    } else {
+      Assert.assertTrue(selector.isNull());
+      Assert.assertTrue(selector.getObject().isNumericNull());
+    }
+  }
 }