lidb 1.2.0__py3-none-any.whl

lidb/qdf/lazy.py

@@ -0,0 +1,174 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/5 21:40
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+from __future__ import annotations
+
+import importlib.util
+import sys
+from functools import lru_cache
+from functools import partial
+from pathlib import Path
+import inspect
+
+import polars as pl
+import logair
+
+from .errors import CalculateError, CompileError, PolarsError, FailError
+from .expr import Expr
+
+# 动态加载模块
+module_name = "udf"
+module_path = Path(__file__).parent / "udf" / "__init__.py"
+spec = importlib.util.spec_from_file_location(module_name, module_path)
+module = importlib.util.module_from_spec(spec)
+sys.modules[module_name] = module
+spec.loader.exec_module(module)
+
+
+@lru_cache(maxsize=512)
+def parse_expr(expr: str) -> Expr:
+    return Expr(expr)
+
+logger = logair.get_logger(__name__)
+
+class LQDF:
+
+    def __init__(self,
+                 data: pl.LazyFrame | pl.DataFrame,
+                 index: tuple[str] = ("date", "time", "asset"),
+                 align: bool = True, ):
+        assert isinstance(data, (pl.LazyFrame, pl.DataFrame)), "data must be a polars DataFrame or LazyFrame"
+        self.data: pl.LazyFrame = data.lazy().cast({pl.Decimal: pl.Float64}).drop_nulls(subset=index)
+        self._index = self.data.select(index).collect()
+        self.dims = [self._index[name].n_unique() for name in index]
+        if align:
+            lev_vals: list[pl.LazyFrame] = [self._index.select(name).unique().lazy() for name in index]
+            full_index = lev_vals[0]
+            for lev_val in lev_vals[1:]:
+                full_index = full_index.join(lev_val, how="cross")
+            self.data = full_index.join(self.data, on=index, how='left')
+        self.data = self.data.sort(index)
+        self.failed = list()
+        self._cols = set(self.data.collect_schema().names())
+        self._expr_cache = dict()  # type: dict[Expr, str]
+        self._cur_expr_cache = dict()
+
+    def __str__(self):
+        return self.data.__str__()
+
+    def __repr__(self):
+        return self.data.__str__()
+
+    def register_udf(self, func: callable, name: str = None):
+        name = name if name is not None else func.__name__
+        setattr(module, name, func)
+
+    def _compile_expr(self, expr: str):
+        """str表达式 -> polars 表达式"""
+        try:
+            expr_parsed = Expr(expr)
+            alias = expr_parsed.alias  # if expr_parsed.alias is not None else str(expr_parsed)
+            if alias in self._cols:
+                return pl.col(alias), alias
+            # 如果该表达式已有对应列，直接复用
+            if expr_parsed in self._expr_cache:
+                expr_pl: pl.Expr = pl.col(self._expr_cache[expr_parsed]).alias(alias)
+                self.data = self.data.with_columns(expr_pl)
+                return pl.col(alias), alias
+            elif expr_parsed in self._cur_expr_cache:
+                expr_pl: pl.Expr = pl.col(self._cur_expr_cache[expr_parsed]).alias(alias)
+                self.data = self.data.with_columns(expr_pl)
+                return pl.col(alias), alias
+
+            def recur_compile(expr_: Expr):
+                """递归编译"""
+                alias_ = expr_.alias
+                if alias_ in self._cols:
+                    # 已存在：直接select数据源
+                    return pl.col(alias_)
+                if expr_ in self._expr_cache:
+                    return pl.col(self._expr_cache[expr_]).alias(alias_)
+                elif expr_ in self._cur_expr_cache:
+                    return pl.col(self._cur_expr_cache[expr_]).alias(alias_)
+                func = getattr(module, expr_.fn_name)
+                _params = sorted(list(inspect.signature(func).parameters.keys()))
+                if "dims" in _params:
+                    func = partial(func, dims=self.dims)
+                args = list()
+                kwargs = dict()
+                for arg in expr_.args:
+                    if isinstance(arg, Expr):
+                        args.append(recur_compile(arg))
+                    elif isinstance(arg, dict):
+                        kwargs.update(arg)
+                    elif isinstance(arg, str):
+                        if arg.lower() == "null":
+                            args.append(None)
+                        else:
+                            args.append(pl.col(arg))
+                    else:
+                        args.append(arg)  # or args.append(pl.lit(arg))
+                try:
+                    expr_pl: pl.Expr = func(*args, **kwargs)
+                    self.data = self.data.with_columns(expr_pl.alias(alias_))
+                    self._cur_expr_cache[expr_] = alias_
+                    return pl.col(alias_)
+                except Exception as e:
+                    raise CompileError(message=f"{expr_.fn_name}({', '.join([str(arg) for arg in args])})\n{e}") from e
+
+            return recur_compile(expr_parsed), alias
+        except (CalculateError, CompileError, PolarsError) as e:
+            raise e
+        except Exception as e:
+            # 所有未处理的错误统一抛出为 CompileError
+            raise CompileError(message=f"[编译器外层]\n{e}") from e
+
+    def sql(self, *exprs: str, ):
+        """
+        表达式查询
+        Parameters
+        ----------
+        exprs: str
+            表达式，比如 "ts_mean(close, 5) as close_ma5"
+        Returns
+        -------
+            polars.LazyFrame
+        """
+        self.failed = list()
+        exprs_select = list()
+        self._cur_expr_cache = {}
+
+        for expr in exprs:
+            try:
+                compiled, alias = self._compile_expr(expr)
+                if compiled is not None:
+                    exprs_select.append(alias)
+            except Exception as e:
+                self.failed.append(FailError(expr, e))
+        if self.failed:
+            logger.warning(f"sql failed num：{len(self.failed)}/{len(exprs)}: \n {self.failed}")
+        self.data = self.data.fill_nan(None)
+        new_expr_cache = dict()
+        try:
+            current_cols = set(self.data.collect_schema().names())
+            # 缓存整理：只保留当前表达式的缓存
+            self._expr_cache.update(self._cur_expr_cache)
+            for k, v in self._expr_cache.items():
+                if v in current_cols:
+                    new_expr_cache[k] = v
+            self._expr_cache = new_expr_cache
+            final_df = self.data.select(*self._index.columns, *exprs_select)
+            return final_df
+            # return final_df
+        except Exception as e:
+            # 缓存整理：只保留当前表达式的缓存
+            for k, v in self._expr_cache.items():
+                if v in self._cols:
+                    new_expr_cache[k] = v
+            self._expr_cache = new_expr_cache
+            raise PolarsError(message=f"LazyFrame.collect() step error:\n{e}") from e

