Fix extractionFns on number-wrapping dimension selectors. (#15761)

When an ExtractionFn is used on top of a numeric column, it wasn't applied to
nulls (nulls are returned as-is). This patch fixes it.

processing/src/main/java/org/apache/druid/segment/DoubleWrappingDimensionSelector.java

@@ -43,12 +43,10 @@ public class DoubleWrappingDimensionSelector extends BaseSingleValueDimensionSel
   @Override
   protected String getValue()
   {
-    if (selector.isNull()) {
+    if (extractionFn == null) {
-      return null;
+      return selector.isNull() ? null : String.valueOf(selector.getDouble());
-    } else if (extractionFn == null) {
-      return String.valueOf(selector.getDouble());
     } else {
-      return extractionFn.apply(selector.getDouble());
+      return selector.isNull() ? extractionFn.apply(null) : extractionFn.apply(selector.getDouble());
     }
   }
 

processing/src/main/java/org/apache/druid/segment/FloatWrappingDimensionSelector.java

@@ -40,12 +40,10 @@ public class FloatWrappingDimensionSelector extends BaseSingleValueDimensionSele
   @Override
   protected String getValue()
   {
-    if (selector.isNull()) {
+    if (extractionFn == null) {
-      return null;
+      return selector.isNull() ? null : String.valueOf(selector.getFloat());
-    } else if (extractionFn == null) {
-      return String.valueOf(selector.getFloat());
     } else {
-      return extractionFn.apply(selector.getFloat());
+      return selector.isNull() ? extractionFn.apply(null) : extractionFn.apply(selector.getFloat());
     }
   }
 

processing/src/main/java/org/apache/druid/segment/LongWrappingDimensionSelector.java

@@ -40,12 +40,10 @@ public class LongWrappingDimensionSelector extends BaseSingleValueDimensionSelec
   @Override
   protected String getValue()
   {
-    if (selector.isNull()) {
+    if (extractionFn == null) {
-      return null;
+      return selector.isNull() ? null : String.valueOf(selector.getLong());
-    } else if (extractionFn == null) {
-      return String.valueOf(selector.getLong());
     } else {
-      return extractionFn.apply(selector.getLong());
+      return selector.isNull() ? extractionFn.apply(null) : extractionFn.apply(selector.getLong());
     }
   }
 

processing/src/test/java/org/apache/druid/segment/WrappingDimensionSelectorTest.java

@@ -20,10 +20,13 @@
 package org.apache.druid.segment;
 
 import org.apache.druid.common.config.NullHandling;
+import org.apache.druid.query.extraction.DimExtractionFn;
 import org.apache.druid.testing.InitializedNullHandlingTest;
 import org.junit.Assert;
 import org.junit.Test;
 
+import javax.annotation.Nullable;
+
 public class WrappingDimensionSelectorTest extends InitializedNullHandlingTest
 {
   @Test
@@ -39,8 +42,10 @@ public class WrappingDimensionSelectorTest extends InitializedNullHandlingTest
     lngSelector.increment();
     if (NullHandling.sqlCompatible()) {
       Assert.assertTrue(lngSelector.isNull());
+      Assert.assertNull(lngWrapSelector.getValue());
     } else {
       Assert.assertEquals(0L, lngSelector.getLong());
+      Assert.assertEquals("0", lngWrapSelector.getValue());
     }
 
     lngSelector.increment();
@@ -69,8 +74,10 @@ public class WrappingDimensionSelectorTest extends InitializedNullHandlingTest
     dblSelector.increment();
     if (NullHandling.sqlCompatible()) {
       Assert.assertTrue(dblSelector.isNull());
+      Assert.assertNull(dblWrapSelector.getValue());
     } else {
       Assert.assertEquals(0d, dblSelector.getDouble(), 0);
+      Assert.assertEquals("0.0", dblWrapSelector.getValue());
     }
 
     dblSelector.increment();
@@ -99,8 +106,10 @@ public class WrappingDimensionSelectorTest extends InitializedNullHandlingTest
     flSelector.increment();
     if (NullHandling.sqlCompatible()) {
       Assert.assertTrue(flSelector.isNull());
+      Assert.assertNull(flWrapSelector.getValue());
     } else {
       Assert.assertEquals(0f, flSelector.getFloat(), 0);
+      Assert.assertEquals("0.0", flWrapSelector.getValue());
     }
 
     flSelector.increment();
@@ -115,4 +124,134 @@ public class WrappingDimensionSelectorTest extends InitializedNullHandlingTest
     Assert.assertEquals(-45.0f, flSelector.getFloat(), 0);
     Assert.assertEquals("-45.0", flWrapSelector.getValue());
   }
+
+  @Test
+  public void testLongWrappingDimensionSelectorExtractionFn()
+  {
+    Long[] vals = new Long[]{24L, null, 50L, 0L, -60L};
+    TestNullableLongColumnSelector lngSelector = new TestNullableLongColumnSelector(vals);
+    final TestExtractionFn extractionFn = new TestExtractionFn();
+
+    LongWrappingDimensionSelector lngWrapSelector = new LongWrappingDimensionSelector(lngSelector, extractionFn);
+    Assert.assertEquals(24L, lngSelector.getLong());
+    Assert.assertEquals("24x", lngWrapSelector.getValue());
+
+    lngSelector.increment();
+    if (NullHandling.sqlCompatible()) {
+      Assert.assertTrue(lngSelector.isNull());
+      Assert.assertEquals("nullx", lngWrapSelector.getValue());
+    } else {
+      Assert.assertEquals(0L, lngSelector.getLong());
+      Assert.assertEquals("0x", lngWrapSelector.getValue());
+    }
+
+    lngSelector.increment();
+    Assert.assertEquals(50L, lngSelector.getLong());
+    Assert.assertEquals("50x", lngWrapSelector.getValue());
+
+    lngSelector.increment();
+    Assert.assertEquals(0L, lngSelector.getLong());
+    Assert.assertEquals("0x", lngWrapSelector.getValue());
+
+    lngSelector.increment();
+    Assert.assertEquals(-60L, lngSelector.getLong());
+    Assert.assertEquals("-60x", lngWrapSelector.getValue());
+  }
+
+  @Test
+  public void testDoubleWrappingDimensionSelectorExtractionFn()
+  {
+    Double[] vals = new Double[]{32.0d, null, 5.0d, 0.0d, -45.0d};
+    TestNullableDoubleColumnSelector dblSelector = new TestNullableDoubleColumnSelector(vals);
+    final TestExtractionFn extractionFn = new TestExtractionFn();
+
+    DoubleWrappingDimensionSelector dblWrapSelector = new DoubleWrappingDimensionSelector(dblSelector, extractionFn);
+    Assert.assertEquals(32.0d, dblSelector.getDouble(), 0);
+    Assert.assertEquals("32.0x", dblWrapSelector.getValue());
+
+    dblSelector.increment();
+    if (NullHandling.sqlCompatible()) {
+      Assert.assertTrue(dblSelector.isNull());
+      Assert.assertEquals("nullx", dblWrapSelector.getValue());
+    } else {
+      Assert.assertEquals(0d, dblSelector.getDouble(), 0);
+      Assert.assertEquals("0.0x", dblWrapSelector.getValue());
+    }
+
+    dblSelector.increment();
+    Assert.assertEquals(5.0d, dblSelector.getDouble(), 0);
+    Assert.assertEquals("5.0x", dblWrapSelector.getValue());
+
+    dblSelector.increment();
+    Assert.assertEquals(0.0d, dblSelector.getDouble(), 0);
+    Assert.assertEquals("0.0x", dblWrapSelector.getValue());
+
+    dblSelector.increment();
+    Assert.assertEquals(-45.0d, dblSelector.getDouble(), 0);
+    Assert.assertEquals("-45.0x", dblWrapSelector.getValue());
+  }
+
+  @Test
+  public void testFloatWrappingDimensionSelectorExtractionFn()
+  {
+    Float[] vals = new Float[]{32.0f, null, 5.0f, 0.0f, -45.0f};
+    TestNullableFloatColumnSelector flSelector = new TestNullableFloatColumnSelector(vals);
+    final TestExtractionFn extractionFn = new TestExtractionFn();
+
+    FloatWrappingDimensionSelector flWrapSelector = new FloatWrappingDimensionSelector(flSelector, extractionFn);
+    Assert.assertEquals(32.0f, flSelector.getFloat(), 0);
+    Assert.assertEquals("32.0x", flWrapSelector.getValue());
+
+    flSelector.increment();
+    if (NullHandling.sqlCompatible()) {
+      Assert.assertTrue(flSelector.isNull());
+      Assert.assertEquals("nullx", flWrapSelector.getValue());
+    } else {
+      Assert.assertEquals(0f, flSelector.getFloat(), 0);
+      Assert.assertEquals("0.0x", flWrapSelector.getValue());
+    }
+
+    flSelector.increment();
+    Assert.assertEquals(5.0f, flSelector.getFloat(), 0);
+    Assert.assertEquals("5.0x", flWrapSelector.getValue());
+
+    flSelector.increment();
+    Assert.assertEquals(0.0f, flSelector.getFloat(), 0);
+    Assert.assertEquals("0.0x", flWrapSelector.getValue());
+
+    flSelector.increment();
+    Assert.assertEquals(-45.0f, flSelector.getFloat(), 0);
+    Assert.assertEquals("-45.0x", flWrapSelector.getValue());
+  }
+
+  /**
+   * Concats {@link String#valueOf(int)} with "x".
+   */
+  private static class TestExtractionFn extends DimExtractionFn
+  {
+    @Override
+    public byte[] getCacheKey()
+    {
+      throw new UnsupportedOperationException();
+    }
+
+    @Nullable
+    @Override
+    public String apply(@Nullable String value)
+    {
+      return value + "x";
+    }
+
+    @Override
+    public boolean preservesOrdering()
+    {
+      return false;
+    }
+
+    @Override
+    public ExtractionType getExtractionType()
+    {
+      return ExtractionType.MANY_TO_ONE;
+    }
+  }
 }