fix pkl file not loading in StaticDataLoader (#1896)

* fix pkl file not loading in StaticDataLoader

* resolve hard code

* resolve hard code

Author: SunsetWolf

examples/data_demo/data_cache_demo.py

@@ -1,8 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 """
-    The motivation of this demo
-    - To show the data modules of Qlib is Serializable, users can dump processed data to disk to avoid duplicated data preprocessing
+The motivation of this demo
+- To show the data modules of Qlib is Serializable, users can dump processed data to disk to avoid duplicated data preprocessing
 """
 
 from copy import deepcopy

examples/data_demo/data_mem_resuse_demo.py

@@ -1,8 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 """
-    The motivation of this demo
-    - To show the data modules of Qlib is Serializable, users can dump processed data to disk to avoid duplicated data preprocessing
+The motivation of this demo
+- To show the data modules of Qlib is Serializable, users can dump processed data to disk to avoid duplicated data preprocessing
 """
 
 from copy import deepcopy

examples/orderbook_data/create_dataset.py

@@ -1,10 +1,10 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 """
-    NOTE:
-    - This scripts is a demo to import example data import Qlib
-    - !!!!!!!!!!!!!!!TODO!!!!!!!!!!!!!!!!!!!:
-        - Its structure is not well designed and very ugly, your contribution is welcome to make importing dataset easier
+NOTE:
+- This scripts is a demo to import example data import Qlib
+- !!!!!!!!!!!!!!!TODO!!!!!!!!!!!!!!!!!!!:
+    - Its structure is not well designed and very ugly, your contribution is welcome to make importing dataset easier
 """
 from datetime import date, datetime as dt
 import os

examples/workflow_by_code.py

@@ -1,7 +1,7 @@
 #  Copyright (c) Microsoft Corporation.
 #  Licensed under the MIT License.
 """
-Qlib provides two kinds of interfaces. 
+Qlib provides two kinds of interfaces.
 (1) Users could define the Quant research workflow by a simple configuration.
 (2) Qlib is designed in a modularized way and supports creating research workflow by code just like building blocks.
 

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
   "matplotlib",
   "jupyter",
   "nbconvert",
+  "pyarrow",
 ]
 
 [project.optional-dependencies]

qlib/contrib/torch.py

@@ -1,9 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 """
-    This module is not a necessary part of Qlib.
-    They are just some tools for convenience
-    It is should not imported into the core part of qlib
+This module is not a necessary part of Qlib.
+They are just some tools for convenience
+It is should not imported into the core part of qlib
 """
 import torch
 import numpy as np

qlib/data/dataset/loader.py

@@ -279,8 +279,11 @@ def _maybe_load_raw_data(self):
             )
             self._data.sort_index(inplace=True)
         elif isinstance(self._config, (str, Path)):
-            with Path(self._config).open("rb") as f:
-                self._data = pickle.load(f)
+            if str(self._config).strip().endswith(".parquet"):
+                self._data = pd.read_parquet(self._config, engine="pyarrow")
+            else:
+                with Path(self._config).open("rb") as f:
+                    self._data = pickle.load(f)
         elif isinstance(self._config, pd.DataFrame):
             self._data = self._config
 

qlib/rl/trainer/trainer.py

@@ -200,7 +200,7 @@ def fit(self, vessel: TrainingVesselBase, ckpt_path: Path | None = None) -> None
 
         if ckpt_path is not None:
             _logger.info("Resuming states from %s", str(ckpt_path))
-            self.load_state_dict(torch.load(ckpt_path))
+            self.load_state_dict(torch.load(ckpt_path, weights_only=False))
         else:
             self.initialize()
 

tests/rl/test_trainer.py

@@ -194,7 +194,7 @@ def test_trainer_checkpoint():
     assert (output_dir / "002.pth").exists()
     assert os.readlink(output_dir / "latest.pth") == str(output_dir / "002.pth")
 
-    trainer.load_state_dict(torch.load(output_dir / "001.pth"))
+    trainer.load_state_dict(torch.load(output_dir / "001.pth", weights_only=False))
     assert trainer.current_iter == 1
     assert trainer.current_episode == 100