From 6707e0d9fa61cf24fb9d981d9ba2cf75a4b160a7 Mon Sep 17 00:00:00 2001 From: Purujit Saha Date: Fri, 18 Oct 2024 21:03:08 +0000 Subject: [PATCH] =?UTF-8?q?Revert=20"Revert=20"Add=20chunk=20offset=20to?= =?UTF-8?q?=20file=20id=20key=20to=20make=20each=20chunk=20have=20a=20uni?= =?UTF-8?q?=E2=80=A6=20(#848)"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 0930648dfaaefa5abb6fe7aa4dcb83423ac1f11e. --- .../streaming/internal/BlobBuilder.java | 2 +- .../ingest/streaming/internal/Flusher.java | 5 ++- .../streaming/internal/ParquetFlusher.java | 36 +++++++++++++++---- .../streaming/internal/ParquetRowBuffer.java | 3 ++ .../net/snowflake/ingest/utils/Constants.java | 1 + .../streaming/internal/RowBufferTest.java | 28 ++++++++++++--- 6 files changed, 62 insertions(+), 13 deletions(-) diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java b/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java index 060af357f..1aa81e1b1 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/BlobBuilder.java @@ -86,7 +86,7 @@ static Blob constructBlobAndMetadata( Flusher flusher = channelsDataPerTable.get(0).createFlusher(); Flusher.SerializationResult serializedChunk = - flusher.serialize(channelsDataPerTable, filePath); + flusher.serialize(channelsDataPerTable, filePath, curDataSize); if (!serializedChunk.channelsMetadataList.isEmpty()) { final byte[] compressedChunkData; diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/Flusher.java b/src/main/java/net/snowflake/ingest/streaming/internal/Flusher.java index 5a426e873..b110ae36a 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/Flusher.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/Flusher.java @@ -20,12 +20,15 @@ public interface Flusher { /** * Serialize buffered rows into the underlying format. * + * @param fullyQualifiedTableName * @param channelsDataPerTable buffered rows * @param filePath file path + * @param chunkStartOffset * @return {@link SerializationResult} * @throws IOException */ - SerializationResult serialize(List> channelsDataPerTable, String filePath) + SerializationResult serialize( + List> channelsDataPerTable, String filePath, long chunkStartOffset) throws IOException; /** Holds result of the buffered rows conversion: channel metadata and stats. */ diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ParquetFlusher.java b/src/main/java/net/snowflake/ingest/streaming/internal/ParquetFlusher.java index e7272d94a..1432e94b3 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ParquetFlusher.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ParquetFlusher.java @@ -16,6 +16,7 @@ import net.snowflake.ingest.utils.Logging; import net.snowflake.ingest.utils.Pair; import net.snowflake.ingest.utils.SFException; +import org.apache.parquet.Preconditions; import org.apache.parquet.column.ParquetProperties; import org.apache.parquet.hadoop.SnowflakeParquetWriter; import org.apache.parquet.schema.MessageType; @@ -52,13 +53,17 @@ public ParquetFlusher( @Override public SerializationResult serialize( - List> channelsDataPerTable, String filePath) + List> channelsDataPerTable, + String filePath, + long chunkStartOffset) throws IOException { - return serializeFromJavaObjects(channelsDataPerTable, filePath); + return serializeFromJavaObjects(channelsDataPerTable, filePath, chunkStartOffset); } private SerializationResult serializeFromJavaObjects( - List> channelsDataPerTable, String filePath) + List> channelsDataPerTable, + String filePath, + long chunkStartOffset) throws IOException { List channelsMetadataList = new ArrayList<>(); long rowCount = 0L; @@ -124,10 +129,7 @@ private SerializationResult serializeFromJavaObjects( } Map metadata = channelsDataPerTable.get(0).getVectors().metadata; - // We insert the filename in the file itself as metadata so that streams can work on replicated - // mixed tables. For a more detailed discussion on the topic see SNOW-561447 and - // http://go/streams-on-replicated-mixed-tables - metadata.put(Constants.PRIMARY_FILE_ID_KEY, StreamingIngestUtils.getShortname(filePath)); + addFileIdToMetadata(filePath, chunkStartOffset, metadata); parquetWriter = new SnowflakeParquetWriter( mergedData, @@ -154,6 +156,26 @@ private SerializationResult serializeFromJavaObjects( parquetWriter.getExtendedMetadataSize()); } + private static void addFileIdToMetadata( + String filePath, long chunkStartOffset, Map metadata) { + // We insert the filename in the file itself as metadata so that streams can work on replicated + // mixed tables. For a more detailed discussion on the topic see SNOW-561447 and + // http://go/streams-on-replicated-mixed-tables + // Using chunk offset as suffix ensures that for interleaved tables, the file + // id key is unique for each chunk. Each chunk is logically a separate Parquet file that happens + // to be bundled together. + if (chunkStartOffset == 0) { + metadata.put(Constants.PRIMARY_FILE_ID_KEY, StreamingIngestUtils.getShortname(filePath)); + } else { + String shortName = StreamingIngestUtils.getShortname(filePath); + final String[] parts = shortName.split("\\."); + Preconditions.checkState(parts.length == 2, "Invalid file name format"); + metadata.put( + Constants.PRIMARY_FILE_ID_KEY, + String.format("%s_%d.%s", parts[0], chunkStartOffset, parts[1])); + } + } + /** * Validates that rows count in metadata matches the row count in Parquet footer and the row count * written by the parquet writer diff --git a/src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java b/src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java index 12a6cc875..a60ec2e6e 100644 --- a/src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java +++ b/src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java @@ -20,9 +20,11 @@ import java.util.Set; import java.util.function.Consumer; import net.snowflake.client.jdbc.internal.google.common.collect.Sets; +import net.snowflake.ingest.connection.RequestBuilder; import net.snowflake.ingest.connection.TelemetryService; import net.snowflake.ingest.streaming.OffsetTokenVerificationFunction; import net.snowflake.ingest.streaming.OpenChannelRequest; +import net.snowflake.ingest.utils.Constants; import net.snowflake.ingest.utils.ErrorCode; import net.snowflake.ingest.utils.IcebergDataTypeParser; import net.snowflake.ingest.utils.SFException; @@ -89,6 +91,7 @@ public void setupSchema(List columns) { if (!clientBufferParameters.getEnableIcebergStreaming()) { metadata.put("sfVer", "1,1"); } + metadata.put(Constants.SDK_VERSION_KEY, RequestBuilder.DEFAULT_VERSION); List parquetTypes = new ArrayList<>(); int id = 1; diff --git a/src/main/java/net/snowflake/ingest/utils/Constants.java b/src/main/java/net/snowflake/ingest/utils/Constants.java index 754c81cff..1294a9a0b 100644 --- a/src/main/java/net/snowflake/ingest/utils/Constants.java +++ b/src/main/java/net/snowflake/ingest/utils/Constants.java @@ -36,6 +36,7 @@ public class Constants { public static final String SNOWFLAKE_OAUTH_TOKEN_ENDPOINT = "/oauth/token-request"; public static final String PRIMARY_FILE_ID_KEY = "primaryFileId"; // Don't change, should match Parquet Scanner + public static final String SDK_VERSION_KEY = "sdkVersion"; public static final long RESPONSE_SUCCESS = 0L; // Don't change, should match server side public static final long RESPONSE_ERR_GENERAL_EXCEPTION_RETRY_REQUEST = 10L; // Don't change, should match server side diff --git a/src/test/java/net/snowflake/ingest/streaming/internal/RowBufferTest.java b/src/test/java/net/snowflake/ingest/streaming/internal/RowBufferTest.java index 4a0e7d4c5..0076e6e29 100644 --- a/src/test/java/net/snowflake/ingest/streaming/internal/RowBufferTest.java +++ b/src/test/java/net/snowflake/ingest/streaming/internal/RowBufferTest.java @@ -27,6 +27,7 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; +import net.snowflake.ingest.connection.RequestBuilder; import net.snowflake.ingest.streaming.InsertValidationResponse; import net.snowflake.ingest.streaming.OpenChannelRequest; import net.snowflake.ingest.utils.Constants; @@ -2028,11 +2029,30 @@ public void testParquetFileNameMetadata() throws IOException { data.setChannelContext(new ChannelFlushContext("name", "db", "schema", "table", 1L, "key", 0L)); ParquetFlusher flusher = (ParquetFlusher) bufferUnderTest.createFlusher(); - Flusher.SerializationResult result = - flusher.serialize(Collections.singletonList(data), filePath); + { + Flusher.SerializationResult result = + flusher.serialize(Collections.singletonList(data), filePath, 0); - BdecParquetReader reader = new BdecParquetReader(result.chunkData.toByteArray()); - Assert.assertEquals(filePath, reader.getKeyValueMetadata().get(Constants.PRIMARY_FILE_ID_KEY)); + BdecParquetReader reader = new BdecParquetReader(result.chunkData.toByteArray()); + Assert.assertEquals( + "testParquetFileNameMetadata.bdec", + reader.getKeyValueMetadata().get(Constants.PRIMARY_FILE_ID_KEY)); + Assert.assertEquals( + RequestBuilder.DEFAULT_VERSION, + reader.getKeyValueMetadata().get(Constants.SDK_VERSION_KEY)); + } + { + Flusher.SerializationResult result = + flusher.serialize(Collections.singletonList(data), filePath, 13); + + BdecParquetReader reader = new BdecParquetReader(result.chunkData.toByteArray()); + Assert.assertEquals( + "testParquetFileNameMetadata_13.bdec", + reader.getKeyValueMetadata().get(Constants.PRIMARY_FILE_ID_KEY)); + Assert.assertEquals( + RequestBuilder.DEFAULT_VERSION, + reader.getKeyValueMetadata().get(Constants.SDK_VERSION_KEY)); + } } @Test