lidb/qdf/lazy2.py

@@ -0,0 +1,161 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/5 21:40
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+from __future__ import annotations
+
+import importlib.util
+import sys
+from functools import lru_cache
+from functools import partial
+from pathlib import Path
+import inspect
+
+import polars as pl
+from polars import selectors as cs
+import logair
+
+from .errors import CalculateError, CompileError, PolarsError, FailError
+from .expr import Expr
+from .expr import build_dependency_graph, topological_sort
+
+# 动态加载模块
+module_name = "udf"
+module_path = Path(__file__).parent / "udf" / "__init__.py"
+spec = importlib.util.spec_from_file_location(module_name, module_path)
+module = importlib.util.module_from_spec(spec)
+sys.modules[module_name] = module
+spec.loader.exec_module(module)
+
+logger = logair.get_logger(__name__)
+
+class LQDF:
+
+    def __init__(self,
+                 data: pl.LazyFrame | pl.DataFrame,
+                 index: tuple[str] = ("date", "time", "asset"),
+                 align: bool = True, ):
+        assert isinstance(data, (pl.LazyFrame, pl.DataFrame)), "data must be a polars DataFrame or LazyFrame"
+        self.data: pl.LazyFrame = data.lazy().cast({pl.Decimal: pl.Float64}).fill_nan(None).drop_nulls(subset=index)
+        self._index = self.data.select(index).collect()
+        self.dims = [self._index[name].n_unique() for name in index]
+        if align:
+            lev_vals: list[pl.LazyFrame] = [self._index.select(name).unique().lazy() for name in index]
+            full_index = lev_vals[0]
+            for lev_val in lev_vals[1:]:
+                full_index = full_index.join(lev_val, how="cross")
+            self.data = full_index.join(self.data, on=index, how='left')
+        self.data = self.data.sort(index)
+        self.failed = list()
+        self._cols = set(self.data.collect_schema().names())
+        self._expr_cache = dict()  # type: dict[Expr, str]
+        self._cur_expr_cache = dict()
+        self._fn_map = dict()
+
+    def __str__(self):
+        return self.data.__str__()
+
+    def __repr__(self):
+        return self.data.__str__()
+
+    def register_udf(self, func: callable, name: str = None):
+        name = name if name is not None else func.__name__
+        setattr(module, name, func)
+
+    def _to_pl_expr(self, e: Expr) -> pl.Expr:
+        """递归编译"""
+        alias = e.alias
+        if alias in self._cols:
+            return pl.col(alias)
+        elif e in self._expr_cache:
+            return pl.col(self._expr_cache[e]).alias(alias)
+        elif e in self._cur_expr_cache:
+            return pl.col(self._cur_expr_cache[e]).alias(alias)
+        else:
+            func = self._fn_map.get(e.fn_name)
+            if func is None:
+                func = getattr(module, e.fn_name)
+                _params = sorted(list(inspect.signature(func).parameters.keys()))
+                if "dims" in _params:
+                    func = partial(func, dims=self.dims)
+                self._fn_map[e.fn_name] = func
+
+            args = list()
+            kwargs = dict()
+            for arg in e.args:
+                if isinstance(arg, Expr):
+                    if arg in self._expr_cache:
+                        args.append(pl.col(self._expr_cache[arg]))
+                    elif arg in self._cur_expr_cache:
+                        args.append(pl.col(self._cur_expr_cache[arg]))
+                    # else:
+                        # 拓扑解析依赖结构出错
+                elif isinstance(arg, dict):
+                    kwargs.update(arg)
+                elif isinstance(arg, str):
+                    if arg.lower() == "null":
+                        args.append(None)
+                    else:
+                        args.append(pl.col(arg))
+                else:
+                    args.append(arg)  # or args.append(pl.lit(arg))
+            try:
+                expr_pl: pl.Expr = func(*args, **kwargs)
+                self._cur_expr_cache[e] = alias
+                return expr_pl.alias(alias)
+            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, ) -> pl.DataFrame:
+        """
+        表达式查询
+        Parameters
+        ----------
+        exprs: str
+            表达式，比如 "ts_mean(close, 5) as close_ma5"
+        Returns
+        -------
+            polars.LazyFrame
+        """
+        self.failed = list()
+        # exprs_select = list()
+        exprs_parsed = list()
+        self._cur_expr_cache = {}
+        for expr in exprs:
+            try:
+                # compiled, alias = self._compile_expr(expr)
+                e_parsed = Expr(expr)
+                exprs_parsed.append(e_parsed)
+            except Exception as e:
+                self.failed.append(FailError(expr, e))
+        if self.failed:
+            logger.warning(f"sql failed num：{len(self.failed)}/{len(exprs)}: \n {self.failed}")
+        graph, indegree, expr_map = build_dependency_graph(exprs_parsed, self._cols)
+        lvls: list[list[Expr]] = topological_sort(graph, indegree, expr_map)
+        for batch_exprs in lvls:
+            batch_exprs = [self._to_pl_expr(e) for e in batch_exprs]
+            self.data = self.data.with_columns(*batch_exprs).fill_nan(None)
+        new_expr_cache = dict()
+        try:
+            current_cols = set(self.data.collect_schema().names())
+            # 缓存整理：只保留当前表达式的缓存
+            self._expr_cache.update(self._cur_expr_cache)
+            for k, v in self._expr_cache.items():
+                if v in current_cols:
+                    new_expr_cache[k] = v
+            self._expr_cache = new_expr_cache
+            final_df = self.data.select(*self._index.columns, *[e.alias for e in exprs_parsed])
+            return final_df
+            # return final_df
+        except Exception as e:
+            # 缓存整理：只保留当前表达式的缓存
+            for k, v in self._expr_cache.items():
+                if v in self._cols:
+                    new_expr_cache[k] = v
+            self._expr_cache = new_expr_cache
+            raise PolarsError(message=f"LazyFrame.collect() step error:\n{e}") from e

