You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by bl...@apache.org on 2020/05/15 15:58:14 UTC
[incubator-iceberg] branch master updated: Add CloseableIterator
(#1011)
This is an automated email from the ASF dual-hosted git repository.
blue pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-iceberg.git
The following commit(s) were added to refs/heads/master by this push:
new 2f9f41d Add CloseableIterator (#1011)
2f9f41d is described below
commit 2f9f41d1acc11dfe85e14d695bb2bc2398928fc9
Author: openinx <op...@gmail.com>
AuthorDate: Fri May 15 23:58:07 2020 +0800
Add CloseableIterator (#1011)
---
.../org/apache/iceberg/io/CloseableIterable.java | 32 ++++++++----
.../org/apache/iceberg/io/CloseableIterator.java | 57 ++++++++++++++++++++++
.../org/apache/iceberg/BaseManifestReader.java | 9 ++--
.../java/org/apache/iceberg/avro/AvroIterable.java | 8 +--
.../org/apache/iceberg/util/ParallelIterable.java | 5 +-
.../org/apache/iceberg/data/TableScanIterable.java | 5 +-
.../iceberg/mr/mapreduce/IcebergInputFormat.java | 28 +++++------
.../java/org/apache/iceberg/orc/OrcIterable.java | 6 ++-
.../apache/iceberg/parquet/ParquetIterable.java | 7 ++-
.../org/apache/iceberg/parquet/ParquetReader.java | 7 ++-
.../iceberg/parquet/VectorizedParquetReader.java | 7 ++-
.../iceberg/spark/source/BaseDataReader.java | 17 +++----
.../apache/iceberg/spark/source/RowDataReader.java | 21 ++++----
13 files changed, 138 insertions(+), 71 deletions(-)
diff --git a/api/src/main/java/org/apache/iceberg/io/CloseableIterable.java b/api/src/main/java/org/apache/iceberg/io/CloseableIterable.java
index 5e47ab2..03c1dd3 100644
--- a/api/src/main/java/org/apache/iceberg/io/CloseableIterable.java
+++ b/api/src/main/java/org/apache/iceberg/io/CloseableIterable.java
@@ -32,6 +32,15 @@ import java.util.function.Predicate;
import org.apache.iceberg.exceptions.RuntimeIOException;
public interface CloseableIterable<T> extends Iterable<T>, Closeable {
+
+ /**
+ * Returns an closeable iterator over elements of type {@code T}.
+ *
+ * @return an {@link CloseableIterator}.
+ */
+ @Override
+ CloseableIterator<T> iterator();
+
static <E> CloseableIterable<E> withNoopClose(E entry) {
return withNoopClose(ImmutableList.of(entry));
}
@@ -43,8 +52,8 @@ public interface CloseableIterable<T> extends Iterable<T>, Closeable {
}
@Override
- public Iterator<E> iterator() {
- return iterable.iterator();
+ public CloseableIterator<E> iterator() {
+ return CloseableIterator.withClose(iterable.iterator());
}
};
}
@@ -61,8 +70,8 @@ public interface CloseableIterable<T> extends Iterable<T>, Closeable {
}
@Override
- public Iterator<E> iterator() {
- return iterable.iterator();
+ public CloseableIterator<E> iterator() {
+ return CloseableIterator.withClose(iterable.iterator());
}
};
}
@@ -81,9 +90,14 @@ public interface CloseableIterable<T> extends Iterable<T>, Closeable {
}
@Override
- public Iterator<O> iterator() {
- return new Iterator<O>() {
- private final Iterator<I> inner = iterable.iterator();
+ public CloseableIterator<O> iterator() {
+ return new CloseableIterator<O>() {
+ private final CloseableIterator<I> inner = iterable.iterator();
+
+ @Override
+ public void close() throws IOException {
+ inner.close();
+ }
@Override
public boolean hasNext() {
@@ -116,13 +130,13 @@ public interface CloseableIterable<T> extends Iterable<T>, Closeable {
}
@Override
- public Iterator<E> iterator() {
+ public CloseableIterator<E> iterator() {
ConcatCloseableIterator<E> iter = new ConcatCloseableIterator<>(inputs);
addCloseable(iter);
return iter;
}
- private static class ConcatCloseableIterator<E> implements Iterator<E>, Closeable {
+ private static class ConcatCloseableIterator<E> implements CloseableIterator<E> {
private final Iterator<CloseableIterable<E>> iterables;
private CloseableIterable<E> currentIterable = null;
private Iterator<E> currentIterator = null;
diff --git a/api/src/main/java/org/apache/iceberg/io/CloseableIterator.java b/api/src/main/java/org/apache/iceberg/io/CloseableIterator.java
new file mode 100644
index 0000000..e2b2e4f
--- /dev/null
+++ b/api/src/main/java/org/apache/iceberg/io/CloseableIterator.java
@@ -0,0 +1,57 @@
+/*
+ * 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.iceberg.io;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Iterator;
+
+public interface CloseableIterator<T> extends Iterator<T>, Closeable {
+
+ static <E> CloseableIterator<E> empty() {
+ return withClose(Collections.emptyIterator());
+ }
+
+ static <E> CloseableIterator<E> withClose(Iterator<E> iterator) {
+ if (iterator instanceof CloseableIterator) {
+ return (CloseableIterator<E>) iterator;
+ }
+
+ return new CloseableIterator<E>() {
+ @Override
+ public void close() throws IOException {
+ if (iterator instanceof Closeable) {
+ ((Closeable) iterator).close();
+ }
+ }
+
+ @Override
+ public boolean hasNext() {
+ return iterator.hasNext();
+ }
+
+ @Override
+ public E next() {
+ return iterator.next();
+ }
+ };
+ }
+}
diff --git a/core/src/main/java/org/apache/iceberg/BaseManifestReader.java b/core/src/main/java/org/apache/iceberg/BaseManifestReader.java
index 1496f01..2d72b62 100644
--- a/core/src/main/java/org/apache/iceberg/BaseManifestReader.java
+++ b/core/src/main/java/org/apache/iceberg/BaseManifestReader.java
@@ -25,7 +25,6 @@ import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import java.io.IOException;
import java.util.Collection;
-import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -39,6 +38,7 @@ import org.apache.iceberg.expressions.InclusiveMetricsEvaluator;
import org.apache.iceberg.expressions.Projections;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.types.Types;
@@ -219,11 +219,12 @@ abstract class BaseManifestReader<T, ThisT> extends CloseableGroup implements Cl
*/
@Override
@SuppressWarnings("unchecked")
- public Iterator<T> iterator() {
+ public CloseableIterator<T> iterator() {
if (dropStats(rowFilter, columns)) {
- return (Iterator<T>) CloseableIterable.transform(liveEntries(), e -> e.file().copyWithoutStats()).iterator();
+ return (CloseableIterator<T>) CloseableIterable.transform(liveEntries(), e -> e.file().copyWithoutStats())
+ .iterator();
} else {
- return (Iterator<T>) CloseableIterable.transform(liveEntries(), e -> e.file().copy()).iterator();
+ return (CloseableIterator<T>) CloseableIterable.transform(liveEntries(), e -> e.file().copy()).iterator();
}
}
diff --git a/core/src/main/java/org/apache/iceberg/avro/AvroIterable.java b/core/src/main/java/org/apache/iceberg/avro/AvroIterable.java
index 8ab1942..ef463bc 100644
--- a/core/src/main/java/org/apache/iceberg/avro/AvroIterable.java
+++ b/core/src/main/java/org/apache/iceberg/avro/AvroIterable.java
@@ -20,7 +20,6 @@
package org.apache.iceberg.avro;
import com.google.common.collect.Maps;
-import java.io.Closeable;
import java.io.IOException;
import java.util.Iterator;
import java.util.Map;
@@ -32,6 +31,7 @@ import org.apache.avro.io.DatumReader;
import org.apache.iceberg.exceptions.RuntimeIOException;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.InputFile;
public class AvroIterable<D> extends CloseableGroup implements CloseableIterable<D> {
@@ -73,7 +73,7 @@ public class AvroIterable<D> extends CloseableGroup implements CloseableIterable
}
@Override
- public Iterator<D> iterator() {
+ public CloseableIterator<D> iterator() {
FileReader<D> fileReader = initMetadata(newFileReader());
if (start != null) {
@@ -86,7 +86,7 @@ public class AvroIterable<D> extends CloseableGroup implements CloseableIterable
return new AvroReuseIterator<>(fileReader);
}
- return fileReader;
+ return CloseableIterator.withClose(fileReader);
}
private DataFileReader<D> newFileReader() {
@@ -173,7 +173,7 @@ public class AvroIterable<D> extends CloseableGroup implements CloseableIterable
}
}
- private static class AvroReuseIterator<D> implements Iterator<D>, Closeable {
+ private static class AvroReuseIterator<D> implements CloseableIterator<D> {
private final FileReader<D> reader;
private D reused = null;
diff --git a/core/src/main/java/org/apache/iceberg/util/ParallelIterable.java b/core/src/main/java/org/apache/iceberg/util/ParallelIterable.java
index b77338a..f9b01b6 100644
--- a/core/src/main/java/org/apache/iceberg/util/ParallelIterable.java
+++ b/core/src/main/java/org/apache/iceberg/util/ParallelIterable.java
@@ -31,6 +31,7 @@ import java.util.concurrent.Future;
import org.apache.iceberg.exceptions.RuntimeIOException;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
public class ParallelIterable<T> extends CloseableGroup implements CloseableIterable<T> {
private final Iterable<? extends Iterable<T>> iterables;
@@ -43,13 +44,13 @@ public class ParallelIterable<T> extends CloseableGroup implements CloseableIter
}
@Override
- public Iterator<T> iterator() {
+ public CloseableIterator<T> iterator() {
ParallelIterator<T> iter = new ParallelIterator<>(iterables, workerPool);
addCloseable(iter);
return iter;
}
- private static class ParallelIterator<T> implements Iterator<T>, Closeable {
+ private static class ParallelIterator<T> implements CloseableIterator<T> {
private final Iterator<Runnable> tasks;
private final ExecutorService workerPool;
private final Future<?>[] taskFutures;
diff --git a/data/src/main/java/org/apache/iceberg/data/TableScanIterable.java b/data/src/main/java/org/apache/iceberg/data/TableScanIterable.java
index a683c01..eec0dde 100644
--- a/data/src/main/java/org/apache/iceberg/data/TableScanIterable.java
+++ b/data/src/main/java/org/apache/iceberg/data/TableScanIterable.java
@@ -44,6 +44,7 @@ import org.apache.iceberg.expressions.Evaluator;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.orc.ORC;
import org.apache.iceberg.parquet.Parquet;
@@ -71,7 +72,7 @@ class TableScanIterable extends CloseableGroup implements CloseableIterable<Reco
}
@Override
- public Iterator<Record> iterator() {
+ public CloseableIterator<Record> iterator() {
ScanIterator iter = new ScanIterator(tasks, caseSensitive);
addCloseable(iter);
return iter;
@@ -128,7 +129,7 @@ class TableScanIterable extends CloseableGroup implements CloseableIterable<Reco
super.close(); // close data files
}
- private class ScanIterator implements Iterator<Record>, Closeable {
+ private class ScanIterator implements CloseableIterator<Record> {
private final Iterator<FileScanTask> tasks;
private final boolean caseSensitive;
private Closeable currentCloseable = null;
diff --git a/mr/src/main/java/org/apache/iceberg/mr/mapreduce/IcebergInputFormat.java b/mr/src/main/java/org/apache/iceberg/mr/mapreduce/IcebergInputFormat.java
index 0d35644..58752d3 100644
--- a/mr/src/main/java/org/apache/iceberg/mr/mapreduce/IcebergInputFormat.java
+++ b/mr/src/main/java/org/apache/iceberg/mr/mapreduce/IcebergInputFormat.java
@@ -20,11 +20,9 @@
package org.apache.iceberg.mr.mapreduce;
import com.google.common.base.Preconditions;
-import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
-import java.io.Closeable;
import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
@@ -69,6 +67,7 @@ import org.apache.iceberg.hadoop.HadoopInputFile;
import org.apache.iceberg.hadoop.HadoopTables;
import org.apache.iceberg.hadoop.Util;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.mr.SerializationUtil;
import org.apache.iceberg.orc.ORC;
@@ -288,8 +287,7 @@ public class IcebergInputFormat<T> extends InputFormat<Void, T> {
private Map<String, Integer> namesToPos;
private Iterator<FileScanTask> tasks;
private T currentRow;
- private Iterator<T> currentIterator;
- private Closeable currentCloseable;
+ private CloseableIterator<T> currentIterator;
@Override
public void initialize(InputSplit split, TaskAttemptContext newContext) {
@@ -315,10 +313,10 @@ public class IcebergInputFormat<T> extends InputFormat<Void, T> {
currentRow = currentIterator.next();
return true;
} else if (tasks.hasNext()) {
- currentCloseable.close();
+ currentIterator.close();
currentIterator = open(tasks.next());
} else {
- currentCloseable.close();
+ currentIterator.close();
return false;
}
}
@@ -347,7 +345,7 @@ public class IcebergInputFormat<T> extends InputFormat<Void, T> {
@Override
public void close() throws IOException {
- currentCloseable.close();
+ currentIterator.close();
}
private static Map<String, Integer> buildNameToPos(Schema expectedSchema) {
@@ -359,25 +357,27 @@ public class IcebergInputFormat<T> extends InputFormat<Void, T> {
return nameToPos;
}
- private Iterator<T> open(FileScanTask currentTask) {
+ private CloseableIterator<T> open(FileScanTask currentTask) {
DataFile file = currentTask.file();
// schema of rows returned by readers
PartitionSpec spec = currentTask.spec();
Set<Integer> idColumns = Sets.intersection(spec.identitySourceIds(), TypeUtil.getProjectedIds(expectedSchema));
boolean hasJoinedPartitionColumns = !idColumns.isEmpty();
+ CloseableIterable<T> iterable;
if (hasJoinedPartitionColumns) {
Schema readDataSchema = TypeUtil.selectNot(expectedSchema, idColumns);
Schema identityPartitionSchema = TypeUtil.select(expectedSchema, idColumns);
- return Iterators.transform(
- open(currentTask, readDataSchema),
+ iterable = CloseableIterable.transform(open(currentTask, readDataSchema),
row -> withIdentityPartitionColumns(row, identityPartitionSchema, spec, file.partition()));
} else {
- return open(currentTask, expectedSchema);
+ iterable = open(currentTask, expectedSchema);
}
+
+ return iterable.iterator();
}
- private Iterator<T> open(FileScanTask currentTask, Schema readSchema) {
+ private CloseableIterable<T> open(FileScanTask currentTask, Schema readSchema) {
DataFile file = currentTask.file();
// TODO we should make use of FileIO to create inputFile
InputFile inputFile = HadoopInputFile.fromLocation(file.path(), context.getConfiguration());
@@ -396,8 +396,8 @@ public class IcebergInputFormat<T> extends InputFormat<Void, T> {
throw new UnsupportedOperationException(
String.format("Cannot read %s file: %s", file.format().name(), file.path()));
}
- currentCloseable = iterable;
- return iterable.iterator();
+
+ return iterable;
}
@SuppressWarnings("unchecked")
diff --git a/orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java b/orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java
index 814c5bc..77eb24c 100644
--- a/orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java
+++ b/orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java
@@ -27,6 +27,7 @@ import org.apache.iceberg.Schema;
import org.apache.iceberg.exceptions.RuntimeIOException;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.InputFile;
import org.apache.orc.Reader;
import org.apache.orc.TypeDescription;
@@ -56,14 +57,15 @@ class OrcIterable<T> extends CloseableGroup implements CloseableIterable<T> {
@SuppressWarnings("unchecked")
@Override
- public Iterator<T> iterator() {
+ public CloseableIterator<T> iterator() {
Reader orcFileReader = ORC.newFileReader(file, config);
addCloseable(orcFileReader);
TypeDescription readOrcSchema = ORCSchemaUtil.buildOrcProjection(schema, orcFileReader.getSchema());
- return new OrcIterator(
+ Iterator<T> iterator = new OrcIterator(
newOrcIterator(file, readOrcSchema, start, length, orcFileReader),
readerFunction.apply(readOrcSchema));
+ return CloseableIterator.withClose(iterator);
}
private static VectorizedRowBatchIterator newOrcIterator(InputFile file,
diff --git a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetIterable.java b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetIterable.java
index 79d702e..5ec5b59 100644
--- a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetIterable.java
+++ b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetIterable.java
@@ -19,13 +19,12 @@
package org.apache.iceberg.parquet;
-import java.io.Closeable;
import java.io.IOException;
-import java.util.Iterator;
import java.util.NoSuchElementException;
import org.apache.iceberg.exceptions.RuntimeIOException;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.parquet.hadoop.ParquetReader;
public class ParquetIterable<T> extends CloseableGroup implements CloseableIterable<T> {
@@ -36,7 +35,7 @@ public class ParquetIterable<T> extends CloseableGroup implements CloseableItera
}
@Override
- public Iterator<T> iterator() {
+ public CloseableIterator<T> iterator() {
try {
ParquetReader<T> reader = builder.build();
addCloseable(reader);
@@ -46,7 +45,7 @@ public class ParquetIterable<T> extends CloseableGroup implements CloseableItera
}
}
- private static class ParquetIterator<T> implements Iterator<T>, Closeable {
+ private static class ParquetIterator<T> implements CloseableIterator<T> {
private final ParquetReader<T> parquet;
private boolean needsAdvance = false;
private boolean hasNext = false;
diff --git a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetReader.java b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetReader.java
index 9ba36b3..17d96c1 100644
--- a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetReader.java
+++ b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetReader.java
@@ -19,9 +19,7 @@
package org.apache.iceberg.parquet;
-import java.io.Closeable;
import java.io.IOException;
-import java.util.Iterator;
import java.util.function.Function;
import org.apache.iceberg.Schema;
import org.apache.iceberg.exceptions.RuntimeIOException;
@@ -29,6 +27,7 @@ import org.apache.iceberg.expressions.Expression;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.InputFile;
import org.apache.parquet.ParquetReadOptions;
import org.apache.parquet.column.page.PageReadStore;
@@ -70,13 +69,13 @@ public class ParquetReader<T> extends CloseableGroup implements CloseableIterabl
}
@Override
- public Iterator<T> iterator() {
+ public CloseableIterator<T> iterator() {
FileIterator<T> iter = new FileIterator<>(init());
addCloseable(iter);
return iter;
}
- private static class FileIterator<T> implements Iterator<T>, Closeable {
+ private static class FileIterator<T> implements CloseableIterator<T> {
private final ParquetFileReader reader;
private final boolean[] shouldSkip;
private final ParquetValueReader<T> model;
diff --git a/parquet/src/main/java/org/apache/iceberg/parquet/VectorizedParquetReader.java b/parquet/src/main/java/org/apache/iceberg/parquet/VectorizedParquetReader.java
index 077464b..c3f87ee 100644
--- a/parquet/src/main/java/org/apache/iceberg/parquet/VectorizedParquetReader.java
+++ b/parquet/src/main/java/org/apache/iceberg/parquet/VectorizedParquetReader.java
@@ -19,9 +19,7 @@
package org.apache.iceberg.parquet;
-import java.io.Closeable;
import java.io.IOException;
-import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
@@ -32,6 +30,7 @@ import org.apache.iceberg.expressions.Expression;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.io.CloseableGroup;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.InputFile;
import org.apache.parquet.ParquetReadOptions;
import org.apache.parquet.column.page.PageReadStore;
@@ -78,13 +77,13 @@ public class VectorizedParquetReader<T> extends CloseableGroup implements Closea
}
@Override
- public Iterator<T> iterator() {
+ public CloseableIterator<T> iterator() {
FileIterator<T> iter = new FileIterator<>(init());
addCloseable(iter);
return iter;
}
- private static class FileIterator<T> implements Iterator<T>, Closeable {
+ private static class FileIterator<T> implements CloseableIterator<T> {
private final ParquetFileReader reader;
private final boolean[] shouldSkip;
private final VectorizedReader<T> model;
diff --git a/spark/src/main/java/org/apache/iceberg/spark/source/BaseDataReader.java b/spark/src/main/java/org/apache/iceberg/spark/source/BaseDataReader.java
index fab6680..d481ab4 100644
--- a/spark/src/main/java/org/apache/iceberg/spark/source/BaseDataReader.java
+++ b/spark/src/main/java/org/apache/iceberg/spark/source/BaseDataReader.java
@@ -22,16 +22,14 @@ package org.apache.iceberg.spark.source;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Iterables;
-import java.io.Closeable;
import java.io.IOException;
-import java.util.Collections;
import java.util.Iterator;
import java.util.Map;
import org.apache.iceberg.CombinedScanTask;
import org.apache.iceberg.FileScanTask;
import org.apache.iceberg.encryption.EncryptedFiles;
import org.apache.iceberg.encryption.EncryptionManager;
-import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.FileIO;
import org.apache.iceberg.io.InputFile;
import org.apache.spark.rdd.InputFileBlockHolder;
@@ -42,14 +40,12 @@ import org.apache.spark.sql.sources.v2.reader.InputPartitionReader;
*
* @param <T> is the Java class returned by this reader whose objects contain one or more rows.
*/
-@SuppressWarnings("checkstyle:VisibilityModifier")
abstract class BaseDataReader<T> implements InputPartitionReader<T> {
private final Iterator<FileScanTask> tasks;
private final FileIO fileIo;
private final Map<String, InputFile> inputFiles;
- private Iterator<T> currentIterator;
- Closeable currentCloseable;
+ private CloseableIterator<T> currentIterator;
private T current = null;
BaseDataReader(CombinedScanTask task, FileIO fileIo, EncryptionManager encryptionManager) {
@@ -64,8 +60,7 @@ abstract class BaseDataReader<T> implements InputPartitionReader<T> {
ImmutableMap.Builder<String, InputFile> inputFileBuilder = ImmutableMap.builder();
decryptedFiles.forEach(decrypted -> inputFileBuilder.put(decrypted.location(), decrypted));
this.inputFiles = inputFileBuilder.build();
- this.currentCloseable = CloseableIterable.empty();
- this.currentIterator = Collections.emptyIterator();
+ this.currentIterator = CloseableIterator.empty();
}
@Override
@@ -75,7 +70,7 @@ abstract class BaseDataReader<T> implements InputPartitionReader<T> {
this.current = currentIterator.next();
return true;
} else if (tasks.hasNext()) {
- this.currentCloseable.close();
+ this.currentIterator.close();
this.currentIterator = open(tasks.next());
} else {
return false;
@@ -88,14 +83,14 @@ abstract class BaseDataReader<T> implements InputPartitionReader<T> {
return current;
}
- abstract Iterator<T> open(FileScanTask task);
+ abstract CloseableIterator<T> open(FileScanTask task);
@Override
public void close() throws IOException {
InputFileBlockHolder.unset();
// close the current iterator
- this.currentCloseable.close();
+ this.currentIterator.close();
// exhaust the task iterator
while (tasks.hasNext()) {
diff --git a/spark/src/main/java/org/apache/iceberg/spark/source/RowDataReader.java b/spark/src/main/java/org/apache/iceberg/spark/source/RowDataReader.java
index 8e240a1..f4ea6c7 100644
--- a/spark/src/main/java/org/apache/iceberg/spark/source/RowDataReader.java
+++ b/spark/src/main/java/org/apache/iceberg/spark/source/RowDataReader.java
@@ -21,12 +21,10 @@ package org.apache.iceberg.spark.source;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
-import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -43,6 +41,7 @@ import org.apache.iceberg.avro.Avro;
import org.apache.iceberg.common.DynMethods;
import org.apache.iceberg.encryption.EncryptionManager;
import org.apache.iceberg.io.CloseableIterable;
+import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.io.FileIO;
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.orc.ORC;
@@ -88,7 +87,7 @@ class RowDataReader extends BaseDataReader<InternalRow> {
}
@Override
- Iterator<InternalRow> open(FileScanTask task) {
+ CloseableIterator<InternalRow> open(FileScanTask task) {
DataFile file = task.file();
// update the current file for Spark's filename() function
@@ -102,7 +101,8 @@ class RowDataReader extends BaseDataReader<InternalRow> {
if (projectsIdentityPartitionColumns) {
if (SUPPORTS_CONSTANTS.contains(file.format())) {
- return open(task, expectedSchema, PartitionUtil.constantsMap(task, RowDataReader::convertConstant));
+ return open(task, expectedSchema, PartitionUtil.constantsMap(task, RowDataReader::convertConstant))
+ .iterator();
}
// schema used to read data files
@@ -115,16 +115,17 @@ class RowDataReader extends BaseDataReader<InternalRow> {
InternalRow partition = convertToRow.apply(file.partition());
joined.withRight(partition);
- return Iterators.transform(
- Iterators.transform(open(task, readSchema, ImmutableMap.of()), joined::withLeft),
+ CloseableIterable<InternalRow> transformedIterable = CloseableIterable.transform(
+ CloseableIterable.transform(open(task, readSchema, ImmutableMap.of()), joined::withLeft),
APPLY_PROJECTION.bind(projection(expectedSchema, joinedSchema))::invoke);
+ return transformedIterable.iterator();
}
// return the base iterator
- return open(task, expectedSchema, ImmutableMap.of());
+ return open(task, expectedSchema, ImmutableMap.of()).iterator();
}
- private Iterator<InternalRow> open(FileScanTask task, Schema readSchema, Map<Integer, ?> idToConstant) {
+ private CloseableIterable<InternalRow> open(FileScanTask task, Schema readSchema, Map<Integer, ?> idToConstant) {
CloseableIterable<InternalRow> iter;
if (task.isDataTask()) {
iter = newDataIterable(task.asDataTask(), readSchema);
@@ -151,9 +152,7 @@ class RowDataReader extends BaseDataReader<InternalRow> {
}
}
- this.currentCloseable = iter;
-
- return iter.iterator();
+ return iter;
}
private CloseableIterable<InternalRow> newAvroIterable(