Destination Snowflake Execute COPY in parallel #10212
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| /* | ||
| * Copyright (c) 2021 Airbyte, Inc., all rights reserved. | ||
| */ | ||
|
|
||
| package io.airbyte.integrations.destination.snowflake; | ||
|
|
||
| import java.util.List; | ||
| import java.util.StringJoiner; | ||
| import java.util.concurrent.CompletableFuture; | ||
| import java.util.concurrent.ExecutorService; | ||
| import java.util.concurrent.Executors; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| interface SnowflakeParallelCopyStreamCopier { | ||
|
|
||
| /** | ||
| * Generates list of staging files. See more | ||
| * https://docs.snowflake.com/en/user-guide/data-load-considerations-load.html#lists-of-files | ||
| */ | ||
| default String generateFilesList(List<String> files) { | ||
| StringJoiner joiner = new StringJoiner(","); | ||
| files.forEach(filename -> joiner.add("'" + filename.substring(filename.lastIndexOf("/") + 1) + "'")); | ||
| return joiner.toString(); | ||
| } | ||
|
|
||
| /** | ||
| * Executes async copying of staging files.This method should block until the copy/upload has | ||
| * completed. | ||
| */ | ||
| default void copyFilesInParallel(List<List<String>> partitions) { | ||
| ExecutorService executorService = Executors.newFixedThreadPool(5); | ||
| List<CompletableFuture<Void>> futures = partitions.stream() | ||
| .map(partition -> CompletableFuture.runAsync(() -> copyIntoStage(partition), executorService)) | ||
| .collect(Collectors.toList()); | ||
|
|
||
| try { | ||
| // wait until all futures ready | ||
| CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join(); | ||
| } catch (Exception e) { | ||
| throw new RuntimeException("Failed to copy files from stage to tmp table {}" + e); | ||
| } finally { | ||
| executorService.shutdown(); | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Copies staging files to the temporary table using <COPY INTO> statement | ||
| */ | ||
| void copyIntoStage(List<String> files); | ||
|
|
||
| /** | ||
| * Generates full bucket/container path to staging files | ||
| */ | ||
| String generateBucketPath(); | ||
|
|
||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -6,23 +6,33 @@ | |
|
|
||
| import com.amazonaws.services.s3.AmazonS3; | ||
| import com.google.common.annotations.VisibleForTesting; | ||
| import com.google.common.collect.Lists; | ||
| import io.airbyte.commons.lang.Exceptions; | ||
| import io.airbyte.db.jdbc.JdbcDatabase; | ||
| import io.airbyte.integrations.destination.ExtendedNameTransformer; | ||
| import io.airbyte.integrations.destination.jdbc.SqlOperations; | ||
| import io.airbyte.integrations.destination.jdbc.copy.s3.S3CopyConfig; | ||
| import io.airbyte.integrations.destination.jdbc.copy.s3.S3StreamCopier; | ||
| import io.airbyte.integrations.destination.s3.S3DestinationConfig; | ||
| import io.airbyte.integrations.destination.s3.util.S3OutputPathHelper; | ||
| import io.airbyte.protocol.models.ConfiguredAirbyteStream; | ||
| import java.sql.SQLException; | ||
| import java.sql.Timestamp; | ||
| import java.time.Instant; | ||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| public class SnowflakeS3StreamCopier extends S3StreamCopier implements SnowflakeParallelCopyStreamCopier { | ||
|
|
||
| private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeS3StreamCopier.class); | ||
|
|
||
| // From https://docs.aws.amazon.com/redshift/latest/dg/t_loading-tables-from-s3.html | ||
| // "Split your load data files so that the files are about equal size, between 1 MB and 1 GB after | ||
| // compression" | ||
| public static final int MAX_PARTS_PER_FILE = 4; | ||
| public static final int MAX_FILES_PER_COPY = 1000; | ||
|
|
||
| public SnowflakeS3StreamCopier(final String stagingFolder, | ||
| final String schema, | ||
|
|
@@ -54,6 +64,7 @@ public SnowflakeS3StreamCopier(final String stagingFolder, | |
| final SqlOperations sqlOperations, | ||
| final Timestamp uploadTime, | ||
| final ConfiguredAirbyteStream configuredAirbyteStream) { | ||
|
|
||
| super(stagingFolder, | ||
| schema, | ||
| client, | ||
|
|
@@ -66,24 +77,47 @@ public SnowflakeS3StreamCopier(final String stagingFolder, | |
| MAX_PARTS_PER_FILE); | ||
| } | ||
|
|
||
| @Override | ||
| public void copyStagingFileToTemporaryTable() throws Exception { | ||
|
There was a problem hiding this comment. is this completely identical to the methods in the GCS copier? if yes, let's extract them to a shared implementation rather than copy-pasting There was a problem hiding this comment. @edgao extracted shared implementation into SnowflakeParallelCopyStreamCopier interface |
||
| List<List<String>> partitions = Lists.partition(new ArrayList<>(stagingWritersByFile.keySet()), MAX_FILES_PER_COPY); | ||
| LOGGER.info("Starting parallel copy to tmp table: {} in destination for stream: {}, schema: {}. Chunks count {}", tmpTableName, streamName, | ||
| schemaName, partitions.size()); | ||
|
|
||
| copyFilesInParallel(partitions); | ||
| LOGGER.info("Copy to tmp table {} in destination for stream {} complete.", tmpTableName, streamName); | ||
| } | ||
|
|
||
| @Override | ||
| public void copyIntoStage(List<String> files) { | ||
| final var copyQuery = String.format( | ||
| "COPY INTO %s.%s FROM '%s' " | ||
| + "CREDENTIALS=(aws_key_id='%s' aws_secret_key='%s') " | ||
| + "file_format = (type = csv field_delimiter = ',' skip_header = 0 FIELD_OPTIONALLY_ENCLOSED_BY = '\"') " | ||
| + "files = (" + generateFilesList(files) + " );", | ||
| schemaName, | ||
| tmpTableName, | ||
| generateBucketPath(), | ||
| s3Config.getAccessKeyId(), | ||
| s3Config.getSecretAccessKey()); | ||
|
|
||
| Exceptions.toRuntime(() -> db.execute(copyQuery)); | ||
| } | ||
|
|
||
| @Override | ||
| public String generateBucketPath() { | ||
| return "s3://" + s3Config.getBucketName() + "/" | ||
| + S3OutputPathHelper.getOutputPrefix(s3Config.getBucketPath(), configuredAirbyteStream.getStream()) + "/"; | ||
| } | ||
|
|
||
| @Override | ||
| public void copyS3CsvFileIntoTable(final JdbcDatabase database, | ||
| final String s3FileLocation, | ||
| final String schema, | ||
| final String tableName, | ||
| final S3DestinationConfig s3Config) | ||
| throws SQLException { | ||
| throw new RuntimeException("Snowflake Stream Copier should not copy individual files without use of a parallel copy"); | ||
|
|
||
| } | ||
|
|
||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can you add javadoc comments to the methods in this interface?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
done