Improving s3 cache strategy.

igorborgest · igorborgest · commit 5b39a170cfd6 · 2020-08-31T13:09:44.000-03:00
diff --git a/awswrangler/_config.py b/awswrangler/_config.py
@@ -29,7 +29,7 @@ class _ConfigArg(NamedTuple):
     "database": _ConfigArg(dtype=str, nullable=True),
     "max_cache_query_inspections": _ConfigArg(dtype=int, nullable=False),
     "max_cache_seconds": _ConfigArg(dtype=int, nullable=False),
-    "s3_read_ahead_size": _ConfigArg(dtype=int, nullable=False, enforced=True),
+    "s3_block_size": _ConfigArg(dtype=int, nullable=False, enforced=True),
 }
 
 
@@ -206,13 +206,13 @@ def max_cache_seconds(self, value: int) -> None:
         self._set_config_value(key="max_cache_seconds", value=value)
 
     @property
-    def s3_read_ahead_size(self) -> int:
-        """Property s3_read_ahead_size."""
-        return cast(int, self["s3_read_ahead_size"])
+    def s3_block_size(self) -> int:
+        """Property s3_block_size."""
+        return cast(int, self["s3_block_size"])
 
-    @s3_read_ahead_size.setter
-    def s3_read_ahead_size(self, value: int) -> None:
-        self._set_config_value(key="s3_read_ahead_size", value=value)
+    @s3_block_size.setter
+    def s3_block_size(self, value: int) -> None:
+        self._set_config_value(key="s3_block_size", value=value)
 
 
 def _inject_config_doc(doc: Optional[str], available_configs: Tuple[str, ...]) -> str:
diff --git a/awswrangler/s3/_fs.py b/awswrangler/s3/_fs.py
@@ -4,20 +4,21 @@
 import io
 import itertools
 import logging
+import math
 import socket
 from contextlib import contextmanager
 from typing import Any, BinaryIO, Dict, Iterator, List, Optional, Set, Tuple, Union, cast
 
 import boto3
-from botocore.exceptions import ClientError
+from botocore.exceptions import ClientError, ReadTimeoutError
 
 from awswrangler import _utils, exceptions
 from awswrangler._config import apply_configs
 from awswrangler.s3._describe import size_objects
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
-_S3_RETRYABLE_ERRORS: Tuple[Any, Any] = (socket.timeout, ConnectionError)
+_S3_RETRYABLE_ERRORS: Tuple[Any, Any, Any] = (socket.timeout, ConnectionError, ReadTimeoutError)
 
 _MIN_WRITE_BLOCK: int = 5_242_880  # 5 MB (5 * 2**20)
 _MIN_PARALLEL_READ_BLOCK: int = 5_242_880  # 5 MB (5 * 2**20)
@@ -178,14 +179,15 @@ class _S3Object:  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
         path: str,
-        s3_read_ahead_size: int,
+        s3_block_size: int,
         mode: str,
         use_threads: bool,
         s3_additional_kwargs: Optional[Dict[str, str]],
         boto3_session: Optional[boto3.Session],
         newline: Optional[str],
         encoding: Optional[str],
     ) -> None:
+        self.closed: bool = False
         self._use_threads = use_threads
         self._newline: str = "\n" if newline is None else newline
         self._encoding: str = "utf-8" if encoding is None else encoding
