Refactor TestBackgroundHiveSplitLoader#drain
This changes loop semantics. Now `drain` is more correct, as it doesn't
try to get splits when source `isFinished()`.
@@ -18,7 +18,6 @@ import com.facebook.presto.hive.metastore.Column;
import com.facebook.presto.hive.metastore.StorageFormat;
import com.facebook.presto.hive.metastore.Table;
import com.facebook.presto.spi.ConnectorSession;
-import com.facebook.presto.spi.ConnectorSplit;
import com.facebook.presto.spi.SchemaTableName;
import com.facebook.presto.spi.predicate.Domain;
import com.facebook.presto.spi.predicate.TupleDomain;
@@ -55,6 +54,7 @@ import static com.facebook.presto.hive.HiveType.HIVE_STRING;
import static com.facebook.presto.hive.HiveUtil.getRegularColumnHandles;
import static com.facebook.presto.spi.predicate.TupleDomain.withColumnDomains;
import static com.facebook.presto.spi.type.VarcharType.VARCHAR;
+import static com.google.common.collect.ImmutableList.toImmutableList;
import static io.airlift.concurrent.Threads.daemonThreadsNamed;
import static io.airlift.slice.Slices.utf8Slice;
import static io.airlift.units.DataSize.Unit.GIGABYTE;
@@ -166,17 +166,21 @@ public class TestBackgroundHiveSplitLoader
private List<String> drain(HiveSplitSource source)
throws Exception
{
- ImmutableList.Builder<String> paths = ImmutableList.builder();
- while (true) {
- List<ConnectorSplit> splits = source.getNextBatch(100).get();
- for (ConnectorSplit connectorSplit : splits) {
- paths.add(((HiveSplit) connectorSplit).getPath());
- }
- if (source.isFinished()) {
- break;
+ return drainSplits(source).stream()
+ .map(HiveSplit::getPath)
+ .collect(toImmutableList());
}
+
+ private List<HiveSplit> drainSplits(HiveSplitSource source)
+ throws Exception
+ {
+ ImmutableList.Builder<HiveSplit> splits = ImmutableList.builder();
+ while (!source.isFinished()) {
+ source.getNextBatch(100).get().stream()
+ .map(HiveSplit.class::cast)
+ .forEach(splits::add);
}
- return paths.build();
+ return splits.build();
}
private static BackgroundHiveSplitLoader backgroundHiveSplitLoader(
