diff --git a/quidem-ut/src/main/java/org/apache/druid/quidem/KttmNestedComponentSupplier.java b/quidem-ut/src/main/java/org/apache/druid/quidem/KttmNestedComponentSupplier.java index 64524b9376c..d3fc758a2c6 100644 --- a/quidem-ut/src/main/java/org/apache/druid/quidem/KttmNestedComponentSupplier.java +++ b/quidem-ut/src/main/java/org/apache/druid/quidem/KttmNestedComponentSupplier.java @@ -59,7 +59,7 @@ public class KttmNestedComponentSupplier extends StandardComponentSupplier public static class Micro extends KttmNestedComponentSupplier { public Micro(TempDirProducer tempDirProducer) { - super(tempDirProducer, Fraction.getFraction(1, 10000), Integer.MAX_VALUE); + super(tempDirProducer, Fraction.getFraction(1, 10000)); } } @@ -68,14 +68,13 @@ public class KttmNestedComponentSupplier extends StandardComponentSupplier public KttmNestedComponentSupplier(TempDirProducer tempDirProducer) { - this(tempDirProducer, Fraction.ONE, Integer.MAX_VALUE); + this(tempDirProducer, Fraction.ONE); } - public KttmNestedComponentSupplier(TempDirProducer tempDirProducer, Fraction fraction, int maxRows) + public KttmNestedComponentSupplier(TempDirProducer tempDirProducer, Fraction fraction) { super(tempDirProducer); this.fraction = fraction; - this.maxRows = maxRows; } @Override diff --git a/quidem-ut/src/main/java/org/apache/druid/quidem/ScaledResoureInputDataSource.java b/quidem-ut/src/main/java/org/apache/druid/quidem/ScaledResoureInputDataSource.java index 962d44837fc..dd646c965d0 100644 --- a/quidem-ut/src/main/java/org/apache/druid/quidem/ScaledResoureInputDataSource.java +++ b/quidem-ut/src/main/java/org/apache/druid/quidem/ScaledResoureInputDataSource.java @@ -41,12 +41,10 @@ public class ScaledResoureInputDataSource extends AbstractInputSource { private final ResourceInputSource resourceInputSource; private final Fraction fraction; - private final int maxRows; - public ScaledResoureInputDataSource(Fraction fraction, int maxRows, ResourceInputSource resourceInputSource) + public ScaledResoureInputDataSource(Fraction fraction, ResourceInputSource resourceInputSource) { this.fraction = fraction; - this.maxRows = maxRows; this.resourceInputSource = resourceInputSource; } @@ -66,7 +64,7 @@ public class ScaledResoureInputDataSource extends AbstractInputSource public InputSourceReader reader(InputRowSchema inputRowSchema, InputFormat inputFormat, File temporaryDirectory) { InputSourceReader reader = resourceInputSource.reader(inputRowSchema, inputFormat, temporaryDirectory); - return new FilteredReader(reader, this::filterPredicate, maxRows); + return new FilteredReader(reader, this::filterPredicate); } public boolean filterPredicate(InputRow inputRow) @@ -76,27 +74,24 @@ public class ScaledResoureInputDataSource extends AbstractInputSource static class FilteredReader implements InputSourceReader { - private InputSourceReader reader; - private Predicate filterPredicate; - private int maxRows; + private final InputSourceReader reader; + private final Predicate filterPredicate; - public FilteredReader(InputSourceReader reader, Predicate filterPredicate, int maxRows) + public FilteredReader(InputSourceReader reader, Predicate filterPredicate) { this.reader = reader; this.filterPredicate = filterPredicate; - this.maxRows = maxRows; } @Override public CloseableIterator read(InputStats inputStats) throws IOException { - FilteringCloseableInputRowIterator filteredIterator = new FilteringCloseableInputRowIterator( + return new FilteringCloseableInputRowIterator( reader.read(inputStats), filterPredicate, NoopRowIngestionMeters.INSTANCE, new ParseExceptionHandler(NoopRowIngestionMeters.INSTANCE, false, 0, 0) ); - return new LimitedCloseableIterator<>(filteredIterator, maxRows); } @Override @@ -105,39 +100,4 @@ public class ScaledResoureInputDataSource extends AbstractInputSource return reader.sample(); } } - - static class LimitedCloseableIterator implements CloseableIterator - { - private final CloseableIterator delegate; - private final int maxRows; - private int count = 0; - - public LimitedCloseableIterator(CloseableIterator delegate, int maxRows) - { - this.delegate = delegate; - this.maxRows = maxRows; - } - - @Override - public boolean hasNext() - { - return count < maxRows && delegate.hasNext(); - } - - @Override - public InputRow next() - { - if (count >= maxRows) { - throw new IllegalStateException("Exceeded maxRows"); - } - count++; - return delegate.next(); - } - - @Override - public void close() throws IOException - { - delegate.close(); - } - } }