DruidInputSource: Fix issues in column projection, timestamp handling.

core/src/main/java/org/apache/druid/data/input/ColumnsFilter.java

@@ -0,0 +1,183 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.data.input;
+
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * Used by some {@link InputSourceReader} implementations in order to know what columns will need to be read out
+ * of the {@link InputRow} objects they create.
+ *
+ * This is meant to be useful as an optimization: if we're reading from a columnar data format, then when a column
+ * isn't going to be needed, we shouldn't read it.
+ *
+ * @see InputSource#reader accepts objects of this class
+ */
+public abstract class ColumnsFilter
+{
+  /**
+   * Accepts all columns.
+   */
+  public static ColumnsFilter all()
+  {
+    return new ExclusionBased(Collections.emptySet());
+  }
+
+  /**
+   * Accepts a specific list of columns.
+   */
+  public static ColumnsFilter inclusionBased(final Set<String> inclusions)
+  {
+    return new InclusionBased(inclusions);
+  }
+
+
+  /**
+   * Accepts all columns, except those on a specific list.
+   */
+  public static ColumnsFilter exclusionBased(final Set<String> exclusions)
+  {
+    return new ExclusionBased(exclusions);
+  }
+
+  /**
+   * Check if a column should be included or not.
+   */
+  public abstract boolean apply(String column);
+
+  /**
+   * Returns a new filter with a particular column added. The returned filter will return true from {@link #apply}
+   * on this column.
+   */
+  public abstract ColumnsFilter plus(String column);
+
+  public static class InclusionBased extends ColumnsFilter
+  {
+    private final Set<String> inclusions;
+
+    private InclusionBased(Set<String> inclusions)
+    {
+      this.inclusions = inclusions;
+    }
+
+    @Override
+    public boolean apply(String column)
+    {
+      return inclusions.contains(column);
+    }
+
+    @Override
+    public ColumnsFilter plus(String column)
+    {
+      if (inclusions.contains(column)) {
+        return this;
+      } else {
+        final Set<String> copy = new HashSet<>(inclusions);
+        copy.add(column);
+        return new InclusionBased(copy);
+      }
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      InclusionBased that = (InclusionBased) o;
+      return Objects.equals(inclusions, that.inclusions);
+    }
+
+    @Override
+    public int hashCode()
+    {
+      return Objects.hash(inclusions);
+    }
+
+    @Override
+    public String toString()
+    {
+      return "ColumnsFilter.InclusionBased{" +
+             "inclusions=" + inclusions +
+             '}';
+    }
+  }
+
+  public static class ExclusionBased extends ColumnsFilter
+  {
+    private final Set<String> exclusions;
+
+    public ExclusionBased(Set<String> exclusions)
+    {
+      this.exclusions = exclusions;
+    }
+
+    @Override
+    public boolean apply(String column)
+    {
+      return !exclusions.contains(column);
+    }
+
+    @Override
+    public ColumnsFilter plus(String column)
+    {
+      if (!exclusions.contains(column)) {
+        return this;
+      } else {
+        final Set<String> copy = new HashSet<>(exclusions);
+        copy.remove(column);
+        return new ExclusionBased(copy);
+      }
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      ExclusionBased that = (ExclusionBased) o;
+      return Objects.equals(exclusions, that.exclusions);
+    }
+
+    @Override
+    public int hashCode()
+    {
+      return Objects.hash(exclusions);
+    }
+
+    @Override
+    public String toString()
+    {
+      return "ColumnsFilter.ExclusionBased{" +
+             "exclusions=" + exclusions +
+             '}';
+    }
+  }
+}

core/src/main/java/org/apache/druid/data/input/InputRowSchema.java

@@ -22,22 +22,24 @@
 import org.apache.druid.data.input.impl.DimensionsSpec;
 import org.apache.druid.data.input.impl.TimestampSpec;
 
-import java.util.List;
-
 /**
  * Schema of {@link InputRow}.
  */
 public class InputRowSchema
 {
   private final TimestampSpec timestampSpec;
   private final DimensionsSpec dimensionsSpec;
-  private final List<String> metricNames;
+  private final ColumnsFilter columnsFilter;
 
-  public InputRowSchema(TimestampSpec timestampSpec, DimensionsSpec dimensionsSpec, List<String> metricNames)
+  public InputRowSchema(
+      final TimestampSpec timestampSpec,
+      final DimensionsSpec dimensionsSpec,
+      final ColumnsFilter columnsFilter
+  )
   {
     this.timestampSpec = timestampSpec;
     this.dimensionsSpec = dimensionsSpec;
-    this.metricNames = metricNames;
+    this.columnsFilter = columnsFilter;
   }
 
   public TimestampSpec getTimestampSpec()
@@ -50,8 +52,17 @@ public DimensionsSpec getDimensionsSpec()
     return dimensionsSpec;
   }
 
