Refactor parallel indexing perfect rollup partitioning

core/src/main/java/org/apache/druid/data/input/HandlingInputRowIterator.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;
+
+import org.apache.druid.java.util.common.parsers.CloseableIterator;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * Decorated {@link CloseableIterator<InputRow>} that can process rows with {@link InputRowHandler}s.
+ */
+public class HandlingInputRowIterator implements CloseableIterator<InputRow>
+{
+  @FunctionalInterface
+  public interface InputRowHandler
+  {
+    /**
+     * @return True if inputRow was successfully handled and no further processing is needed
+     */
+    boolean handle(InputRow inputRow);
+  }
+
+  private final CloseableIterator<InputRow> delegate;
+  private final List<InputRowHandler> inputRowHandlers;
+
+  /**
+   * @param inputRowIterator Source of {@link InputRow}s
+   * @param inputRowHandlers Before yielding the next {@link InputRow}, each {@link InputRowHandler} is sequentially
+   *                         applied to the {@link InputRow} until one of them returns true or all of the handlers are
+   *                         applied.
+   */
+  public HandlingInputRowIterator(
+      CloseableIterator<InputRow> inputRowIterator,
+      List<InputRowHandler> inputRowHandlers
+  )
+  {
+    this.delegate = inputRowIterator;
+    this.inputRowHandlers = inputRowHandlers;
+  }
+
+  @Override
+  public boolean hasNext()
+  {
+    return delegate.hasNext();
+  }
+
+  /**
+   * @return Next {@link InputRow} or null if row was successfully handled by an {@link InputRowHandler}.
+   */
+  @Override
+  @Nullable
+  public InputRow next()
+  {
+    InputRow inputRow = delegate.next();
+
+    // NOTE: This loop invokes a virtual call per input row, which may have significant overhead for large inputs
+    // (e.g. InputSourceProcessor). If performance suffers, this implementation or the clients will need to change.
+    for (InputRowHandler inputRowHandler : inputRowHandlers) {
+      if (inputRowHandler.handle(inputRow)) {
+        return null;
+      }
+    }
+
+    return inputRow;
+  }
+
+  @Override
+  public void close() throws IOException
+  {
+    delegate.close();
+  }
+}
+

core/src/main/java/org/apache/druid/indexer/partitions/DynamicPartitionsSpec.java

@@ -81,6 +81,12 @@ public boolean needsDeterminePartitions(boolean useForHadoopTask)
     return false;
   }
 
+  @Override
+  public String getForceGuaranteedRollupIncompatiblityReason()
+  {
+    return NAME + " partitions unsupported";
+  }
+
   @Override
   public boolean equals(Object o)
   {

core/src/main/java/org/apache/druid/indexer/partitions/HashedPartitionsSpec.java

@@ -37,6 +37,8 @@ public class HashedPartitionsSpec implements DimensionBasedPartitionsSpec
   @VisibleForTesting
   static final String NUM_SHARDS = "numShards";
 
+  private static final String FORCE_GUARANTEED_ROLLUP_COMPATIBLE = "";
+
   @Nullable
   private final Integer maxRowsPerSegment;
   @Nullable
@@ -149,6 +151,12 @@ public List<String> getPartitionDimensions()
     return partitionDimensions;
   }
 
+  @Override
+  public String getForceGuaranteedRollupIncompatiblityReason()
+  {
+    return getNumShards() == null ? NUM_SHARDS + " must be specified" : FORCE_GUARANTEED_ROLLUP_COMPATIBLE;
+  }
+
   @Override
   public boolean equals(Object o)
   {

core/src/main/java/org/apache/druid/indexer/partitions/PartitionsSpec.java

@@ -19,6 +19,7 @@
 
 package org.apache.druid.indexer.partitions;
 
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonSubTypes;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 
@@ -75,4 +76,29 @@ static Integer resolveHistoricalNullIfNeeded(@Nullable Integer val)
   {
     return isEffectivelyNull(val) ? null : val;
   }
+
+  /**
+   * @return True if this partitionSpec's type is compatible with forceGuaranteedRollup=true.
+   */
+  @JsonIgnore
+  default boolean isForceGuaranteedRollupCompatibleType()
+  {
+    return !(this instanceof DynamicPartitionsSpec);
+  }
+
+  /**
+   * @return True if this partitionSpec's property values are compatible with forceGuaranteedRollup=true.
+   */
+  @JsonIgnore
+  default boolean isForceGuaranteedRollupCompatible()
+  {
+    return getForceGuaranteedRollupIncompatiblityReason().isEmpty();
+  }
+
+  /**
+   * @return Message describing why this partitionSpec is incompatible with forceGuaranteedRollup=true. Empty string if
+   * the partitionSpec is compatible.
+   */
+  @JsonIgnore
+  String getForceGuaranteedRollupIncompatiblityReason();
 }

core/src/main/java/org/apache/druid/indexer/partitions/SingleDimensionPartitionsSpec.java

