Merge pull request #493 from bxdd/optimize_resam_data

optimize performance of resam data in rule_strategy & exchange

qlib/data/dataset/handler.py

@@ -197,7 +197,7 @@ class DataHandler(Serializable):
         -------
         pd.DataFrame.
         """
-        from .storage import HasingStockStorage
+        from .storage import BaseHandlerStorage
 
         data_storage = self._data
         if isinstance(data_storage, pd.DataFrame):
@@ -211,10 +211,17 @@ class DataHandler(Serializable):
                 # Fetch column  first will be more friendly to SepDataFrame
                 data_df = fetch_df_by_col(data_df, col_set)
                 data_df = fetch_df_by_index(data_df, selector, level, fetch_orig=self.fetch_orig)
-        elif isinstance(data_storage, HasingStockStorage):
-            if proc_func is not None:
-                raise ValueError("proc_func is not supported by the HasingStockStorage")
-            data_df = data_storage.fetch(selector=selector, level=level, col_set=col_set, fetch_orig=self.fetch_orig)
+        elif isinstance(data_storage, BaseHandlerStorage):
+            if not data_storage.is_proc_func_supported():
+                if proc_func is not None:
+                    raise ValueError(f"proc_func is not supported by the storage {type(data_storage)}")
+                data_df = data_storage.fetch(
+                    selector=selector, level=level, col_set=col_set, fetch_orig=self.fetch_orig
+                )
+            else:
+                data_df = data_storage.fetch(
+                    selector=selector, level=level, col_set=col_set, fetch_orig=self.fetch_orig, proc_func=proc_func
+                )
         else:
             raise TypeError(f"data_storage should be pd.DataFrame|HasingStockStorage, not {type(data_storage)}")
 
@@ -522,7 +529,7 @@ class DataHandlerLP(DataHandler):
         -------
         pd.DataFrame:
         """
-        from .storage import HasingStockStorage
+        from .storage import BaseHandlerStorage
 
         data_storage = self._get_df_by_key(data_key)
         if isinstance(data_storage, pd.DataFrame):
@@ -537,10 +544,17 @@ class DataHandlerLP(DataHandler):
                 data_df = fetch_df_by_col(data_df, col_set)
                 data_df = fetch_df_by_index(data_df, selector, level, fetch_orig=self.fetch_orig)
 
-        elif isinstance(data_storage, HasingStockStorage):
-            if proc_func is not None:
-                raise ValueError("proc_func is not supported by the HasingStockStorage")
-            data_df = data_storage.fetch(selector=selector, level=level, col_set=col_set, fetch_orig=self.fetch_orig)
+        elif isinstance(data_storage, BaseHandlerStorage):
+            if not data_storage.is_proc_func_supported():
+                if proc_func is not None:
+                    raise ValueError(f"proc_func is not supported by the storage {type(data_storage)}")
+                data_df = data_storage.fetch(
+                    selector=selector, level=level, col_set=col_set, fetch_orig=self.fetch_orig
+                )
+            else:
+                data_df = data_storage.fetch(
+                    selector=selector, level=level, col_set=col_set, fetch_orig=self.fetch_orig, proc_func=proc_func
+                )
         else:
             raise TypeError(f"data_storage should be pd.DataFrame|HasingStockStorage, not {type(data_storage)}")
 

qlib/data/dataset/storage.py

@@ -14,6 +14,7 @@ class BaseHandlerStorage:
         level: Union[str, int] = "datetime",
         col_set: Union[str, List[str]] = DataHandler.CS_ALL,
         fetch_orig: bool = True,
+        proc_func: Callable = None,
         **kwargs,
     ) -> pd.DataFrame:
         """fetch data from the data storage
@@ -24,6 +25,7 @@ class BaseHandlerStorage:
             describe how to select data by index
         level : Union[str, int]
             which index level to select the data
+            - if level is None, apply selector to df directly
         col_set : Union[str, List[str]]
             - if isinstance(col_set, str):
                 select a set of meaningful columns.(e.g. features, columns)
@@ -33,15 +35,24 @@ class BaseHandlerStorage:
                 select several sets of meaningful columns, the returned data has multiple level
         fetch_orig : bool
             Return the original data instead of copy if possible.
+        proc_func: Callable
+            please refer to the doc of DataHandler.fetch
 
+        Returns
+        -------
+        pd.DataFrame
+            the dataframe fetched
         """
-
         raise NotImplementedError("fetch is method not implemented!")
 
     @staticmethod
     def from_df(df: pd.DataFrame):
         raise NotImplementedError("from_df method is not implemented!")
 
+    def is_proc_func_supported(self):
+        """whether the arg `proc_func` in `fetch` method is supported."""
+        raise NotImplementedError("is_proc_func_supported method is not implemented!")
+
 
 class HasingStockStorage(BaseHandlerStorage):
     def __init__(self, df):
@@ -105,3 +116,7 @@ class HasingStockStorage(BaseHandlerStorage):
             return fetch_stock_df_list[0]
         else:
             return pd.concat(fetch_stock_df_list, sort=False, copy=~fetch_orig)
+
+    def is_proc_func_supported(self):
+        """the arg `proc_func` in `fetch` method is not supported in HasingStockStorage"""
+        return False