chamikaramj commented on a change in pull request #12223:
URL: https://github.com/apache/beam/pull/12223#discussion_r461190881
##########
File path:
sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
##########
@@ -235,12 +284,151 @@ public ReadFiles withAvroDataModel(GenericData model) {
return toBuilder().setAvroDataModel(model).build();
}
+ public ReadFiles withSplit() {
+ return toBuilder().setSplit(true).build();
+ }
+
@Override
public PCollection<GenericRecord> expand(PCollection<FileIO.ReadableFile>
input) {
checkNotNull(getSchema(), "Schema can not be null");
- return input
- .apply(ParDo.of(new ReadFn(getAvroDataModel())))
- .setCoder(AvroCoder.of(getSchema()));
+ if (getSplit() != null) {
+ return input
+ .apply(ParDo.of(new SplitReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ } else {
+ return input
+ .apply(ParDo.of(new ReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ }
+ }
+
+ @DoFn.BoundedPerElement
+ static class SplitReadFn extends DoFn<FileIO.ReadableFile, GenericRecord> {
+ private Class<? extends GenericData> modelClass;
+ private static final Logger LOG =
LoggerFactory.getLogger(SplitReadFn.class);
+ ReadSupport<GenericRecord> readSupport;
+
+ SplitReadFn(GenericData model) {
+ this.modelClass = model != null ? model.getClass() : null;
+ }
+
+ private static <K, V> Map<K, Set<V>> toSetMultiMap(Map<K, V> map) {
+ Map<K, Set<V>> setMultiMap = new HashMap<K, Set<V>>();
+ for (Map.Entry<K, V> entry : map.entrySet()) {
+ Set<V> set = new HashSet<V>();
+ set.add(entry.getValue());
+ setMultiMap.put(entry.getKey(), Collections.unmodifiableSet(set));
+ }
+ return Collections.unmodifiableMap(setMultiMap);
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element FileIO.ReadableFile file,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<GenericRecord> outputReceiver)
+ throws Exception {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ ResourceId filename = file.getMetadata().resourceId();
+ throw new RuntimeException(String.format("File has to be seekable:
%s", filename));
+ }
+
+ SeekableByteChannel seekableByteChannel = file.openSeekable();
+ ReadSupport<GenericRecord> readSupport;
+ InputFile inputFile = new BeamParquetInputFile(seekableByteChannel);
+ Configuration conf = setConf();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ readSupport = new AvroReadSupport<GenericRecord>(model);
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ Filter filter = checkNotNull(options.getRecordFilter(), "filter");
+ conf = ((HadoopReadOptions) options).getConf();
+ for (String property : options.getPropertyNames()) {
+ conf.set(property, options.getProperty(property));
+ }
+ FileMetaData parquetFileMetadata =
reader.getFooter().getFileMetaData();
+ MessageType fileSchema = parquetFileMetadata.getSchema();
+ Map<String, String> fileMetadata =
parquetFileMetadata.getKeyValueMetaData();
+
+ ReadSupport.ReadContext readContext =
+ readSupport.init(new InitContext(conf,
toSetMultiMap(fileMetadata), fileSchema));
+ ColumnIOFactory columnIOFactory = new
ColumnIOFactory(parquetFileMetadata.getCreatedBy());
+ MessageType requestedSchema = readContext.getRequestedSchema();
+ RecordMaterializer<GenericRecord> recordConverter =
+ readSupport.prepareForRead(conf, fileMetadata, fileSchema,
readContext);
+ boolean strictTypeChecking = options.isEnabled(STRICT_TYPE_CHECKING,
true);
+ boolean filterRecords = options.useRecordFilter();
+ reader.setRequestedSchema(requestedSchema);
+ MessageColumnIO columnIO =
+ columnIOFactory.getColumnIO(requestedSchema, fileSchema,
strictTypeChecking);
+ for (int i = 0; i < reader.getRowGroups().size(); i++) {
+ if (i < tracker.currentRestriction().getFrom()) {
+ reader.skipNextRowGroup();
+ continue;
+ }
+ if (tracker.tryClaim((long) i)) {
+ PageReadStore pages = reader.readNextRowGroup();
+ i += 1;
+ RecordReader<GenericRecord> recordReader =
+ columnIO.getRecordReader(
+ pages, recordConverter, filterRecords ? filter :
FilterCompat.NOOP);
+ GenericRecord read;
+ long current = 0;
+ long totalRows = pages.getRowCount();
+ while (current < totalRows) {
+ read = recordReader.read();
+ outputReceiver.output(read);
+ current += 1;
+ }
+ } else {
+ break;
+ }
+ }
+ }
+
+ private Configuration setConf() throws Exception {
+ Configuration conf = new Configuration();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ if (model != null
+ && (model.getClass() == GenericData.class || model.getClass() ==
SpecificData.class)) {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, true);
+ } else {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, false);
+ }
+ return conf;
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element FileIO.ReadableFile
file) throws Exception {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ ResourceId filename = file.getMetadata().resourceId();
+ throw new RuntimeException(String.format("File has to be seekable:
%s", filename));
+ }
+ SeekableByteChannel seekableByteChannel = file.openSeekable();
+ InputFile inputFile = new BeamParquetInputFile(seekableByteChannel);
+ Configuration conf = setConf();
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ return new OffsetRange(0, reader.getRowGroups().size());
+ }
+
+ @SplitRestriction
+ public void split(@Restriction OffsetRange restriction,
OutputReceiver<OffsetRange> out) {
Review comment:
Usual pattern for files is to initial split into 64MB blocks and dynamic
split further as needed. I believe this is the tried and tested method for
existing file-based sources. This might require changing to a tracker that use
byte positions (0 to size of file) instead of a tracker that goes from 0 to
number of blocks.
##########
File path:
sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
##########
@@ -235,12 +277,195 @@ public ReadFiles withAvroDataModel(GenericData model) {
return toBuilder().setAvroDataModel(model).build();
}
+ public ReadFiles withSplit() {
+ return toBuilder().setSplit(true).build();
+ }
+
@Override
public PCollection<GenericRecord> expand(PCollection<FileIO.ReadableFile>
input) {
checkNotNull(getSchema(), "Schema can not be null");
- return input
- .apply(ParDo.of(new ReadFn(getAvroDataModel())))
- .setCoder(AvroCoder.of(getSchema()));
+ if (!getSplit()) {
+ return input
+ .apply(ParDo.of(new SplitReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ } else {
+ return input
+ .apply(ParDo.of(new ReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ }
+ }
+
+ @DoFn.BoundedPerElement
+ static class SplitReadFn extends DoFn<FileIO.ReadableFile, GenericRecord> {
+ private Class<? extends GenericData> modelClass;
+ private static final Logger LOG =
LoggerFactory.getLogger(SplitReadFn.class);
+ ReadSupport<GenericRecord> readSupport;
+
+ SplitReadFn(GenericData model) {
+ this.modelClass = model != null ? model.getClass() : null;
+ }
+
+ private static <K, V> Map<K, Set<V>> toSetMultiMap(Map<K, V> map) {
+ Map<K, Set<V>> setMultiMap = new HashMap<K, Set<V>>();
+ for (Map.Entry<K, V> entry : map.entrySet()) {
+ Set<V> set = new HashSet<V>();
+ set.add(entry.getValue());
+ setMultiMap.put(entry.getKey(), Collections.unmodifiableSet(set));
+ }
+ return Collections.unmodifiableMap(setMultiMap);
+ }
+
+ private InputFile getInputFile(FileIO.ReadableFile file) throws
IOException {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ throw new RuntimeException(
+ String.format("File has to be seekable: %s",
file.getMetadata().resourceId()));
+ }
+ return new BeamParquetInputFile(file.openSeekable());
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element FileIO.ReadableFile file,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<GenericRecord> outputReceiver)
+ throws Exception {
+ ReadSupport<GenericRecord> readSupport;
+ InputFile inputFile = getInputFile(file);
+ Configuration conf = setConf();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ readSupport = new AvroReadSupport<GenericRecord>(model);
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ Filter filter = checkNotNull(options.getRecordFilter(), "filter");
+ conf = ((HadoopReadOptions) options).getConf();
+ for (String property : options.getPropertyNames()) {
+ conf.set(property, options.getProperty(property));
+ }
+ FileMetaData parquetFileMetadata =
reader.getFooter().getFileMetaData();
+ MessageType fileSchema = parquetFileMetadata.getSchema();
+ Map<String, String> fileMetadata =
parquetFileMetadata.getKeyValueMetaData();
+
+ ReadSupport.ReadContext readContext =
+ readSupport.init(new InitContext(conf,
toSetMultiMap(fileMetadata), fileSchema));
+ ColumnIOFactory columnIOFactory = new
ColumnIOFactory(parquetFileMetadata.getCreatedBy());
+ MessageType requestedSchema = readContext.getRequestedSchema();
+ RecordMaterializer<GenericRecord> recordConverter =
+ readSupport.prepareForRead(conf, fileMetadata, fileSchema,
readContext);
+ boolean strictTypeChecking = options.isEnabled(STRICT_TYPE_CHECKING,
true);
+ boolean filterRecords = options.useRecordFilter();
+ reader.setRequestedSchema(requestedSchema);
+ MessageColumnIO columnIO =
+ columnIOFactory.getColumnIO(requestedSchema, fileSchema,
strictTypeChecking);
+ long currentBlock = tracker.currentRestriction().getFrom();
+ for (int i = 0; i < currentBlock; i++) {
+ reader.skipNextRowGroup();
+ }
+ while (tracker.tryClaim(currentBlock)) {
+
+ LOG.info("reading block" + currentBlock);
+ PageReadStore pages = reader.readNextRowGroup();
+ currentBlock += 1;
+ RecordReader<GenericRecord> recordReader =
+ columnIO.getRecordReader(
+ pages, recordConverter, filterRecords ? filter :
FilterCompat.NOOP);
+ long currentRow = 0;
+ long totalRows = pages.getRowCount();
+ while (currentRow < totalRows) {
+ try {
+ GenericRecord record;
+ currentRow += 1;
+ try {
+ record = recordReader.read();
+ } catch (RecordMaterializer.RecordMaterializationException e) {
+ LOG.debug("skipping a corrupt record");
+ continue;
+ }
+ if (recordReader.shouldSkipCurrentRecord()) {
+ // this record is being filtered via the filter2 package
+ LOG.debug("skipping record");
+ continue;
+ }
+ if (record == null) {
+ // only happens with FilteredRecordReader at end of block
+ LOG.debug("filtered record reader reached end of block");
+ break;
+ }
+ outputReceiver.output(record);
+ } catch (RuntimeException e) {
+ throw new ParquetDecodingException(format("Can not read value at
%d in block %d in file %s", currentRow, currentBlock, file.toString()), e);
+ }
+ }
+ LOG.info("finish read " + currentRow + " rows");
+ }
+
+
+
+ }
+
+ private Configuration setConf() throws Exception {
+ Configuration conf = new Configuration();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ if (model != null
+ && (model.getClass() == GenericData.class || model.getClass() ==
SpecificData.class)) {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, true);
+ } else {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, false);
+ }
+ return conf;
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element FileIO.ReadableFile
file) throws Exception {
+ InputFile inputFile = getInputFile(file);
+ Configuration conf = setConf();
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ return new OffsetRange(0, reader.getRowGroups().size());
Review comment:
Does this mean that we will be incrementing progress (obtained by the
runner through restriction tracker) only when a new block is read ? This might
be too coarse grained.
##########
File path:
sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
##########
@@ -235,12 +283,164 @@ public ReadFiles withAvroDataModel(GenericData model) {
return toBuilder().setAvroDataModel(model).build();
}
+ public ReadFiles withSplit() {
+ return toBuilder().setSplit(true).build();
+ }
+
@Override
public PCollection<GenericRecord> expand(PCollection<FileIO.ReadableFile>
input) {
checkNotNull(getSchema(), "Schema can not be null");
- return input
- .apply(ParDo.of(new ReadFn(getAvroDataModel())))
- .setCoder(AvroCoder.of(getSchema()));
+ if (!getSplit()) {
+ return input
+ .apply(ParDo.of(new SplitReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ } else {
+ return input
+ .apply(ParDo.of(new ReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ }
+ }
+
+ @DoFn.BoundedPerElement
+ static class SplitReadFn extends DoFn<FileIO.ReadableFile, GenericRecord> {
+ private Class<? extends GenericData> modelClass;
+ private static final Logger LOG =
LoggerFactory.getLogger(SplitReadFn.class);
+ ReadSupport<GenericRecord> readSupport;
+
+ SplitReadFn(GenericData model) {
+ this.modelClass = model != null ? model.getClass() : null;
+ }
+
+ private static <K, V> Map<K, Set<V>> toSetMultiMap(Map<K, V> map) {
+ Map<K, Set<V>> setMultiMap = new HashMap<K, Set<V>>();
+ for (Map.Entry<K, V> entry : map.entrySet()) {
+ Set<V> set = new HashSet<V>();
+ set.add(entry.getValue());
+ setMultiMap.put(entry.getKey(), Collections.unmodifiableSet(set));
+ }
+ return Collections.unmodifiableMap(setMultiMap);
+ }
+
+ private InputFile getInputFile(FileIO.ReadableFile file) throws
IOException {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ throw new RuntimeException(
+ String.format("File has to be seekable: %s",
file.getMetadata().resourceId()));
+ }
+ return new BeamParquetInputFile(file.openSeekable());
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element FileIO.ReadableFile file,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<GenericRecord> outputReceiver)
+ throws Exception {
+ ReadSupport<GenericRecord> readSupport;
+ InputFile inputFile = getInputFile(file);
+ Configuration conf = setConf();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ readSupport = new AvroReadSupport<GenericRecord>(model);
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ Filter filter = checkNotNull(options.getRecordFilter(), "filter");
+ conf = ((HadoopReadOptions) options).getConf();
+ for (String property : options.getPropertyNames()) {
+ conf.set(property, options.getProperty(property));
+ }
+ FileMetaData parquetFileMetadata =
reader.getFooter().getFileMetaData();
+ MessageType fileSchema = parquetFileMetadata.getSchema();
+ Map<String, String> fileMetadata =
parquetFileMetadata.getKeyValueMetaData();
+
+ ReadSupport.ReadContext readContext =
+ readSupport.init(new InitContext(conf,
toSetMultiMap(fileMetadata), fileSchema));
+ ColumnIOFactory columnIOFactory = new
ColumnIOFactory(parquetFileMetadata.getCreatedBy());
+ MessageType requestedSchema = readContext.getRequestedSchema();
+ RecordMaterializer<GenericRecord> recordConverter =
+ readSupport.prepareForRead(conf, fileMetadata, fileSchema,
readContext);
+ boolean strictTypeChecking = options.isEnabled(STRICT_TYPE_CHECKING,
true);
+ boolean filterRecords = options.useRecordFilter();
+ reader.setRequestedSchema(requestedSchema);
+ MessageColumnIO columnIO =
+ columnIOFactory.getColumnIO(requestedSchema, fileSchema,
strictTypeChecking);
+ long currentBlock = tracker.currentRestriction().getFrom();
+ for (int i = 0; i < currentBlock; i++) {
+ reader.skipNextRowGroup();
+ }
+ while (tracker.tryClaim(currentBlock)) {
+ PageReadStore pages = reader.readNextRowGroup();
+ currentBlock += 1;
+ RecordReader<GenericRecord> recordReader =
+ columnIO.getRecordReader(
+ pages, recordConverter, filterRecords ? filter :
FilterCompat.NOOP);
+ GenericRecord read;
+ long currentRow = 0;
+ long totalRows = pages.getRowCount();
+ while (currentRow < totalRows) {
+ outputReceiver.output(recordReader.read());
+ currentRow += 1;
+ }
+ }
+ }
+
+ private Configuration setConf() throws Exception {
+ Configuration conf = new Configuration();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ if (model != null
+ && (model.getClass() == GenericData.class || model.getClass() ==
SpecificData.class)) {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, true);
+ } else {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, false);
+ }
+ return conf;
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element FileIO.ReadableFile
file) throws Exception {
+ InputFile inputFile = getInputFile(file);
+ Configuration conf = setConf();
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ return new OffsetRange(0, reader.getRowGroups().size());
+ }
+
+ @SplitRestriction
+ public void split(@Restriction OffsetRange restriction,
OutputReceiver<OffsetRange> out) {
+ for (OffsetRange range : restriction.split(1, 0)) {
+ out.output(range);
+ }
+ }
+
+ @NewTracker
+ public OffsetRangeTracker newTracker(@Restriction OffsetRange
restriction) {
+ return new OffsetRangeTracker(restriction);
+ }
+
+ @GetSize
+ public double getSize(@Element FileIO.ReadableFile file, @Restriction
OffsetRange restriction)
+ throws Exception {
+ InputFile inputFile = getInputFile(file);
+ Configuration conf = setConf();
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ if (restriction == null) {
+ return reader.getRecordCount();
+ } else {
+ long start = restriction.getFrom();
+ long end = restriction.getTo();
+ List<BlockMetaData> blocks = reader.getRowGroups();
+ double size = 0;
+ for (long i = start; i < end; i++) {
+ size += blocks.get((int) i).getRowCount();
+ }
+ return size;
Review comment:
If you use byte sizes for tracking implementation here will end up being
trivial.
##########
File path:
sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
##########
@@ -235,12 +284,151 @@ public ReadFiles withAvroDataModel(GenericData model) {
return toBuilder().setAvroDataModel(model).build();
}
+ public ReadFiles withSplit() {
+ return toBuilder().setSplit(true).build();
+ }
+
@Override
public PCollection<GenericRecord> expand(PCollection<FileIO.ReadableFile>
input) {
checkNotNull(getSchema(), "Schema can not be null");
- return input
- .apply(ParDo.of(new ReadFn(getAvroDataModel())))
- .setCoder(AvroCoder.of(getSchema()));
+ if (getSplit() != null) {
+ return input
+ .apply(ParDo.of(new SplitReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ } else {
+ return input
+ .apply(ParDo.of(new ReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ }
+ }
+
+ @DoFn.BoundedPerElement
+ static class SplitReadFn extends DoFn<FileIO.ReadableFile, GenericRecord> {
+ private Class<? extends GenericData> modelClass;
+ private static final Logger LOG =
LoggerFactory.getLogger(SplitReadFn.class);
+ ReadSupport<GenericRecord> readSupport;
+
+ SplitReadFn(GenericData model) {
+ this.modelClass = model != null ? model.getClass() : null;
+ }
+
+ private static <K, V> Map<K, Set<V>> toSetMultiMap(Map<K, V> map) {
+ Map<K, Set<V>> setMultiMap = new HashMap<K, Set<V>>();
+ for (Map.Entry<K, V> entry : map.entrySet()) {
+ Set<V> set = new HashSet<V>();
+ set.add(entry.getValue());
+ setMultiMap.put(entry.getKey(), Collections.unmodifiableSet(set));
+ }
+ return Collections.unmodifiableMap(setMultiMap);
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element FileIO.ReadableFile file,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<GenericRecord> outputReceiver)
+ throws Exception {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ ResourceId filename = file.getMetadata().resourceId();
+ throw new RuntimeException(String.format("File has to be seekable:
%s", filename));
+ }
+
+ SeekableByteChannel seekableByteChannel = file.openSeekable();
+ ReadSupport<GenericRecord> readSupport;
+ InputFile inputFile = new BeamParquetInputFile(seekableByteChannel);
+ Configuration conf = setConf();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ readSupport = new AvroReadSupport<GenericRecord>(model);
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ Filter filter = checkNotNull(options.getRecordFilter(), "filter");
+ conf = ((HadoopReadOptions) options).getConf();
+ for (String property : options.getPropertyNames()) {
+ conf.set(property, options.getProperty(property));
+ }
+ FileMetaData parquetFileMetadata =
reader.getFooter().getFileMetaData();
+ MessageType fileSchema = parquetFileMetadata.getSchema();
+ Map<String, String> fileMetadata =
parquetFileMetadata.getKeyValueMetaData();
+
+ ReadSupport.ReadContext readContext =
+ readSupport.init(new InitContext(conf,
toSetMultiMap(fileMetadata), fileSchema));
+ ColumnIOFactory columnIOFactory = new
ColumnIOFactory(parquetFileMetadata.getCreatedBy());
+ MessageType requestedSchema = readContext.getRequestedSchema();
+ RecordMaterializer<GenericRecord> recordConverter =
+ readSupport.prepareForRead(conf, fileMetadata, fileSchema,
readContext);
+ boolean strictTypeChecking = options.isEnabled(STRICT_TYPE_CHECKING,
true);
+ boolean filterRecords = options.useRecordFilter();
+ reader.setRequestedSchema(requestedSchema);
+ MessageColumnIO columnIO =
+ columnIOFactory.getColumnIO(requestedSchema, fileSchema,
strictTypeChecking);
+ for (int i = 0; i < reader.getRowGroups().size(); i++) {
+ if (i < tracker.currentRestriction().getFrom()) {
+ reader.skipNextRowGroup();
+ continue;
+ }
+ if (tracker.tryClaim((long) i)) {
+ PageReadStore pages = reader.readNextRowGroup();
+ i += 1;
+ RecordReader<GenericRecord> recordReader =
+ columnIO.getRecordReader(
+ pages, recordConverter, filterRecords ? filter :
FilterCompat.NOOP);
+ GenericRecord read;
+ long current = 0;
+ long totalRows = pages.getRowCount();
+ while (current < totalRows) {
+ read = recordReader.read();
+ outputReceiver.output(read);
+ current += 1;
+ }
+ } else {
+ break;
+ }
+ }
+ }
+
+ private Configuration setConf() throws Exception {
+ Configuration conf = new Configuration();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ if (model != null
+ && (model.getClass() == GenericData.class || model.getClass() ==
SpecificData.class)) {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, true);
+ } else {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, false);
+ }
+ return conf;
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element FileIO.ReadableFile
file) throws Exception {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ ResourceId filename = file.getMetadata().resourceId();
+ throw new RuntimeException(String.format("File has to be seekable:
%s", filename));
+ }
+ SeekableByteChannel seekableByteChannel = file.openSeekable();
+ InputFile inputFile = new BeamParquetInputFile(seekableByteChannel);
+ Configuration conf = setConf();
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ return new OffsetRange(0, reader.getRowGroups().size());
+ }
+
+ @SplitRestriction
+ public void split(@Restriction OffsetRange restriction,
OutputReceiver<OffsetRange> out) {
Review comment:
@boyuanzz do you think we should add an abstract block-based SDF
implementation/utility that will be shared between multiple source
implementations (for example, Parquet, Avro) similar to [1] for bounded sources
?
[1]
https://github.com/apache/beam/blob/master/sdks/java/core/src/main/java/org/apache/beam/sdk/io/BlockBasedSource.java
##########
File path:
sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
##########
@@ -235,12 +284,151 @@ public ReadFiles withAvroDataModel(GenericData model) {
return toBuilder().setAvroDataModel(model).build();
}
+ public ReadFiles withSplit() {
+ return toBuilder().setSplit(true).build();
+ }
+
@Override
public PCollection<GenericRecord> expand(PCollection<FileIO.ReadableFile>
input) {
checkNotNull(getSchema(), "Schema can not be null");
- return input
- .apply(ParDo.of(new ReadFn(getAvroDataModel())))
- .setCoder(AvroCoder.of(getSchema()));
+ if (getSplit() != null) {
+ return input
+ .apply(ParDo.of(new SplitReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ } else {
+ return input
+ .apply(ParDo.of(new ReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ }
+ }
+
+ @DoFn.BoundedPerElement
+ static class SplitReadFn extends DoFn<FileIO.ReadableFile, GenericRecord> {
+ private Class<? extends GenericData> modelClass;
+ private static final Logger LOG =
LoggerFactory.getLogger(SplitReadFn.class);
+ ReadSupport<GenericRecord> readSupport;
+
+ SplitReadFn(GenericData model) {
+ this.modelClass = model != null ? model.getClass() : null;
+ }
+
+ private static <K, V> Map<K, Set<V>> toSetMultiMap(Map<K, V> map) {
+ Map<K, Set<V>> setMultiMap = new HashMap<K, Set<V>>();
+ for (Map.Entry<K, V> entry : map.entrySet()) {
+ Set<V> set = new HashSet<V>();
+ set.add(entry.getValue());
+ setMultiMap.put(entry.getKey(), Collections.unmodifiableSet(set));
+ }
+ return Collections.unmodifiableMap(setMultiMap);
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element FileIO.ReadableFile file,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<GenericRecord> outputReceiver)
+ throws Exception {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ ResourceId filename = file.getMetadata().resourceId();
+ throw new RuntimeException(String.format("File has to be seekable:
%s", filename));
+ }
+
+ SeekableByteChannel seekableByteChannel = file.openSeekable();
+ ReadSupport<GenericRecord> readSupport;
+ InputFile inputFile = new BeamParquetInputFile(seekableByteChannel);
+ Configuration conf = setConf();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ readSupport = new AvroReadSupport<GenericRecord>(model);
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ Filter filter = checkNotNull(options.getRecordFilter(), "filter");
+ conf = ((HadoopReadOptions) options).getConf();
+ for (String property : options.getPropertyNames()) {
+ conf.set(property, options.getProperty(property));
+ }
+ FileMetaData parquetFileMetadata =
reader.getFooter().getFileMetaData();
+ MessageType fileSchema = parquetFileMetadata.getSchema();
+ Map<String, String> fileMetadata =
parquetFileMetadata.getKeyValueMetaData();
+
+ ReadSupport.ReadContext readContext =
+ readSupport.init(new InitContext(conf,
toSetMultiMap(fileMetadata), fileSchema));
+ ColumnIOFactory columnIOFactory = new
ColumnIOFactory(parquetFileMetadata.getCreatedBy());
+ MessageType requestedSchema = readContext.getRequestedSchema();
+ RecordMaterializer<GenericRecord> recordConverter =
+ readSupport.prepareForRead(conf, fileMetadata, fileSchema,
readContext);
+ boolean strictTypeChecking = options.isEnabled(STRICT_TYPE_CHECKING,
true);
+ boolean filterRecords = options.useRecordFilter();
+ reader.setRequestedSchema(requestedSchema);
+ MessageColumnIO columnIO =
+ columnIOFactory.getColumnIO(requestedSchema, fileSchema,
strictTypeChecking);
+ for (int i = 0; i < reader.getRowGroups().size(); i++) {
+ if (i < tracker.currentRestriction().getFrom()) {
+ reader.skipNextRowGroup();
+ continue;
+ }
+ if (tracker.tryClaim((long) i)) {
+ PageReadStore pages = reader.readNextRowGroup();
+ i += 1;
+ RecordReader<GenericRecord> recordReader =
+ columnIO.getRecordReader(
+ pages, recordConverter, filterRecords ? filter :
FilterCompat.NOOP);
+ GenericRecord read;
+ long current = 0;
+ long totalRows = pages.getRowCount();
+ while (current < totalRows) {
+ read = recordReader.read();
+ outputReceiver.output(read);
+ current += 1;
+ }
+ } else {
+ break;
+ }
+ }
+ }
+
+ private Configuration setConf() throws Exception {
+ Configuration conf = new Configuration();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ if (model != null
+ && (model.getClass() == GenericData.class || model.getClass() ==
SpecificData.class)) {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, true);
+ } else {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, false);
+ }
+ return conf;
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element FileIO.ReadableFile
file) throws Exception {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ ResourceId filename = file.getMetadata().resourceId();
+ throw new RuntimeException(String.format("File has to be seekable:
%s", filename));
+ }
+ SeekableByteChannel seekableByteChannel = file.openSeekable();
+ InputFile inputFile = new BeamParquetInputFile(seekableByteChannel);
+ Configuration conf = setConf();
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ return new OffsetRange(0, reader.getRowGroups().size());
+ }
+
+ @SplitRestriction
+ public void split(@Restriction OffsetRange restriction,
OutputReceiver<OffsetRange> out) {
Review comment:
May be we should add a BlockBasedTracker that extends OffsetRangeTracker
that goes from 0 to size of file and can split at block boundaries but can
provide better progress signals (based on element estimates).
##########
File path:
sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java
##########
@@ -235,12 +284,151 @@ public ReadFiles withAvroDataModel(GenericData model) {
return toBuilder().setAvroDataModel(model).build();
}
+ public ReadFiles withSplit() {
+ return toBuilder().setSplit(true).build();
+ }
+
@Override
public PCollection<GenericRecord> expand(PCollection<FileIO.ReadableFile>
input) {
checkNotNull(getSchema(), "Schema can not be null");
- return input
- .apply(ParDo.of(new ReadFn(getAvroDataModel())))
- .setCoder(AvroCoder.of(getSchema()));
+ if (getSplit() != null) {
+ return input
+ .apply(ParDo.of(new SplitReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ } else {
+ return input
+ .apply(ParDo.of(new ReadFn(getAvroDataModel())))
+ .setCoder(AvroCoder.of(getSchema()));
+ }
+ }
+
+ @DoFn.BoundedPerElement
+ static class SplitReadFn extends DoFn<FileIO.ReadableFile, GenericRecord> {
+ private Class<? extends GenericData> modelClass;
+ private static final Logger LOG =
LoggerFactory.getLogger(SplitReadFn.class);
+ ReadSupport<GenericRecord> readSupport;
+
+ SplitReadFn(GenericData model) {
+ this.modelClass = model != null ? model.getClass() : null;
+ }
+
+ private static <K, V> Map<K, Set<V>> toSetMultiMap(Map<K, V> map) {
+ Map<K, Set<V>> setMultiMap = new HashMap<K, Set<V>>();
+ for (Map.Entry<K, V> entry : map.entrySet()) {
+ Set<V> set = new HashSet<V>();
+ set.add(entry.getValue());
+ setMultiMap.put(entry.getKey(), Collections.unmodifiableSet(set));
+ }
+ return Collections.unmodifiableMap(setMultiMap);
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element FileIO.ReadableFile file,
+ RestrictionTracker<OffsetRange, Long> tracker,
+ OutputReceiver<GenericRecord> outputReceiver)
+ throws Exception {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ ResourceId filename = file.getMetadata().resourceId();
+ throw new RuntimeException(String.format("File has to be seekable:
%s", filename));
+ }
+
+ SeekableByteChannel seekableByteChannel = file.openSeekable();
+ ReadSupport<GenericRecord> readSupport;
+ InputFile inputFile = new BeamParquetInputFile(seekableByteChannel);
+ Configuration conf = setConf();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ readSupport = new AvroReadSupport<GenericRecord>(model);
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ Filter filter = checkNotNull(options.getRecordFilter(), "filter");
+ conf = ((HadoopReadOptions) options).getConf();
+ for (String property : options.getPropertyNames()) {
+ conf.set(property, options.getProperty(property));
+ }
+ FileMetaData parquetFileMetadata =
reader.getFooter().getFileMetaData();
+ MessageType fileSchema = parquetFileMetadata.getSchema();
+ Map<String, String> fileMetadata =
parquetFileMetadata.getKeyValueMetaData();
+
+ ReadSupport.ReadContext readContext =
+ readSupport.init(new InitContext(conf,
toSetMultiMap(fileMetadata), fileSchema));
+ ColumnIOFactory columnIOFactory = new
ColumnIOFactory(parquetFileMetadata.getCreatedBy());
+ MessageType requestedSchema = readContext.getRequestedSchema();
+ RecordMaterializer<GenericRecord> recordConverter =
+ readSupport.prepareForRead(conf, fileMetadata, fileSchema,
readContext);
+ boolean strictTypeChecking = options.isEnabled(STRICT_TYPE_CHECKING,
true);
+ boolean filterRecords = options.useRecordFilter();
+ reader.setRequestedSchema(requestedSchema);
+ MessageColumnIO columnIO =
+ columnIOFactory.getColumnIO(requestedSchema, fileSchema,
strictTypeChecking);
+ for (int i = 0; i < reader.getRowGroups().size(); i++) {
+ if (i < tracker.currentRestriction().getFrom()) {
+ reader.skipNextRowGroup();
+ continue;
+ }
+ if (tracker.tryClaim((long) i)) {
+ PageReadStore pages = reader.readNextRowGroup();
+ i += 1;
+ RecordReader<GenericRecord> recordReader =
+ columnIO.getRecordReader(
+ pages, recordConverter, filterRecords ? filter :
FilterCompat.NOOP);
+ GenericRecord read;
+ long current = 0;
+ long totalRows = pages.getRowCount();
+ while (current < totalRows) {
+ read = recordReader.read();
+ outputReceiver.output(read);
+ current += 1;
+ }
+ } else {
+ break;
+ }
+ }
+ }
+
+ private Configuration setConf() throws Exception {
+ Configuration conf = new Configuration();
+ GenericData model = null;
+ if (modelClass != null) {
+ model = (GenericData) modelClass.getMethod("get").invoke(null);
+ }
+ if (model != null
+ && (model.getClass() == GenericData.class || model.getClass() ==
SpecificData.class)) {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, true);
+ } else {
+ conf.setBoolean(AvroReadSupport.AVRO_COMPATIBILITY, false);
+ }
+ return conf;
+ }
+
+ @GetInitialRestriction
+ public OffsetRange getInitialRestriction(@Element FileIO.ReadableFile
file) throws Exception {
+ if (!file.getMetadata().isReadSeekEfficient()) {
+ ResourceId filename = file.getMetadata().resourceId();
+ throw new RuntimeException(String.format("File has to be seekable:
%s", filename));
+ }
+ SeekableByteChannel seekableByteChannel = file.openSeekable();
+ InputFile inputFile = new BeamParquetInputFile(seekableByteChannel);
+ Configuration conf = setConf();
+ ParquetReadOptions options = HadoopReadOptions.builder(conf).build();
+ ParquetFileReader reader = ParquetFileReader.open(inputFile, options);
+ return new OffsetRange(0, reader.getRowGroups().size());
+ }
+
+ @SplitRestriction
+ public void split(@Restriction OffsetRange restriction,
OutputReceiver<OffsetRange> out) {
Review comment:
Yes, I think we can build a BlockBasedTracker that uses number of blocks
for restriction but aware of byte positions. This will allow it to better
perform initial splitting and provide a better progress signal based on
estimated element positions (by overriding getProgress()).
We do something like this for progress in existing block-based source:
https://github.com/apache/beam/blob/master/sdks/java/core/src/main/java/org/apache/beam/sdk/io/BlockBasedSource.java#L241
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
