lidb 2.0.20__py3-none-any.whl

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

@@ -0,0 +1,45 @@
+# Copyright (c) ZhangYundi.
+# Licensed under the MIT License. 
+# Created on 2025/7/17 14:40
+# Description:
+
+from pathlib import Path
+from dynaconf import Dynaconf
+import logair
+import os
+
+
+USERHOME = Path("~").expanduser() # 用户家目录
+NAME = "lidb"
+DB_PATH = USERHOME / NAME
+CONFIG_PATH = USERHOME / ".config" / NAME / "settings.toml"
+
+logger = logair.get_logger(NAME)
+
+
+if not CONFIG_PATH.exists():
+    try:
+        CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
+    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\n[POLARS]\nmax_threads=32\n'
+        f.write(template_content)
+    logger.info(f"Settings file created: {CONFIG_PATH}")
+
+def get_settings():
+    try:
+        return Dynaconf(settings_files=[CONFIG_PATH])
+    except Exception as e:
+        logger.error(f"Read settings file failed: {e}")
+        return {}
+
+# 读取配置文件覆盖
+_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

@@ -0,0 +1,111 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Copyright (c) 2025 ZhangYundi
+Licensed under the MIT License.
+Created on 2024/11/6 下午7:25
+Email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+import re
+from pathlib import Path
+from urllib.parse import unquote
+
+import polars as pl
+import sqlparse
+
+
+def format_sql(sql_content):
+    """将sql语句进行规范化，并去除sql中的注释，输入和输出均为字符串"""
+    parse_str = sqlparse.format(sql_content, reindent=True, strip_comments=True)
+    return parse_str
+
+
+def extract_temp_tables(with_clause):
+    """从WITH子句中提取临时表名，输出为列表"""
+    temp_tables = re.findall(r'\b(\w+)\s*as\s*\(', with_clause, re.IGNORECASE)
+    return temp_tables
+
+
+def extract_table_names_from_sql(sql_query):
+    """从sql中提取对应的表名称，输出为列表"""
+    table_names = set()
+    # 解析SQL语句
+    parsed = sqlparse.parse(sql_query)
+    # 正则表达式模式，用于匹配表名
+    table_name_pattern = r'\bFROM\s+([^\s\(\)\,]+)|\bJOIN\s+([^\s\(\)\,]+)'
+
+    # 用于存储WITH子句中的临时表名
+    remove_with_name = []
+
+    # 遍历解析后的语句块
+    for statement in parsed:
+        # 转换为字符串
+        statement_str = str(statement)  # .lower()
+
+        # 将字符串中的特殊语法置空
+        statement_str = re.sub(r'(substring|extract)\s*\(((.|\s)*?)\)', '', statement_str)
+
+        # 查找匹配的表名
+        matches = re.findall(table_name_pattern, statement_str, re.IGNORECASE)
+
+        for match in matches:
+            # 提取非空的表名部分
+            for name in match:
+                if name:
+                    # 对于可能包含命名空间的情况，只保留最后一部分作为表名
+                    table_name = name.split('.')[-1]
+                    # 去除表名中的特殊符号
+                    table_name = re.sub(r'("|`|\'|;)', '', table_name)
+                    table_names.add(table_name)
+
+        # 处理特殊的WITH语句
+        if 'with' in statement_str:
+            remove_with_name = extract_temp_tables(statement_str)
+    # 移除多余的表名
+    if remove_with_name:
+        table_names = list(set(table_names) - set(remove_with_name))
+
+    return table_names
+
+
+def parse_hive_partition_structure(root_path: Path | str, file_pattern: str = "*.parquet") -> pl.DataFrame:
+    """
+    通用Hive分区结构解析器
+
+    Args:
+        root_path: 根路径 (如 /data)
+        file_pattern: 文件匹配模式 (默认 "*.parquet")
+
+    Returns:
+        polars.DataFrame
+    """
+    if isinstance(root_path, str):
+        root_path = Path(root_path)
+
+    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)
+
+        # 收集分区信息
+        partition_dict = {}
+        for part in relative_path.parts[:-1]:  # 排除文件名
+            if '=' in part:
+                key, value = part.split('=', 1)
+                value = unquote(value)
+
+                partition_dict[key] = value
+
+        # 记录分区组合
+        combination = tuple(partition_dict.items())
+        partition_combinations.add(combination)
+
+    # 转换为普通dict
+    res = [dict(combo) for combo in partition_combinations]
+
+    return pl.DataFrame(res)