@@ -194,11 +196,13 @@ def __init__(
         if mode not in {"rb", "wb", "r", "w"}:
             raise NotImplementedError("File mode must be {'rb', 'wb', 'r', 'w'}, not %s" % mode)
         self._mode: str = "rb" if mode is None else mode
-        self._s3_read_ahead_size: int = s3_read_ahead_size
+        if s3_block_size < 2:
+            raise exceptions.InvalidArgumentValue("s3_block_size MUST > 1")
+        self._s3_block_size: int = s3_block_size
+        self._s3_half_block_size: int = s3_block_size // 2
         self._s3_additional_kwargs: Dict[str, str] = {} if s3_additional_kwargs is None else s3_additional_kwargs
         self._client: boto3.client = _utils.client(service_name="s3", session=self._boto3_session)
         self._loc: int = 0
-        self.closed: bool = False
 
         if self.readable() is True:
             self._cache: bytes = b""
@@ -209,6 +213,7 @@ def __init__(
                 raise exceptions.InvalidArgumentValue(f"S3 object w/o defined size: {path}")
             self._size: int = size
             _logger.debug("self._size: %s", self._size)
+            _logger.debug("self._s3_block_size: %s", self._s3_block_size)
         elif self.writable() is True:
             self._mpu: Dict[str, Any] = {}
             self._buffer: io.BytesIO = io.BytesIO()
@@ -289,16 +294,60 @@ def _fetch_range_proxy(self, start: int, end: int) -> bytes:
             )
 
     def _fetch(self, start: int, end: int) -> None:
-        if end > self._size:
-            end = self._size
+        end = self._size if end > self._size else end
+        start = 0 if start < 0 else start
+
+        if start >= self._start and end <= self._end:
+            return None  # Does not require download
 
-        if start < self._start or end > self._end:
+        if end - start >= self._s3_block_size:  # Fetching length greater than cache length
+            self._cache = self._fetch_range_proxy(start, end)
             self._start = start
-            if ((end - start) < self._s3_read_ahead_size) and (end < self._size):
-                self._end = start + self._s3_read_ahead_size
-            else:
-                self._end = end
-            self._cache = self._fetch_range_proxy(self._start, self._end)
+            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) / 2))
+        new_block_start: int = mid - self._s3_half_block_size
+        new_block_end: int = mid + self._s3_half_block_size
+        _logger.debug("new_block_start: %s / new_block_end: %s / mid: %s", new_block_start, new_block_end, mid)
+        if new_block_start < 0 and new_block_end > self._size:  # both ends overflowing
+            new_block_start = 0
+            new_block_end = self._size
+        elif new_block_end > self._size:  # right overflow
+            new_block_start = new_block_start - (new_block_end - self._size)
+            new_block_start = 0 if new_block_start < 0 else new_block_start
+            new_block_end = self._size
+        elif new_block_start < 0:  # left overflow
+            new_block_end = new_block_end + (0 - new_block_start)
+            new_block_end = self._size if new_block_end > self._size else new_block_end
+            new_block_start = 0
+        _logger.debug(
+            "new_block_start: %s / new_block_end: %s/ self._start: %s / self._end: %s",
+            new_block_start,
+            new_block_end,
+            self._start,
+            self._end,
+        )
+
+        # Calculating missing bytes in cache
+        if (new_block_start < self._start and new_block_end > self._end) or (
+            new_block_start > self._end and new_block_end < self._start
+        ):  # Full block download
+            self._cache = self._fetch_range_proxy(new_block_start, new_block_end)
+        elif new_block_end > self._end:
+            prune_diff: int = new_block_start - self._start
+            self._cache = self._cache[prune_diff:] + self._fetch_range_proxy(self._end, new_block_end)
+        elif new_block_start < self._start:
+            prune_diff = new_block_end - self._end
+            self._cache = self._cache[:-prune_diff] + self._fetch_range_proxy(new_block_start, self._start)
+        else:
+            raise RuntimeError("Wrangler's cache calculation error.")
+        self._start = new_block_start
+        self._end = new_block_end
+
+        return None
 
     def read(self, length: int = -1) -> Union[bytes, str]:
         """Return cached data and fetch on demand chunks."""
@@ -313,12 +362,11 @@ def read(self, length: int = -1) -> Union[bytes, str]:
         self._fetch(self._loc, self._loc + length)
         out: bytes = self._cache[self._loc - self._start : self._loc - self._start + length]
         self._loc += len(out)
-
         return out
 
     def readline(self, length: int = -1) -> Union[bytes, str]:
         """Read until the next line terminator."""
-        self._fetch(self._loc, self._loc + self._s3_read_ahead_size)
+        self._fetch(self._loc, self._loc + self._s3_block_size)
         while True:
             found: int = self._cache[self._loc - self._start :].find(self._newline.encode(encoding=self._encoding))
 
@@ -329,7 +377,7 @@ def readline(self, length: int = -1) -> Union[bytes, str]:
             if self._end >= self._size:
                 return self.read(length)
 
-            self._fetch(self._loc, self._end + self._s3_read_ahead_size)
+            self._fetch(self._loc, self._end + self._s3_half_block_size)
 
     def readlines(self) -> List[Union[bytes, str]]:
         """Return all lines as list."""