@@ -39,10 +39,12 @@
  */
 public class SingleDimensionPartitionsSpec implements DimensionBasedPartitionsSpec
 {
-  static final String NAME = "single_dim";
+  public static final String NAME = "single_dim";
   static final String OLD_NAME = "dimension";  // for backward compatibility
 
+  private static final String PARITION_DIMENSION = "partitionDimension";
   private static final String MAX_PARTITION_SIZE = "maxPartitionSize";
+  private static final String FORCE_GUARANTEED_ROLLUP_COMPATIBLE = "";
 
   private final Integer targetRowsPerSegment;
   private final Integer maxRowsPerSegment;
@@ -56,7 +58,7 @@ public class SingleDimensionPartitionsSpec implements DimensionBasedPartitionsSp
   public SingleDimensionPartitionsSpec(
       @JsonProperty(DimensionBasedPartitionsSpec.TARGET_ROWS_PER_SEGMENT) @Nullable Integer targetRowsPerSegment,
       @JsonProperty(PartitionsSpec.MAX_ROWS_PER_SEGMENT) @Nullable Integer maxRowsPerSegment,
-      @JsonProperty("partitionDimension") @Nullable String partitionDimension,
+      @JsonProperty(PARITION_DIMENSION) @Nullable String partitionDimension,
       @JsonProperty("assumeGrouped") boolean assumeGrouped,  // false by default
 
       // Deprecated properties preserved for backward compatibility:
@@ -170,6 +172,12 @@ public List<String> getPartitionDimensions()
     return partitionDimension == null ? Collections.emptyList() : Collections.singletonList(partitionDimension);
   }
 
+  @Override
+  public String getForceGuaranteedRollupIncompatiblityReason()
+  {
+    return NAME + " partitions unsupported";
+  }
+
   @Override
   public boolean needsDeterminePartitions(boolean useForHadoopTask)
   {

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

@@ -0,0 +1,179 @@
+/*
+ * 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 org.apache.druid.java.util.common.CloseableIterators;
+import org.apache.druid.java.util.common.parsers.CloseableIterator;
+import org.easymock.EasyMock;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.runners.Enclosed;
+import org.junit.runner.RunWith;
+
+import javax.annotation.Nullable;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+
+@RunWith(Enclosed.class)
+public class HandlingInputRowIteratorTest
+{
+  public static class AbsentRowTest
+  {
+    private static final CloseableIterator<InputRow> EMPTY_ITERATOR = CloseableIterators.withEmptyBaggage(
+        new Iterator<InputRow>()
+        {
+          @Override
+          public boolean hasNext()
+          {
+            return false;
+          }
+
+          @Nullable
+          @Override
+          public InputRow next()
+          {
+            throw new NoSuchElementException();
+          }
+        }
+    );
+
+    private HandlingInputRowIterator target;
+
+    @Before
+    public void setup()
+    {
+      target = new HandlingInputRowIterator(EMPTY_ITERATOR, Collections.emptyList());
+    }
+
+    @Test
+    public void doesNotHaveNext()
+    {
+      Assert.assertFalse(target.hasNext());
+    }
+
+    @Test(expected = NoSuchElementException.class)
+    public void throwsExceptionWhenYieldingNext()
+    {
+      target.next();
+    }
+  }
+
+  public static class PresentRowTest
+  {
+    private static final InputRow INPUT_ROW1 = EasyMock.mock(InputRow.class);
+    private static final InputRow INPUT_ROW2 = EasyMock.mock(InputRow.class);
+    private static final List<InputRow> INPUT_ROWS = Arrays.asList(INPUT_ROW1, INPUT_ROW2);
+
+    private TestInputRowHandler successfulHandler;
+    private TestInputRowHandler unsuccessfulHandler;
+
+    @Before
+    public void setup()
+    {
+      successfulHandler = new TestInputRowHandler(true);
+      unsuccessfulHandler = new TestInputRowHandler(false);
+    }
+
+    @Test
+    public void hasNext()
+    {
+      HandlingInputRowIterator target = createInputRowIterator(unsuccessfulHandler, unsuccessfulHandler);
+      Assert.assertTrue(target.hasNext());
+      Assert.assertFalse(unsuccessfulHandler.invoked);
+    }
+
+    @Test
+    public void yieldsNextIfUnhandled()
+    {
+      HandlingInputRowIterator target = createInputRowIterator(unsuccessfulHandler, unsuccessfulHandler);
+      Assert.assertEquals(INPUT_ROW1, target.next());
+      Assert.assertTrue(unsuccessfulHandler.invoked);
+    }
+
+    @Test
+    public void yieldsNullIfHandledByFirst()
+    {
+      HandlingInputRowIterator target = createInputRowIterator(successfulHandler, unsuccessfulHandler);
+      Assert.assertNull(target.next());
+      Assert.assertTrue(successfulHandler.invoked);
+      Assert.assertFalse(unsuccessfulHandler.invoked);
+    }
+
+    @Test
+    public void yieldsNullIfHandledBySecond()
+    {
+      HandlingInputRowIterator target = createInputRowIterator(unsuccessfulHandler, successfulHandler);
+      Assert.assertNull(target.next());
+      Assert.assertTrue(unsuccessfulHandler.invoked);
+      Assert.assertTrue(successfulHandler.invoked);
+    }
+
+    private static HandlingInputRowIterator createInputRowIterator(
+        HandlingInputRowIterator.InputRowHandler firstHandler,
+        HandlingInputRowIterator.InputRowHandler secondHandler
+    )
+    {
+      CloseableIterator<InputRow> iterator = CloseableIterators.withEmptyBaggage(
+          new Iterator<InputRow>()
+          {
+            private final Iterator<InputRow> delegate = INPUT_ROWS.iterator();
+
+            @Override
+            public boolean hasNext()
+            {
+              return delegate.hasNext();
+            }
+
+            @Nullable
+            @Override
+            public InputRow next()
+            {
+              return delegate.next();
+            }
+          }
+      );
+
+      return new HandlingInputRowIterator(iterator, Arrays.asList(firstHandler, secondHandler));
+    }
+
+    private static class TestInputRowHandler implements HandlingInputRowIterator.InputRowHandler
+    {
+      boolean invoked = false;
+
+      private final boolean successful;
+
+      TestInputRowHandler(boolean successful)
+      {
+        this.successful = successful;
+      }
+
+      @Override
+      public boolean handle(InputRow inputRow)
+      {
+        invoked = true;
+        return successful;
+      }
+    }
+  }
+}