🎉 Destination snowflake: reduce memory consumption (#10297)

* Avoid redundant adapter construction * Remove unused logger * Avoid redundant creation of buffer map * Decrease max batch byte size to 128 mb * Format code * Move data adapter to an instance variable * Bump version * Bump version in seed
airbytehq · Feb 15, 2022 · 6301cfa · 6301cfa
1 parent ed276f4
commit 6301cfa
Show file tree

Hide file tree

Showing 10 changed files with 54 additions and 77 deletions.
diff --git a/airbyte-config/init/src/main/resources/seed/destination_definitions.yaml b/airbyte-config/init/src/main/resources/seed/destination_definitions.yaml
@@ -185,7 +185,7 @@
 - name: Snowflake
   destinationDefinitionId: 424892c4-daac-4491-b35d-c6688ba547ba
   dockerRepository: airbyte/destination-snowflake
-  dockerImageTag: 0.4.9
+  dockerImageTag: 0.4.10
   documentationUrl: https://docs.airbyte.io/integrations/destinations/snowflake
   icon: snowflake.svg
 - name: MariaDB ColumnStore

diff --git a/airbyte-config/init/src/main/resources/seed/destination_specs.yaml b/airbyte-config/init/src/main/resources/seed/destination_specs.yaml
@@ -3817,7 +3817,7 @@
     supported_destination_sync_modes:
     - "overwrite"
     - "append"
-- dockerImage: "airbyte/destination-snowflake:0.4.9"
+- dockerImage: "airbyte/destination-snowflake:0.4.10"
   spec:
     documentationUrl: "https://docs.airbyte.io/integrations/destinations/snowflake"
     connectionSpecification:
@@ -3899,32 +3899,30 @@
           description: "The loading method used to send data to Snowflake."
           order: 8
           oneOf:
-          - title: "[Recommended] Internal Staging"
+          - title: "Select another option"
             additionalProperties: false
-            description: "Writes large batches of records to a file, uploads the file\
-              \ to Snowflake, then uses <pre>COPY INTO table</pre> to upload the file.\
-              \ Recommended for large production workloads for better speed and scalability."
+            description: "Select another option"
             required:
             - "method"
             properties:
               method:
                 type: "string"
                 enum:
-                - "Internal Staging"
-                default: "Internal Staging"
-          - title: "Standard Inserts"
+                - "Standard"
+                default: "Standard"
+          - title: "[Recommended] Internal Staging"
             additionalProperties: false
-            description: "Uses <pre>INSERT</pre> statements to send batches of records\
-              \ to Snowflake. Easiest (no setup) but not recommended for large production\
-              \ workloads due to slow speed."
+            description: "Writes large batches of records to a file, uploads the file\
+              \ to Snowflake, then uses <pre>COPY INTO table</pre> to upload the file.\
+              \ Recommended for large production workloads for better speed and scalability."
             required:
             - "method"
             properties:
               method:
                 type: "string"
                 enum:
-                - "Standard"
-                default: "Standard"
+                - "Internal Staging"
+                default: "Internal Staging"
           - title: "AWS S3 Staging"
             additionalProperties: false
             description: "Writes large batches of records to a file, uploads the file\

diff --git a/.../io/airbyte/integrations/destination/buffered_stream_consumer/BufferedStreamConsumer.java b/.../io/airbyte/integrations/destination/buffered_stream_consumer/BufferedStreamConsumer.java
@@ -23,10 +23,8 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -81,13 +79,13 @@ public class BufferedStreamConsumer extends FailureTrackingAirbyteMessageConsume
   private final RecordWriter recordWriter;
   private final CheckedConsumer<Boolean, Exception> onClose;
   private final Set<AirbyteStreamNameNamespacePair> streamNames;
-  private final List<AirbyteMessage> buffer;
   private final ConfiguredAirbyteCatalog catalog;
   private final CheckedFunction<JsonNode, Boolean, Exception> isValidRecord;
-  private final Map<AirbyteStreamNameNamespacePair, Long> pairToIgnoredRecordCount;
+  private final Map<AirbyteStreamNameNamespacePair, Long> streamToIgnoredRecordCount;
   private final Consumer<AirbyteMessage> outputRecordCollector;
   private final long maxQueueSizeInBytes;
   private long bufferSizeInBytes;
+  private Map<AirbyteStreamNameNamespacePair, List<AirbyteRecordMessage>> streamBuffer;
 
   private boolean hasStarted;
   private boolean hasClosed;
@@ -112,9 +110,9 @@ public BufferedStreamConsumer(final Consumer<AirbyteMessage> outputRecordCollect
     this.catalog = catalog;
     this.streamNames = AirbyteStreamNameNamespacePair.fromConfiguredCatalog(catalog);
     this.isValidRecord = isValidRecord;
-    this.buffer = new ArrayList<>(10_000);
     this.bufferSizeInBytes = 0;
-    this.pairToIgnoredRecordCount = new HashMap<>();
+    this.streamToIgnoredRecordCount = new HashMap<>();
+    this.streamBuffer = new HashMap<>();
   }
 
   @Override
@@ -123,7 +121,7 @@ protected void startTracked() throws Exception {
     Preconditions.checkState(!hasStarted, "Consumer has already been started.");
     hasStarted = true;
 
-    pairToIgnoredRecordCount.clear();
+    streamToIgnoredRecordCount.clear();
     LOGGER.info("{} started.", BufferedStreamConsumer.class);
 
     onStart.call();
@@ -141,7 +139,7 @@ protected void acceptTracked(final AirbyteMessage message) throws Exception {
       }
 
       if (!isValidRecord.apply(message.getRecord().getData())) {
-        pairToIgnoredRecordCount.put(stream, pairToIgnoredRecordCount.getOrDefault(stream, 0L) + 1L);
+        streamToIgnoredRecordCount.put(stream, streamToIgnoredRecordCount.getOrDefault(stream, 0L) + 1L);
         return;
       }
 
@@ -151,15 +149,12 @@ protected void acceptTracked(final AirbyteMessage message) throws Exception {
       final long messageSizeInBytes = ByteUtils.getSizeInBytesForUTF8CharSet(Jsons.serialize(recordMessage.getData()));
       if (bufferSizeInBytes + messageSizeInBytes > maxQueueSizeInBytes) {
         LOGGER.info("Flushing buffer...");
-        AirbyteSentry.executeWithTracing("FlushBuffer",
-            this::flushQueueToDestination,
-            Map.of("stream", stream.getName(),
-                "namespace", Objects.requireNonNullElse(stream.getNamespace(), "null"),
-                "bufferSizeInBytes", bufferSizeInBytes));
+        flushQueueToDestination(bufferSizeInBytes);
         bufferSizeInBytes = 0;
       }
 
-      buffer.add(message);
+      final List<AirbyteRecordMessage> bufferedRecords = streamBuffer.computeIfAbsent(stream, k -> new ArrayList<>());
+      bufferedRecords.add(message.getRecord());
       bufferSizeInBytes += messageSizeInBytes;
 
     } else if (message.getType() == Type.STATE) {
@@ -170,16 +165,13 @@ protected void acceptTracked(final AirbyteMessage message) throws Exception {
 
   }
 
-  private void flushQueueToDestination() throws Exception {
-    final Map<AirbyteStreamNameNamespacePair, List<AirbyteRecordMessage>> recordsByStream = buffer.stream()
-        .map(AirbyteMessage::getRecord)
-        .collect(Collectors.groupingBy(AirbyteStreamNameNamespacePair::fromRecordMessage));
-
-    buffer.clear();
-
-    for (final Map.Entry<AirbyteStreamNameNamespacePair, List<AirbyteRecordMessage>> entry : recordsByStream.entrySet()) {
-      recordWriter.accept(entry.getKey(), entry.getValue());
-    }
+  private void flushQueueToDestination(long bufferSizeInBytes) throws Exception {
+    AirbyteSentry.executeWithTracing("FlushBuffer", () -> {
+      for (final Map.Entry<AirbyteStreamNameNamespacePair, List<AirbyteRecordMessage>> entry : streamBuffer.entrySet()) {
+        recordWriter.accept(entry.getKey(), entry.getValue());
+      }
+    }, Map.of("bufferSizeInBytes", bufferSizeInBytes));
+    streamBuffer = new HashMap<>();
 
     if (pendingState != null) {
       lastFlushedState = pendingState;
@@ -199,13 +191,13 @@ protected void close(final boolean hasFailed) throws Exception {
     Preconditions.checkState(!hasClosed, "Has already closed.");
     hasClosed = true;
 
-    pairToIgnoredRecordCount
+    streamToIgnoredRecordCount
         .forEach((pair, count) -> LOGGER.warn("A total of {} record(s) of data from stream {} were invalid and were ignored.", count, pair));
     if (hasFailed) {
       LOGGER.error("executing on failed close procedure.");
     } else {
       LOGGER.info("executing on success close procedure.");
-      flushQueueToDestination();
+      flushQueueToDestination(bufferSizeInBytes);
     }
 
     try {

diff --git a/.../destination-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/DataAdapter.java b/.../destination-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/DataAdapter.java
@@ -8,13 +8,9 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import java.util.function.Function;
 import java.util.function.Predicate;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 public class DataAdapter {
 
-  private static final Logger LOGGER = LoggerFactory.getLogger(DataAdapter.class);
-
   private final Predicate<JsonNode> filterValueNode;
   private final Function<JsonNode, JsonNode> valueNodeAdapter;
 

diff --git a/...nation-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/JdbcSqlOperations.java b/...nation-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/JdbcSqlOperations.java
@@ -19,18 +19,28 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.UUID;
 import org.apache.commons.csv.CSVFormat;
 import org.apache.commons.csv.CSVPrinter;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
+@SuppressWarnings("OptionalUsedAsFieldOrParameterType")
 public abstract class JdbcSqlOperations implements SqlOperations {
 
-  private static final Logger LOGGER = LoggerFactory.getLogger(JdbcSqlOperations.class);
   protected static final String SHOW_SCHEMAS = "show schemas;";
   protected static final String NAME = "name";
 
+  // this adapter modifies record message before inserting them to the destination
+  protected final Optional<DataAdapter> dataAdapter;
+
+  protected JdbcSqlOperations() {
+    this.dataAdapter = Optional.empty();
+  }
+
+  protected JdbcSqlOperations(final DataAdapter dataAdapter) {
+    this.dataAdapter = Optional.of(dataAdapter);
+  }
+
   @Override
   public void createSchemaIfNotExists(final JdbcDatabase database, final String schemaName) throws Exception {
     if (!isSchemaExists(database, schemaName)) {
@@ -63,21 +73,14 @@ public String createTableQuery(final JdbcDatabase database, final String schemaN
   }
 
   protected void writeBatchToFile(final File tmpFile, final List<AirbyteRecordMessage> records) throws Exception {
-    PrintWriter writer = null;
-    try {
-      writer = new PrintWriter(tmpFile, StandardCharsets.UTF_8);
-      final var csvPrinter = new CSVPrinter(writer, CSVFormat.DEFAULT);
-
+    try (final PrintWriter writer = new PrintWriter(tmpFile, StandardCharsets.UTF_8);
+        final CSVPrinter csvPrinter = new CSVPrinter(writer, CSVFormat.DEFAULT)) {
       for (final AirbyteRecordMessage record : records) {
         final var uuid = UUID.randomUUID().toString();
         final var jsonData = Jsons.serialize(formatData(record.getData()));
         final var emittedAt = Timestamp.from(Instant.ofEpochMilli(record.getEmittedAt()));
         csvPrinter.printRecord(uuid, jsonData, emittedAt);
       }
-    } finally {
-      if (writer != null) {
-        writer.close();
-      }
     }
   }
 
@@ -137,7 +140,7 @@ public final void insertRecords(final JdbcDatabase database,
       throws Exception {
     AirbyteSentry.executeWithTracing("InsertRecords",
         () -> {
-          records.forEach(airbyteRecordMessage -> getDataAdapter().adapt(airbyteRecordMessage.getData()));
+          dataAdapter.ifPresent(adapter -> records.forEach(airbyteRecordMessage -> adapter.adapt(airbyteRecordMessage.getData())));
           insertRecordsInternal(database, records, schemaName, tableName);
         },
         Map.of("schema", Objects.requireNonNullElse(schemaName, "null"), "table", tableName, "recordCount", records.size()));
@@ -149,8 +152,4 @@ protected abstract void insertRecordsInternal(JdbcDatabase database,
                                                 String tableName)
       throws Exception;
 
-  protected DataAdapter getDataAdapter() {
-    return new DataAdapter(j -> false, c -> c);
-  }
-
 }
diff --git a/...res/src/main/java/io/airbyte/integrations/destination/postgres/PostgresSqlOperations.java b/...res/src/main/java/io/airbyte/integrations/destination/postgres/PostgresSqlOperations.java
@@ -5,7 +5,6 @@
 package io.airbyte.integrations.destination.postgres;
 
 import io.airbyte.db.jdbc.JdbcDatabase;
-import io.airbyte.integrations.destination.jdbc.DataAdapter;
 import io.airbyte.integrations.destination.jdbc.JdbcSqlOperations;
 import io.airbyte.protocol.models.AirbyteRecordMessage;
 import java.io.BufferedReader;
@@ -17,12 +16,12 @@
 import java.util.List;
 import org.postgresql.copy.CopyManager;
 import org.postgresql.core.BaseConnection;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 public class PostgresSqlOperations extends JdbcSqlOperations {
 
-  private static final Logger LOGGER = LoggerFactory.getLogger(PostgresSqlOperations.class);
+  public PostgresSqlOperations() {
+    super(new PostgresDataAdapter());
+  }
 
   @Override
   public void insertRecordsInternal(final JdbcDatabase database,
@@ -58,9 +57,4 @@ public void insertRecordsInternal(final JdbcDatabase database,
     });
   }
 
-  @Override
-  protected DataAdapter getDataAdapter() {
-    return new PostgresDataAdapter();
-  }
-
 }
diff --git a/airbyte-integrations/connectors/destination-snowflake/Dockerfile b/airbyte-integrations/connectors/destination-snowflake/Dockerfile
@@ -18,8 +18,8 @@ COPY build/distributions/${APPLICATION}*.tar ${APPLICATION}.tar
 
 RUN tar xf ${APPLICATION}.tar --strip-components=1
 
-ENV APPLICATION_VERSION 0.4.9
+ENV APPLICATION_VERSION 0.4.10
 ENV ENABLE_SENTRY true
 
-LABEL io.airbyte.version=0.4.9
+LABEL io.airbyte.version=0.4.10
 LABEL io.airbyte.name=airbyte/destination-snowflake
diff --git a/...ake/src/main/java/io/airbyte/integrations/destination/snowflake/SnowflakeDestination.java b/...ake/src/main/java/io/airbyte/integrations/destination/snowflake/SnowflakeDestination.java
@@ -10,13 +10,9 @@
 import io.airbyte.integrations.base.IntegrationRunner;
 import io.airbyte.integrations.destination.jdbc.copy.SwitchingDestination;
 import java.util.Map;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 public class SnowflakeDestination extends SwitchingDestination<SnowflakeDestination.DestinationType> {
 
-  private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeDestination.class);
-
   enum DestinationType {
     COPY_S3,
     COPY_GCS,

diff --git a/...o/airbyte/integrations/destination/snowflake/SnowflakeInternalStagingConsumerFactory.java b/...o/airbyte/integrations/destination/snowflake/SnowflakeInternalStagingConsumerFactory.java
@@ -46,7 +46,7 @@ public class SnowflakeInternalStagingConsumerFactory {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeInternalStagingConsumerFactory.class);
 
-  private static final long MAX_BATCH_SIZE_BYTES = 1024 * 1024 * 1024 / 4; // 256mb
+  private static final long MAX_BATCH_SIZE_BYTES = 128 * 1024 * 1024; // 128mb
   private final String CURRENT_SYNC_PATH = UUID.randomUUID().toString();
 
   public AirbyteMessageConsumer create(final Consumer<AirbyteMessage> outputRecordCollector,

diff --git a/docs/integrations/destinations/snowflake.md b/docs/integrations/destinations/snowflake.md
@@ -217,6 +217,8 @@ Finally, you need to add read/write permissions to your bucket with that email.
 
 | Version | Date       | Pull Request | Subject |
 |:--------|:-----------| :-----       | :------ |
+| 0.4.10  | 2022-02-14 | [\#10297](https://github.com/airbytehq/airbyte/pull/10297) | Halve the record buffer size to reduce memory consumption. |
+| 0.4.9   | 2022-02-14 | [\#10256](https://github.com/airbytehq/airbyte/pull/10256) | Add `ExitOnOutOfMemoryError` JVM flag. |
 | 0.4.8   | 2022-02-01 | [\#9959](https://github.com/airbytehq/airbyte/pull/9959) | Fix null pointer exception from buffered stream consumer. |
 | 0.4.7   | 2022-01-29 | [\#9745](https://github.com/airbytehq/airbyte/pull/9745) | Integrate with Sentry. |
 | 0.4.6   | 2022-01-28 | [#9623](https://github.com/airbytehq/airbyte/pull/9623) | Add jdbc_url_params support for optional JDBC parameters |