Merge pull request #2715 from navis/stringformat-null-handling

stringFormat extractionFn should be able to return null on null values (Fix for #2706)

docs/content/querying/dimensionspecs.md

@@ -346,10 +346,10 @@ For example, `'/druid/prod/historical'` is transformed to `'the dru'` as regular
 Returns the dimension value formatted according to the given format string.
 
 ```json
-{ "type" : "stringFormat", "format" : <sprintf_expression> }
+{ "type" : "stringFormat", "format" : <sprintf_expression>, "nullHandling" : <optional attribute for handling null value> }
 ```
 
-For example if you want to concat "[" and "]" before and after the actual dimension value, you need to specify "[%s]" as format string.
+For example if you want to concat "[" and "]" before and after the actual dimension value, you need to specify "[%s]" as format string. "nullHandling" can be one of `nullString`, `emptyString` or `returnNull`. With "[%s]" format, each configuration will result `[null]`, `[]`, `null`. Default is `nullString`.
 
 ### Filtered DimensionSpecs
 

processing/src/main/java/io/druid/query/extraction/StringFormatExtractionFn.java

@@ -21,6 +21,7 @@ package io.druid.query.extraction;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonValue;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.metamx.common.StringUtils;
@@ -32,15 +33,42 @@ import java.nio.ByteBuffer;
  */
 public class StringFormatExtractionFn extends DimExtractionFn
 {
+  public static enum NullHandling
+  {
+    NULLSTRING,
+    EMPTYSTRING,
+    RETURNNULL;
+
+    @JsonCreator
+    public static NullHandling forValue(String value)
+    {
+      return value == null ? NULLSTRING : NullHandling.valueOf(value.toUpperCase());
+    }
+
+    @JsonValue
+    public String toValue()
+    {
+      return name().toLowerCase();
+    }
+  }
+
   private final String format;
+  private final NullHandling nullHandling;
 
   @JsonCreator
   public StringFormatExtractionFn(
-      @JsonProperty("format") String format
+      @JsonProperty("format") String format,
+      @JsonProperty("nullHandling") NullHandling nullHandling
   )
   {
     Preconditions.checkArgument(!Strings.isNullOrEmpty(format), "format string should not be empty");
     this.format = format;
+    this.nullHandling = nullHandling == null ? NullHandling.NULLSTRING : nullHandling;
+  }
+
+  public StringFormatExtractionFn(String format)
+  {
+    this(format, NullHandling.NULLSTRING);
   }
 
   @JsonProperty
@@ -49,12 +77,19 @@ public class StringFormatExtractionFn extends DimExtractionFn
     return format;
   }
 
+  @JsonProperty
+  public NullHandling getNullHandling()
+  {
+    return nullHandling;
+  }
+
   @Override
   public byte[] getCacheKey()
   {
     byte[] bytes = StringUtils.toUtf8(format);
-    return ByteBuffer.allocate(1 + bytes.length)
+    return ByteBuffer.allocate(2 + bytes.length)
                      .put(ExtractionCacheHelper.CACHE_TYPE_ID_STRING_FORMAT)
+                     .put((byte) nullHandling.ordinal())
                      .put(bytes)
                      .array();
   }
@@ -62,7 +97,15 @@ public class StringFormatExtractionFn extends DimExtractionFn
   @Override
   public String apply(String value)
   {
-    return String.format(format, value);
+    if (value == null) {
+      if (nullHandling == NullHandling.RETURNNULL) {
+        return null;
+      }
+      if (nullHandling == NullHandling.EMPTYSTRING) {
+        value = "";
+      }
+    }
+    return Strings.emptyToNull(String.format(format, value));
   }
 
   @Override
@@ -89,13 +132,26 @@ public class StringFormatExtractionFn extends DimExtractionFn
 
     StringFormatExtractionFn that = (StringFormatExtractionFn) o;
 
+    if (nullHandling != that.nullHandling) {
+      return false;
+    }
     return format.equals(that.format);
-
   }
 
   @Override
   public int hashCode()
   {
-    return format.hashCode();
+    int result = format.hashCode();
+    result = 31 * result + nullHandling.ordinal();
+    return result;
+  }
+
+  @Override
+  public String toString()
+  {
+    return "StringFormatExtractionFn{" +
+           "format='" + format + '\'' +
+           ", nullHandling=" + nullHandling +
+           '}';
   }
 }

processing/src/test/java/io/druid/query/extraction/StringFormatExtractionFnTest.java

@@ -19,6 +19,7 @@
 
 package io.druid.query.extraction;
 
+import com.fasterxml.jackson.databind.JsonMappingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.druid.jackson.DefaultObjectMapper;
 import org.junit.Assert;
@@ -39,18 +40,54 @@ public class StringFormatExtractionFnTest
   }
 
   @Test
-  public void testApplyNull() throws Exception
+  public void testApplyNull1() throws Exception
   {
-    StringFormatExtractionFn fn = new StringFormatExtractionFn("[%s]");
     String test = null;
-    Assert.assertEquals("[null]", fn.apply(test));
+    Assert.assertEquals("[null]", format("[%s]", "nullString").apply(test));
+    Assert.assertEquals("[]", format("[%s]", "emptyString").apply(test));
+    Assert.assertNull(format("[%s]", "returnNull").apply(test));
+  }
+
+  @Test
+  public void testApplyNull2() throws Exception
+  {
+    String test = null;
+    Assert.assertEquals("null", format("%s", "nullString").apply(test));
+    Assert.assertNull(format("%s", "emptyString").apply(test));
+    Assert.assertNull(format("%s", "returnNull").apply(test));
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testInvalidOption1() throws Exception
+  {
+    new StringFormatExtractionFn("");
   }
 
   @Test
   public void testSerde() throws Exception
+  {
+    validateSerde("{ \"type\" : \"stringFormat\", \"format\" : \"[%s]\" }");
+    validateSerde(
+        "{ \"type\" : \"stringFormat\", \"format\" : \"[%s]\", \"nullHandling\" : \"returnNull\" }"
+    );
+  }
+
+  @Test(expected = JsonMappingException.class)
+  public void testInvalidOption2() throws Exception
+  {
+    validateSerde(
+        "{ \"type\" : \"stringFormat\", \"format\" : \"[%s]\", \"nullHandling\" : \"invalid\" }"
+    );
+  }
+
+  public StringFormatExtractionFn format(String format, String nullHandling)
+  {
+    return new StringFormatExtractionFn(format, StringFormatExtractionFn.NullHandling.forValue(nullHandling));
+  }
+
+  private void validateSerde(String json) throws java.io.IOException
   {
     final ObjectMapper objectMapper = new DefaultObjectMapper();
-    final String json = "{ \"type\" : \"stringFormat\", \"format\" : \"[%s]\" }";
     StringFormatExtractionFn extractionFn = (StringFormatExtractionFn) objectMapper.readValue(json, ExtractionFn.class);
 
     Assert.assertEquals("[%s]", extractionFn.getFormat());