@@ -472,7 +520,7 @@ def open_s3_object(
     mode: str,
     use_threads: bool = False,
     s3_additional_kwargs: Optional[Dict[str, str]] = None,
-    s3_read_ahead_size: int = 4_194_304,  # 4 MB (4 * 2**20)
+    s3_block_size: int = 4_194_304,  # 4 MB (4 * 2**20)
     boto3_session: Optional[boto3.Session] = None,
     newline: Optional[str] = "\n",
     encoding: Optional[str] = "utf-8",
@@ -483,7 +531,7 @@ def open_s3_object(
     try:
         s3obj = _S3Object(
             path=path,
-            s3_read_ahead_size=s3_read_ahead_size,
+            s3_block_size=s3_block_size,
             mode=mode,
             use_threads=use_threads,
             s3_additional_kwargs=s3_additional_kwargs,
@@ -494,7 +542,13 @@ def open_s3_object(
         if "b" in mode:  # binary
             yield s3obj
         else:  # text
-            text_s3obj = io.TextIOWrapper(cast(BinaryIO, s3obj), encoding=encoding, newline=newline)
+            text_s3obj = io.TextIOWrapper(
+                buffer=cast(BinaryIO, s3obj),
+                encoding=encoding,
+                newline=newline,
+                line_buffering=False,
+                write_through=False,
+            )
             yield text_s3obj
     finally:
         if text_s3obj is not None and text_s3obj.closed is False:
diff --git a/awswrangler/s3/_read_parquet.py b/awswrangler/s3/_read_parquet.py
@@ -40,7 +40,7 @@ def _read_parquet_metadata_file(
         path=path,
         mode="rb",
         use_threads=use_threads,
-        s3_read_ahead_size=1_048_576,  # 1 MB (1 * 2**20)
+        s3_block_size=1_048_576,  # 1 MB (1 * 2**20)
         s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,
     ) as f:
@@ -257,7 +257,7 @@ def _read_parquet_chunked(
             path=path,
             mode="rb",
             use_threads=use_threads,
-            s3_read_ahead_size=10_485_760,  # 10 MB (10 * 2**20)
+            s3_block_size=10_485_760,  # 10 MB (10 * 2**20)
             s3_additional_kwargs=s3_additional_kwargs,
             boto3_session=boto3_session,
         ) as f:
@@ -319,7 +319,7 @@ def _read_parquet_file(
         path=path,
         mode="rb",
         use_threads=use_threads,
-        s3_read_ahead_size=134_217_728,  # 128 MB (128 * 2**20)
+        s3_block_size=134_217_728,  # 128 MB (128 * 2**20)
         s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,
     ) as f:
@@ -339,7 +339,7 @@ def _count_row_groups(
         path=path,
         mode="rb",
         use_threads=use_threads,
-        s3_read_ahead_size=1_048_576,  # 1 MB (1 * 2**20)
+        s3_block_size=1_048_576,  # 1 MB (1 * 2**20)
         s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,
     ) as f:
@@ -361,7 +361,7 @@ def _read_parquet_row_group(
         path=path,
         mode="rb",
         use_threads=use_threads,
-        s3_read_ahead_size=10_485_760,  # 10 MB (10 * 2**20)
+        s3_block_size=10_485_760,  # 10 MB (10 * 2**20)
         s3_additional_kwargs=s3_additional_kwargs,
         boto3_session=boto3_session,
     ) as f:
diff --git a/awswrangler/s3/_read_text.py b/awswrangler/s3/_read_text.py
@@ -51,7 +51,7 @@ def _read_text_chunked(
         with open_s3_object(
             path=path,
             mode=mode,
-            s3_read_ahead_size=10_485_760,  # 10 MB (10 * 2**20)
+            s3_block_size=10_485_760,  # 10 MB (10 * 2**20)
             encoding=encoding,
             use_threads=use_threads,
             s3_additional_kwargs=s3_additional_kwargs,
@@ -78,7 +78,7 @@ def _read_text_file(
         path=path,
         mode=mode,
         use_threads=use_threads,
-        s3_read_ahead_size=134_217_728,  # 128 MB (128 * 2**20)
+        s3_block_size=134_217_728,  # 128 MB (128 * 2**20)
         encoding=encoding,
         s3_additional_kwargs=s3_additional_kwargs,
         newline=newline,
diff --git a/pytest.ini b/pytest.ini
@@ -3,6 +3,6 @@ log_cli=False
 filterwarnings =
     ignore::DeprecationWarning
 addopts =
-    --log-cli-format "[%(asctime)s][%(levelname)s][%(name)s][%(funcName)s] %(message)s"
+    --log-cli-format "[%(asctime)s][%(levelname)s][%(name)s][%(funcName)s][%(thread)d] %(message)s"
     --verbose
     --capture=sys
diff --git a/tests/test_config.py b/tests/test_config.py
@@ -24,11 +24,11 @@ def test_basics(path, glue_database, glue_table):
 
     # Testing configured s3 block size
     size = 1 * 2 ** 20  # 1 MB
-    wr.config.s3_read_ahead_size = size
+    wr.config.s3_block_size = size
     with open_s3_object(path, mode="wb") as s3obj:
         s3obj.write(b"foo")
     with open_s3_object(path, mode="rb") as s3obj:
-        assert s3obj._s3_read_ahead_size == size
+        assert s3obj._s3_block_size == size
 
     # Resetting all configs
     wr.config.reset()
diff --git a/tests/test_fs.py b/tests/test_fs.py
@@ -34,7 +34,7 @@ def test_read_full(path, mode, use_threads):
     bucket, key = wr._utils.parse_path(path)
     text = "AHDG*AWY&GD*A&WGd*AWgd87AGWD*GA*G*g*AGˆˆ&ÂDTW&ˆˆD&ÂTW7ˆˆTAWˆˆDAW&ˆˆAWGDIUHWOD#N"
     client_s3.put_object(Body=text, Bucket=bucket, Key=key)
-    with open_s3_object(path, mode=mode, s3_read_ahead_size=100, newline="\n", use_threads=use_threads) as s3obj:
+    with open_s3_object(path, mode=mode, s3_block_size=100, newline="\n", use_threads=use_threads) as s3obj:
         if mode == "r":
             assert s3obj.read() == text
         else:
@@ -52,7 +52,7 @@ def test_read_chunked(path, mode, block_size, use_threads):
     bucket, key = wr._utils.parse_path(path)
     text = "0123456789"
     client_s3.put_object(Body=text, Bucket=bucket, Key=key)
-    with open_s3_object(path, mode=mode, s3_read_ahead_size=block_size, newline="\n", use_threads=use_threads) as s3obj:
+    with open_s3_object(path, mode=mode, s3_block_size=block_size, newline="\n", use_threads=use_threads) as s3obj:
         if mode == "r":
             for i in range(3):
                 assert s3obj.read(1) == text[i]
@@ -67,22 +67,23 @@ def test_read_chunked(path, mode, block_size, use_threads):
 
 @pytest.mark.parametrize("use_threads", [True, False])
 @pytest.mark.parametrize("mode", ["r", "rb"])
-@pytest.mark.parametrize("block_size", [1, 2, 3, 10, 23, 48, 65, 100])
+@pytest.mark.parametrize("block_size", [2, 3, 10, 23, 48, 65, 100])
 def test_read_line(path, mode, block_size, use_threads):
     client_s3 = boto3.client("s3")
     path = f"{path}0.txt"
     bucket, key = wr._utils.parse_path(path)
     text = "0\n11\n22222\n33333333333333\n44444444444444444444444444444444444444444444\n55555"
     expected = ["0\n", "11\n", "22222\n", "33333333333333\n", "44444444444444444444444444444444444444444444\n", "55555"]
     client_s3.put_object(Body=text, Bucket=bucket, Key=key)
-    with open_s3_object(path, mode=mode, s3_read_ahead_size=block_size, newline="\n", use_threads=use_threads) as s3obj:
+    with open_s3_object(path, mode=mode, s3_block_size=block_size, newline="\n", use_threads=use_threads) as s3obj:
         for i, line in enumerate(s3obj):
             if mode == "r":
                 assert line == expected[i]
             else:
                 assert line == expected[i].encode("utf-8")
         s3obj.seek(0)
         lines = s3obj.readlines()
+        print(lines)
         if mode == "r":
             assert lines == expected
         else:
@@ -136,11 +137,7 @@ def test_additional_kwargs(path, kms_key_id, s3_additional_kwargs, use_threads):
     with open_s3_object(path, mode="w", s3_additional_kwargs=s3_additional_kwargs, use_threads=use_threads) as s3obj:
         s3obj.write("foo")
     with open_s3_object(
-        path,
-        mode="r",
-        s3_read_ahead_size=10_000_000,
-        s3_additional_kwargs=s3_additional_kwargs,
-        use_threads=use_threads,
+        path, mode="r", s3_block_size=10_000_000, s3_additional_kwargs=s3_additional_kwargs, use_threads=use_threads,
     ) as s3obj:
         assert s3obj.read() == "foo"
     desc = wr.s3.describe_objects([path])[path]
@@ -160,3 +157,20 @@ def test_pyarrow(path, glue_table, glue_database):
     ensure_data_types(df2, has_list=True)
     assert df2.shape == (3, 19)
     assert df.iint8.sum() == df2.iint8.sum()
+
+
+@pytest.mark.parametrize("use_threads", [True, False])
+@pytest.mark.parametrize("block_size", [2, 3, 5, 8, 9, 15])
+@pytest.mark.parametrize("text", ["012345678", "0123456789"])
+def test_cache(path, use_threads, block_size, text):
+    client_s3 = boto3.client("s3")
+    path = f"{path}0.txt"
+    bucket, key = wr._utils.parse_path(path)
+    client_s3.put_object(Body=text, Bucket=bucket, Key=key)
+    with open_s3_object(path, mode="rb", s3_block_size=block_size, use_threads=use_threads) as s3obj:
+        for i in range(len(text)):
+            value = s3obj.read(1)
+            print(value)
+            assert value == text[i].encode("utf-8")
+            assert len(s3obj._cache) in (block_size, block_size - 1, len(text))
+    assert s3obj._cache == b""
