lidb 1.2.0__py3-none-any.whl

lidb/__init__.py

@@ -0,0 +1,30 @@
+# Copyright (c) ZhangYundi.
+# Licensed under the MIT License. 
+
+from .init import (
+    NAME,
+    DB_PATH,
+    CONFIG_PATH,
+    get_settings,
+)
+
+from .database import (
+    sql,
+    put,
+    has,
+    tb_path,
+    read_mysql,
+    write_mysql,
+    execute_mysql,
+    read_ck,
+    scan,
+)
+
+from .table import Table, TableMode
+from .dataset import Dataset, DataLoader
+from .qdf import from_polars, Expr
+from .svc import DataService, D
+
+from .parse import parse_hive_partition_structure
+
+__version__ = "1.2.0"

lidb/database.py

@@ -0,0 +1,234 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Copyright (c) 2025 ZhangYundi
+Licensed under the MIT License.
+Created on 2024/7/1 09:44
+Email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+import re
+from pathlib import Path
+from typing import Literal
+
+import pymysql
+
+from .parse import extract_table_names_from_sql
+from .init import DB_PATH, logger, get_settings
+import urllib
+import polars as pl
+
+# ======================== 本地数据库 catdb ========================
+def tb_path(tb_name: str) -> Path:
+    """
+    返回指定表名 完整的本地路径
+    Parameters
+    ----------
+    tb_name: str
+       表名，路径写法: a/b/c
+    Returns
+    -------
+    pathlib.Path
+        full_abs_path: pathlib.Path
+        完整的本地绝对路径 $DB_PATH/a/b/c
+    """
+    return Path(DB_PATH, tb_name)
+
+
+def put(df, tb_name: str, partitions: list[str] | None = None):
+    """
+    将一个DataFrame写入到指定名称的表格目录中，支持分区存储。
+
+    该函数负责将给定的DataFrame（df）根据提供的表名（tb_name）写入到本地文件系统中。
+    如果指定了分区（partitions），则会按照这些分区列将数据分割存储。如果目录不存在，会自动创建目录。
+
+    Parameters
+    ----------
+    df: polars.DataFrame
+    tb_name: str
+        表的名称，用于确定存储数据的目录
+    partitions: list[str] | None
+        指定用于分区的列名列表。如果未提供，则不进行分区。
+
+    Returns
+    -------
+
+    """
+    if df is None:
+        logger.warning(f"put failed: input data is None.")
+        return
+    if df.is_empty():
+        logger.warning(f"put failed: input data is empty.")
+        return
+    tbpath = tb_path(tb_name)
+    if not tbpath.exists():
+        tbpath.mkdir(parents=True, exist_ok=True)
+    if partitions is not None:
+        df.write_parquet(tbpath, partition_by=partitions)
+    else:
+        df.write_parquet(tbpath / "data.parquet")
+
+def has(tb_name: str) -> bool:
+    """
+    判定给定的表名是否存在
+    Parameters
+    ----------
+    tb_name: str
+
+    Returns
+    -------
+
+    """
+    return tb_path(tb_name).exists()
+
+def sql(query: str, ):
+    """
+    sql 查询，从本地paquet文件中查询数据
+
+    Parameters
+    ----------
+    query: str
+        sql查询语句
+    Returns
+    -------
+
+    """
+    import polars as pl
+
+    tbs = extract_table_names_from_sql(query)
+    convertor = dict()
+    for tb in tbs:
+        db_path = tb_path(tb)
+        format_tb = f"read_parquet('{db_path}/**/*.parquet')"
+        convertor[tb] = format_tb
+    pattern = re.compile("|".join(re.escape(k) for k in convertor.keys()))
+    new_query = pattern.sub(lambda m: convertor[m.group(0)], query)
+    return pl.sql(new_query)
+
+def scan(tb: str,) -> pl.LazyFrame:
+    """polars.scan_parquet"""
+    tb = tb_path(tb)
+    return pl.scan_parquet(tb)
+
+def read_mysql(query: str, db_conf: str = "DATABASES.mysql"):
+    """
+    从MySQL数据库中读取数据。
+    Parameters
+    ----------
+    query: str
+        查询语句
+    db_conf: str
+        对应的配置 $DB_PATH/conf/settings.toml
+    Returns
+    -------
+    polars.DataFrame
+    """
+    import polars as pl
+    try:
+        db_setting = get_settings().get(db_conf, {})
+        required_keys = ['user', 'password', 'url', 'db']
+        missing_keys = [key for key in required_keys if key not in db_setting]
+        if missing_keys:
+            raise KeyError(f"Missing required keys in database config: {missing_keys}")
+
+        user = urllib.parse.quote_plus(db_setting['user'])
+        password = urllib.parse.quote_plus(db_setting['password'])
+        uri = f"mysql://{user}:{password}@{db_setting['url']}/{db_setting['db']}"
+        return pl.read_database_uri(query, uri)
+
+    except KeyError as e:
+        raise RuntimeError("Database configuration error: missing required fields.") from e
+    except Exception as e:
+        raise RuntimeError(f"Failed to execute MySQL query: {e}") from e
+
+def write_mysql(df: pl.DataFrame,
+                remote_tb: str,
+                db_conf: str,
+                if_table_exists: Literal["append", "replace", "fail"]="append"):
+    """将 polars.DataFrame 写入mysql"""
+    try:
+        db_setting = get_settings().get(db_conf, {})
+        required_keys = ['user', 'password', 'url', 'db']
+        missing_keys = [key for key in required_keys if key not in db_setting]
+        if missing_keys:
+            raise KeyError(f"Missing required keys in database config: {missing_keys}")
+
+        user = urllib.parse.quote_plus(db_setting['user'])
+        password = urllib.parse.quote_plus(db_setting['password'])
+        uri = f"mysql+pymysql://{user}:{password}@{db_setting['url']}/{db_setting['db']}"
+        return df.write_database(remote_tb,
+                                 connection=uri,
+                                 if_table_exists=if_table_exists)
+
+    except KeyError as e:
+        raise RuntimeError("Database configuration error: missing required fields.") from e
+    except Exception as e:
+        raise RuntimeError(f"Failed to write MySQL: {e}") from e
+
+def execute_mysql(sql: str, db_conf: str):
+    """执行mysql语句"""
+    try:
+        db_setting = get_settings().get(db_conf, {})
+        required_keys = ['user', 'password', 'url', 'db']
+        missing_keys = [key for key in required_keys if key not in db_setting]
+        if missing_keys:
+            raise KeyError(f"Missing required keys in database config: {missing_keys}")
+
+        user = urllib.parse.quote_plus(db_setting['user'])
+        password = urllib.parse.quote_plus(db_setting['password'])
+        url = urllib.parse.quote_plus(db_setting["url"])
+        host, port = url.split(":")
+
+    except KeyError as e:
+        raise RuntimeError("Database configuration error: missing required fields.") from e
+    except Exception as e:
+        raise RuntimeError(f"Failed to parse config: {e}") from e
+
+    connection = pymysql.connect(
+        host=host,
+        port=port,
+        user=user,
+        password=password,
+        database=db_setting['db']  # or extract from connection string
+    )
+    try:
+        with connection.cursor() as cursor:
+            cursor.execute(sql)
+        connection.commit()
+    except Exception as e:
+        raise RuntimeError(f"Failed to execute MySQL: {e}") from e
+    finally:
+        connection.close()
+
+
+def read_ck(query: str, db_conf: str = "DATABASES.ck"):
+    """
+    从Clickhouse集群读取数据。
+    Parameters
+    ----------
+    query: str
+        查询语句
+    db_conf: str
+        对应的配置 $DB_PATH/conf/settings.toml
+    Returns
+    -------
+    polars.DataFrame
+    """
+    import clickhouse_df
+    try:
+        db_setting = get_settings().get(db_conf, {})
+        required_keys = ['user', 'password', 'urls']
+        missing_keys = [key for key in required_keys if key not in db_setting]
+        if missing_keys:
+            raise KeyError(f"Missing required keys in database config: {missing_keys}")
+
+        user = urllib.parse.quote_plus(db_setting['user'])
+        password = urllib.parse.quote_plus(db_setting['password'])
+
+        with clickhouse_df.connect(db_setting['urls'], user=user, password=password):
+            return clickhouse_df.to_polars(query)
+
+    except KeyError as e:
+        raise RuntimeError("Database configuration error: missing required fields.") from e
+    except Exception as e:
+        raise RuntimeError(f"Failed to execute ClickHouse query: {e}") from e