lidb/qdf/__init__.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/5 21:40
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+from __future__ import annotations
+
+from .qdf import QDF
+from .lazy import LQDF
+from .expr import Expr
+from typing import TYPE_CHECKING
+from pathlib import Path
+from ..dataset import scan
+
+if TYPE_CHECKING:
+    import polars as pl
+
+
+
+def from_polars(df: pl.DataFrame | pl.LazyFrame | Path | str, index: tuple[str] = ("date", "time", "asset"), align: bool = False, ) -> QDF:
+    """polars dataframe 转为 表达式数据库"""
+    if isinstance(df, (Path, str)):
+        df = scan(df)
+    return QDF(df, index, align,)
+
+# def to_lazy(df: pl.DataFrame | pl.LazyFrame | Path | str, index: tuple[str] = ("date", "time", "asset"), align: bool = False, ) -> LQDF:
+#     """polars dataframe/lazy frame/table path 转为 表达式数据库"""
+#     if isinstance(df, (Path, str)):
+#         df = scan(df)
+#     return LQDF(df, index, align,)
+

lidb/qdf/errors.py

@@ -0,0 +1,65 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/5/16 10:47
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+
+from dataclasses import dataclass
+
+@dataclass
+class ParseError(Exception):
+    message: str
+
+    def __str__(self):
+        return self.message
+
+    def __repr__(self):
+        return self.__str__()
+
+@dataclass
+class CalculateError(Exception):
+    message: str
+
+    def __str__(self):
+        return self.message
+
+    def __repr__(self):
+        return self.__str__()
+
+@dataclass
+class CompileError(Exception):
+    message: str
+
+    def __str__(self):
+        return self.message
+
+    def __repr__(self):
+        return self.__str__()
+
+@dataclass
+class PolarsError(Exception):
+    message: str
+
+    def __str__(self):
+        return self.message
+
+    def __repr__(self):
+        return self.__str__()
+
+@dataclass
+class FailError:
+    expr: str
+    error: Exception
+
+    def __str__(self):
+        return f"""
+[失败表达式]: {self.expr}
+[错误类型]: {self.error.__class__.__name__}
+[错误信息]: \n{self.error}
+"""
+
+    def __repr__(self):
+        return self.__str__()

lidb/qdf/expr.py

