lidb 1.2.0__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.2.0"
+__version__ = "2.0.6"

lidb/dataset.py

@@ -5,21 +5,25 @@
 
 from __future__ import annotations
 
+import shutil
 from collections import defaultdict
 from enum import Enum
 from functools import partial
 from typing import Callable, Literal
 
 import logair
+import pandas as pd
 import polars as pl
 import polars.selectors as cs
 import xcals
 import ygo
+from varname import varname
 
 from .database import put, tb_path, scan, DB_PATH
 from .parse import parse_hive_partition_structure
-from .qdf import QDF, from_polars
+import inspect
 
+DEFAULT_DS_PATH = DB_PATH / "datasets"
 
 class InstrumentType(Enum):
     STOCK = "Stock"  # 股票
@@ -30,7 +34,7 @@ class InstrumentType(Enum):
 def complete_data(fn, date, save_path, partitions):
     logger = logair.get_logger(__name__)
     try:
-        data = fn(date=date)
+        data = fn()
         if data is None:
             # 保存数据的逻辑在fn中实现了
             return
@@ -44,59 +48,107 @@ def complete_data(fn, date, save_path, partitions):
         cols = data.columns
         if "date" not in cols:
             data = data.with_columns(pl.lit(date).alias("date")).select("date", *cols)
-
+        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}")
-        logger.warning(e)
+        logger.error(f"{save_path}: Error when complete data for {date}\n", exc_info=e)
 
 
 class Dataset:
 
     def __init__(self,
-                 fn: Callable[..., pl.DataFrame],
-                 tb: str,
+                 *depends: Dataset,
+                 fn: Callable[..., pl.DataFrame | pl.LazyFrame],
+                 tb: str = "",
                  update_time: str = "",
+                 window: str = "1d",
                  partitions: list[str] = None,
-                 by_asset: bool = True,
-                 by_time: bool = False):
+                 is_hft: bool = False,
+                 data_name: str = "",
+                 frame: int = 1):
         """
 
         Parameters
         ----------
+        depends: Dataset
+            底层依赖数据集
         fn: str
-            数据集计算函数
+            数据集计算函数。如果要用到底层依赖数据集，则必须显示定义形参 `depend`
         tb: str
-            数据集保存表格
+            数据集保存表格, 如果没有指定，默认 {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]
-            分区
-        by_asset: bool
-            是否按照标的进行分区，默认 True
-        by_time: bool
-            是否按照标的进行分区，默认 False
+            分区: 如果指定为 None, 则自动从 fn 参数推断，如果不需要分区，应该将其设定为空列表: []
+        is_hft: bool
+            是否是高频数据，如果是，则会按照asset进行分区存储，默认 False
+            hft定义为：时间步长 < 1min
+        window: str
+            配合depends使用，在取depends时，会回看window周期，最小单位为`d`。不足 `d` 的会往上取整为`1d`
+        data_name: str
+            数据名，默认为空，会自动推断，如果指定了，则使用指定名
+        frame: int
+            用于自动推断 数据名
         """
+        self._depends = list(depends)
+        self._name = ""
         self.fn = fn
         self.fn_params_sig = ygo.fn_signature_params(fn)
-        self._by_asset = by_asset
-        self._by_time = by_time
-        self._append_partitions = ["asset", "date"] if by_asset else ["date", ]
-        if by_time:
-            self._append_partitions.append("time")
+        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 = 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)
+        self._window = window
+        self._days = window_td.days
+        if window_td.seconds > 0:
+            self._days += 1
+        # 检测是否高频数据：如果是高频数据，则按照标的进行分区，高频的定义为时间差 < 60s
+        self._append_partitions = ["asset", "date"] if is_hft else ["date", ]
         if partitions is not None:
             partitions = [k for k in partitions if k not in self._append_partitions]
             partitions = [*partitions, *self._append_partitions]
         else:
-            partitions = self._append_partitions
+            # partitions = self._append_partitions
+            partitions = [k for k in self.fn_params_sig if k not in self._append_partitions and k != "depend"]
+            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 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)
+        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)
 
-        self.tb = tb
-        self.save_path = tb_path(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:
@@ -106,20 +158,35 @@ 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)
-        ds = Dataset(fn=fn,
+        """赋值时也会同步更新底层依赖数据集的同名参数"""
+
+        fn = ygo.delay(self.fn)(*fn_args, **fn_kwargs)
+        ds = Dataset(*self._depends,
+                     fn=fn,
                      tb=self.tb,
                      partitions=self.partitions,
-                     by_asset=self._by_asset,
-                     by_time=self._by_time,
-                     update_time=self.update_time)
+                     update_time=self.update_time,
+                     is_hft=self._is_hft,
+                     window=self._window,
+                     frame=self._frame+1)
+        ds.data_name = self.data_name
         return ds
 
