Skip to content
This repository was archived by the owner on Nov 28, 2025. It is now read-only.

fix tfrecord compress on hdfs #108

Closed
bigbear2017 wants to merge 1 commit into from
Closed

fix tfrecord compress on hdfs #108

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions hadoop/src/main/java/org/tensorflow/hadoop/io/TFRecordFileOutputFormat.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,20 +32,29 @@ public class TFRecordFileOutputFormat extends FileOutputFormat<BytesWritable, Nu
@Override public RecordWriter<BytesWritable, NullWritable> getRecordWriter(

@fhoering fhoering Mar 14, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

relevant imports are missing:
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.GzipCodec;
import org.apache.hadoop.util.ReflectionUtils;

import java.io.DataOutputStream;

TaskAttemptContext context) throws IOException, InterruptedException {
Configuration conf = context.getConfiguration();
Path file = getDefaultWorkFile(context, "");
FileSystem fs = file.getFileSystem(conf);

boolean isCompressed = getCompressOutput(context);
CompressionCodec codec = null;
String extension = "";
if (isCompressed) {
Class<? extends CompressionCodec> codecClass = getOutputCompressorClass(context, GzipCodec.class);
codec = ReflectionUtils.newInstance(codecClass, conf);
extension = codec.getDefaultExtension();
}
int bufferSize = TFRecordIOConf.getBufferSize(conf);
final FSDataOutputStream fsdos = fs.create(file, true, bufferSize);
Path file = getDefaultWorkFile(context, extension);
FileSystem fs = file.getFileSystem(conf);
final DataOutputStream fsdos = isCompressed ?
new DataOutputStream(codec.createOutputStream(fs.create(file, true, bufferSize))) :
fs.create(file, true, bufferSize);

@fhoering fhoering Mar 14, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fs.create(file, true, bufferSize) could be factorized in a variable

final TFRecordWriter writer = new TFRecordWriter(fsdos);
return new RecordWriter<BytesWritable, NullWritable>() {
@Override public void write(BytesWritable key, NullWritable value)
throws IOException, InterruptedException {
throws IOException, InterruptedException {
writer.write(key.getBytes(), 0, key.getLength());
}

@Override public void close(TaskAttemptContext context)
throws IOException, InterruptedException {
throws IOException, InterruptedException {
fsdos.close();
}
};
Expand Down