RowAdapter: Add a default implementation for timestampFunction. (apac…

…he#11885) Enables simpler implementations for adapters that want to treat the timestamp as "just another column".
natsumehu · Nov 8, 2021 · a5bd0b8 · a5bd0b8
1 parent 7237dc8
commit a5bd0b8
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 46 deletions.
diff --git a/processing/src/main/java/org/apache/druid/query/InlineDataSource.java b/processing/src/main/java/org/apache/druid/query/InlineDataSource.java
@@ -27,7 +27,6 @@
 import com.google.common.collect.Lists;
 import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.segment.RowAdapter;
-import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
 
@@ -37,8 +36,6 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
-import java.util.function.Function;
-import java.util.function.ToLongFunction;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -196,30 +193,13 @@ public RowSignature getRowSignature()
 
   public RowAdapter<Object[]> rowAdapter()
   {
-    return new RowAdapter<Object[]>()
-    {
-      @Override
-      public ToLongFunction<Object[]> timestampFunction()
-      {
-        final int columnNumber = signature.indexOf(ColumnHolder.TIME_COLUMN_NAME);
-
-        if (columnNumber >= 0) {
-          return row -> (long) row[columnNumber];
-        } else {
-          return row -> 0L;
-        }
-      }
+    return columnName -> {
+      final int columnNumber = signature.indexOf(columnName);
 
-      @Override
-      public Function<Object[], Object> columnFunction(String columnName)
-      {
-        final int columnNumber = signature.indexOf(columnName);
-
-        if (columnNumber >= 0) {
-          return row -> row[columnNumber];
-        } else {
-          return row -> null;
-        }
+      if (columnNumber >= 0) {
+        return row -> row[columnNumber];
+      } else {
+        return row -> null;
       }
     };
   }

diff --git a/processing/src/main/java/org/apache/druid/segment/RowAdapter.java b/processing/src/main/java/org/apache/druid/segment/RowAdapter.java
@@ -19,6 +19,8 @@
 
 package org.apache.druid.segment;
 
+import org.apache.druid.segment.column.ColumnHolder;
+
 import java.util.function.Function;
 import java.util.function.ToLongFunction;
 
@@ -30,8 +32,23 @@ public interface RowAdapter<RowType>
 {
   /**
    * Returns a function that retrieves timestamps from rows.
+   *
+   * The default implementation delegates to {@link #columnFunction} and expects it to already contain long-typed
+   * values or nulls. Nulls, if present, will be converted to zeroes.
    */
-  ToLongFunction<RowType> timestampFunction();
+  default ToLongFunction<RowType> timestampFunction()
+  {
+    final Function<RowType, Object> timeColumnFunction = columnFunction(ColumnHolder.TIME_COLUMN_NAME);
+    return row -> {
+      final Object obj = timeColumnFunction.apply(row);
+
+      if (obj == null) {
+        return 0L;
+      } else {
+        return (long) obj;
+      }
+    };
+  }
 
   /**
    * Returns a function that retrieves the value for column "columnName" from rows.

diff --git a/processing/src/test/java/org/apache/druid/query/scan/ScanQueryResultOrderingTest.java b/processing/src/test/java/org/apache/druid/query/scan/ScanQueryResultOrderingTest.java
@@ -56,8 +56,6 @@
 import java.util.List;
 import java.util.Set;
 import java.util.TreeSet;
-import java.util.function.Function;
-import java.util.function.ToLongFunction;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -72,24 +70,13 @@ public class ScanQueryResultOrderingTest
   private static final String DATASOURCE = "datasource";
   private static final String ID_COLUMN = "id";
 
-  private static final RowAdapter<Object[]> ROW_ADAPTER = new RowAdapter<Object[]>()
-  {
-    @Override
-    public ToLongFunction<Object[]> timestampFunction()
-    {
+  private static final RowAdapter<Object[]> ROW_ADAPTER = columnName -> {
+    if (ID_COLUMN.equals(columnName)) {
+      return row -> row[1];
+    } else if (ColumnHolder.TIME_COLUMN_NAME.equals(columnName)) {
       return row -> ((DateTime) row[0]).getMillis();
-    }
-
-    @Override
-    public Function<Object[], Object> columnFunction(String columnName)
-    {
-      if (ID_COLUMN.equals(columnName)) {
-        return row -> row[1];
-      } else if (ColumnHolder.TIME_COLUMN_NAME.equals(columnName)) {
-        return timestampFunction()::applyAsLong;
-      } else {
-        return row -> null;
-      }
+    } else {
+      return row -> null;
     }
   };