Fix ZeroRowFileCreator to use DataSink instead of OutputStream interface

prestodb · May 20, 2021 · a6d5270 · a6d5270
1 parent 5eb36ae
commit a6d5270
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 8 deletions.
diff --git a/presto-common/src/main/java/com/facebook/presto/common/io/DataSink.java b/presto-common/src/main/java/com/facebook/presto/common/io/DataSink.java
@@ -13,10 +13,12 @@
  */
 package com.facebook.presto.common.io;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.util.List;
 
 public interface DataSink
+        extends Closeable
 {
     /**
      * Number of bytes written to this sink so far.
@@ -38,6 +40,7 @@ void write(List<DataOutput> outputData)
     /**
      * File is complete
      */
+    @Override
     void close()
             throws IOException;
 }
diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveZeroRowFileCreator.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveZeroRowFileCreator.java
@@ -14,27 +14,32 @@
 package com.facebook.presto.hive;
 
 import com.facebook.airlift.log.Logger;
+import com.facebook.presto.common.io.DataOutput;
+import com.facebook.presto.common.io.DataSink;
 import com.facebook.presto.hive.datasink.DataSinkFactory;
 import com.facebook.presto.hive.metastore.StorageFormat;
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.PrestoException;
+import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
+import io.airlift.slice.Slices;
+import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.hive.ql.exec.FileSinkOperator.RecordWriter;
 import org.apache.hadoop.mapred.JobConf;
 
 import javax.inject.Inject;
 
 import java.io.IOException;
-import java.io.OutputStream;
 import java.io.UncheckedIOException;
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Properties;
 
 import static com.facebook.airlift.concurrent.MoreFutures.getFutureValue;
+import static com.facebook.presto.common.io.DataOutput.createDataOutput;
 import static com.facebook.presto.hive.HiveCompressionCodec.NONE;
 import static com.facebook.presto.hive.HiveErrorCode.HIVE_WRITER_CLOSE_ERROR;
 import static com.facebook.presto.hive.HiveWriteUtils.initializeSerializer;
@@ -79,7 +84,7 @@ public void createFiles(ConnectorSession session, HdfsContext hdfsContext, Path
         List<ListenableFuture<?>> commitFutures = new ArrayList<>();
 
         for (String fileName : fileNames) {
-            commitFutures.add(executor.submit(() -> createFile(hdfsContext, new Path(destinationDirectory, fileName), fileContent)));
+            commitFutures.add(executor.submit(() -> createFile(hdfsContext, new Path(destinationDirectory, fileName), fileContent, session)));
         }
 
         ListenableFuture<?> listenableFutureAggregate = whenAllSucceed(commitFutures).call(() -> null, directExecutor());
@@ -144,10 +149,14 @@ private byte[] generateZeroRowFile(
         }
     }
 
-    private void createFile(HdfsContext hdfsContext, Path path, byte[] content)
+    private void createFile(HdfsContext hdfsContext, Path path, byte[] content, ConnectorSession session)
     {
-        try (OutputStream outputStream = hdfsEnvironment.getFileSystem(hdfsContext, path).create(path, false)) {
-            outputStream.write(content);
+        try {
+            FileSystem fs = hdfsEnvironment.getFileSystem(hdfsContext, path);
+            try (DataSink dataSink = dataSinkFactory.createDataSink(session, fs, path)) {
+                DataOutput dataOutput = createDataOutput(Slices.wrappedBuffer(content));
+                dataSink.write(ImmutableList.of(dataOutput));
+            }
         }
         catch (IOException e) {
             throw new PrestoException(HIVE_WRITER_CLOSE_ERROR, "Error write zero-row file to Hive", e);

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/pagefile/PageFileWriterFactory.java b/presto-hive/src/main/java/com/facebook/presto/hive/pagefile/PageFileWriterFactory.java
@@ -120,10 +120,8 @@ public static void createEmptyPageFile(
             FileSystem fileSystem,
             Path path)
     {
-        try {
-            DataSink dataSink = dataSinkFactory.createDataSink(session, fileSystem, path);
+        try (DataSink dataSink = dataSinkFactory.createDataSink(session, fileSystem, path)) {
             dataSink.write(ImmutableList.of(createEmptyPageFileFooterOutput()));
-            dataSink.close();
         }
         catch (IOException e) {
             throw new PrestoException(HIVE_WRITER_OPEN_ERROR, "Error creating empty pagefile", e);

diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoSparkQueryRunner.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoSparkQueryRunner.java
@@ -102,6 +102,18 @@ public void testTableWrite()
                 45000);
     }
 
+    @Test
+    public void testZeroFileCreatorForBucketedTable()
+    {
+        assertUpdate(
+                getSession(),
+                format("CREATE TABLE hive.hive_test.test_hive_orders_bucketed_join_zero_file WITH (bucketed_by=array['orderkey'], bucket_count=8) AS " +
+                        "SELECT orderkey, custkey, orderstatus, totalprice, orderdate, orderpriority, clerk, shippriority, comment " +
+                        "FROM orders_bucketed " +
+                        "WHERE orderkey = 1"),
+                1);
+    }
+
     @Test
     public void testBucketedTableWriteSimple()
     {