quantnodes 3.0.0__py3-none-any.whl

QuantNodes/operator_node/__init__.py

@@ -0,0 +1,28 @@
+# coding=utf-8
+"""
+OperatorNode 操作节点模块
+
+提供数据操作、SQL 构建和转换的节点类型：
+
+- OperatorNode: 操作节点基类
+- ChainOperator: 链式操作节点
+- SQLBuilderNode: SQL 构建节点
+- TableQueryNode: 表查询执行节点
+- TransformNode: 数据转换节点
+- SQLBuilder: SQL 构建工具类
+"""
+
+from QuantNodes.operator_node.base import OperatorNode, ChainOperator
+from QuantNodes.operator_node.sql_builder import SQLBuilderNode
+from QuantNodes.operator_node.query_node import TableQueryNode
+from QuantNodes.operator_node.transform import TransformNode
+from QuantNodes.operator_node.sql_utils import SQLBuilder
+
+__all__ = [
+    'OperatorNode',
+    'ChainOperator',
+    'SQLBuilderNode',
+    'TableQueryNode',
+    'TransformNode',
+    'SQLBuilder',
+]

QuantNodes/operator_node/base.py

@@ -0,0 +1,97 @@
+# coding=utf-8
+"""
+OperatorNode 基类模块
+
+提供操作节点的基础架构，继承自 BaseNode。
+用于数据操作、SQL 构建和转换。
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+from QuantNodes.core.node import BaseNode
+
+
+class OperatorNode(BaseNode, ABC):
+    """
+    操作节点基类
+
+    所有数据操作节点都继承自此类，提供统一的操作接口。
+    支持链式调用，可以像构建 SQL 一样组合操作。
+
+    Subclasses must implement:
+        _execute_operation(): 执行具体操作
+
+    Examples:
+        >>> query_node = TableQueryNode(table="users")
+        >>> result = query_node.execute()
+    """
+
+    def __init__(self, name: str = None, config: Dict[str, Any] = None, **kwargs):
+        super().__init__(name=name or self.__class__.__name__, config=config, **kwargs)
+        self._chained_result: Any = None
+
+    @abstractmethod
+    def _execute_operation(self, input_data: Any = None, **kwargs) -> Any:
+        """
+        执行具体操作
+
+        Args:
+            input_data: 输入数据
+            **kwargs: 执行参数
+
+        Returns:
+            操作结果
+        """
+        pass
+
+    def _execute(self, input_data: Any = None, **kwargs) -> Any:
+        """执行操作"""
+        return self._execute_operation(input_data, **kwargs)
+
+    def then(self, other: 'OperatorNode') -> 'ChainOperator':
+        """
+        链式调用：将操作链接到另一个操作
+
+        Args:
+            other: 下一个操作节点
+
+        Returns:
+            ChainOperator 包含两个操作的链接
+        """
+        return ChainOperator([self, other])
+
+    def __rshift__(self, other: 'OperatorNode') -> 'ChainOperator':
+        """重载 >> 运算符支持链式调用"""
+        return self.then(other)
+
+
+class ChainOperator(OperatorNode):
+    """
+    链式操作节点
+
+    将多个 OperatorNode 链接在一起执行。
+    """
+
+    def __init__(
+        self, operators: List[OperatorNode], name: str = None,
+        config: Dict[str, Any] = None,
+    ):
+        super().__init__(name=name or "Chain", config=config)
+        self.operators = operators
+
+    def _execute_operation(self, input_data: Any = None, **kwargs) -> Any:
+        """依次执行每个操作"""
+        result = input_data
+        for op in self.operators:
+            result = op.execute(result, **kwargs)
+        return result
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {"operators": [op.serialize() for op in self.operators]}
+
+    @classmethod
+    def _from_dict_impl(cls, data: Dict[str, Any]) -> 'ChainOperator':
+        operators = [BaseNode.deserialize(op_data) for op_data in data["operators"]]
+        return ChainOperator(operators=operators, name=data.get("name"))

QuantNodes/operator_node/query_node.py

@@ -0,0 +1,129 @@
+# coding=utf-8
+"""
+TableQueryNode - 表查询执行节点
+
+执行 SQL 查询并返回 DataFrame 结果。
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from QuantNodes.core.node import BaseNode
+from QuantNodes.operator_node.base import OperatorNode
+from QuantNodes.operator_node.sql_builder import SQLBuilderNode
+
+
+class TableQueryNode(OperatorNode):
+    """
+    表查询执行节点
+
+    执行 SQL 查询或使用 SQLBuilderNode 构建查询，返回 DataFrame。
+
+    Examples:
+        >>> # 直接执行 SQL
+        >>> node = TableQueryNode(db_node=clickhouse_node, sql="SELECT * FROM users")
+        >>> df = node.execute()
+        >>>
+        >>> # 使用链式调用
+        >>> node = TableQueryNode(db_node=clickhouse_node)
+        >>> df = (node
+        ...     .from_table("users")
+        ...     .select(["id", "name"])
+        ...     .where(["active = 1"])
+        ...     .execute())
+    """
+
+    def __init__(
+        self,
+        db_node: Optional[BaseNode] = None,
+        sql: Optional[str] = None,
+        builder: Optional[SQLBuilderNode] = None,
+        name: str = None,
+        config: Dict[str, Any] = None,
+        **kwargs
+    ):
+        """
+        Args:
+            db_node: 数据库节点 (DatabaseNode 实例)
+            sql: 直接执行的 SQL 语句
+            builder: SQLBuilderNode 用于构建查询
+            name: 节点名称
+            config: 额外配置
+            **kwargs: 额外参数
+        """
+        super().__init__(name=name or "TableQuery", config=config, **kwargs)
+        self._db_node = db_node
+        self._sql = sql
+        self._builder = builder
+
+        self._table: Optional[str] = None
+        self._columns: Optional[List[str]] = None
+        self._where: List[str] = []
+        self._group_by: List[str] = []
+        self._order_by: List[str] = []
+        self._limit: Optional[int] = None
+
+    def from_table(self, table: str) -> 'TableQueryNode':
+        """设置查询表"""
+        self._table = table
+        return self
+
+    def select(self, columns: List[str]) -> 'TableQueryNode':
+        """选择列"""
+        self._columns = columns
+        return self
+
+    def where(self, conditions: List[str]) -> 'TableQueryNode':
+        """添加 WHERE 条件"""
+        self._where.extend(conditions)
+        return self
+
+    def group_by(self, columns: List[str]) -> 'TableQueryNode':
+        """添加 GROUP BY"""
+        self._group_by.extend(columns)
+        return self
+
+    def order_by(self, columns: List[str]) -> 'TableQueryNode':
+        """添加 ORDER BY"""
+        self._order_by.extend(columns)
+        return self
+
+    def limit(self, n: int) -> 'TableQueryNode':
+        """添加 LIMIT"""
+        self._limit = n
+        return self
+
+    def _execute_operation(self, input_data: Any = None, **kwargs) -> Any:
+        """执行查询"""
+        if self._db_node is None:
+            raise ValueError("db_node is required for TableQueryNode")
+
+        if self._sql:
+            sql = self._sql
+        elif self._builder:
+            sql = self._builder.to_sql()
+        else:
+            sql = self._build_sql()
+
+        return self._db_node.execute(sql, **kwargs)
+
+    def _build_sql(self) -> str:
+        """构建 SQL 语句"""
+        if self._table is None:
+            raise ValueError("table or sql must be set")
+
+        builder = SQLBuilderNode(table=self._table, columns=self._columns or ['*'])
+
+        if self._where:
+            builder.where(self._where)
+        if self._group_by:
+            builder.group_by(self._group_by)
+        if self._order_by:
+            builder.order_by(self._order_by)
+        if self._limit:
+            builder.limit(self._limit)
+
+        return builder.to_sql()
+
+    def __repr__(self) -> str:
+        return f"<TableQueryNode table='{self._table}'>"

