SNOW-922991: fix file chunks size for files over 80GB (#1777)

Author: sfc-gh-abenzaoui

DESCRIPTION.md

@@ -12,6 +12,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
   - Added support for `use_logical_type` in `write_pandas`.
   - Removed dependencies on pycryptodomex and oscrypto. All connections now go through OpenSSL via the cryptography library, which was already a dependency.
+  - Fixed issue with ingesting files over 80 GB to S3.
   - Added the `backoff_policy` argument to `snowflake.connector.connect` allowing for configurable backoff policy between retries of failed requests. See available implementations in the `backoff_policies` module.
   - Added the `socket_timeout` argument to `snowflake.connector.connect` specifying socket read and connect timeout.
   - Fixed `login_timeout` and `network_timeout` behaviour. Retries of login and network requests are now properly halted after these timeouts expire.

src/snowflake/connector/constants.py

@@ -318,6 +318,14 @@ class IterUnit(Enum):
     TABLE_UNIT = "table"
 
 
+# Amazon S3 multipart upload limits
+# https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html
+S3_DEFAULT_CHUNK_SIZE = 8 * 1024**2
+S3_MAX_OBJECT_SIZE = 5 * 1024**4
+S3_MAX_PART_SIZE = 5 * 1024**3
+S3_MIN_PART_SIZE = 5 * 1024**2
+S3_MAX_PARTS = 10000
+
 S3_CHUNK_SIZE = 8388608  # boto3 default
 AZURE_CHUNK_SIZE = 4 * megabyte
 

src/snowflake/connector/file_transfer_agent.py

@@ -7,6 +7,7 @@
 
 import binascii
 import glob
+import math
 import mimetypes
 import os
 import sys
@@ -21,12 +22,17 @@
 from .azure_storage_client import SnowflakeAzureRestClient
 from .compat import GET_CWD, IS_WINDOWS
 from .constants import (
+    AZURE_CHUNK_SIZE,
     AZURE_FS,
     CMD_TYPE_DOWNLOAD,
     CMD_TYPE_UPLOAD,
     GCS_FS,
     LOCAL_FS,
+    S3_DEFAULT_CHUNK_SIZE,
     S3_FS,
+    S3_MAX_OBJECT_SIZE,
+    S3_MAX_PARTS,
+    S3_MIN_PART_SIZE,
     ResultStatus,
     megabyte,
 )
@@ -183,6 +189,28 @@ def _update_progress(
     return progress == 1.0
 
 
+def _chunk_size_calculator(file_size: int) -> int:
+    # S3 has limitation on the num of parts to be uploaded, this helper method recalculate the num of parts
+    if file_size > S3_MAX_OBJECT_SIZE:
+        # check if we don't exceed the allowed S3 max file size 5 TiB
+        raise ValueError(
+            f"File size {file_size} exceeds the maximum file size {S3_MAX_OBJECT_SIZE} allowed in S3."
+        )
+
+    # num_parts = math.ceil(file_size / default_chunk_size)
+    # if num_parts is greater than the allowed S3_MAX_PARTS, we update our chunk_size, otherwise we use the default one
+    calculated_chunk_size = (
+        max(math.ceil(file_size / S3_MAX_PARTS), S3_MIN_PART_SIZE)
+        if math.ceil(file_size / S3_DEFAULT_CHUNK_SIZE) > S3_MAX_PARTS
+        else S3_DEFAULT_CHUNK_SIZE
+    )
+    if calculated_chunk_size != S3_DEFAULT_CHUNK_SIZE:
+        logger.debug(
+            f"Setting chunksize to {calculated_chunk_size} instead of the default {S3_DEFAULT_CHUNK_SIZE}."
+        )
+    return calculated_chunk_size
+
+
 def percent(seen_so_far: int, size: float) -> float:
     return 1.0 if seen_so_far >= size or size <= 0 else float(seen_so_far / size)
 
@@ -630,8 +658,6 @@ def function_and_callback_wrapper(
     def _create_file_transfer_client(
         self, meta: SnowflakeFileMeta
     ) -> SnowflakeStorageClient:
-        from .constants import AZURE_CHUNK_SIZE, S3_CHUNK_SIZE
-
         if self._stage_location_type == LOCAL_FS:
             return SnowflakeLocalStorageClient(
                 meta,
@@ -651,7 +677,7 @@ def _create_file_transfer_client(
                 meta,
                 self._credentials,
                 self._stage_info,
-                S3_CHUNK_SIZE,
+                _chunk_size_calculator(meta.src_file_size),
                 use_accelerate_endpoint=self._use_accelerate_endpoint,
                 use_s3_regional_url=self._use_s3_regional_url,
             )

test/unit/test_compute_chunk_size.py

@@ -0,0 +1,38 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+import pytest
+
+pytestmark = pytest.mark.skipolddriver
+
+
+def test_check_chunk_size():
+    from snowflake.connector.constants import (
+        S3_MAX_OBJECT_SIZE,
+        S3_MAX_PART_SIZE,
+        S3_MIN_PART_SIZE,
+    )
+    from snowflake.connector.file_transfer_agent import _chunk_size_calculator
+
+    expected_chunk_size = 8 * 1024**2
+    sample_file_size_2gb = 2 * 1024**3
+    sample_file_size_under_5tb = 4.9 * 1024**4
+    sample_file_size_6tb = 6 * 1024**4
+    sample_chunk_size_4mb = 4 * 1024**2
+
+    chunk_size_1 = _chunk_size_calculator(sample_file_size_2gb)
+    assert chunk_size_1 == expected_chunk_size
+
+    chunk_size_2 = _chunk_size_calculator(int(sample_file_size_under_5tb))
+    assert chunk_size_2 <= S3_MAX_PART_SIZE
+
+    with pytest.raises(ValueError) as exc:
+        _chunk_size_calculator(sample_file_size_6tb)
+    assert (
+        f"File size {sample_file_size_6tb} exceeds the maximum file size {S3_MAX_OBJECT_SIZE} allowed in S3."
+        in str(exc)
+    )
+
+    chunk_size_1 = _chunk_size_calculator(sample_chunk_size_4mb)
+    assert chunk_size_1 >= S3_MIN_PART_SIZE