Add support for random prefix with deletion vector in Delta Lake

ebyhr · ebyhr · commit b5e2b2437e2c · 2024-11-07T22:07:17.000+09:00
diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMergeSink.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMergeSink.java
@@ -130,6 +130,7 @@ public class DeltaLakeMergeSink
     private final ParquetReaderOptions parquetReaderOptions;
     private final boolean deletionVectorEnabled;
     private final Map<String, DeletionVectorEntry> deletionVectors;
+    private final int randomPrefixLength;
 
     @Nullable
     private DeltaLakeCdfPageSink cdfPageSink;
@@ -153,7 +154,8 @@ public DeltaLakeMergeSink(
             ParquetReaderOptions parquetReaderOptions,
             FileFormatDataSourceStats fileFormatDataSourceStats,
             boolean deletionVectorEnabled,
-            Map<String, DeletionVectorEntry> deletionVectors)
+            Map<String, DeletionVectorEntry> deletionVectors,
+            int randomPrefixLength)
     {
         this.typeOperators = requireNonNull(typeOperators, "typeOperators is null");
         this.session = requireNonNull(session, "session is null");
@@ -181,6 +183,7 @@ public DeltaLakeMergeSink(
         this.parquetReaderOptions = requireNonNull(parquetReaderOptions, "parquetReaderOptions is null");
         this.deletionVectorEnabled = deletionVectorEnabled;
         this.deletionVectors = ImmutableMap.copyOf(requireNonNull(deletionVectors, "deletionVectors is null"));
+        this.randomPrefixLength = randomPrefixLength;
         dataColumnsIndices = new int[tableColumnCount];
         dataAndRowIdColumnsIndices = new int[tableColumnCount + 1];
         for (int i = 0; i < tableColumnCount; i++) {
@@ -408,7 +411,7 @@ private Slice writeDeletionVector(
 
         DeletionVectorEntry deletionVectorEntry;
         try {
-            deletionVectorEntry = writeDeletionVectors(fileSystem, rootTableLocation, deletedRows);
+            deletionVectorEntry = writeDeletionVectors(fileSystem, rootTableLocation, deletedRows, randomPrefixLength);
         }
         catch (IOException e) {
             throw new TrinoException(DELTA_LAKE_BAD_WRITE, "Unable to write deletion vector file", e);
diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakePageSinkProvider.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakePageSinkProvider.java
@@ -54,6 +54,7 @@
 import static io.trino.plugin.deltalake.DeltaLakeParquetSchemas.createParquetSchemaMapping;
 import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.changeDataFeedEnabled;
 import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.extractSchema;
+import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.getRandomPrefixLength;
 import static io.trino.plugin.deltalake.transactionlog.DeltaLakeSchemaSupport.isDeletionVectorEnabled;
 import static io.trino.spi.type.VarcharType.VARCHAR;
 import static java.util.Objects.requireNonNull;
@@ -199,7 +200,8 @@ public ConnectorMergeSink createMergeSink(ConnectorTransactionHandle transaction
                 parquetReaderOptions,
                 fileFormatDataSourceStats,
                 isDeletionVectorEnabled(tableHandle.metadataEntry(), tableHandle.protocolEntry()),
-                merge.deletionVectors());
+                merge.deletionVectors(),
+                getRandomPrefixLength(tableHandle.metadataEntry()));
     }
 
     private DeltaLakeCdfPageSink createCdfPageSink(
diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/delete/DeletionVectors.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/delete/DeletionVectors.java
@@ -28,6 +28,7 @@
 import java.nio.ByteBuffer;
 import java.util.OptionalInt;
 import java.util.UUID;
+import java.util.concurrent.ThreadLocalRandom;
 import java.util.zip.CRC32;
 import java.util.zip.Checksum;
 
@@ -75,12 +76,18 @@ public static RoaringBitmapArray readDeletionVectors(TrinoFileSystem fileSystem,
     public static DeletionVectorEntry writeDeletionVectors(
             TrinoFileSystem fileSystem,
             Location location,
-            RoaringBitmapArray deletedRows)
+            RoaringBitmapArray deletedRows,
+            int randomPrefixLength)
             throws IOException
     {
         UUID uuid = randomUUID();
-        String deletionVectorFilename = "deletion_vector_" + uuid + ".bin";
         String pathOrInlineDv = encodeUUID(uuid);
+        String deletionVectorFilename = "deletion_vector_" + uuid + ".bin";
+        if (randomPrefixLength > 0) {
+            String randomPrefix = randomPrefix(randomPrefixLength);
+            pathOrInlineDv = randomPrefix + pathOrInlineDv;
+            location = location.appendPath(randomPrefix);
+        }
         int sizeInBytes = MAGIC_NUMBER_BYTE_SIZE + BIT_MAP_COUNT_BYTE_SIZE + BIT_MAP_KEY_BYTE_SIZE + deletedRows.serializedSizeInBytes();
         long cardinality = deletedRows.cardinality();
 
@@ -100,6 +107,16 @@ public static DeletionVectorEntry writeDeletionVectors(
         return new DeletionVectorEntry(UUID_MARKER, pathOrInlineDv, offset, sizeInBytes, cardinality);
     }
 
+    private static String randomPrefix(int length)
+    {
+        String alphanumeric = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
+        StringBuilder prefix = new StringBuilder(length);
+        for (int i = 0; i < length; i++) {
+            prefix.append(alphanumeric.charAt(ThreadLocalRandom.current().nextInt(alphanumeric.length())));
+        }
+        return prefix.toString();
+    }
+
     private static byte[] serializeAsByteArray(RoaringBitmapArray bitmaps, int sizeInBytes)
     {
         ByteBuffer buffer = ByteBuffer.allocate(sizeInBytes).order(LITTLE_ENDIAN);
diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeSchemaSupport.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/transactionlog/DeltaLakeSchemaSupport.java
@@ -204,6 +204,18 @@ public static boolean isDeletionVectorEnabled(MetadataEntry metadataEntry, Proto
         return parseBoolean(metadataEntry.getConfiguration().get(DELETION_VECTORS_CONFIGURATION_KEY));
     }
 
+    public static int getRandomPrefixLength(MetadataEntry metadataEntry)
+    {
+        boolean randomizeFilePrefixes = parseBoolean(metadataEntry.getConfiguration().get("delta.randomizeFilePrefixes"));
+        if (randomizeFilePrefixes) {
+            // 2 is the default value in Delta Lake
+            int randomPrefixLength = Integer.parseInt(metadataEntry.getConfiguration().getOrDefault("delta.randomPrefixLength", "2"));
+            checkArgument(randomPrefixLength >= 0, "randomPrefixLength must be >= 0: %s", randomPrefixLength);
+            return randomPrefixLength;
+        }
+        return 0;
+    }
+
     public static List<String> enabledUniversalFormats(MetadataEntry metadataEntry)
     {
         String formats = metadataEntry.getConfiguration().get(UNIVERSAL_FORMAT_CONFIGURATION_KEY);
diff --git a/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeBasic.java b/plugin/trino-delta-lake/src/test/java/io/trino/plugin/deltalake/TestDeltaLakeBasic.java
@@ -1243,6 +1243,26 @@ public void testDeletionVectorsCheckPoint()
         assertUpdate("DROP TABLE " + tableName);
     }
 
+    @Test
+    public void testDeletionVectorsRandomPrefix()
+            throws Exception
+    {
+        String tableName = "deletion_vectors_random_prefix" + randomNameSuffix();
+
+        Path tableLocation = catalogDir.resolve(tableName);
+        copyDirectoryContents(new File(Resources.getResource("deltalake/deletion_vector_random_prefix").toURI()).toPath(), tableLocation);
+        assertUpdate("CALL system.register_table('%s', '%s', '%s')".formatted(getSession().getSchema().orElseThrow(), tableName, tableLocation.toUri()));
+
+        assertUpdate("INSERT INTO " + tableName + " VALUES (1, 10), (2, 20), (3, 30)", 3);
+        assertUpdate("DELETE FROM " + tableName + " WHERE a = 1", 1);
+        assertQuery("SELECT * FROM " + tableName, "VALUES (2, 20), (3, 30)");
+
+        assertThat(fileNamesIn(new URI(tableLocation.toString()).getPath(), true))
+                .anyMatch(path -> path.matches(tableLocation + "/[a-zA-Z0-9_]{3}/deletion_vector_[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}.bin"));
+
+        assertUpdate("DROP TABLE " + tableName);
+    }
+
     @Test
     public void testUnsupportedVacuumDeletionVectors()
             throws Exception
diff --git a/plugin/trino-delta-lake/src/test/resources/deltalake/deletion_vector_random_prefix/README.md b/plugin/trino-delta-lake/src/test/resources/deltalake/deletion_vector_random_prefix/README.md
@@ -0,0 +1,12 @@
+Data generated using OSS Delta Lake 3.2.0:
+
+```sql
+CREATE TABLE tpch.test_dv_random_prefix 
+(a int, b int) 
+USING delta 
+TBLPROPERTIES(
+  'delta.enableDeletionVectors' = true, 
+  'delta.randomizeFilePrefixes' = true,
+  'delta.randomPrefixLength' = 3
+);
+```
diff --git a/plugin/trino-delta-lake/src/test/resources/deltalake/deletion_vector_random_prefix/_delta_log/00000000000000000000.json b/plugin/trino-delta-lake/src/test/resources/deltalake/deletion_vector_random_prefix/_delta_log/00000000000000000000.json
@@ -0,0 +1,3 @@
+{"commitInfo":{"timestamp":1730875806173,"operation":"CREATE TABLE","operationParameters":{"partitionBy":"[]","clusterBy":"[]","description":null,"isManaged":"true","properties":"{\"delta.randomPrefixLength\":\"3\",\"delta.randomizeFilePrefixes\":\"true\",\"delta.enableDeletionVectors\":\"true\"}"},"isolationLevel":"Serializable","isBlindAppend":true,"operationMetrics":{},"engineInfo":"Apache-Spark/3.5.0 Delta-Lake/3.2.0","txnId":"6723f2fb-7f22-402a-ac00-f2ffd838136e"}}
+{"metaData":{"id":"20f880b6-d9f7-4cd2-a21d-1a33b0b445f1","format":{"provider":"parquet","options":{}},"schemaString":"{\"type\":\"struct\",\"fields\":[{\"name\":\"a\",\"type\":\"integer\",\"nullable\":true,\"metadata\":{}},{\"name\":\"b\",\"type\":\"integer\",\"nullable\":true,\"metadata\":{}}]}","partitionColumns":[],"configuration":{"delta.randomPrefixLength":"3","delta.randomizeFilePrefixes":"true","delta.enableDeletionVectors":"true"},"createdTime":1730875806156}}
+{"protocol":{"minReaderVersion":3,"minWriterVersion":7,"readerFeatures":["deletionVectors"],"writerFeatures":["deletionVectors"]}}
diff --git a/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeDeleteCompatibility.java b/testing/trino-product-tests/src/main/java/io/trino/tests/product/deltalake/TestDeltaLakeDeleteCompatibility.java
@@ -307,9 +307,15 @@ public void testDeletionVectorsWithRandomPrefix()
                 "LOCATION 's3://" + bucketName + "/databricks-compatibility-test-" + tableName + "' " +
                 "TBLPROPERTIES ('delta.enableDeletionVectors' = true, 'delta.randomizeFilePrefixes' = true)");
         try {
-            onDelta().executeQuery("INSERT INTO default." + tableName + " VALUES (1, 11), (2, 22)");
+            onDelta().executeQuery("INSERT INTO default." + tableName + " VALUES (1, 11), (2, 22), (3, 33)");
+
             onDelta().executeQuery("DELETE FROM default." + tableName + " WHERE a = 2");
+            assertThat(onDelta().executeQuery("SELECT * FROM default." + tableName))
+                    .containsOnly(row(1, 11), row(3, 33));
+            assertThat(onTrino().executeQuery("SELECT * FROM delta.default." + tableName))
+                    .containsOnly(row(1, 11), row(3, 33));
 
+            onTrino().executeQuery("DELETE FROM delta.default." + tableName + " WHERE a = 3");
             assertThat(onDelta().executeQuery("SELECT * FROM default." + tableName))
                     .containsOnly(row(1, 11));
             assertThat(onTrino().executeQuery("SELECT * FROM delta.default." + tableName))

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+{"commitInfo":{"timestamp":1730875806173,"operation":"CREATE TABLE","operationParameters":{"partitionBy":"[]","clusterBy":"[]","description":null,"isManaged":"true","properties":"{\"delta.randomPrefixLength\":\"3\",\"delta.randomizeFilePrefixes\":\"true\",\"delta.enableDeletionVectors\":\"true\"}"},"isolationLevel":"Serializable","isBlindAppend":true,"operationMetrics":{},"engineInfo":"Apache-Spark/3.5.0 Delta-Lake/3.2.0","txnId":"6723f2fb-7f22-402a-ac00-f2ffd838136e"}}`
	`2`	`+{"metaData":{"id":"20f880b6-d9f7-4cd2-a21d-1a33b0b445f1","format":{"provider":"parquet","options":{}},"schemaString":"{\"type\":\"struct\",\"fields\":[{\"name\":\"a\",\"type\":\"integer\",\"nullable\":true,\"metadata\":{}},{\"name\":\"b\",\"type\":\"integer\",\"nullable\":true,\"metadata\":{}}]}","partitionColumns":[],"configuration":{"delta.randomPrefixLength":"3","delta.randomizeFilePrefixes":"true","delta.enableDeletionVectors":"true"},"createdTime":1730875806156}}`
	`3`	`+{"protocol":{"minReaderVersion":3,"minWriterVersion":7,"readerFeatures":["deletionVectors"],"writerFeatures":["deletionVectors"]}}`