From a5bd0b8cc023b771ef95f1e46c19961f50e05900 Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Mon, 8 Nov 2021 10:25:13 -0800
Subject: [PATCH] RowAdapter: Add a default implementation for
 timestampFunction. (#11885)

Enables simpler implementations for adapters that want to treat the
timestamp as "just another column".
---
 .../apache/druid/query/InlineDataSource.java  | 32 ++++---------------
 .../org/apache/druid/segment/RowAdapter.java  | 19 ++++++++++-
 .../scan/ScanQueryResultOrderingTest.java     | 25 ++++-----------
 3 files changed, 30 insertions(+), 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
index 4206732cfc6..163bb986f91 100644
--- 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.base.Preconditions;
 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.Collections;
 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 class InlineDataSource implements DataSource
 
   public RowAdapter<Object[]> rowAdapter()
   {
-    return new RowAdapter<Object[]>()
-    {
-      @Override
-      public ToLongFunction<Object[]> timestampFunction()
-      {
-        final int columnNumber = signature.indexOf(ColumnHolder.TIME_COLUMN_NAME);
+    return columnName -> {
+      final int columnNumber = signature.indexOf(columnName);
 
-        if (columnNumber >= 0) {
-          return row -> (long) row[columnNumber];
-        } else {
-          return row -> 0L;
-        }
-      }
-
-      @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
index 603d349bd4d..e1317e714f8 100644
--- 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
index 002d80957e8..4433fc0a614 100644
--- 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.Comparator;
 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;
     }
   };