feat(pems_data/cache): helpers to get/set DataFrames

use the DataFrame serialization functions as mutators

Author: thekaveman

pems_data/src/pems_data/cache.py

@@ -2,8 +2,11 @@
 import os
 from typing import Any, Callable
 
+import pandas as pd
 import redis
 
+from pems_data.serialization import arrow_bytes_to_df, df_to_arrow_bytes
+
 logger = logging.getLogger(__name__)
 
 
@@ -76,6 +79,10 @@ def get(self, key: str, mutate_func: Callable[[Any], Any] = None) -> Any:
         logger.warning(f"cache unavailable to get: {key}")
         return None
 
+    def get_df(self, key: str) -> pd.DataFrame:
+        """Get a `pandas.DataFrame` from the cache, or None if the key doesn't exist."""
+        return self.get(key, mutate_func=arrow_bytes_to_df)
+
     def set(self, key: str, value: Any, mutate_func: Callable[[Any], Any] = None) -> None:
         """Set a value in the cache.
 
@@ -92,3 +99,7 @@ def set(self, key: str, value: Any, mutate_func: Callable[[Any], Any] = None) ->
             self.c.set(key, value)
         else:
             logger.warning(f"cache unavailable to set: {key}")
+
+    def set_df(self, key: str, value: pd.DataFrame) -> None:
+        """Set a `pandas.DataFrame` in the cache."""
+        self.set(key, value, mutate_func=df_to_arrow_bytes)

tests/pytest/pems_data/conftest.py

@@ -0,0 +1,7 @@
+import pandas as pd
+import pytest
+
+
+@pytest.fixture
+def df() -> pd.DataFrame:
+    return pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})

tests/pytest/pems_data/test_cache.py

@@ -1,5 +1,6 @@
-import redis
+import pandas as pd
 import pytest
+import redis
 
 from pems_data.cache import Cache, redis_connection
 
@@ -107,7 +108,7 @@ def test_get(self, cache: Cache, mock_redis_connection, spy_connect):
         spy_connect.assert_called_once()
         mock_redis_connection.get.assert_called_once_with("test-key")
 
-    def test_get_mutate(self, cache: Cache, mock_redis_connection, spy_connect):
+    def test_get__mutate(self, cache: Cache, mock_redis_connection, spy_connect):
         expected = 2
         mock_redis_connection.get.return_value = 1
 
@@ -123,10 +124,60 @@ def test_set(self, cache: Cache, mock_redis_connection, spy_connect):
         spy_connect.assert_called_once()
         mock_redis_connection.set.assert_called_once_with("test-key", "test-value")
 
-    def test_set_mutate(self, cache: Cache, mock_redis_connection, spy_connect):
+    def test_set__mutate(self, cache: Cache, mock_redis_connection, spy_connect):
         expected = 2
 
         cache.set("test-key", 1, lambda v: v + 1)
 
         spy_connect.assert_called_once()
         mock_redis_connection.set.assert_called_once_with("test-key", expected)
+
+    def test_get_df(self, cache: Cache, mock_redis_connection, mocker, spy_connect):
+        # Mock arrow_bytes_to_df to return a DataFrame
+        df = mocker.Mock(spec=pd.DataFrame)
+        arrow_bytes = b"arrow-bytes"
+        mock_redis_connection.get.return_value = arrow_bytes
+        mock_arrow_bytes_to_df = mocker.patch("pems_data.cache.arrow_bytes_to_df", return_value=df)
+
+        result = cache.get_df("df-key")
+
+        spy_connect.assert_called_once()
+        mock_redis_connection.get.assert_called_once_with("df-key")
+        mock_arrow_bytes_to_df.assert_called_once_with(arrow_bytes)
+        assert result == df
+
+    def test_set_df(self, cache: Cache, mock_redis_connection, spy_connect, df):
+        """Test setting a DataFrame in the cache"""
+        cache.set_df("test-key", df)
+
+        spy_connect.assert_called_once()
+        mock_redis_connection.set.assert_called_once()
+        # Verify first arg is the key
+        assert mock_redis_connection.set.call_args[0][0] == "test-key"
+        # Verify second arg is bytes (arrow serialized)
+        assert isinstance(mock_redis_connection.set.call_args[0][1], bytes)
+
+    def test_set_df__empty_df(self, cache: Cache, mock_redis_connection, spy_connect):
+        empty_df = pd.DataFrame()
+        cache.set_df("test-key", empty_df)
+
+        spy_connect.assert_called_once()
+        mock_redis_connection.set.assert_called_once()
+        # Verify empty DataFrame is handled
+        assert isinstance(mock_redis_connection.set.call_args[0][1], bytes)
+
+    def test_set_df__roundtrip(self, cache: Cache, mock_redis_connection, spy_connect, df, mocker):
+        # Setup mock to return the serialized value on get
+        def mock_set(key, value):
+            mock_redis_connection.get.return_value = value
+
+        mock_redis_connection.set.side_effect = mock_set
+
+        # Set the DataFrame
+        cache.set_df("test-key", df)
+
+        # Get it back
+        result = cache.get_df("test-key")
+
+        pd.testing.assert_frame_equal(result, df)
+        spy_connect.assert_has_calls([mocker.call(), mocker.call()])

tests/pytest/pems_data/test_serialization.py

@@ -5,11 +5,6 @@
 from pems_data.serialization import arrow_bytes_to_df, df_to_arrow_bytes
 
 
-@pytest.fixture
-def df():
-    return pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
-
-
 @pytest.fixture
 def arrow_bytes(df):
     # convert df to actual arrow bytes for testing