+    def alias(self, new_name: str):
+        self._name = new_name
+        return self
+
     def get_value(self, date, eager: bool = True, **constraints):
         """
         取值: 不保证未来数据
@@ -135,6 +202,7 @@ class Dataset:
         -------
 
         """
+        logger = logair.get_logger(f"{__name__}.{self.__class__.__name__}")
         _constraints = {k: v for k, v in constraints.items() if k in self.partitions}
         _limits = {k: v for k, v in constraints.items() if k not in self.partitions}
         search_path = self.save_path
@@ -144,9 +212,22 @@ class Dataset:
             search_path = search_path / f"{k}={v}"
         search_path = search_path / f"date={date}"
 
+        # 处理空文件
+        for file_path in search_path.rglob("*.parquet"):
+            if file_path.stat().st_size == 0:
+                # 删除
+                logger.warning(f"{file_path}: Deleting empty file.")
+                file_path.unlink()
+
         if not self.is_empty(search_path):
             lf = scan(search_path).cast({"date": pl.Utf8})
-            schema = lf.collect_schema()
+            try:
+                schema = lf.collect_schema()
+            except:
+                logger.warning(f"{search_path}: Failed to collect schema.")
+                # 删除该文件夹
+                shutil.rmtree(search_path)
+                return self.get_value(date=date, eager=eager, **constraints)
             _limits = {k: v for k, v in constraints.items() if schema.get(k) is not None}
             lf = lf.filter(date=date, **_limits)
             if not eager:
@@ -156,7 +237,10 @@ class Dataset:
                 return data
         fn = self.fn
         save_path = self.save_path
-
+        if self._is_depend:
+            fn = partial(fn, depend=self._get_depends(date,))
+        else:
+            fn = partial(fn, date=date)
         if self._type_asset:
             if "asset" in _constraints:
                 fn = ygo.delay(self.fn)(asset=_constraints["asset"])
@@ -169,7 +253,6 @@ class Dataset:
                     params[k] = v
                     save_path = save_path / f"{k}={v}"
             fn = ygo.delay(self.fn)(**params)
-        logger = logair.get_logger(__name__)
 
         today = xcals.today()
         now = xcals.now()
@@ -198,7 +281,7 @@ class Dataset:
     def get_history(self,
                     dateList: list[str],
                     n_jobs: int = 5,
-                    backend: Literal["threading", "multiprocessing", "loky"] = "loky",
+                    backend: Literal["threading", "multiprocessing", "loky"] = "threading",
                     eager: bool = True,
                     rep_asset: str = "000001",  # 默认 000001
                     **constraints):
@@ -227,6 +310,14 @@ class Dataset:
             missing_dates = set(dateList).difference(set(exist_dates))
             missing_dates = sorted(list(missing_dates))
         if missing_dates:
+            # 先逐个补齐 depends
+            _end_date = max(missing_dates)
+            _beg_date = min(missing_dates)
+            if self._days > 1:
+                _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)
             fn = self.fn
             save_path = self.save_path
 
@@ -250,12 +341,17 @@ class Dataset:
                 except:
                     pass
                 for date in missing_dates:
-                    go.submit(complete_data, job_name=f"Completing {info_path}")(
-                        fn=fn,
-                        date=date,
-                        save_path=save_path,
-                        partitions=self._append_partitions,
-                    )
+                    if self._is_depend:
+                        fn = partial(fn, depend=self._get_depends(date))
+                    else:
+                        fn = partial(fn, date=date)
+                    go.submit(complete_data,
+                              job_name=f"Completing",
+                              postfix=info_path,
+                              leave=False)(fn=fn,
+                                           date=date,
+                                           save_path=save_path,
+                                           partitions=self._append_partitions, )
                 go.do()
         data = scan(search_path, ).cast({"date": pl.Utf8}).filter(pl.col("date").is_in(dateList), **constraints)
         data = data.sort("date")
@@ -263,6 +359,26 @@ class Dataset:
             return data.collect()
         return data
 