@@ -0,0 +1,370 @@
+# -*- coding: utf-8 -*-
+"""
+---------------------------------------------
+Created on 2025/3/3 19:52
+@author: ZhangYundi
+@email: yundi.xxii@outlook.com
+---------------------------------------------
+"""
+
+from __future__ import annotations
+
+import re
+import warnings
+from dataclasses import dataclass
+
+from lark import Lark, Transformer, v_args
+
+from .errors import ParseError
+from collections import defaultdict, deque
+
+
+# 基类
+class Token:
+    pass
+
+
+@dataclass
+class OperatorToken(Token):
+    """算子类型token"""
+    value: str
+
+
+@dataclass
+class OperandToken(Token):
+    """运算对象token"""
+    value: str | float | int
+
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    grammar = """
+        start: expr
+        ?expr: ternary_expr
+        ?ternary_expr: or_expr
+            | or_expr "?" or_expr ":" ternary_expr -> ternary
+        ?or_expr: and_expr
+            | or_expr "|" and_expr -> or_
+        ?and_expr: comp_expr
+            | and_expr "&" comp_expr -> and_
+        ?comp_expr: eq_expr
+            | comp_expr "<" eq_expr -> lt
+            | comp_expr ">" eq_expr -> gt
+            | comp_expr "<=" eq_expr -> le
+            | comp_expr ">=" eq_expr -> ge
+        ?eq_expr: arith_expr
+            | eq_expr "==" arith_expr -> eq
+            | eq_expr "!=" arith_expr -> neq
+        ?arith_expr: term
+            | arith_expr "+" term -> add
+            | arith_expr "-" term -> sub
+        ?term: pow_expr
+            | term "*" pow_expr -> mul
+            | term "/" pow_expr -> div
+            | term "//" pow_expr -> floordiv // 取整
+            | term "%" pow_expr -> mod      // 求余
+        ?pow_expr: factor
+            | factor "**" pow_expr -> pow
+        ?factor: atom
+            | "-" factor -> neg
+            | "!" factor -> not_
+            | "~" factor -> not_
+        ?atom: function
+            | NAME
+            | NUMBER
+            | FLOAT
+            | "(" expr ")"
+            | implicit_mul // 隐式乘法
+            | attribute_access // 新增：属性访问
+        implicit_mul: (NUMBER | FLOAT) NAME -> implicit_mul  // 隐式乘法
+        attribute_access: atom "." NAME -> attribute_access // 新增：属性访问
+        function: NAME "(" expr_list ")" -> function
+        // expr_list: expr ("," expr)*
+        keyword_arg: NAME "=" expr -> keyword_arg  // 关键字参数
+        expr_list: (expr | keyword_arg) ("," (expr | keyword_arg))*  // 支持关键字参数
+        NAME: /[a-zA-Z_$,][a-zA-Z0-9_$]*/
+        NUMBER: /\\d+/  // regex for numbers
+        FLOAT: /\\d+\\.\\d+([eE][+-]?\\d+)?/ | /\\d+[eE][+-]?\\d+/  // 支持科学计数法
+        %import common.WS
+        %ignore WS
+    """
+
+
+class ExprParser(Transformer):
+    @v_args(inline=True)
+    def ternary(self, a, b, c):
+        return Expr.new("if_", [a, b, c])
+
+    def attribute_access(self, items):
+        return ".".join(items)
+
+    def keyword_arg(self, item):
+        k, v = item
+        return {k: v}
+
+    def NAME(self, name):
+        return str(name)
+
+    def NUMBER(self, number):  # new transformer for numbers
+        return int(number)
+
+    def FLOAT(self, number):
+        return float(number)
+
+    def add(self, items):
+        return Expr.new("add", items)
+
+    def sub(self, items):
+        return Expr.new("sub", items)
+
+    def mul(self, items):
+        return Expr.new("mul", items)
+
+    def div(self, items):
+        return Expr.new("div", items)
+
+    def floordiv(self, items):
+        return Expr.new("floordiv", items)
+
+    def mod(self, items):
+        return Expr.new("mod", items)
+
+    def pow(self, items):
+        return Expr.new("pow", items)
+
+    def neg(self, items):
+        item = items[0]
+        if isinstance(item, (int, float)):
+            return -item
+        return Expr.new("neg", items)
+
+    def not_(self, item):
+        return Expr.new("not_", item)
+
+    def and_(self, items):
+        return Expr.new("and_", items)
+
+    def or_(self, items):
+        return Expr.new("or_", items)
+
+    def eq(self, items):
+        return Expr.new("eq", items)
+
+    def neq(self, items):
+        return Expr.new("neq", items)
+
+    def lt(self, items):
+        return Expr.new("lt", items)
+
+    def gt(self, items):
+        return Expr.new("gt", items)
+
+    def le(self, items):
+        return Expr.new("le", items)
+
+    def ge(self, items):
+        return Expr.new("ge", items)
+
+    def function(self, items):
+        name = items.pop(0)
+        return Expr.new(name, items[0])
+
+    def implicit_mul(self, items):
+        return Expr.new("mul", items)
+
+    def expr_list(self, items):
+        return items
+
+
+parser = Lark(grammar, parser='lalr', transformer=ExprParser())
+
+
+def parse_expr(expression: str) -> Expr:
+    return parser.parse(expression).children[0]
+
+
+class Expr:
+
+    def __init__(self, expr: str | None = None):
+
+        self.fn_name: str | None = ""
+        self.args: list | None = None
+        self.alias: str | None = None
+        if expr:
+            try:
+                self._parse(expr)
+            except Exception as e:
+                raise ParseError(f"{expr}\n{e}")
+
+    @classmethod
+    def new(cls, fn_name: str | None, args: list | None, alias: str | None = None):
+        expr = cls()
+        expr.fn_name = fn_name
+        expr.args = args
+        expr.alias = alias if alias is not None else str(expr)
+        return expr
+
+    def __hash__(self):
+        return hash(str(self).strip())
+
+    def __eq__(self, other):
+        return isinstance(other, Expr) and str(self).strip() == str(other).strip()
+
+    def to_rpn(self) -> list[Token]:
+        """生成逆波兰表达式: (后缀表达式: 运算符在后)"""
+        rpn = list()
+
+        # 递归遍历子表达式
+        def _traverse(node: Expr):
+
+            if node.args is not None:
+                for child in node.args:
+                    if isinstance(child, Expr):
+                        _traverse(child)
+                    else:
+                        rpn.append(OperandToken(child))
+            rpn.append(OperatorToken(node.fn_name))
+
+        _traverse(self)
+
+        return rpn
+
+    def __str__(self):
+        unary_map = {"neg": "-", "not_": "!"}
+        binary_map = {"add": "+",
+                      "mul": "*",
+                      "div": "/",
+                      "sub": "-",
+                      "floordiv": "//",
+                      "mod": "%",
+                      "pow": "**",
+                      "and_": "&",
+                      "or_": "|",
+                      "gt": ">",
+                      "gte": ">=",
+                      "lt": "<",
+                      "lte": "<=",
+                      "eq": "==",
+                      "neq": "!=",
+                      }
+        if self.fn_name is None:
+            return str(self.args[0])
+        if self.fn_name == "if_":
+            cond, body, orelse = self.args
+            return f"{cond}?{body}:{orelse}"
+        elif self.fn_name in ("neg", "not_"):
+            return f"{unary_map.get(self.fn_name)}{self.args[0]}"
+        elif self.fn_name in binary_map:
+            return f"({binary_map.get(self.fn_name).join([str(arg) for arg in self.args])})"
+        else:
+            return f"{self.fn_name}({', '.join([str(arg) for arg in self.args])})"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def _parse(self, expr):
+        """
+        解析表达式
+        """
+        convertor = {
+            'if(': 'if_(',
+            'not(': 'not_(',
+            'and(': 'and_(',
+            'or(': 'or_(',
+            '$': '',
+            "\n": '',
+            "!": "~",
+            ",": ", ",
+        }
+        for old, new in convertor.items():
+            expr = expr.replace(old, new)
+        new_expr = expr
+        match = re.search(r'(?i)(.+?)\s+AS\s+(\w+)', new_expr)
+        alias = None
+        if match:
+            new_expr = match.group(1).strip()
+            alias = match.group(2).strip()
+
+        expr_ = parse_expr(new_expr)
+        self.alias = alias if alias is not None else str(expr_)
+        if not isinstance(expr_, Expr):
+            self.args = [expr_]
+        else:
+            self.fn_name, self.args = expr_.fn_name, expr_.args
+
+    @property
+    def n_args(self) -> int:
+        """返回表达式的参数个数"""
+        return len(self.args)
+
+    @property
+    def depth(self) -> int:
+        """返回表达式的嵌套深度"""
+        _depth = 1
+        _depths = [0]
+        for arg in self.args:
+            if isinstance(arg, Expr):
+                _depths.append(arg.depth)
+        return _depth + max(_depths)
+
+def build_dependency_graph(exprs: list[Expr], avail_cols: set[str]) -> tuple[dict, dict, dict]:
+    """
+    构建表达式的依赖图
+    Parameters
+    ----------
+    exprs: list[Expr]
+    avail_cols: set[str]
+
+    Returns
+    -------
+    tuple[dict, dict]
+        - graph: {expr_alias: [依赖的 expr_alias]}
+        - indegree: {expr_alias: indegree}
+        - expr_map: {expr_alias: Expr}
+
+    """
+
+    graph = defaultdict(list)
+    indegree = defaultdict(int)
+    expr_map = {}
+
+    def collect_deps(e: Expr):
+        alias = e.alias
+        expr_map[alias] = e
+        for arg in e.args:
+            if isinstance(arg, Expr):
+                graph[arg.alias].append(alias)
+                indegree[alias] += 1
+                collect_deps(arg)
+            elif isinstance(arg, str) and arg.lower() != "null":
+                if arg not in avail_cols:
+                    graph[arg].append(alias)
+                    indegree[alias] += 1
+    for expr in exprs:
+        # parsed = parse_expr(expr) if isinstance(expr, str) else expr
+        collect_deps(expr)
+
+    for alias in expr_map:
+        if alias not in indegree:
+            indegree[alias] = 0
+
+    return graph, indegree, expr_map
+
+def topological_sort(graph: dict, indegree: dict, expr_map: dict) -> list[list[str]]:
+    """
+    返回按照层级划分的表达式执行顺序(每层内部没有依赖)
+    """
+    queue = deque([k for k in expr_map if indegree[k] == 0])
+    levels = []
+    while queue:
+        level_size = len(queue)
+        current_level = list()
+        for _ in range(level_size):
+            node = queue.popleft()
+            current_level.append(expr_map[node])
+            for neighbor in graph[node]:
+                indegree[neighbor] -= 1
+                if indegree[neighbor] == 0:
+                    queue.append(neighbor)
+        levels.append(current_level)
+    return levels