Merge pull request #194 from tomwhite/vcf-writer-override

Changes to make it possible to subclass VCF and BCF input formats and…
HadoopGenomics · Apr 23, 2018 · 9f974cf · 9f974cf
2 parents dae6c1c + b1f4d81
commit 9f974cf
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 10 deletions.
diff --git a/src/main/java/org/seqdoop/hadoop_bam/BCFRecordWriter.java b/src/main/java/org/seqdoop/hadoop_bam/BCFRecordWriter.java
@@ -28,6 +28,7 @@
 
 import htsjdk.samtools.util.BlockCompressedOutputStream;
 
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.mapreduce.RecordWriter;
 import org.apache.hadoop.mapreduce.TaskAttemptContext;
@@ -76,13 +77,20 @@ public BCFRecordWriter(
 	{
 		init(
 			output.getFileSystem(ctx.getConfiguration()).create(output),
-			header, writeHeader);
+			header, writeHeader, ctx);
 	}
 	public BCFRecordWriter(
 			OutputStream output, VCFHeader header, boolean writeHeader)
 		throws IOException
 	{
-		init(output, header, writeHeader);
+		init(output, header, writeHeader, null);
+	}
+	public BCFRecordWriter(
+			OutputStream output, VCFHeader header, boolean writeHeader,
+			TaskAttemptContext ctx)
+			throws IOException
+	{
+		init(output, header, writeHeader, ctx);
 	}
 
 	// Working around not being able to call a constructor other than as the
@@ -94,25 +102,31 @@ private void init(
 	{
 		init(
 			output.getFileSystem(ctx.getConfiguration()).create(output),
-			header, writeHeader);
+			header, writeHeader, ctx);
 	}
 	private void init(
-			OutputStream output, VCFHeader header, final boolean writeHeader)
+			OutputStream output, VCFHeader header, final boolean writeHeader,
+			TaskAttemptContext ctx)
 		throws IOException
 	{
 		final BCFStoppableOutputStream stopOut =
 			new BCFStoppableOutputStream(!writeHeader, output);
 
-		writer = new VariantContextWriterBuilder().clearOptions()
-				.setOption(Options.FORCE_BCF)
-				.setOutputBCFStream(stopOut).build();
+		writer = createVariantContextWriter(ctx == null ? null : ctx.getConfiguration(), stopOut);
 
 		writer.writeHeader(header);
 		stopOut.stopped = false;
 
 		setInputHeader(header);
 	}
 
+	protected VariantContextWriter createVariantContextWriter(Configuration conf,
+			OutputStream out) {
+		return new VariantContextWriterBuilder().clearOptions()
+				.setOption(Options.FORCE_BCF)
+				.setOutputBCFStream(out).build();
+	}
+
 	@Override public void close(TaskAttemptContext ctx) throws IOException {
 		writer.close();
 	}

diff --git a/src/main/java/org/seqdoop/hadoop_bam/KeyIgnoringBCFRecordWriter.java b/src/main/java/org/seqdoop/hadoop_bam/KeyIgnoringBCFRecordWriter.java
@@ -49,12 +49,24 @@ public KeyIgnoringBCFRecordWriter(
 	{
 		super(output, header, writeHeader, ctx);
 	}
+	/**
+	 * @deprecated This constructor has no {@link TaskAttemptContext} so it is not
+	 * possible to pass configuration properties to the writer.
+	 */
+	@Deprecated
 	public KeyIgnoringBCFRecordWriter(
 			OutputStream output, VCFHeader header, boolean writeHeader)
 		throws IOException
 	{
 		super(output, header, writeHeader);
 	}
+	public KeyIgnoringBCFRecordWriter(
+			OutputStream output, VCFHeader header, boolean writeHeader,
+			TaskAttemptContext ctx)
+			throws IOException
+	{
+		super(output, header, writeHeader, ctx);
+	}
 
 	@Override public void write(K ignored, VariantContextWritable vc) {
 		writeRecord(vc.get());

diff --git a/src/main/java/org/seqdoop/hadoop_bam/KeyIgnoringVCFOutputFormat.java b/src/main/java/org/seqdoop/hadoop_bam/KeyIgnoringVCFOutputFormat.java
@@ -132,7 +132,7 @@ public RecordWriter<K,VariantContextWritable> getRecordWriter(
 		}
 	}
 
-	private RecordWriter<K,VariantContextWritable> getRecordWriter(
+	public RecordWriter<K,VariantContextWritable> getRecordWriter(
 			TaskAttemptContext ctx, OutputStream outputStream)
 			throws IOException
 	{
@@ -144,8 +144,8 @@ private RecordWriter<K,VariantContextWritable> getRecordWriter(
 				WRITE_HEADER_PROPERTY, true);
 
 		switch (format) {
-			case BCF: return new KeyIgnoringBCFRecordWriter<K>(outputStream,header,wh);
-			case VCF: return new KeyIgnoringVCFRecordWriter<K>(outputStream,header,wh);
+			case BCF: return new KeyIgnoringBCFRecordWriter<K>(outputStream,header,wh,ctx);
+			case VCF: return new KeyIgnoringVCFRecordWriter<K>(outputStream,header,wh,ctx);
 			default: assert false; return null;
 		}
 	}

diff --git a/src/main/java/org/seqdoop/hadoop_bam/KeyIgnoringVCFRecordWriter.java b/src/main/java/org/seqdoop/hadoop_bam/KeyIgnoringVCFRecordWriter.java
@@ -49,12 +49,25 @@ public KeyIgnoringVCFRecordWriter(
 	{
 		super(output, header, writeHeader, ctx);
 	}
+	/**
+	 * @deprecated This constructor has no {@link TaskAttemptContext} so it is not
+	 * possible to pass configuration properties to the writer.
+	 */
+	@Deprecated
 	public KeyIgnoringVCFRecordWriter(
 			OutputStream output, VCFHeader header, boolean writeHeader)
 		throws IOException
 	{
 		super(output, header, writeHeader);
 	}
+	public KeyIgnoringVCFRecordWriter(
+			OutputStream output, VCFHeader header, boolean writeHeader,
+			TaskAttemptContext ctx)
+			throws IOException
+	{
+		super(output, header, writeHeader, ctx);
+	}
+
 
 	@Override public void write(K ignored, VariantContextWritable vc) {
 		writeRecord(vc.get());

diff --git a/src/main/java/org/seqdoop/hadoop_bam/VCFRecordWriter.java b/src/main/java/org/seqdoop/hadoop_bam/VCFRecordWriter.java
@@ -91,6 +91,14 @@ public VCFRecordWriter(
 		init(output, header, writeHeader, null);
 	}
 
+	public VCFRecordWriter(
+			OutputStream output, VCFHeader header, boolean writeHeader,
+			TaskAttemptContext ctx)
+			throws IOException
+	{
+		init(output, header, writeHeader, ctx);
+	}
+
 	// Working around not being able to call a constructor other than as the
 	// first statement...
 	private void init(