Reduce filesystem calls for reading delta deletion vector

Author: raunaqmorarka

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/delete/DeletionVectors.java

@@ -17,12 +17,12 @@
 import io.delta.kernel.internal.deletionvectors.Base85Codec;
 import io.trino.filesystem.Location;
 import io.trino.filesystem.TrinoFileSystem;
+import io.trino.filesystem.TrinoInput;
 import io.trino.filesystem.TrinoInputFile;
 import io.trino.plugin.deltalake.transactionlog.DeletionVectorEntry;
 import io.trino.spi.TrinoException;
 import org.roaringbitmap.RoaringBitmap;
 
-import java.io.DataInputStream;
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -32,7 +32,7 @@
 import java.util.zip.Checksum;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.base.Preconditions.checkState;
+import static io.airlift.slice.SizeOf.SIZE_OF_INT;
 import static io.delta.kernel.internal.deletionvectors.Base85Codec.decodeUUID;
 import static io.delta.kernel.internal.deletionvectors.Base85Codec.encodeUUID;
 import static io.trino.plugin.deltalake.DeltaLakeErrorCode.DELTA_LAKE_INVALID_SCHEMA;
@@ -63,7 +63,7 @@ public static RoaringBitmapArray readDeletionVectors(TrinoFileSystem fileSystem,
     {
         if (deletionVector.storageType().equals(UUID_MARKER)) {
             TrinoInputFile inputFile = fileSystem.newInputFile(location.appendPath(toFileName(deletionVector.pathOrInlineDv())));
-            byte[] buffer = readDeletionVector(inputFile, deletionVector.offset().orElseThrow(), deletionVector.sizeInBytes());
+            ByteBuffer buffer = readDeletionVector(inputFile, deletionVector.offset().orElseThrow(), deletionVector.sizeInBytes());
             return deserializeDeletionVectors(buffer);
         }
         if (deletionVector.storageType().equals(INLINE_MARKER) || deletionVector.storageType().equals(PATH_MARKER)) {
@@ -125,38 +125,40 @@ public static String toFileName(String pathOrInlineDv)
         return "%sdeletion_vector_%s.bin".formatted(prefix, uuid);
     }
 
-    public static byte[] readDeletionVector(TrinoInputFile inputFile, int offset, int expectedSize)
+    public static ByteBuffer readDeletionVector(TrinoInputFile inputFile, int offset, int expectedSize)
             throws IOException
     {
-        byte[] bytes = new byte[expectedSize];
-        try (DataInputStream inputStream = new DataInputStream(inputFile.newStream())) {
-            checkState(inputStream.skip(offset) == offset);
-            int actualSize = inputStream.readInt();
+        try (TrinoInput input = inputFile.newInput()) {
+            ByteBuffer buffer = input.readFully(offset, SIZE_OF_INT + expectedSize + SIZE_OF_INT).toByteBuffer();
+            int actualSize = buffer.getInt(0);
             if (actualSize != expectedSize) {
                 throw new TrinoException(DELTA_LAKE_INVALID_SCHEMA, "The size of deletion vector %s expects %s but got %s".formatted(inputFile.location(), expectedSize, actualSize));
             }
-            inputStream.readFully(bytes);
-            int checksum = inputStream.readInt();
-            if (calculateChecksum(bytes) != checksum) {
+            int checksum = buffer.getInt(SIZE_OF_INT + expectedSize);
+            if (calculateChecksum(buffer.array(), buffer.arrayOffset() + SIZE_OF_INT, expectedSize) != checksum) {
                 throw new TrinoException(DELTA_LAKE_INVALID_SCHEMA, "Checksum mismatch for deletion vector: " + inputFile.location());
             }
+            return buffer.slice(SIZE_OF_INT, expectedSize).order(LITTLE_ENDIAN);
         }
-        return bytes;
     }
 
     private static int calculateChecksum(byte[] data)
+    {
+        return calculateChecksum(data, 0, data.length);
+    }
+
+    private static int calculateChecksum(byte[] data, int offset, int length)
     {
         // Delta Lake allows integer overflow intentionally because it's fine from checksum perspective
         // https://github.com/delta-io/delta/blob/039a29abb4abc72ac5912651679233dc983398d6/spark/src/main/scala/org/apache/spark/sql/delta/storage/dv/DeletionVectorStore.scala#L115
         Checksum crc = new CRC32();
-        crc.update(data);
+        crc.update(data, offset, length);
         return (int) crc.getValue();
     }
 
-    private static RoaringBitmapArray deserializeDeletionVectors(byte[] bytes)
+    private static RoaringBitmapArray deserializeDeletionVectors(ByteBuffer buffer)
             throws IOException
     {
-        ByteBuffer buffer = ByteBuffer.wrap(bytes).order(LITTLE_ENDIAN);
         checkArgument(buffer.order() == LITTLE_ENDIAN, "Byte order must be little endian: %s", buffer.order());
         int magicNumber = buffer.getInt();
         if (magicNumber == PORTABLE_ROARING_BITMAP_MAGIC_NUMBER) {

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeAlluxioCacheFileOperations.java

@@ -247,14 +247,10 @@ public void testCacheDeletionVectorsFileOperations()
                         .add(new CacheOperation("Alluxio.readCached", "data", 0, 796))
                         .add(new CacheOperation("Input.readFully", "data", 0, 796))
                         .add(new CacheOperation("Alluxio.writeCache", "data", 0, 796))
-                        .add(new CacheOperation("Alluxio.readCached", "deletion_vector", 5, 34))
-                        .add(new CacheOperation("Alluxio.readCached", "deletion_vector", 1, 4))
-                        .add(new CacheOperation("Alluxio.readCached", "deletion_vector", 39, 4))
-                        .add(new CacheOperation("Alluxio.readExternalStream", "deletion_vector", 38, 1))
-                        .add(new CacheOperation("Alluxio.readExternalStream", "deletion_vector", 1, 4))
-                        .add(new CacheOperation("InputFile.newStream", "deletion_vector"))
+                        .add(new CacheOperation("Alluxio.readCached", "deletion_vector", 1, 42))
+                        .add(new CacheOperation("Input.readFully", "deletion_vector", 0, 43))
                         .add(new CacheOperation("InputFile.length", "deletion_vector"))
-                        .addCopies(new CacheOperation("Alluxio.writeCache", "deletion_vector", 0, 43), 2)
+                        .add(new CacheOperation("Alluxio.writeCache", "deletion_vector", 0, 43))
                         .add(new CacheOperation("InputFile.newStream", "_last_checkpoint"))
                         .build());
         assertFileSystemAccesses(
@@ -268,9 +264,7 @@ public void testCacheDeletionVectorsFileOperations()
                         .add(new CacheOperation("InputFile.length", "00000000000000000002.json"))
                         .add(new CacheOperation("InputFile.length", "00000000000000000003.json"))
                         .add(new CacheOperation("Alluxio.readCached", "data", 0, 796))
-                        .add(new CacheOperation("Alluxio.readCached", "deletion_vector", 5, 34))
-                        .add(new CacheOperation("Alluxio.readCached", "deletion_vector", 1, 4))
-                        .add(new CacheOperation("Alluxio.readCached", "deletion_vector", 39, 4))
+                        .add(new CacheOperation("Alluxio.readCached", "deletion_vector", 1, 42))
                         .add(new CacheOperation("InputFile.length", "deletion_vector"))
                         .add(new CacheOperation("InputFile.newStream", "extended_stats.json"))
                         .add(new CacheOperation("InputFile.newStream", "extendeded_stats.json"))

plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeFileOperations.java

@@ -941,7 +941,7 @@ public void testDeletionVectors()
                         .add(new FileOperation(TRANSACTION_LOG_JSON, "00000000000000000002.json", "InputFile.newStream"))
                         .add(new FileOperation(TRANSACTION_LOG_JSON, "00000000000000000003.json", "InputFile.newStream"))
                         .add(new FileOperation(LAST_CHECKPOINT, "_last_checkpoint", "InputFile.newStream"))
-                        .add(new FileOperation(DELETION_VECTOR, "deletion_vector_a52eda8c-0a57-4636-814b-9c165388f7ca.bin", "InputFile.newStream"))
+                        .add(new FileOperation(DELETION_VECTOR, "deletion_vector_a52eda8c-0a57-4636-814b-9c165388f7ca.bin", "InputFile.newInput"))
                         .add(new FileOperation(DATA, "no partition", "InputFile.newInput"))
                         .build());
         assertFileSystemAccesses(
@@ -954,7 +954,7 @@ public void testDeletionVectors()
                         .add(new FileOperation(STARBURST_EXTENDED_STATS_JSON, "extendeded_stats.json", "InputFile.newStream"))
                         .add(new FileOperation(TRINO_EXTENDED_STATS_JSON, "extended_stats.json", "InputFile.newStream"))
                         .add(new FileOperation(LAST_CHECKPOINT, "_last_checkpoint", "InputFile.newStream"))
-                        .add(new FileOperation(DELETION_VECTOR, "deletion_vector_a52eda8c-0a57-4636-814b-9c165388f7ca.bin", "InputFile.newStream"))
+                        .add(new FileOperation(DELETION_VECTOR, "deletion_vector_a52eda8c-0a57-4636-814b-9c165388f7ca.bin", "InputFile.newInput"))
                         .add(new FileOperation(DATA, "no partition", "InputFile.newInput"))
                         .build());