lidb 1.3.6__py3-none-any.whl → 2.0.6__py3-none-any.whl

lidb/__init__.py

@@ -22,9 +22,10 @@ from .database import (
 
 from .table import Table, TableMode
 from .dataset import Dataset, DataLoader
+from .decorator import dataset
 from .qdf import from_polars, Expr
 from .svc import DataService, D
 
 from .parse import parse_hive_partition_structure
 
-__version__ = "1.3.6"
+__version__ = "2.0.6"

lidb/dataset.py

@@ -21,6 +21,7 @@ from varname import varname
 
 from .database import put, tb_path, scan, DB_PATH
 from .parse import parse_hive_partition_structure
+import inspect
 
 DEFAULT_DS_PATH = DB_PATH / "datasets"
 
@@ -50,6 +51,9 @@ def complete_data(fn, date, save_path, partitions):
         else:
             data = data.cast({"date": pl.Utf8})
         data = data.filter(date=date)
+        if "time" in data.columns:
+            if data["time"].n_unique() < 2:
+                data = data.drop("time")
         put(data, save_path, partitions=partitions)
     except Exception as e:
         logger.error(f"{save_path}: Error when complete data for {date}\n", exc_info=e)
@@ -65,6 +69,7 @@ class Dataset:
                  window: str = "1d",
                  partitions: list[str] = None,
                  is_hft: bool = False,
+                 data_name: str = "",
                  frame: int = 1):
         """
 
@@ -75,16 +80,22 @@ class Dataset:
         fn: str
             数据集计算函数。如果要用到底层依赖数据集，则必须显示定义形参 `depend`
         tb: str
-            数据集保存表格, 如果没有指定，默认 {lidb.DB_PATH}/datasets/
+            数据集保存表格, 如果没有指定，默认 {lidb.DB_PATH}/datasets/<module>
         update_time: str
             更新时间: 默认没有-实时更新，也就是可以取到当天值
+            更新时间只允许三种情况：
+            - 1. 盘前时间点：比如 08:00:00, 09:00:00, 09:15:00 ...
+            - 2. 盘中时间点：归为实时更新，使用空值 ""
+            - 3. 盘后时间点：比如 15:00:00, 16:30:00, 20:00:00 ...
         partitions: list[str]
-            分区
+            分区: 如果指定为 None, 则自动从 fn 参数推断，如果不需要分区，应该将其设定为空列表: []
         is_hft: bool
             是否是高频数据，如果是，则会按照asset进行分区存储，默认 False
             hft定义为：时间步长 < 1min
         window: str
             配合depends使用，在取depends时，会回看window周期，最小单位为`d`。不足 `d` 的会往上取整为`1d`
+        data_name: str
+            数据名，默认为空，会自动推断，如果指定了，则使用指定名
         frame: int
             用于自动推断 数据名
         """
@@ -95,13 +106,19 @@ class Dataset:
         self._is_depend = "depend" in self.fn_params_sig and len(self._depends) > 0
         self._is_hft = is_hft
         self._frame = frame
-        self.data_name = ""
-        try:
-            self.data_name = varname(frame, strict=False)
-        except Exception as e:
-            pass
+        self.data_name = data_name
+        if not self.data_name:
+            try:
+                self.data_name = varname(frame, strict=False)
+            except Exception as e:
+                pass
         if self.data_name:
             self.data_name = self.data_name.replace('ds_', '')
+        fn_params = ygo.fn_params(self.fn)
+        self.fn_params = {k: v for (k, v) in fn_params}
+        # 更新底层依赖数据集的同名参数
+        self._update_depends()
+
         if pd.Timedelta(window).days < 1:
             window = "1d"
         window_td = pd.Timedelta(window)
@@ -120,17 +137,18 @@ class Dataset:
             partitions = [*partitions, *self._append_partitions]
         self.partitions = partitions
         self._type_asset = "asset" in self.fn_params_sig
+        if "09:30:00" < update_time < "15:00:00":
+            update_time = ""
         self.update_time = update_time
         # 根据底层依赖调整update_time
-        if self._depends:
+        if update_time and self._depends:
             dep_ut = [ds.update_time for ds in self._depends]
             dep_ut.append(update_time)
             self.update_time = max(dep_ut)