-  public List<String> getMetricNames()
+  /**
+   * A {@link ColumnsFilter} that can filter down the list of columns that must be read after flattening.
+   *
+   * Logically, Druid applies ingestion spec components in a particular order: first flattenSpec (if any), then
+   * timestampSpec, then transformSpec, and finally dimensionsSpec and metricsSpec.
+   *
+   * If a flattenSpec is provided, this method returns a filter that should be applied after flattening. So, it will
+   * be based on what needs to pass between the flattenSpec and everything beyond it.
+   */
+  public ColumnsFilter getColumnsFilter()
   {
-    return metricNames;
+    return columnsFilter;
   }
 }

core/src/main/java/org/apache/druid/data/input/InputSource.java

@@ -78,5 +78,9 @@ public interface InputSource
    * @param inputFormat        to parse data. It can be null if {@link #needsFormat()} = true
    * @param temporaryDirectory to store temp data. It will be cleaned up automatically once the task is finished.
    */
-  InputSourceReader reader(InputRowSchema inputRowSchema, @Nullable InputFormat inputFormat, File temporaryDirectory);
+  InputSourceReader reader(
+      InputRowSchema inputRowSchema,
+      @Nullable InputFormat inputFormat,
+      File temporaryDirectory
+  );
 }

core/src/main/java/org/apache/druid/data/input/MapBasedInputRow.java

@@ -25,8 +25,10 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 /**
+ *
  */
 @PublicApi
 public class MapBasedInputRow extends MapBasedRow implements InputRow
@@ -59,6 +61,28 @@ public List<String> getDimensions()
     return dimensions;
   }
 
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    if (!super.equals(o)) {
+      return false;
+    }
+    MapBasedInputRow that = (MapBasedInputRow) o;
+    return Objects.equals(dimensions, that.dimensions);
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(super.hashCode(), dimensions);
+  }
+
   @Override
   public String toString()
   {

core/src/test/java/org/apache/druid/data/input/FirehoseFactoryToInputSourceAdaptorTest.java

@@ -37,7 +37,6 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.stream.Stream;
@@ -71,7 +70,7 @@ public void testUnimplementedInputFormat() throws IOException
         new InputRowSchema(
             inputRowParser.getParseSpec().getTimestampSpec(),
             inputRowParser.getParseSpec().getDimensionsSpec(),
-            Collections.emptyList()
+            ColumnsFilter.all()
         ),
         null,
         null

core/src/test/java/org/apache/druid/data/input/impl/ColumnsFilterTest.java

@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.data.input.impl;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import nl.jqno.equalsverifier.EqualsVerifier;
+import org.apache.druid.data.input.ColumnsFilter;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class ColumnsFilterTest
+{
+  private static final List<String> COLUMNS = ImmutableList.of("a", "b", "c");
+
+  @Test
+  public void testAll()
+  {
+    Assert.assertEquals(
+        ImmutableList.of("a", "b", "c"),
+        apply(ColumnsFilter.all(), COLUMNS)
+    );
+  }
+
+  @Test
+  public void testInclusionBased()
+  {
+    Assert.assertEquals(
+        ImmutableList.of("b"),
+        apply(ColumnsFilter.inclusionBased(ImmutableSet.of("b")), COLUMNS)
+    );
+  }
+
+  @Test
+  public void testInclusionBasedPlus()
+  {
+    Assert.assertEquals(
+        ColumnsFilter.inclusionBased(ImmutableSet.of("a", "b", "c")),
+        ColumnsFilter.inclusionBased(ImmutableSet.of("b", "c")).plus("a").plus("c")
+    );
+  }
+
+  @Test
+  public void testExclusionBased()
+  {
+    Assert.assertEquals(
+        ImmutableList.of("a", "c"),
+        apply(ColumnsFilter.exclusionBased(ImmutableSet.of("b")), COLUMNS)
+    );
+  }
+
+  @Test
+  public void testExclusionBasedPlus()
+  {
+    Assert.assertEquals(
+        ColumnsFilter.exclusionBased(ImmutableSet.of("b")),
+        ColumnsFilter.exclusionBased(ImmutableSet.of("b", "c")).plus("a").plus("c")
+    );
+  }
+
+  @Test
+  public void testEquals()
+  {
+    EqualsVerifier.forClass(ColumnsFilter.InclusionBased.class).usingGetClass().verify();
+    EqualsVerifier.forClass(ColumnsFilter.ExclusionBased.class).usingGetClass().verify();
+  }
+
+  private List<String> apply(ColumnsFilter columnsFilter, List<String> columns)
+  {
+    return columns.stream().filter(columnsFilter::apply).collect(Collectors.toList());
+  }
+}