QuantNodes/operator_node/sql_builder.py

@@ -0,0 +1,125 @@
+# coding=utf-8
+"""
+SQLBuilderNode - SQL 构建节点
+
+基于 SQLUtils 构建的 SQL 生成节点。
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from QuantNodes.operator_node.base import OperatorNode
+
+
+class SQLBuilderNode(OperatorNode):
+    """
+    SQL 构建节点
+
+    构建 SQL 查询语句，支持链式调用。
+
+    Examples:
+        >>> builder = SQLBuilderNode(table="users")
+        >>> sql = builder.select(["id", "name"]).where(["active = 1"]).execute()
+        >>> print(sql)
+        SELECT id, name FROM users WHERE active = 1
+    """
+
+    def __init__(
+        self,
+        table: str,
+        columns: Optional[List[str]] = None,
+        name: str = None,
+        config: Dict[str, Any] = None,
+        **kwargs
+    ):
+        """
+        Args:
+            table: 表名 (格式: db.table 或 table)
+            columns: 要选择的列，None 表示 *
+            name: 节点名称
+            config: 额外配置
+            **kwargs: 额外参数
+        """
+        super().__init__(name=name or "SQLBuilder", config=config, **kwargs)
+        self._table = table
+        self._columns = columns or ['*']
+        self._where: List[str] = []
+        self._group_by: List[str] = []
+        self._order_by: List[str] = []
+        self._limit: Optional[int] = None
+        self._having: List[str] = []
+        self._joins: List[Dict[str, Any]] = []
+        self._sample: Optional[str] = None
+
+    def select(self, columns: List[str]) -> 'SQLBuilderNode':
+        """选择列"""
+        self._columns = columns
+        return self
+
+    def where(self, conditions: List[str]) -> 'SQLBuilderNode':
+        """添加 WHERE 条件"""
+        self._where.extend(conditions)
+        return self
+
+    def group_by(self, columns: List[str]) -> 'SQLBuilderNode':
+        """添加 GROUP BY"""
+        self._group_by.extend(columns)
+        return self
+
+    def having(self, conditions: List[str]) -> 'SQLBuilderNode':
+        """添加 HAVING 条件"""
+        self._having.extend(conditions)
+        return self
+
+    def order_by(self, columns: List[str]) -> 'SQLBuilderNode':
+        """添加 ORDER BY"""
+        self._order_by.extend(columns)
+        return self
+
+    def limit(self, n: int) -> 'SQLBuilderNode':
+        """添加 LIMIT"""
+        self._limit = n
+        return self
+
+    def join(self, join_type: str, table: str, condition: str) -> 'SQLBuilderNode':
+        """添加 JOIN"""
+        self._joins.append({
+            'type': join_type,
+            'table': table,
+            'condition': condition
+        })
+        return self
+
+    def sample(self, ratio: str) -> 'SQLBuilderNode':
+        """添加 SAMPLE"""
+        self._sample = ratio
+        return self
+
+    def _execute_operation(self, input_data: Any = None, **kwargs) -> str:
+        """生成 SQL 语句"""
+        from QuantNodes.operator_node.sql_utils import SQLBuilder
+
+        db_table = self._table
+
+        sql = SQLBuilder.create_select_sql(
+            DB_TABLE=db_table,
+            cols=self._columns if self._columns != ['*'] else ['*'],
+            sample=self._sample,
+            array_join=None,
+            join=self._joins[0] if self._joins else None,
+            prewhere=None,
+            where=self._where if self._where else None,
+            having=self._having if self._having else None,
+            group_by=self._group_by if self._group_by else None,
+            order_by=self._order_by if self._order_by else None,
+            limit_by=None,
+            limit=self._limit
+        )
+        return sql
+
+    def to_sql(self) -> str:
+        """返回 SQL 语句（execute 的别名）"""
+        return self.execute()
+
+    def __repr__(self) -> str:
+        return f"<SQLBuilderNode table='{self._table}'>"

QuantNodes/operator_node/sql_utils.py

@@ -0,0 +1,172 @@
+# coding=utf-8
+"""
+SQL utilities
+
+SQL building utilities for generating SQL queries.
+"""
+
+
+class TableEngineCreator(object):
+    @staticmethod
+    def _assemble_cols_2_clause(prefix, cols, default=''):
+        if cols is None:
+            return default
+        else:
+            cols_str = ','.join(cols)
+            return f"{prefix} ( {cols_str} ) "
+
+    @classmethod
+    def ReplacingMergeTree_creator(cls, DB_TABLE, cols_def, order_by_cols,
+                                   sample_by_cols=None,
+                                   ON_CLUSTER='', partition_by_cols=None, primary_by_cols=None):
+
+        order_by_cols_str = ','.join(order_by_cols)
+        ORDER_BY_CLAUSE = f"ORDER BY ( {order_by_cols_str} )"
+
+        SAMPLE_CLAUSE = cls._assemble_cols_2_clause('SAMPLE BY', sample_by_cols, default='')
+
+        PRIMARY_BY_CLAUSE = cls._assemble_cols_2_clause('PRIMARY BY', primary_by_cols, default='')
+
+        PARTITION_by_CLAUSE = cls._assemble_cols_2_clause('PARTITION BY', partition_by_cols, default='')  # noqa: E501 (long var name)
+
+        return cls.raw_create_ReplacingMergeTree_table_sql(
+            DB_TABLE, cols_def, ORDER_BY_CLAUSE,
+            PRIMARY_BY_CLAUSE=PRIMARY_BY_CLAUSE,
+            SAMPLE_CLAUSE=SAMPLE_CLAUSE,
+            ENGINE_TYPE='ReplacingMergeTree', ON_CLUSTER=ON_CLUSTER,
+            PARTITION_by_CLAUSE=PARTITION_by_CLAUSE,
+        )
+
+    @staticmethod
+    def raw_create_ReplacingMergeTree_table_sql(
+        DB_TABLE, cols_def, ORDER_BY_CLAUSE,
+        PRIMARY_BY_CLAUSE='', SAMPLE_CLAUSE='',
+        ENGINE_TYPE='ReplacingMergeTree', ON_CLUSTER='',
+        PARTITION_by_CLAUSE='', TTL='',
+    ):
+        maid_body = (
+            f"CREATE TABLE IF NOT EXISTS {DB_TABLE} {ON_CLUSTER} "
+            f"( {cols_def} ) ENGINE = {ENGINE_TYPE}"
+        )
+
+        settings = "SETTINGS index_granularity = 8192"
+        conds = f"{PARTITION_by_CLAUSE} {ORDER_BY_CLAUSE} {PRIMARY_BY_CLAUSE} {SAMPLE_CLAUSE} {TTL}"
+
+        base = f"{maid_body} {conds}  {settings}"
+        return base
+
+
+class SQLBuilder:
+    """SQL building utility class"""
+
+    @staticmethod
+    def _assemble_sample(sample):
+        if sample is None:
+            return ""
+        elif isinstance(sample, (int, float)):
+            if sample < 1:
+                return f"SAMPLE {sample}"
+            else:
+                return f"SAMPLE {sample}"
+        return ""
+
+    @staticmethod
+    def _assemble_array_join(array_join_list):
+        if array_join_list is None:
+            return ""
+        clauses = []
+        for arr in array_join_list:
+            clauses.append(f"ARRAY JOIN {arr}")
+        return " ".join(clauses)
+
+    @staticmethod
+    def _assemble_join(join_info):
+        if join_info is None:
+            return ""
+        join_type = join_info.get('type', '')
+        using = join_info.get('USING', '')
+        return f"{join_type} USING ( {using} )"
+
+    @staticmethod
+    def _assemble_where_like(where_list, prefix='WHERE'):
+        if where_list is None:
+            return ""
+        clauses = []
+        for where in where_list:
+            clauses.append(where)
+        where_clause = ' AND '.join(clauses)
+        return f"{prefix} {where_clause}"
+
+    @staticmethod
+    def _assemble_group_by(group_by_list):
+        if group_by_list is None:
+            return ""
+        cols_str = ','.join(group_by_list)
+        return f"GROUP BY ( {cols_str} )"
+
+    @staticmethod
+    def _assemble_order_by(order_by_list):
+        if order_by_list is None:
+            return ""
+        cols_str = ','.join(order_by_list)
+        return f"ORDER BY ( {cols_str} )"
+
+    @staticmethod
+    def _assemble_limit_by(limit_by):
+        if limit_by is None:
+            return ""
+        n = limit_by.get('N', 10)
+        cols = limit_by.get('limit_by_cols', [])
+        cols_str = ','.join(cols)
+        return f"LIMIT {n} BY {cols_str}"
+
+    @staticmethod
+    def _assemble_limit(limit):
+        if limit is None:
+            return ""
+        return f"LIMIT {limit}"
+
+    @classmethod
+    def raw_create_select_sql(cls, SELECT_CLAUSE, DB_TABLE, SAMPLE_CLAUSE='', ARRAY_JOIN_CLAUSE='',
+                           JOIN_CLAUSE='', PREWHERE_CLAUSE='', WHERE_CLAUSE='',
+                           GROUP_BY_CLAUSE='', HAVING_CLAUSE='', ORDER_BY_CLAUSE='',
+                           LIMIT_N_CLAUSE='', LIMIT_CLAUSE=''):
+        if DB_TABLE.lower().startswith('select '):
+            DB_TABLE = f"( {DB_TABLE} )"
+        main_body = f"SELECT {SELECT_CLAUSE} FROM {DB_TABLE} {SAMPLE_CLAUSE}"
+        join = f"{ARRAY_JOIN_CLAUSE} {JOIN_CLAUSE}"
+        where_conditions = f"{PREWHERE_CLAUSE} {WHERE_CLAUSE} {GROUP_BY_CLAUSE} {HAVING_CLAUSE} "
+        order_limit = f"{ORDER_BY_CLAUSE} {LIMIT_N_CLAUSE} {LIMIT_CLAUSE}"
+        sql = f"{main_body} {join} {where_conditions} {order_limit}"
+        return sql
+
+    @classmethod
+    def create_select_sql(
+        cls, DB_TABLE: str, cols: list,
+        sample: (int, float, None) = None,
+        array_join: (list, None) = None, join: (dict, None) = None,
+        prewhere: (list, None) = None, where: (list, None) = None,
+        having: (list, None) = None, group_by: (list, None) = None,
+        order_by: (list, None) = None, limit_by: (dict, None) = None,
+        limit: (int, None) = None,
+    ) -> str:
+        SELECT_CLAUSE = ','.join(cols)
+        SAMPLE_CLAUSE = cls._assemble_sample(sample=sample)
+        ARRAY_JOIN_CLAUSE = cls._assemble_array_join(array_join_list=array_join)
+        JOIN_CLAUSE = cls._assemble_join(join)
+        PREWHERE_CLAUSE = cls._assemble_where_like(prewhere, prefix='PREWHERE')
+        WHERE_CLAUSE = cls._assemble_where_like(where, prefix='WHERE')
+        HAVING_CLAUSE = cls._assemble_where_like(having, prefix='HAVING')
+        GROUP_BY_CLAUSE = cls._assemble_group_by(group_by)
+        ORDER_BY_CLAUSE = cls._assemble_order_by(order_by)
+        LIMIT_N_CLAUSE = cls._assemble_limit_by(limit_by)
+        LIMIT_CLAUSE = cls._assemble_limit(limit)
+
+        return cls.raw_create_select_sql(
+            SELECT_CLAUSE, DB_TABLE, SAMPLE_CLAUSE, ARRAY_JOIN_CLAUSE, JOIN_CLAUSE,
+            PREWHERE_CLAUSE, WHERE_CLAUSE, GROUP_BY_CLAUSE, HAVING_CLAUSE,
+            ORDER_BY_CLAUSE, LIMIT_N_CLAUSE, LIMIT_CLAUSE,
+        )
+
+
+__all__ = ['SQLBuilder', 'TableEngineCreator']

QuantNodes/operator_node/transform.py

@@ -0,0 +1,130 @@
+# coding=utf-8
+"""
+TransformNode - 数据转换节点
+
+提供数据转换操作，如选择列、过滤、聚合等。
+"""
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Union
+
+import pandas as pd
+
+from QuantNodes.operator_node.base import OperatorNode
+
+
+class TransformNode(OperatorNode):
+    """
+    数据转换节点
+
+    对 DataFrame 进行各种转换操作，支持链式调用。
+
+    Examples:
+        >>> # 选择列
+        >>> node = TransformNode().select(["col1", "col2"])
+        >>> result = node.execute(df)
+        >>>
+        >>> # 过滤行
+        >>> node = TransformNode().filter(lambda df: df["value"] > 0)
+        >>> result = node.execute(df)
+        >>>
+        >>> # 聚合
+        >>> node = TransformNode().aggregate(group_by=["category"], agg={"value": "sum"})
+        >>> result = node.execute(df)
+    """
+
+    def __init__(
+        self,
+        name: str = None,
+        config: Dict[str, Any] = None,
+        **kwargs
+    ):
+        super().__init__(name=name or "Transform", config=config, **kwargs)
+        self._operations: List[Callable] = []
+
+    def select(self, columns: List[str]) -> 'TransformNode':
+        """选择列"""
+        def op(df: pd.DataFrame) -> pd.DataFrame:
+            return df[columns]
+        self._operations.append(op)
+        return self
+
+    def drop(self, columns: List[str]) -> 'TransformNode':
+        """删除列"""
+        def op(df: pd.DataFrame) -> pd.DataFrame:
+            return df.drop(columns=columns)
+        self._operations.append(op)
+        return self
+
+    def filter(self, condition: Union[str, Callable]) -> 'TransformNode':
+        """过滤行"""
+        if isinstance(condition, str):
+            def op(df: pd.DataFrame) -> pd.DataFrame:
+                return df.query(condition)
+        else:
+            def op(df: pd.DataFrame) -> pd.DataFrame:
+                return df[condition(df)]
+        self._operations.append(op)
+        return self
+
+    def aggregate(
+        self,
+        group_by: List[str],
+        agg: Dict[str, Union[str, List[str]]]
+    ) -> 'TransformNode':
+        """聚合操作"""
+        def op(df: pd.DataFrame) -> pd.DataFrame:
+            return df.groupby(group_by).agg(agg).reset_index()
+        self._operations.append(op)
+        return self
+
+    def sort_by(self, columns: Union[str, List[str]], ascending: bool = True) -> 'TransformNode':
+        """排序"""
+        def op(df: pd.DataFrame) -> pd.DataFrame:
+            return df.sort_values(by=columns, ascending=ascending)
+        self._operations.append(op)
+        return self
+
+    def rename(self, columns: Dict[str, str]) -> 'TransformNode':
+        """重命名列"""
+        def op(df: pd.DataFrame) -> pd.DataFrame:
+            return df.rename(columns=columns)
+        self._operations.append(op)
+        return self
+
+    def fillna(self, value: Any) -> 'TransformNode':
+        """填充空值"""
+        def op(df: pd.DataFrame) -> pd.DataFrame:
+            return df.fillna(value)
+        self._operations.append(op)
+        return self
+
+    def apply(self, func: Callable, **kwargs) -> 'TransformNode':
+        """应用自定义函数"""
+        def op(df: pd.DataFrame) -> pd.DataFrame:
+            return df.apply(func, **kwargs)
+        self._operations.append(op)
+        return self
+
+    def then(self, other: 'TransformNode') -> 'TransformNode':
+        """链式调用"""
+        combined = TransformNode(name=f"{self.name}_then_{other.name}")
+        combined._operations = self._operations + other._operations
+        return combined
+
+    def _execute_operation(self, input_data: Any = None, **kwargs) -> Any:
+        """执行转换"""
+        if input_data is None:
+            raise ValueError("input_data (DataFrame) is required")
+
+        if not isinstance(input_data, pd.DataFrame):
+            raise TypeError(f"Expected DataFrame, got {type(input_data)}")
+
+        result = input_data
+        for op in self._operations:
+            result = op(result)
+
+        return result
+
+    def __repr__(self) -> str:
+        return f"<TransformNode operations={len(self._operations)}>"