lidb/qdf/qdf.py

@@ -0,0 +1,161 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/5 21:40
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+from __future__ import annotations
+
+import importlib.util
+import inspect
+import sys
+from functools import partial
+from pathlib import Path
+
+import logair
+import polars as pl
+from tqdm.auto import tqdm
+
+from .errors import CompileError, FailError
+from .expr import Expr
+from .expr import build_dependency_graph, topological_sort
+
+# 动态加载模块
+module_name = "udf"
+module_path = Path(__file__).parent / "udf" / "__init__.py"
+spec = importlib.util.spec_from_file_location(module_name, module_path)
+module = importlib.util.module_from_spec(spec)
+sys.modules[module_name] = module
+spec.loader.exec_module(module)
+
+logger = logair.get_logger(__name__)
+
+
+class QDF:
+
+    def __init__(self,
+                 data: pl.LazyFrame | pl.DataFrame,
+                 index: tuple[str] = ("date", "time", "asset"),
+                 align: bool = True, ):
+        assert isinstance(data, (pl.LazyFrame, pl.DataFrame)), "data must be a polars DataFrame or LazyFrame"
+        self.data: pl.LazyFrame = (data
+                                   .lazy()
+                                   .cast({pl.Decimal: pl.Float64})
+                                   .fill_nan(None)
+                                   .drop_nulls(subset=index)
+                                   .sort(index))
+        self.data: pl.DataFrame = self.data.collect()
+
+        self._index: pl.DataFrame = self.data.select(index)
+        self.dims = [self._index[name].n_unique() for name in index]
+
+        index_greater_than_one = [index[i] for i, dim in enumerate(self.dims) if dim > 1]
+        index_one = [index[i] for i, dim in enumerate(self.dims) if dim == 1]
+
+        if align:
+            lev_vals: list[pl.DataFrame] = [self._index.select(name).unique(maintain_order=True) for name in
+                                            index_greater_than_one]
+            full_index = lev_vals[0]
+            for lev_val in lev_vals[1:]:
+                full_index = full_index.join(lev_val, how="cross", maintain_order="left")
+            self.data = full_index.join(self.data, on=index_greater_than_one, how='left', maintain_order="left")
+            if index_one:
+                self.data = self.data.with_columns(*[pl.lit(self._index[name][0]).alias(name) for name in index_one])
+
+        self.failed = list()
+        self._cols = set(self.data.columns)
+        self._expr_cache = dict()  # {expr: pl.col(expr.alias)}
+        self._cur_batch_expr_cache = dict()
+        # self._alias_expr_map = defaultdict(pl.Expr) # {alias: pl.col(ref:alias)}
+        self._fn_map = dict()
+
+    def __str__(self):
+        return self.data.__str__()
+
+    def __repr__(self):
+        return self.data.__str__()
+
+    def register_udf(self, func: callable, name: str = None):
+        name = name if name is not None else func.__name__
+        setattr(module, name, func)
+
+    def _to_pl_expr(self, e: Expr) -> pl.Expr:
+        alias = e.alias
+        if alias in self._cols:
+            return pl.col(alias)
+        elif e in self._expr_cache:
+            return self._expr_cache[e].alias(alias)
+        else:
+            func = self._fn_map.get(e.fn_name)
+            if func is None:
+                func = getattr(module, e.fn_name)
+                _params = sorted(list(inspect.signature(func).parameters.keys()))
+                if "dims" in _params:
+                    func = partial(func, dims=self.dims)
+                self._fn_map[e.fn_name] = func
+            args = list()
+            kwargs = dict()
+            for arg in e.args:
+                if isinstance(arg, Expr):
+                    args.append(self._expr_cache[arg])
+                    # else:
+                    # 拓扑解析依赖结构出错
+                elif isinstance(arg, dict):
+                    kwargs.update(arg)
+                elif isinstance(arg, str):
+                    if arg.lower() == "null":
+                        args.append(None)
+                    else:
+                        args.append(pl.col(arg))
+                else:
+                    args.append(arg)  # or args.append(pl.lit(arg))
+            try:
+                expr_pl: pl.Expr = func(*args, **kwargs)
+                self._cur_batch_expr_cache[e] = pl.col(alias)
+                return expr_pl.alias(alias)
+            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:
+        """
+        表达式查询
+        Parameters
+        ----------
+        exprs: str
+            表达式，比如 "ts_mean(close, 5) as close_ma5"
+        show_progress: bool
+            是否展示进度条
+        Returns
+        -------
+            polars.DataFrame
+        """
+        self.failed = list()
+        exprs_parsed = list()
+        for expr in exprs:
+            try:
+                e_parsed = Expr(expr)
+                exprs_parsed.append(e_parsed)
+            except Exception as e:
+                self.failed.append(FailError(expr, e))
+        if self.failed:
+            logger.warning(f"sql failed num：{len(self.failed)}/{len(exprs)}: \n {self.failed}")
+        graph, indegree, expr_map = build_dependency_graph(exprs_parsed, self._cols)
+        lvls: list[list[Expr]] = topological_sort(graph, indegree, expr_map)
+        pbar = None
+        lvl_num = len(lvls)
+        if show_progress:
+            pbar = tqdm(total=lvl_num, desc=f"{len(exprs)}")
+        for i, batch_exprs in enumerate(lvls):
+            if show_progress:
+                pbar.set_postfix_str(f"level-{i + 1}:{len(batch_exprs)}")
+            batch_exprs = [self._to_pl_expr(e) for e in batch_exprs]
+            self.data = self.data.with_columns(*batch_exprs).fill_nan(None)
+            self._cols = set(self.data.columns)
+            self._expr_cache.update(self._cur_batch_expr_cache)
+            if show_progress:
+                pbar.update(1)
+
+        final_df = self.data.select(*self._index.columns, *[e.alias for e in exprs_parsed])
+        return final_df

