From 00cb0a2900a4cc8248532327aba521add9cc91b4 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Fri, 26 Jan 2024 06:26:13 -0800 Subject: [PATCH] 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. --- .../DoubleWrappingDimensionSelector.java | 8 +- .../FloatWrappingDimensionSelector.java | 8 +- .../LongWrappingDimensionSelector.java | 8 +- .../WrappingDimensionSelectorTest.java | 139 ++++++++++++++++++ 4 files changed, 148 insertions(+), 15 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/segment/DoubleWrappingDimensionSelector.java b/processing/src/main/java/org/apache/druid/segment/DoubleWrappingDimensionSelector.java index 4f1dac049f6..a4dbcf6fbdf 100644 --- a/processing/src/main/java/org/apache/druid/segment/DoubleWrappingDimensionSelector.java +++ b/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()) { - return null; - } else if (extractionFn == null) { - return String.valueOf(selector.getDouble()); + if (extractionFn == null) { + return selector.isNull() ? null : String.valueOf(selector.getDouble()); } else { - return extractionFn.apply(selector.getDouble()); + return selector.isNull() ? extractionFn.apply(null) : extractionFn.apply(selector.getDouble()); } } diff --git a/processing/src/main/java/org/apache/druid/segment/FloatWrappingDimensionSelector.java b/processing/src/main/java/org/apache/druid/segment/FloatWrappingDimensionSelector.java index ec5f958dab9..6455dc8e227 100644 --- a/processing/src/main/java/org/apache/druid/segment/FloatWrappingDimensionSelector.java +++ b/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()) { - return null; - } else if (extractionFn == null) { - return String.valueOf(selector.getFloat()); + if (extractionFn == null) { + return selector.isNull() ? null : String.valueOf(selector.getFloat()); } else { - return extractionFn.apply(selector.getFloat()); + return selector.isNull() ? extractionFn.apply(null) : extractionFn.apply(selector.getFloat()); } } diff --git a/processing/src/main/java/org/apache/druid/segment/LongWrappingDimensionSelector.java b/processing/src/main/java/org/apache/druid/segment/LongWrappingDimensionSelector.java index de25d2ebd87..37e4ec79de2 100644 --- a/processing/src/main/java/org/apache/druid/segment/LongWrappingDimensionSelector.java +++ b/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()) { - return null; - } else if (extractionFn == null) { - return String.valueOf(selector.getLong()); + if (extractionFn == null) { + return selector.isNull() ? null : String.valueOf(selector.getLong()); } else { - return extractionFn.apply(selector.getLong()); + return selector.isNull() ? extractionFn.apply(null) : extractionFn.apply(selector.getLong()); } } diff --git a/processing/src/test/java/org/apache/druid/segment/WrappingDimensionSelectorTest.java b/processing/src/test/java/org/apache/druid/segment/WrappingDimensionSelectorTest.java index 8ba3425807d..fd4ebf69b30 100644 --- a/processing/src/test/java/org/apache/druid/segment/WrappingDimensionSelectorTest.java +++ b/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; + } + } }