lidb/dataset.py

@@ -0,0 +1,442 @@
+# Copyright (c) ZhangYundi.
+# Licensed under the MIT License. 
+# Created on 2025/10/27 14:13
+# Description:
+
+from __future__ import annotations
+
+from collections import defaultdict
+from enum import Enum
+from functools import partial
+from typing import Callable, Literal
+
+import logair
+import polars as pl
+import polars.selectors as cs
+import xcals
+import ygo
+
+from .database import put, tb_path, scan, DB_PATH
+from .parse import parse_hive_partition_structure
+from .qdf import QDF, from_polars
+
+
+class InstrumentType(Enum):
+    STOCK = "Stock"  # 股票
+    ETF = "ETF"  #
+    CB = "ConvertibleBond"  # 可转债
+
+
+def complete_data(fn, date, save_path, partitions):
+    logger = logair.get_logger(__name__)
+    try:
+        data = fn(date=date)
+        if data is None:
+            # 保存数据的逻辑在fn中实现了
+            return
+        # 剔除以 `_` 开头的列
+        data = data.select(~cs.starts_with("_"))
+        if not isinstance(data, (pl.DataFrame, pl.LazyFrame)):
+            logger.error(f"{save_path}: Result of dataset.fn must be polars.DataFrame or polars.LazyFrame.")
+            return
+        if isinstance(data, pl.LazyFrame):
+            data = data.collect()
+        cols = data.columns
+        if "date" not in cols:
+            data = data.with_columns(pl.lit(date).alias("date")).select("date", *cols)
+
+        put(data, save_path, partitions=partitions)
+    except Exception as e:
+        logger.error(f"{save_path}: Error when complete data for {date}")
+        logger.warning(e)
+
+
+class Dataset:
+
+    def __init__(self,
+                 fn: Callable[..., pl.DataFrame],
+                 tb: str,
+                 update_time: str = "",
+                 partitions: list[str] = None,
+                 by_asset: bool = True,
+                 by_time: bool = False):
+        """
+
+        Parameters
+        ----------
+        fn: str
+            数据集计算函数
+        tb: str
+            数据集保存表格
+        update_time: str
+            更新时间: 默认没有-实时更新，也就是可以取到当天值
+        partitions: list[str]
+            分区
+        by_asset: bool
+            是否按照标的进行分区，默认 True
+        by_time: bool
+            是否按照标的进行分区，默认 False
+        """
+        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")
+        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
+        self.partitions = partitions
+        self._type_asset = "asset" in self.fn_params_sig
+        self.update_time = update_time
+
+        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:
+                v = self.fn_params[k]
+                if isinstance(v, (list, tuple)) and not isinstance(v, str):
+                    v = sorted(v)
+                self.constraints[k] = v
+                self.save_path = self.save_path / f"{k}={v}"
+
+    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,
+                     tb=self.tb,
+                     partitions=self.partitions,
+                     by_asset=self._by_asset,
+                     by_time=self._by_time,
+                     update_time=self.update_time)
+        return ds
+
+    def get_value(self, date, eager: bool = True, **constraints):
+        """
+        取值: 不保证未来数据
+        Parameters
+        ----------
+        date: str
+            取值日期
+        eager: bool
+        constraints: dict
+            取值的过滤条件
+
+        Returns
+        -------
+
+        """
+        _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
+        for k, v in _constraints.items():
+            if isinstance(v, (list, tuple)) and not isinstance(v, str):
+                v = sorted(v)
+            search_path = search_path / f"{k}={v}"
+        search_path = search_path / f"date={date}"
+
+        if not self.is_empty(search_path):
+            lf = scan(search_path).cast({"date": pl.Utf8})
+            schema = lf.collect_schema()
+            _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:
+                return lf
+            data = lf.collect()
+            if not data.is_empty():
+                return data
+        fn = self.fn
+        save_path = self.save_path
+
+        if self._type_asset:
+            if "asset" in _constraints:
+                fn = ygo.delay(self.fn)(asset=_constraints["asset"])
+        if len(self.constraints) < len(self.partitions) - len(self._append_partitions):
+            # 如果分区指定的字段没有在Dataset定义中指定，需要在get_value中指定
+            params = dict()
+            for k in self.partitions[:-len(self._append_partitions)]:
+                if k not in self.constraints:
+                    v = constraints[k]
+                    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()
+        if (date > today) or (date == today and now < self.update_time):
+            logger.warning(f"{self.tb}: {date} is not ready, waiting for {self.update_time}")
+            return
+        complete_data(fn, date, save_path, self._append_partitions)
+
+        lf = scan(search_path).cast({"date": pl.Utf8})
+        schema = lf.collect_schema()
+        _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:
+            return lf
+        return lf.collect()
+
+    def get_pit(self, date: str, query_time: str, eager: bool = True, **contraints):
+        """取值：如果取值时间早于更新时间，则返回上一天的值"""
+        if not self.update_time:
+            return self.get_value(date, **contraints)
+        val_date = date
+        if query_time < self.update_time:
+            val_date = xcals.shift_tradeday(date, -1)
+        return self.get_value(val_date, eager=eager, **contraints).with_columns(date=pl.lit(date), )
+
+    def get_history(self,
+                    dateList: list[str],
+                    n_jobs: int = 5,
+                    backend: Literal["threading", "multiprocessing", "loky"] = "loky",
+                    eager: bool = True,
+                    rep_asset: str = "000001",  # 默认 000001
+                    **constraints):
+        """获取历史值: 不保证未来数据"""
+        _constraints = {k: v for k, v in constraints.items() if k in self.partitions}
+        search_path = self.save_path
+        for k, v in _constraints.items():
+            if isinstance(v, (list, tuple)) and not isinstance(v, str):
+                v = sorted(v)
+            search_path = search_path / f"{k}={v}"
+        if self.is_empty(search_path):
+            # 需要补全全部数据
+            missing_dates = dateList
+        else:
+            if not self._type_asset:
+                _search_path = self.save_path
+                for k, v in _constraints.items():
+                    if k != "asset":
+                        _search_path = _search_path / f"{k}={v}"
+                    else:
+                        _search_path = _search_path / f"asset={rep_asset}"
+                hive_info = parse_hive_partition_structure(_search_path)
+            else:
+                hive_info = parse_hive_partition_structure(search_path)
+            exist_dates = hive_info["date"].to_list()
+            missing_dates = set(dateList).difference(set(exist_dates))
+            missing_dates = sorted(list(missing_dates))
+        if missing_dates:
+            fn = self.fn
+            save_path = self.save_path
+
+            if self._type_asset:
+                if "asset" in _constraints:
+                    fn = ygo.delay(self.fn)(asset=_constraints["asset"])
+
+            if len(self.constraints) < len(self.partitions) - len(self._append_partitions):
+                params = dict()
+                for k in self.partitions[:-len(self._append_partitions)]:
+                    if k not in self.constraints:
+                        v = constraints[k]
+                        params[k] = v
+                        save_path = save_path / f"{k}={v}"
+                fn = ygo.delay(self.fn)(**params)
+
+            with ygo.pool(n_jobs=n_jobs, backend=backend) as go:
+                info_path = self.save_path
+                try:
+                    info_path = info_path.relative_to(DB_PATH)
+                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,
+                    )
+                go.do()
+        data = scan(search_path, ).cast({"date": pl.Utf8}).filter(pl.col("date").is_in(dateList), **constraints)
+        data = data.sort("date")
+        if eager:
+            return data.collect()
+        return data
+
+
+def loader(data_name: str,
+           ds: Dataset,
+           date_list: list[str],
+           prev_date_list: list[str],
+           prev_date_mapping: dict[str, str],
+           time: str,
+           **constraints) -> pl.LazyFrame:
+    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:
+        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 time < ds.update_time:
+        lf = lf.with_columns(date=pl.col("date").replace(prev_date_mapping))
+    return data_name, lf
+
+
+def load_ds(ds_conf: dict[str, list[Dataset]],
+            beg_date: str,
+            end_date: str,
+            time: str,
+            n_jobs: int = 7,
+            backend: Literal["threading", "multiprocessing", "loky"] = "threading",
+            eager: bool = False,
+            **constraints) -> dict[str, pl.DataFrame | pl.LazyFrame]:
+    """
+    加载数据集
+    Parameters
+    ----------
+    ds_conf: dict[str, list[Dataset]]
+        数据集配置: key-data_name, value-list[Dataset]
+    beg_date: str
+        开始日期
+    end_date: str
+        结束日期
+    time: str
+        取值时间
+    n_jobs: int
+        并发数量
+    backend: str
+    eager: bool
+        是否返回 DataFrame
+        - True: 返回DataFrame
+        - False: 返回LazyFrame
+    constraints
+        限制条件，比如 asset='000001'
+    Returns
+    -------
+    dict[str, polars.DataFrame | polars.LazyFrame]
+        - key: data_name
+        - value: polars.DataFrame
+
+    """
+    if beg_date > end_date:
+        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_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:
+        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)
+        for name, lf in go.do():
+            results[name].append(lf)
+    index = ("date", "time", "asset")
+    LFs = {
+        name: (pl.concat(lfList, how="align")
+               .sort(index)
+               .select(*index,
+                       cs.exclude(index))
+               )
+        for name, lfList in results.items()}
+    if not eager:
+        return LFs
+    return {
+        name: lf.collect()
+        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
+
+    def get(self,
+            ds_list: list[Dataset],
+            beg_date: str,
+            end_date: str,
+            n_jobs: int = 11,
+            backend: Literal["threading", "multiprocessing", "loky"] = "threading",
+            **constraints):
+        """
+        添加数据集
+        Parameters
+        ----------
+        ds_list: list[Dataset]
+        beg_date: str
+        end_date: str
+        n_jobs: int
+        backend: str
+        constraints
+
+        Returns
+        -------
+
+        """
+        lf = load_ds(ds_conf={self._name: ds_list},
+                     beg_date=beg_date,
+                     end_date=end_date,
+                     n_jobs=n_jobs,
+                     backend=backend,
+                     eager=False,
+                     **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()
+
+    @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()
+        return self._df
+
+    def sql(self, *exprs: str) -> pl.DataFrame:
+        return self._db.sql(*exprs)