Merge pull request #736 from Kanaries/codex/fix-data-parser-cache-leak-attributed

[codex] Fix data parser cache leaks with attribution

Author: ObservedObserver

pygwalker/data_parsers/base.py

@@ -1,6 +1,6 @@
 from typing import Generic, Dict, List, Any, Optional
 from typing_extensions import Literal
-from functools import lru_cache
+from functools import cached_property, lru_cache
 from datetime import datetime, date
 from datetime import timedelta
 import abc
@@ -142,16 +142,14 @@ def __init__(
         self.infer_number_to_dimension = infer_number_to_dimension
         self.other_params = other_params
 
-    @property
-    @lru_cache()
+    @cached_property
     def field_metas(self) -> List[Dict[str, str]]:
         duckdb.register("pygwalker_mid_table", self._duckdb_df)
         result = duckdb.query("SELECT * FROM pygwalker_mid_table LIMIT 1")
         data = result.fetchone()
         return get_data_meta_type(dict(zip(result.columns, data))) if data else []
 
-    @property
-    @lru_cache()
+    @cached_property
     def raw_fields(self) -> List[Dict[str, str]]:
         return [self._infer_prop(col, self.field_specs) for _, col in enumerate(self._example_df.columns)]
 

pygwalker/data_parsers/cloud_dataset_parser.py

@@ -1,5 +1,5 @@
 from typing import Any, Dict, List, Optional
-from functools import lru_cache
+from functools import cached_property
 from decimal import Decimal
 import logging
 import io
@@ -41,14 +41,12 @@ def _get_example_pandas_df(self) -> pd.DataFrame:
                 example_df[column] = example_df[column].astype(float)
         return example_df
 
-    @property
-    @lru_cache()
+    @cached_property
     def field_metas(self) -> List[Dict[str, str]]:
         data = self._get_all_datas(1)
         return get_data_meta_type(data[0]) if data else []
 
-    @property
-    @lru_cache()
+    @cached_property
     def raw_fields(self) -> List[Dict[str, str]]:
         pandas_parser = PandasDataFrameDataParser(
             self.example_pandas_df,

pygwalker/data_parsers/database_parser.py

@@ -1,5 +1,5 @@
 from typing import Any, Dict, List, Optional
-from functools import lru_cache
+from functools import cached_property
 from decimal import Decimal
 import logging
 import json
@@ -179,14 +179,12 @@ def _format_sql(self, sql: str) -> str:
     def placeholder_table_name(self) -> str:
         return "___pygwalker_temp_view_name___"
 
-    @property
-    @lru_cache()
+    @cached_property
     def field_metas(self) -> List[Dict[str, str]]:
         data = self._get_datas_by_sql(f"SELECT * FROM {self.placeholder_table_name} LIMIT 1")
         return get_data_meta_type(data[0]) if data else []
 
-    @property
-    @lru_cache()
+    @cached_property
     def raw_fields(self) -> List[Dict[str, str]]:
         pandas_parser = PandasDataFrameDataParser(
             self.example_pandas_df,

pygwalker/data_parsers/spark_parser.py

@@ -1,5 +1,5 @@
 from typing import Any, Dict, List, Optional
-from functools import lru_cache
+from functools import cached_property
 import logging
 import io
 
@@ -41,8 +41,7 @@ def __init__(
         self.infer_number_to_dimension = infer_number_to_dimension
         self.other_params = other_params
 
-    @property
-    @lru_cache()
+    @cached_property
     def raw_fields(self) -> List[Dict[str, str]]:
         pandas_parser = PandasDataFrameDataParser(
             self.example_pandas_df,
@@ -53,8 +52,7 @@ def raw_fields(self) -> List[Dict[str, str]]:
         )
         return pandas_parser.raw_fields
 
-    @property
-    @lru_cache()
+    @cached_property
     def field_metas(self) -> List[Dict[str, str]]:
         data = self.get_datas_by_sql("SELECT * FROM pygwalker_mid_table LIMIT 1")
         return get_data_meta_type(data[0]) if data else []

tests/test_data_parsers.py

@@ -1,6 +1,9 @@
+import gc
 import os.path
+import pickle
 import subprocess
 import sys
+import weakref
 
 from sqlalchemy import create_engine
 import pandas as pd
@@ -11,6 +14,7 @@
 from pygwalker.services.data_parsers import get_dataset_hash, get_parser
 from pygwalker.data_parsers.database_parser import Connector, DatabaseDataParser, text
 from pygwalker.data_parsers.database_parser import _check_view_sql
+from pygwalker.data_parsers.pandas_parser import PandasDataFrameDataParser
 from pygwalker.errors import ViewSqlSameColumnError
 
 datas = [
@@ -94,6 +98,29 @@ def test_get_parser_reports_supported_inputs_for_unsupported_dataset():
     assert "cloud dataset id string" in message
 
 
+@pytest.mark.parametrize("cached_property_name", ["raw_fields", "field_metas"])
+def test_pandas_parser_cached_properties_do_not_retain_parser_or_dataframe(cached_property_name):
+    def create_refs():
+        df = pd.DataFrame({"city": ["London", "Tokyo"], "value": [1, 2]})
+        parser = PandasDataFrameDataParser(df, [], True, True, {})
+        getattr(parser, cached_property_name)
+        return weakref.ref(parser), weakref.ref(df)
+
+    parser_ref, df_ref = create_refs()
+    gc.collect()
+
+    assert parser_ref() is None
+    assert df_ref() is None
+
+
+def test_pandas_parser_remains_picklable_after_cached_properties_are_loaded():
+    parser = PandasDataFrameDataParser(pd.DataFrame({"city": ["London"], "value": [1]}), [], True, True, {})
+    parser.raw_fields
+    parser.field_metas
+
+    pickle.dumps(parser)
+
+
 @pytest.mark.parametrize(
     "module_name",
     [