+    def _get_depends(self, date: str) -> pl.LazyFrame | None:
+        # 获取依赖数据集数据
+        if not self._depends:
+            return None
+        end_date = date
+        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, ],
+            "show_progress": False,
+            "eager": False,
+            "process_time": False,  # 不处理时间
+        }
+        res = load_ds(**params)
+        return res["depend"]
+
 
 def loader(data_name: str,
            ds: Dataset,
@@ -270,35 +386,80 @@ 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 include_time:
-        lf = lf.filter(time=time)
-    else:
-        lf = lf.with_columns(time=pl.lit(time))
+    if process_time and time:
+        if include_time:
+            lf = lf.filter(time=time)
+        else:
+            lf = lf.with_columns(time=pl.lit(time))
     if time < ds.update_time:
         lf = lf.with_columns(date=pl.col("date").replace(prev_date_mapping))
+    keep = {"date", "time", "asset"}
+    if ds._name:
+        columns = lf.collect_schema().names()
+        rename_cols = set(columns).difference(keep)
+        if len(rename_cols) > 1:
+            lf = lf.rename({k: f"{ds._name}.{k}" for k in rename_cols})
+        else:
+            lf = lf.rename({k: ds._name for k in rename_cols})
     return data_name, lf
 
 
 def load_ds(ds_conf: dict[str, list[Dataset]],
             beg_date: str,
             end_date: str,
-            time: str,
+            times: list[str],
             n_jobs: int = 7,
             backend: Literal["threading", "multiprocessing", "loky"] = "threading",
+            show_progress: bool = True,
             eager: bool = False,
+            process_time: bool = True,
             **constraints) -> dict[str, pl.DataFrame | pl.LazyFrame]:
     """
     加载数据集
@@ -310,15 +471,20 @@ def load_ds(ds_conf: dict[str, list[Dataset]],
         开始日期
     end_date: str
         结束日期
-    time: str
+    times: list[str]
         取值时间
     n_jobs: int
         并发数量
     backend: str
+    show_progress: bool
     eager: bool
         是否返回 DataFrame
         - True: 返回DataFrame
         - False: 返回LazyFrame
+    process_time: bool
+        是否处理源数据的时间: 根据实参 time. 用于应对不同场景
+        场景1：依赖因子不处理，底层数据是什么就返回什么
+        场景2：zoo.load 用来加载测试日内不同时间点的数据，就应该处理
     constraints
         限制条件，比如 asset='000001'
     Returns
@@ -332,31 +498,79 @@ def load_ds(ds_conf: dict[str, list[Dataset]],
         raise ValueError("beg_date must be less than end_date")
     date_list = xcals.get_tradingdays(beg_date, end_date)
     beg_date, end_date = date_list[0], date_list[-1]
-    prev_date_list = xcals.get_tradingdays(xcals.shift_tradeday(beg_date, -1), xcals.shift_tradeday(end_date, -1))
+    prev_date_list = xcals.get_tradingdays(xcals.shift_tradeday(beg_date, -1),
+                                           xcals.shift_tradeday(end_date, -1))
     prev_date_mapping = {prev_date: date_list[i] for i, prev_date in enumerate(prev_date_list)}
     results = defaultdict(list)
-    with ygo.pool(n_jobs=n_jobs, backend=backend) as go:
+    index = ("date", "time", "asset")
+    _index = ("date", "asset")
+    with ygo.pool(n_jobs=n_jobs,
+                  backend=backend,
+                  show_progress=show_progress) as go:
         for data_name, ds_list in ds_conf.items():
             for ds in ds_list:
-                go.submit(loader,
-                          job_name="Loading",
-                          postfix=data_name)(data_name=data_name,
-                                             ds=ds,
-                                             date_list=date_list,
-                                             prev_date_list=prev_date_list,
-                                             prev_date_mapping=prev_date_mapping,
-                                             time=time,
-                                             **constraints)
+                _data_name = f"{data_name}:{ds.tb}"
+                for time in times:
+                    go.submit(loader,
+                              job_name="Loading",
+                              postfix=data_name, )(data_name=_data_name,
+                                                   ds=ds,
+                                                   date_list=date_list,
+                                                   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)
-    index = ("date", "time", "asset")
-    LFs = {
+    # _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_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 results.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 {
@@ -364,21 +578,20 @@ def load_ds(ds_conf: dict[str, list[Dataset]],
         for name, lf in LFs.items()
     }
 
-
 class DataLoader:
 
     def __init__(self, name: str):
         self._name = name
-        self._lf: pl.LazyFrame = None
-        self._df: pl.DataFrame = None
         self._index: tuple[str] = ("date", "time", "asset")
-        self._db: QDF = None
-        self._one: pl.DataFrame = None
+        self._df: pl.LazyFrame | pl.DataFrame = None
+        # self._db: QDF = None
 
     def get(self,
             ds_list: list[Dataset],
             beg_date: str,
             end_date: str,
+            times: list[str],
+            eager: bool = False,
             n_jobs: int = 11,
             backend: Literal["threading", "multiprocessing", "loky"] = "threading",
             **constraints):
@@ -389,6 +602,9 @@ class DataLoader:
         ds_list: list[Dataset]
         beg_date: str
         end_date: str
+        times: list[str]
+            加载的时间列表
+        eager: bool
         n_jobs: int
         backend: str
         constraints
@@ -402,41 +618,25 @@ class DataLoader:
                      end_date=end_date,
                      n_jobs=n_jobs,
                      backend=backend,
-                     eager=False,
+                     times=times,
+                     eager=eager,
+                     process_time=True,
                      **constraints)
-        self._lf = lf.get(self._name)
-        self._df = None
-        self._db = from_polars(self._lf, self._index, align=True)
-        dateList = xcals.get_tradingdays(beg_date, end_date)
-        _data_name = f"{self._name}(one_day)"
-        self._one = load_ds(ds_conf={_data_name: ds_list},
-                            beg_date=dateList[0],
-                            end_date=dateList[0],
-                            n_jobs=n_jobs,
-                            backend=backend,
-                            eager=False,
-                            **constraints).get(_data_name).collect()
+        self._df = lf[self._name]
 
     @property
     def name(self) -> str:
         return self._name
 
     @property
-    def one_day(self) -> pl.DataFrame:
-        return self._one
-
-    @property
-    def schema(self) -> pl.Schema:
-        return self._one.schema
-
-    @property
-    def columns(self) -> list[str]:
-        return self._one.columns
-
-    def collect(self) -> pl.DataFrame:
-        if self._df is None:
-            self._df = self._lf.collect()
+    def data(self) -> pl.DataFrame | None:
+        """返回全量数据"""
+        if isinstance(self._df, pl.LazyFrame):
+            self._df = self._df.collect()
         return self._df
 
-    def sql(self, *exprs: str) -> pl.DataFrame:
-        return self._db.sql(*exprs)
+    def add_data(self, df: pl.DataFrame | pl.LazyFrame):
+        """添加dataframe, index 保持为原有的 _df.index"""
+        if isinstance(df, pl.LazyFrame):
+            df = df.collect()
+        self._df = pl.concat([self._df, df], how="align").sort(self._index)

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/init.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 from dynaconf import Dynaconf
 import logair
+import os
 
 
 USERHOME = Path("~").expanduser() # 用户家目录
@@ -22,8 +23,7 @@ if not CONFIG_PATH.exists():
     except Exception as e:
         logger.error(f"Failed to create settings file: {e}")
     with open(CONFIG_PATH, "w") as f:
-        template_content = f'[global]\npath="{DB_PATH}"\n'
-    with open(CONFIG_PATH, "w") as f:
+        template_content = f'[GLOBAL]\npath="{DB_PATH}"\n\n[POLARS]\nmax_threads=32\n'
         f.write(template_content)
     logger.info(f"Settings file created: {CONFIG_PATH}")
 
@@ -38,5 +38,8 @@ def get_settings():
 _settiings = get_settings()
 if _settiings is not None:
     setting_db_path = _settiings.get(f"global.path", "")
+    # 配置 polars
+    setting_polars_threads = _settiings.get("polars.max_threads", 32)
+    os.environ["POLARS_MAX_THREADS"] = str(setting_polars_threads)
     if setting_db_path:
         DB_PATH = Path(setting_db_path)

lidb/parse.py

@@ -86,6 +86,10 @@ def parse_hive_partition_structure(root_path: Path | str, file_pattern: str = "*
     partition_combinations = set()
 
     for file_path in root_path.rglob(file_pattern):
+        if file_path.stat().st_size == 0:
+            # 删除
+            file_path.unlink()
+            continue
         relative_path = file_path.relative_to(root_path)
 
         # 收集分区信息

lidb/qdf/qdf.py

@@ -118,7 +118,7 @@ class QDF:
             except Exception as error:
                 raise CompileError(message=f"{e.fn_name}({', '.join([str(arg) for arg in args])})\n{error}") from error
 
-    def sql(self, *exprs: str, show_progress: bool = False) -> pl.DataFrame:
+    def sql(self, *exprs: str, show_progress: bool = False, leave: bool = False) -> pl.DataFrame:
         """
         表达式查询
         Parameters
@@ -127,6 +127,8 @@ class QDF:
             表达式，比如 "ts_mean(close, 5) as close_ma5"
         show_progress: bool
             是否展示进度条
+        leave: bool
+            是否保留进度条
         Returns
         -------
             polars.DataFrame
@@ -146,7 +148,7 @@ class QDF:
         pbar = None
         lvl_num = len(lvls)
         if show_progress:
-            pbar = tqdm(total=lvl_num, desc=f"{len(exprs)}")
+            pbar = tqdm(total=lvl_num, desc=f"{len(exprs)}", leave=leave)
         for i, batch_exprs in enumerate(lvls):
             if show_progress:
                 pbar.set_postfix_str(f"level-{i + 1}:{len(batch_exprs)}")

lidb/table.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import sys
 from collections.abc import Callable
 from enum import Enum
 
@@ -24,8 +25,6 @@ class TableMode(Enum):
     F = "full" # 全量更新
     I = "increment" # 增量更新
 
-
-
 class Table:
 
     def __init__(self,
@@ -58,10 +57,10 @@ class Table:
         """获取数据并且保存数据"""
         data = ygo.delay(self.fn)(this=self)()
         if data is None:
-            self.logger.error("No data.")
+            self.logger.error(f"{self.tb}: No data.")
             return
         if data.is_empty():
-            self.logger.warning("No data.")
+            self.logger.warning(f"{self.tb}: No data.")
             return
         if self.mode == TableMode.I:
             time_uuid = uuid.uuid1()
@@ -77,33 +76,43 @@ class Table:
     def update(self, verbose: bool = False):
         """更新最新数据: 全量更新, 覆盖旧数据"""
         self.verbose = verbose
+        if self._need_update(date=xcals.today()):
+            self._log("Updating.", "info")
+            self._do_job()
+
+    def _need_update(self, date: str) -> bool:
+        """是否需要更新"""
         existed = self._data_dir.exists()
         if not existed:
             self._data_dir.mkdir(parents=True, exist_ok=True)
-            self._log("Creating new data.", "info")
-            self._do_job()
+            return True
         else:
             modified_time = self.modified_time
             if modified_time is not None:
                 modified_datetime = modified_time.strftime("%Y-%m-%d %H:%M:%S")
                 modified_d, modified_t = modified_datetime.split(" ")
-                if self._updated(data_date=modified_d, data_time=modified_t):
-                    return
-            self._log("Updating.", "info")
-            self._do_job()
-            self._log("Updated.", "info")
+                if self._updated(date, data_date=modified_d, data_time=modified_t):
+                    return False
+            return True
 
-    def get_value(self, eager: bool = True) -> pl.DataFrame | pl.LazyFrame:
+    def get_value(self, date: str, eager: bool = True) -> pl.DataFrame | pl.LazyFrame:
         """获取数据"""
-        self.update(verbose=True)
+        # self.update(verbose=True)
+        if not date:
+            date = xcals.today()
+        self.verbose = True
+        if self._need_update(date):
+            self._log("Update first plz.", "warning")
+            sys.exit()
+
         df = scan(self._data_dir)
         if eager:
             return df.collect()
         return df
 
-    def _updated(self, data_date: str, data_time: str) -> bool:
-        """判断是否需要更新数据"""
-        recent_tradeday = xcals.get_recent_tradeday()
+    def _updated(self, date: str, data_date: str, data_time: str) -> bool:
+        """判断是否已经更新数据"""
+        recent_tradeday = xcals.get_recent_tradeday(date)
         prev_tradeday = xcals.shift_tradeday(recent_tradeday, -1)
         now = xcals.now()
         latest_update_date = recent_tradeday if now >= self.update_time else prev_tradeday

lidb-2.0.6.dist-info/METADATA

@@ -0,0 +1,282 @@
+Metadata-Version: 2.4
+Name: lidb
+Version: 2.0.6
+Summary: Light database for quantor
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: dynaconf>=3.2.11
+Requires-Dist: polars>=1.31.0
+Requires-Dist: sqlparse>=0.5.3
+Requires-Dist: logair>=1.0.8
+Requires-Dist: clickhouse-df>=0.1.5
+Requires-Dist: connectorx>=0.4.3
+Requires-Dist: pymysql>=1.1.2
+Requires-Dist: xcals>=0.0.4
+Requires-Dist: ygo>=1.2.9
+Requires-Dist: lark>=1.3.1
+Requires-Dist: numpy>=2.3.1
+Requires-Dist: tqdm>=4.67.1
+Requires-Dist: varname>=0.15.1
+
+## lidb
+
+### 项目简介
+lidb 是一个基于 Polars 的数据管理和分析库，专为金融量化研究设计。它提供了高效的数据存储、查询和表达式计算功能，支持多种时间序列和横截面数据分析操作。
+
+### 功能特性
+- **多数据源支持**: 本地 Parquet 存储、MySQL、ClickHouse 等数据库连接
+- **高效数据存储**: 基于 Parquet 格式的分区存储机制
+- **SQL 查询接口**: 支持标准 SQL 语法进行数据查询
+- **表达式计算引擎**: 提供丰富的 UDF 函数库，包括时间序列、横截面、维度等分析函数
+- **数据集管理**: 自动化数据补全、历史数据加载和 PIT(Point-in-Time)数据处理
+- **数据服务**: 异步加载数据，用于数据密集型任务的数据加载(如大量标的的高频数据)
+
+### 安装
+```bash
+pip install -U lidb
+```
+
+### 快速开始
+
+#### 基础数据操作
+```python
+import lidb
+import polars as pl
+
+df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+
+# 写入数据
+lidb.put(df, "my_table")
+
+# sql 查询
+res = lidb.sql("select * from my_table;")
+```
+
+#### 数据集使用
+```python
+import lidb
+from lidb import Dataset, dataset
+import polars as pl
+
+# 定义一个tick级别的高频数据集: 高频成交量
+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")
+
+ds_hft_vol = Dataset(fn=hft_vol, 
+                     tb="path/to/hft_vol", 
+                     partitions=["num"], # 默认值 None, 会自动识别 num 
+                     update_time="", # 实时更新
+                     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`
+除了 `Dataset` 类用于管理复杂的、可分区的历史数据集之外，lidb 还提供了一个更轻量级的 `Table` 类。
+它适用于那些不需要复杂分区逻辑，且通常以单一文件形式存储的表格数据。`Table` 类同样支持基于更新时间的自动化数据管理和加载。
+##### 特性
+- **简化数据管理**: 专为单表数据设计，无需复杂的分区结构。
+- **灵活更新策略**: 
+  - **全量更新(`TableMode.F`)**: 每次更新时覆盖旧数据，仅保留最新的数据文件（0.parquet）。
+  - **增量更新(`TableMode.I`)**: 每次更新时生成一个新的带时间戳的文件（<uuid>.parquet），保留历史版本。
+- **自动更新检查**: 根据设定的 `update_time` 和文件修改时间，自动判断是否需要更新数据。
+ 
+##### 使用示例
+```python
+from lidb import Table, TableMode
+import polars as pl
+
+# 1. 定义一个数据获取函数
+def fetch_latest_stock_list() -> pl.DataFrame:
+    # 模拟从某个API或数据库获取最新的股票列表
+    import time
+    time.sleep(1) # 模拟网络延迟
+    return pl.DataFrame({
+        "symbol": ["AAPL", "GOOGL", "MSFT"],
+        "name": ["Apple Inc.", "Alphabet Inc.", "Microsoft Corp."],
+        "sector": ["Technology", "Communication Services", "Technology"]
+    })
+
+# 2. 创建 Table 实例
+# 假设此表每天上午9点更新
+stock_list_table = Table(
+    fn=fetch_latest_stock_list,
+    tb="stock_list",
+    update_time="09:00:00",
+    mode=TableMode.F # 使用全量更新模式
+)
+
+# 3. 更新数据 (可选，get_value 会自动检查并提示更新)
+# stock_list_table.update(verbose=True)
+
+# 4. 获取数据
+# 如果数据过期，get_value 会打印警告并退出，提示先调用 update()
+df = stock_list_table.get_value(date="2023-10-27")
+print(df)
+```
+
+
+#### 表达式计算
+```python
+import lidb
+
+date = "2025-05-15"
+quote_query = f"select * from quote where date = '{date}'"
+quote = lidb.read_ck(quote_query, db_conf="databases.ck")
+
+qdf = lidb.from_polars(quote)
+
+# 使用 QDF 进行表达式计算
+res = qdf.sql(
+    "ts_mean(close, 5) as c_m5", 
+    "cs_rank(volume) as vol_rank", 
+)
+```
+
+#### 数据服务
+lidb 提供了一个名为 `D` 的全局 `DataService` 实例。
+用于在后台线程中预加载数据并缓存，从而提升数据密集型任务的性能。
+这对于需要提前准备大量数据的应用非常有用，例如回测系统或实时数据处理流水线。
+##### 启动数据服务
+你可以通过调用 `D.start()` 方法来启动数据服务，指定一个数据加载函数、需要加载的键列表以及迭代配置。
+```python
+from lidb import D
+import polars as pl
+
+# 定义一个模拟的数据加载函数
+def mock_data_loader(key: str, iterables: list[str]) -> pl.DataFrame:
+  # 模拟耗时操作
+  import time
+  time.sleep(1)
+
+  # 返回简单的 DataFrame 示例
+  return pl.DataFrame({
+    "key": [key],
+    "value": [sum(len(s) for s in iterables)]
+  })
+
+# 启动数据服务
+D.start(
+  fn=mock_data_loader,
+  keys=["2023-01-01", "2023-01-02", "2023-01-03"],
+  iter_conf={"data_source_a": ["a", "b"], "data_source_b": ["x", "y"]},
+  max_cache_size=3
+)
+```
+##### 消费数据
+一旦数据服务启动，你就可以通过 `D.do()` 来消费已加载的数据。
+这个方法接受一个消费者函数作为参数，每当有新数据可用时就会被调用。
+```python
+def data_consumer(data_package: dict):
+    print(f"Consumed data for key: {data_package['key']}")
+    for name, df in data_package['data'].items():
+        print(f"  Data from {name}:")
+        print(df)
+
+# 开始消费数据
+D.do(consumer=data_consumer, wait_secs=1)
+```
+##### 停止数据服务
+当你需要停止数据服务时，你可以调用 `D.stop()` 方法。
+##### 完整示例
+以下是一个完整的示例，演示了如何使用 D 进行异步数据加载与消费:
+```python
+import lidb
+from lidb import D
+import polars as pl
+import time
+
+def fetch_market_data(key: str, iterables: list[str]) -> pl.DataFrame:
+    # 模拟网络请求或复杂计算
+    time.sleep(0.5)
+    return pl.DataFrame({
+        "date": [key],
+        "symbol_count": [len(iterables)],
+        "total_volume": [sum(ord(c) for s in iterables for c in s)]  # Dummy volume
+    })
+
+# 启动服务
+D.start(
+    fn=fetch_market_data,
+    keys=["2023-01-01", "2023-01-02", "2023-01-03"],
+    iter_conf={"symbols": ["AAPL", "GOOGL", "MSFT"]},
+    max_cache_size=2
+)
+
+# 消费者函数
+def handle_data(data_package: dict):
+    print(f"\nReceived data for {data_package['key']}:")
+    print(data_package['data']['market_data'])
+
+# 启动消费过程
+try:
+    D.do(consumer=handle_data, wait_secs=1)
+except KeyboardInterrupt:
+    print("\nShutting down data service...")
+finally:
+    D.stop()
+```
+
+### 核心模块
+
+#### 数据库操作(`database.py`)
+- `put`: 将 `polars.DataFrame` 写入指定表
+- `sql`: 执行 `SQL` 查询
+- `has`: 检查表是否存在
+- `read_mysql`,`write_mysql`: mysql 数据读写
+- `read_ck`: clickhouse 数据读取
+
+#### 数据服务(`svc/data.py`)
+- `DataService`: 数据服务管理
+- `D`: `DataService` 全局实例
+
+#### 数据集管理(`dataset.py`)
+- `Dataset`: 数据集定义和管理
+- `DataLoader`： 数据加载器
+
+#### 表达式计算(`qdf/`)
+- `QDF`: 表达式数据库
+- `Expr`: 表达式解析器
+- `UDF 函数库`:
+    - `base_udf`: 基础运算函数
+    - `ts_udf`: 时间序列函数
+    - `cs_udf`: 横截面函数
+    - `d_udf`: 日期维度函数
+    - `itd_udf`: 日内函数
+
+#### 配置管理(`init.py`)
+- 自动创建配置文件
+- 支持自定义数据存储路径
+- `polars` 线程配置
+#### 配置说明
+首次运行会在 `~/.config/lidb/settings.toml` 创建配置文件:
+```toml
+[GLOBAL]
+path = "~/lidb"  # 数据存储路径
+
+[POLARS]
+max_threads = 32  # Polars 最大线程数
+```
+
+### 许可证
+本项目采用 MIT 许可证, 请在项目根目录下查看
+
+### 联系方式
+Zhangyundi - yundi.xxii@outlook.com

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

@@ -1,15 +1,16 @@
-lidb/__init__.py,sha256=-EUd5pO1o7cBs__gvTsxquuHRBpiVn7mVGkL9miyc2k,504
+lidb/__init__.py,sha256=WuGdkD4QzcCkIG3zbXupaXJV0b3o8gvaMGhs6MhVa_c,536
 lidb/database.py,sha256=DnPXRXvUO6g0kuMo3LPl6eKo_HbD3JNW1qzoaJ14Sgo,7533
-lidb/dataset.py,sha256=j3yFtokbNILVhjV-etAJunnbgfxYAu68Dkr2cgtCYSc,15766
-lidb/init.py,sha256=jLHpeL5mIM4YjdMYAndZlDilMiKXJMr_51Ke3ZSJWCM,1170
-lidb/parse.py,sha256=f7vfj6Nguw1WzUVEUb7fs2Oh-_2YQzB_atJhm3WGC28,3379
-lidb/table.py,sha256=-85U2N1ECDtZTTCJtgOM8XBKyueIgBmYRF5DocPvkh8,4167
+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
 lidb/qdf/__init__.py,sha256=gYiSxijoPQZmbgATQX4GsutjolPpN82Kea0eQz6zGyg,1037
 lidb/qdf/errors.py,sha256=lJhhjDRdQOOKUFGlLQ9ELK4AexXBwYQSYus_V-kc5K8,1180
 lidb/qdf/expr.py,sha256=kBzXwjL_PVsJUL9FIHJ2W_G_OVRqFR-kS2mUHTt9thM,10412
 lidb/qdf/lazy.py,sha256=I08IvSkSC84qJkgtZ7nwvG_4UH07jaHBKRp7qQnwqbs,6937
 lidb/qdf/lazy2.py,sha256=ADKQaxmo-BlndhLY-idWCFypZF1icxKNHNMWEfmWy-Q,6294
-lidb/qdf/qdf.py,sha256=tfPnnQvh8uQZT4aOqJi6bDyDoJwLObvQrFeM2Ilz6vM,6236
+lidb/qdf/qdf.py,sha256=UWG9G1GI0YdG4dMz5uTV731ETEcZelHqnb0QUGrmHPM,6324
 lidb/qdf/udf/__init__.py,sha256=yIySmkWjtJ-Lj_PMP5O4EnXGDjMAPQL40NmFCekKXBw,313
 lidb/qdf/udf/base_udf.py,sha256=ZjRF2UIrZFgznbm1gxFpdf4V92oO84IaakLeeSNF44U,3444
 lidb/qdf/udf/cs_udf.py,sha256=qlBZd2c1enIdGp_DrNyQWzH3cth4ZpLBIE1hGZuJXbA,3528
@@ -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.2.0.dist-info/METADATA,sha256=fj1SvELa0jivjl6dcyut8IHbE7V00h5o6mGJkZa04S0,506
-lidb-1.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lidb-1.2.0.dist-info/top_level.txt,sha256=NgXJNwt6ld6oLXtW1vOPaEh-VO5R0JEX_KmGIJR4ueE,5
-lidb-1.2.0.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,,

lidb-1.2.0.dist-info/METADATA

@@ -1,18 +0,0 @@
-Metadata-Version: 2.4
-Name: lidb
-Version: 1.2.0
-Summary: Light database for quantor
-Requires-Python: >=3.12
-Description-Content-Type: text/markdown
-Requires-Dist: dynaconf>=3.2.11
-Requires-Dist: polars>=1.31.0
-Requires-Dist: sqlparse>=0.5.3
-Requires-Dist: logair>=1.0.1
-Requires-Dist: clickhouse-df>=0.1.5
-Requires-Dist: connectorx>=0.4.3
-Requires-Dist: pymysql>=1.1.2
-Requires-Dist: xcals>=0.0.4
-Requires-Dist: ygo>=1.2.8
-Requires-Dist: lark>=1.3.1
-Requires-Dist: numpy>=2.3.1
-Requires-Dist: tqdm>=4.67.1