Add tests for _one_shot_download. #376

Author: igorborgest

awswrangler/s3/_fs.py

@@ -193,11 +193,13 @@ def __init__(
             raise NotImplementedError("File mode must be {'rb', 'wb', 'r', 'w'}, not %s" % mode)
         self._mode: str = "rb" if mode is None else mode
         self._one_shot_download: bool = False
-        if s3_block_size == 1:
-            raise exceptions.InvalidArgumentValue("s3_block_size MUST > 1 to define a valid size or "
-                                                  "< 1 to avoid blocks and always execute one shot downloads.")
-        elif s3_block_size < 1:
-            _logger.debug(f"s3_block_size of %d, enabling one_shot_download.", s3_block_size)
+        if 0 < s3_block_size < 3:
+            raise exceptions.InvalidArgumentValue(
+                "s3_block_size MUST > 2 to define a valid size or "
+                "< 1 to avoid blocks and always execute one shot downloads."
+            )
+        if s3_block_size <= 0:
+            _logger.debug("s3_block_size of %d, enabling one_shot_download.", s3_block_size)
             self._one_shot_download = True
         self._s3_block_size: int = s3_block_size
         self._s3_half_block_size: int = s3_block_size // 2
@@ -303,18 +305,18 @@ def _fetch(self, start: int, end: int) -> None:
         if start >= self._start and end <= self._end:
             return None  # Does not require download
 
-        if end - start >= self._s3_block_size:  # Fetching length greater than cache length
-            self._cache = self._fetch_range_proxy(start, end)
-            self._start = start
-            self._end = end
-            return None
-
         if self._one_shot_download:
             self._start = 0
             self._end = self._size
             self._cache = self._fetch_range_proxy(self._start, self._end)
             return None
 
+        if end - start >= self._s3_block_size:  # Fetching length greater than cache length
+            self._cache = self._fetch_range_proxy(start, end)
+            self._start = start
+            self._end = end
+            return None
+
         # Calculating block START and END positions
         _logger.debug("Downloading: %s (start) / %s (end)", start, end)
         mid: int = int(math.ceil((start + (end - 1)) / 2))

awswrangler/s3/_read_parquet.py

@@ -314,11 +314,12 @@ def _read_parquet_file(
     s3_additional_kwargs: Optional[Dict[str, str]],
     use_threads: bool,
 ) -> pa.Table:
+    s3_block_size: int = 20_971_520 if columns else -1  # One shot for a full read otherwise 20 MB (20 * 2**20)
     with open_s3_object(
         path=path,
         mode="rb",
         use_threads=use_threads,
-        s3_block_size=-1,  # One shot download
+        s3_block_size=s3_block_size,
         s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,
     ) as f:

tests/test_fs.py

@@ -14,7 +14,7 @@
 
 
 @pytest.mark.parametrize("use_threads", [True, False])
-@pytest.mark.parametrize("block_size", list(range(3, 10)))
+@pytest.mark.parametrize("block_size", list(range(3, 10)) + [-1])
 @pytest.mark.parametrize("length", list(range(1, 10)))
 @pytest.mark.parametrize(
     "seq", [(9, 5, 4, 9, 2, 6), (9, 2, 6, 1, 5, 9), (1, 5, 9, 1, 3, 9), (1, 3, 9, 1, 2, 3), (1, 2, 3, 9, 5, 4)]
@@ -35,7 +35,9 @@ def test_read(path, use_threads, block_size, seq, length):
                 assert data[0:1] == text[i].encode("utf-8")
                 assert data == f.read(length)
                 logger.debug(s3obj._cache)
-                if length > block_size:
+                if block_size < 1:
+                    assert len(s3obj._cache) == s3obj._size
+                elif length > block_size:
                     assert block_size <= len(s3obj._cache) <= length
                 else:
                     assert len(s3obj._cache) == block_size