lidb/qdf/udf/__init__.py

@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/4 20:20
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+
+from .base_udf import *
+from .cs_udf import *
+from .ts_udf import *
+from .d_udf import *
+from .itd_udf import *

lidb/qdf/udf/base_udf.py

@@ -0,0 +1,146 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/4 20:28
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+
+import polars as pl
+import math
+
+"""
+基本算子：一元算子、二元算子、三元算子 以及 polars 支持的表达式(剔除数据泄露的)
+"""
+# ======================== 一元算子 ========================
+
+def not_(expr: pl.Expr): return ~expr
+
+
+def neg(expr: pl.Expr): return -expr
+
+
+def abs(expr: pl.Expr): return expr.abs()
+
+
+def log(expr: pl.Expr, base=math.e): return expr.log(base=base)
+
+
+def sqrt(expr: pl.Expr): return expr.sqrt()
+
+
+def square(expr: pl.Expr): return expr ** 2
+
+
+def cube(expr: pl.Expr): return expr ** 3
+
+
+def cbrt(expr: pl.Expr): return expr ** (1 / 3)
+
+
+def sin(expr: pl.Expr): return expr.sin()
+
+def sinh(expr: pl.Expr): return expr.sinh()
+
+def arcsin(expr: pl.Expr): return expr.arcsin()
+
+def arcsinh(expr: pl.Expr): return expr.arcsinh()
+
+
+def cos(expr: pl.Expr): return expr.cos()
+
+def cosh(expr: pl.Expr): return expr.cosh()
+
+def arccos(expr: pl.Expr): return expr.arccos()
+
+def arccosh(expr: pl.Expr): return expr.arccosh()
+
+def tan(expr: pl.Expr): return expr.tan()
+
+def tanh(expr: pl.Expr): return expr.tanh()
+
+def arctan(expr: pl.Expr): return expr.arctan()
+
+def arctanh(expr: pl.Expr): return expr.arctanh()
+
+
+def sign(expr: pl.Expr): return expr.sign()
+
+
+def sigmoid(expr: pl.Expr): return 1 / (1 + (-expr).exp())
+
+
+# def all(expr: pl.Expr, ignore_nulls: bool = True): return expr.all(ignore_nulls=ignore_nulls)
+
+
+# def any(expr: pl.Expr, ignore_nulls: bool = True): return expr.any(ignore_nulls=ignore_nulls)
+
+def cot(expr: pl.Expr): return expr.cot()
+
+def degrees(expr: pl.Expr): return expr.degrees()
+
+def exp(expr: pl.Expr): return expr.exp()
+
+def log1p(expr: pl.Expr): return expr.log1p()
+
+def clip(expr: pl.Expr, lower_bound, upper_bound): return expr.clip(lower_bound, upper_bound)
+
+# ======================== 二元算子 ========================
+def add(left: pl.Expr, right: pl.Expr): return left + right
+
+
+def sub(left: pl.Expr, right: pl.Expr): return left - right
+
+
+def mul(left: pl.Expr, right: pl.Expr): return left * right
+
+
+def div(left: pl.Expr, right: pl.Expr): return left / right
+
+
+def floordiv(left: pl.Expr, right: pl.Expr): return left // right
+
+
+def mod(left: pl.Expr, right: pl.Expr): return left % right
+
+
+def lt(left: pl.Expr, right: pl.Expr): return left < right
+
+
+def le(left: pl.Expr, right: pl.Expr): return left <= right
+
+
+def gt(left: pl.Expr, right: pl.Expr): return left > right
+
+
+def ge(left: pl.Expr, right: pl.Expr): return left >= right
+
+
+def eq(left: pl.Expr, right: pl.Expr): return left == right
+
+
+def neq(left: pl.Expr, right: pl.Expr): return left != right
+
+def and_(left: pl.Expr, right: pl.Expr): return left & right
+
+def or_(left: pl.Expr, right: pl.Expr): return left | right
+
+def max(*exprs: pl.Expr):
+    return pl.max_horizontal(*exprs)
+
+def min(*exprs: pl.Expr):
+    return pl.min_horizontal(*exprs)
+
+def sum(*exprs: pl.Expr): return pl.sum_horizontal(*exprs)
+
+# ======================== 三元 ========================
+def if_(cond: pl.Expr, body: pl.Expr, or_else: pl.Expr):
+    return pl.when(cond).then(body).otherwise(or_else)
+
+def fib(high: pl.Expr, low: pl.Expr, ratio: float = 0.618):
+    """
+    计算裴波那契回调比率
+    ratio: 0.236 | 0.382 | 0.618 等黄金分割比例
+    """
+    return low + (high - low) * ratio