-
-        self.tb = tb if tb else DEFAULT_DS_PATH / f"{self.data_name}"
+        mod = inspect.getmodule(fn)
+        self.tb = tb if tb else DEFAULT_DS_PATH / mod.__name__ /f"{self.data_name}"
         self.save_path = tb_path(self.tb)
-        fn_params = ygo.fn_params(self.fn)
-        self.fn_params = {k: v for (k, v) in fn_params}
+
         self.constraints = dict()
         for k in self.partitions[:-len(self._append_partitions)]:
             if k in self.fn_params:
@@ -140,12 +158,20 @@ class Dataset:
                 self.constraints[k] = v
                 self.save_path = self.save_path / f"{k}={v}"
 
+    def _update_depends(self):
+        new_deps = list()
+        for dep in self._depends:
+            new_dep = dep(**self.fn_params)
+            new_deps.append(new_dep)
+        self._depends = new_deps
+
     def is_empty(self, path) -> bool:
         return not any(path.rglob("*.parquet"))
 
     def __call__(self, *fn_args, **fn_kwargs):
-        # self.fn =
-        fn = partial(self.fn, *fn_args, **fn_kwargs)
+        """赋值时也会同步更新底层依赖数据集的同名参数"""
+
+        fn = ygo.delay(self.fn)(*fn_args, **fn_kwargs)
         ds = Dataset(*self._depends,
                      fn=fn,
                      tb=self.tb,
@@ -212,7 +238,7 @@ class Dataset:
         fn = self.fn
         save_path = self.save_path
         if self._is_depend:
-            fn = partial(fn, depend=self._get_depends(date))
+            fn = partial(fn, depend=self._get_depends(date,))
         else:
             fn = partial(fn, date=date)
         if self._type_asset:
@@ -288,7 +314,7 @@ class Dataset:
             _end_date = max(missing_dates)
             _beg_date = min(missing_dates)
             if self._days > 1:
-                _beg_date = xcals.shift_tradeday(_beg_date, -self._days)
+                _beg_date = xcals.shift_tradeday(_beg_date, -(self._days-1))
             _depend_dates = xcals.get_tradingdays(_beg_date, _end_date)
             for depend in self._depends:
                 depend.get_history(_depend_dates, eager=False)
@@ -333,19 +359,22 @@ class Dataset:
             return data.collect()
         return data
 
-    def _get_depends(self, date: str) -> pl.DataFrame | None:
+    def _get_depends(self, date: str) -> pl.LazyFrame | None:
         # 获取依赖数据集数据
         if not self._depends:
             return None
         end_date = date
-        beg_date = xcals.shift_tradeday(date, -self._days)
+        beg_date = date
+        if self._days > 1:
+            beg_date = xcals.shift_tradeday(beg_date, -(self._days-1))
         params = {
             "ds_conf": dict(depend=self._depends),
             "beg_date": beg_date,
             "end_date": end_date,
-            "times": [self.update_time],
+            "times": [self.update_time, ],
             "show_progress": False,
-            "eager": True,
+            "eager": False,
+            "process_time": False,  # 不处理时间
         }
         res = load_ds(**params)
         return res["depend"]
@@ -357,20 +386,54 @@ def loader(data_name: str,
            prev_date_list: list[str],
            prev_date_mapping: dict[str, str],
            time: str,
+           process_time: bool,
            **constraints) -> pl.LazyFrame:
-    if time < ds.update_time:
-        if len(prev_date_list) > 1:
-            lf = ds.get_history(prev_date_list, eager=False, **constraints)
+    """
+    Parameters
+    ----------
+    data_name
+    ds
+    date_list
+    prev_date_list
+    prev_date_mapping
+    time
+    process_time: bool
+        是否处理源数据的时间: 根据实参 time. 用于应对不同场景
+        场景1：依赖因子不处理，底层数据是什么就返回什么
+        场景2：zoo.load 用来加载测试日内不同时间点的数据，就应该处理
+    constraints
+
+    Returns
+    -------
+
+    """
+    if time:
+        if time < ds.update_time:
+            if len(prev_date_list) > 1:
+                lf = ds.get_history(prev_date_list, eager=False, **constraints)
+            else:
+                lf = ds.get_value(prev_date_list[0], eager=False, **constraints)
         else:
-            lf = ds.get_value(prev_date_list[0], eager=False, **constraints)
+            if len(date_list) > 1:
+                lf = ds.get_history(date_list, eager=False, **constraints)
+            else:
+                lf = ds.get_value(date_list[0], eager=False, **constraints)
     else:
-        if len(date_list) > 1:
-            lf = ds.get_history(date_list, eager=False, **constraints)
+        if ds.update_time > "09:30:00":
+            # 盘后因子：取上一天的值
+            if len(prev_date_list) > 1:
+                lf = ds.get_history(prev_date_list, eager=False, **constraints)
+            else:
+                lf = ds.get_value(prev_date_list[0], eager=False, **constraints)
         else:
-            lf = ds.get_value(date_list[0], eager=False, **constraints)
+            if len(date_list) > 1:
+                lf = ds.get_history(date_list, eager=False, **constraints)
+            else:
+                lf = ds.get_value(date_list[0], eager=False, **constraints)
+
     schema = lf.collect_schema()
     include_time = schema.get("time") is not None
-    if time:
+    if process_time and time:
         if include_time:
             lf = lf.filter(time=time)
         else:
@@ -396,6 +459,7 @@ def load_ds(ds_conf: dict[str, list[Dataset]],
             backend: Literal["threading", "multiprocessing", "loky"] = "threading",
             show_progress: bool = True,
             eager: bool = False,
+            process_time: bool = True,
             **constraints) -> dict[str, pl.DataFrame | pl.LazyFrame]:
     """
     加载数据集
@@ -417,6 +481,10 @@ def load_ds(ds_conf: dict[str, list[Dataset]],
         是否返回 DataFrame
         - True: 返回DataFrame
         - False: 返回LazyFrame
+    process_time: bool
+        是否处理源数据的时间: 根据实参 time. 用于应对不同场景
+        场景1：依赖因子不处理，底层数据是什么就返回什么
+        场景2：zoo.load 用来加载测试日内不同时间点的数据，就应该处理
     constraints
         限制条件，比如 asset='000001'
     Returns
@@ -435,6 +503,7 @@ def load_ds(ds_conf: dict[str, list[Dataset]],
     prev_date_mapping = {prev_date: date_list[i] for i, prev_date in enumerate(prev_date_list)}
     results = defaultdict(list)
     index = ("date", "time", "asset")
+    _index = ("date", "asset")
     with ygo.pool(n_jobs=n_jobs,
                   backend=backend,
                   show_progress=show_progress) as go:
@@ -450,27 +519,58 @@ def load_ds(ds_conf: dict[str, list[Dataset]],
                                                    prev_date_list=prev_date_list,
                                                    prev_date_mapping=prev_date_mapping,
                                                    time=time,
+                                                   process_time=process_time,
                                                    **constraints)
         for name, lf in go.do():
             results[name].append(lf)
-    _LFs = {
-        name: (pl.concat(lfList, )
-               .select(*index,
-                       cs.exclude(index))
-               )
-        for name, lfList in results.items()}
-    LFs = defaultdict(list)
-    for name, lf in _LFs.items():
+    # _LFs = {
+    # name: (pl.concat(lfList, )
+    # .select(*index,
+    # cs.exclude(index))
+    # )
+    # for name, lfList in results.items()}
+    _LFs_with_time = {}
+    _LFs_without_time = {}
+    for name, lfList in results.items():
+        lf = pl.concat(lfList)
+        # print(lf)
+        if "time" not in lf.collect_schema().names():
+            _LFs_without_time[name] = lf
+        else:
+            _LFs_with_time[name] = lf
+    LFs_with_time = defaultdict(list)
+    LFs_without_time = defaultdict(list)
+    for name, lf in _LFs_with_time.items():
         dn, _ = name.split(":")
-        LFs[dn].append(lf)
-    LFs = {
+        LFs_with_time[dn].append(lf)
+    for name, lf in _LFs_without_time.items():
+        dn, _ = name.split(":")
+        LFs_without_time[dn].append(lf)
+    LFs_with_time = {
         name: (pl.concat(lfList, how="align")
                .sort(index)
                .select(*index,
                        cs.exclude(index))
                )
-        for name, lfList in LFs.items()}
-
+        for name, lfList in LFs_with_time.items()}
+    LFs_without_time = {
+        name: (pl.concat(lfList, how="align")
+               .sort(_index)
+               .select(*_index,
+                       cs.exclude(_index))
+               )
+        for name, lfList in LFs_without_time.items()}
+    dns = list(LFs_with_time.keys()) if LFs_with_time else list(LFs_without_time.keys())
+    LFs = dict()
+    for dn in dns:
+        _lf_with_time = LFs_with_time.get(dn)
+        _lf_without_time = LFs_without_time.get(dn)
+        if _lf_with_time is not None:
+            LFs[dn] = _lf_with_time
+            if _lf_without_time is not None:
+                LFs[dn] = LFs[dn].join(_lf_without_time, on=["date", "asset"], how="left")
+        else:
+            LFs[dn] = _lf_without_time
     if not eager:
         return LFs
     return {
@@ -478,7 +578,6 @@ def load_ds(ds_conf: dict[str, list[Dataset]],
         for name, lf in LFs.items()
     }
 
-
 class DataLoader:
 
     def __init__(self, name: str):
@@ -521,6 +620,7 @@ class DataLoader:
                      backend=backend,
                      times=times,
                      eager=eager,
+                     process_time=True,
                      **constraints)
         self._df = lf[self._name]
 

lidb/decorator.py

@@ -0,0 +1,50 @@
+# Copyright (c) ZhangYundi.
+# Licensed under the MIT License. 
+# Created on 2025/12/31 10:58
+# Description:
+
+from .dataset import Dataset
+from typing import Callable, TypeVar, cast
+
+F = TypeVar('F', bound=Callable)
+
+def dataset(*depends: Dataset,
+            tb: str = "",
+            update_time: str = "",
+            window: str = "1d",
+            partitions: list[str] = None,
+            is_hft: bool = False) -> Callable[[F], Dataset]:
+    """
+    装饰器：将函数转换为Dataset对象
+
+    Parameters
+    ----------
+    depends: Dataset
+        底层依赖数据集
+    tb: str
+        数据集保存表格, 如果没有指定，默认 {DEFAULT_DS_PATH}/
+    update_time: str
+        更新时间: 默认没有-实时更新，也就是可以取到当天值
+    window: str
+        配合depends使用，在取depends时，会回看window周期，最小单位为`d`。不足 `d` 的会往上取整为`1d`
+    partitions: list[str]
+        分区: 如果指定为 None, 则自动从 fn 参数推断，如果不需要分区，应该将其设定为空列表: []
+    is_hft: bool
+        是否是高频数据，如果是，则会按照asset进行分区存储，默认 False
+        hft定义为：时间步长 < 1min
+    """
+    def decorator(fn: F):
+        # 创建Dataset实例
+        ds = Dataset(
+            *depends,
+            fn=fn,
+            tb=tb,
+            update_time=update_time,
+            window=window,
+            partitions=partitions,
+            is_hft=is_hft,
+            data_name=fn.__name__,
+            frame=1
+        )
+        return ds
+    return decorator

{lidb-1.3.6.dist-info → lidb-2.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lidb
-Version: 1.3.6
+Version: 2.0.6
 Summary: Light database for quantor
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown
@@ -55,7 +55,7 @@ res = lidb.sql("select * from my_table;")
 #### 数据集使用
 ```python
 import lidb
-from lidb import Dataset
+from lidb import Dataset, dataset
 import polars as pl
 
 # 定义一个tick级别的高频数据集: 高频成交量
@@ -68,13 +68,24 @@ def hft_vol(date: str, num: int) -> pl.DataFrame | pl.LazyFrame | None:
 
 ds_hft_vol = Dataset(fn=hft_vol, 
                      tb="path/to/hft_vol", 
-                     partitions=["num"], 
+                     partitions=["num"], # 默认值 None, 会自动识别 num 
                      update_time="", # 实时更新
-                     by_asset=True, # 根据asset_id进行分区
+                     is_hft=True, # 根据asset_id进行分区
                     )(num=20)
 
 # 获取历史数据
 history_data = ds_hft_vol.get_history(["2023-01-01", "2023-01-02", ...])
+
+# 更加便捷的创建数据集方式：通过dataset装饰器
+@dataset()
+def hft_vol(date: str, num: int) -> pl.DataFrame | pl.LazyFrame | None:
+    # 假设上游tick行情表在clickhouse
+    quote_query = f"select * from quote where date = '{date}'"
+    quote = lidb.read_ck(quote_query, db_conf="databases.ck")
+    # 特征计算: 比如过去20根tick的成交量总和, 使用表达式引擎计算
+    return lidb.from_polars(quote).sql(f"itd_sum(volume, {num}) as vol_s20")
+
+hft_vol.get_value("2025-05-15")
 ```
 
 #### `Table`
@@ -239,7 +250,6 @@ finally:
 #### 数据集管理(`dataset.py`)
 - `Dataset`: 数据集定义和管理
 - `DataLoader`： 数据加载器
-- `zoo`: alpha因子数据管理
 
 #### 表达式计算(`qdf/`)
 - `QDF`: 表达式数据库

{lidb-1.3.6.dist-info → lidb-2.0.6.dist-info}/RECORD

@@ -1,6 +1,7 @@
-lidb/__init__.py,sha256=tcPBOs0UltwU3tCqjIBHrklITYQcAeWkdrV3_SYCu1I,505
+lidb/__init__.py,sha256=WuGdkD4QzcCkIG3zbXupaXJV0b3o8gvaMGhs6MhVa_c,536
 lidb/database.py,sha256=DnPXRXvUO6g0kuMo3LPl6eKo_HbD3JNW1qzoaJ14Sgo,7533
-lidb/dataset.py,sha256=hC2D2uJ7xV6yRB-j9TOYxb0aqZQME_5_BbXND2NPBK0,20254
+lidb/dataset.py,sha256=rZGUmvRwaIdynWbTFF-D1fPE1NyAbhDLVxJ3J0y1MYo,24363
+lidb/decorator.py,sha256=bFnUPcJED6F95nBxHq1a8j5pM2JF9rjFtNvxIQUs9_I,1605
 lidb/init.py,sha256=N_PiBGZO3hKUhQQYzly3GKHgSf4eJVO7xyxjX-chUpQ,1327
 lidb/parse.py,sha256=6awnc14OK7XBkkSrAJFOCZOQ0JUHmm6yDI9F3kkLwcQ,3494
 lidb/table.py,sha256=NeqOU0EJU3DA0yz-1T2GVLpKASu1_1fdOLK3yxf7DtA,4494
@@ -18,7 +19,7 @@ lidb/qdf/udf/itd_udf.py,sha256=O_OOdSTEaeCoqjtlKnpvNF-_10QoamJL_tw2xEZCYVw,6747
 lidb/qdf/udf/ts_udf.py,sha256=Ag6-ffhmIugkA-st2QY-GP4hclQZcRG8SB-bVa7k5cc,5674
 lidb/svc/__init__.py,sha256=9vQo7gCm5LRgWSiq_UU2hlbwvXi0FlGYt2UDVZixx_U,141
 lidb/svc/data.py,sha256=tLOI_YylnsVejyqv9l-KgPetkPO0QzybOf1PEeFSZNI,4380
-lidb-1.3.6.dist-info/METADATA,sha256=0f7wFU6CZwD_jiqmJjzc_HNCx48mKA24_JBUREiEfSs,8558
-lidb-1.3.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lidb-1.3.6.dist-info/top_level.txt,sha256=NgXJNwt6ld6oLXtW1vOPaEh-VO5R0JEX_KmGIJR4ueE,5
-lidb-1.3.6.dist-info/RECORD,,
+lidb-2.0.6.dist-info/METADATA,sha256=ldndXJNXi7y_k1rh5fRPbBVF4a97LqRykzW2gEk8lEM,9087
+lidb-2.0.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lidb-2.0.6.dist-info/top_level.txt,sha256=NgXJNwt6ld6oLXtW1vOPaEh-VO5R0JEX_KmGIJR4ueE,5
+lidb-2.0.6.dist-info/RECORD,,