hexdag 0.5.0.dev1__py3-none-any.whl

hexdag_plugins/hexdag_etl/hexdag_etl/nodes/pandas_transform.py

@@ -0,0 +1,563 @@
+"""Pandas transform node with multi-operation support for ETL pipelines."""
+
+import asyncio
+import importlib
+from collections.abc import Callable
+from dataclasses import asdict
+from typing import Any
+
+import pandas as pd
+from hexdag.builtin.nodes.base_node_factory import BaseNodeFactory
+from hexdag.core.domain.dag import NodeSpec
+from hexdag.core.registry import node
+from hexdag.core.registry.models import NodeSubtype
+from pydantic import BaseModel
+
+
+class PandasOperation(BaseModel):
+    """Single pandas operation configuration."""
+
+    type: str = "transform"
+    """Operation type: 'transform', 'map', 'filter', 'assign'"""
+
+    method: str | None = None
+    """Pandas method path (e.g., 'pandas.DataFrame.groupby', 'pandas.merge')"""
+
+    args: list[Any] | None = None
+    """Positional arguments for the operation"""
+
+    kwargs: dict[str, Any] | None = None
+    """Keyword arguments for the operation"""
+
+    columns: dict[str, str] | None = None
+    """Column mappings (for 'map' or 'rename' operations)"""
+
+    condition: str | None = None
+    """Filter condition expression (for 'filter' operations)"""
+
+
+@node(name="pandas_transform_node", subtype=NodeSubtype.FUNCTION, namespace="etl")
+class PandasTransformNode(BaseNodeFactory):
+    """Node factory for multi-operation pandas transforms.
+
+    Executes a sequence of pandas operations on DataFrames, supporting:
+    - Chained transformations
+    - Multiple input DataFrames
+    - Artifact storage integration
+    - Complex data cleaning and enrichment
+
+    Examples
+    --------
+    YAML pipeline::
+
+        - kind: pandas_transform_node
+          metadata:
+            name: clean_and_aggregate
+          spec:
+            input_artifacts:
+              - slot: raw_customers
+                key: customers_v1
+              - slot: raw_transactions
+                key: transactions_v1
+            operations:
+              # Operation 1: Join DataFrames
+              - type: transform
+                method: pandas.merge
+                args:
+                  - {{input_artifacts[0]}}
+                  - {{input_artifacts[1]}}
+                kwargs:
+                  on: customer_id
+                  how: left
+
+              # Operation 2: Drop missing values
+              - type: transform
+                method: pandas.DataFrame.dropna
+                kwargs:
+                  subset: [customer_id, amount]
+
+              # Operation 3: Calculate new column
+              - type: transform
+                method: pandas.DataFrame.assign
+                kwargs:
+                  revenue_tier: |
+                    lambda df: pd.cut(
+                      df['amount'],
+                      bins=[0, 100, 500, float('inf')],
+                      labels=['Low', 'Medium', 'High']
+                    )
+
+              # Operation 4: Rename columns
+              - type: map
+                columns:
+                  transaction_id: txn_id
+                  customer_id: cust_id
+                  amount: total_amount
+
+              # Operation 5: Filter rows
+              - type: filter
+                condition: "{{ df['amount'] > 0 }}"
+
+              # Operation 6: Group and aggregate
+              - type: transform
+                method: pandas.DataFrame.groupby
+                args:
+                  - customer_id
+                kwargs:
+                  as_index: false
+
+              # Operation 7: Calculate aggregations
+              - type: transform
+                method: pandas.DataFrame.agg
+                kwargs:
+                  amount: ['count', 'sum', 'mean']
+                  customer_id: 'count'
+
+            output_artifact:
+              slot: enriched_customers
+              key: enriched_v1
+              format: parquet
+              compression: snappy
+    """
+
+    def __call__(
+        self,
+        name: str,
+        operations: list[dict[str, Any]],
+        input_artifacts: list[dict[str, Any]] | None = None,
+        output_artifact: dict[str, Any] | None = None,
+        deps: list[str] | None = None,
+        **kwargs: Any,
+    ) -> NodeSpec:
+        """Create a pandas transform node specification.
+
+        Parameters
+        ----------
+        name : str
+            Node name
+        operations : list[dict]
+            List of pandas operation configurations
+        input_artifacts : list[dict], optional
+            Artifact references for input DataFrames
+        output_artifact : dict, optional
+            Artifact configuration for output DataFrame
+        deps : list[str], optional
+            Dependency node names
+        **kwargs : Any
+            Additional node parameters
+
+        Returns
+        -------
+        NodeSpec
+            Node specification ready for execution
+        """
+        # Convert operation dicts to Pydantic models for validation
+        operation_models = [PandasOperation(**op) for op in operations]
+
+        # Create wrapped function
+        wrapped_fn = self._create_transform_function(name, operation_models, input_artifacts, output_artifact)
+
+        # Define input schema
+        if input_artifacts:
+            input_schema = {"input_data": dict, "**ports": dict}
+        else:
+            input_schema = {"input_data": dict, "**ports": dict}
+
+        # Define output schema
+        output_schema = {"output": dict}
+
+        input_model = self.create_pydantic_model(f"{name}Input", input_schema)
+        output_model = self.create_pydantic_model(f"{name}Output", output_schema)
+
+        # Store parameters
+        node_params = {
+            "operations": operations,
+            "input_artifacts": input_artifacts,
+            "output_artifact": output_artifact,
+            **kwargs,
+        }
+
+        return NodeSpec(
+            name=name,
+            fn=wrapped_fn,
+            in_model=input_model,
+            out_model=output_model,
+            deps=frozenset(deps or []),
+            params=node_params,
+        )
+
+    def _create_transform_function(
+        self,
+        name: str,
+        operations: list[PandasOperation],
+        input_artifacts: list[dict[str, Any]] | None,
+        output_artifact: dict[str, Any] | None,
+    ) -> Callable[..., dict[str, Any]]:
+        """Create the wrapped transformation function.
+
+        Parameters
+        ----------
+        name : str
+            Node name
+        operations : list[PandasOperation]
+            Operations to execute
+        input_artifacts : list[dict], optional
+            Input artifact references
+        output_artifact : dict, optional
+            Output artifact configuration
+
+        Returns
+        -------
+        Callable
+            Async function that executes the transformation
+        """
+
+        async def wrapped_fn(input_data: Any, **ports: Any) -> dict[str, Any]:
+            """Execute pandas transformation operations."""
+            # Initialize result DataFrame
+            df = None
+
+            # Load input artifacts if specified
+            if input_artifacts:
+                artifact_store = ports.get("artifact_store")
+                if not artifact_store:
+                    raise ValueError("artifact_store port required when using input_artifacts")
+
+                loaded_dfs = []
+                for artifact_ref in input_artifacts:
+                    slot = artifact_ref.get("slot")
+                    key = artifact_ref.get("key")
+                    format = artifact_ref.get("format")
+
+                    if not slot or not key:
+                        raise ValueError(f"Invalid artifact reference: {artifact_ref}")
+
+                    # Load from artifact store
+                    df_loaded = await artifact_store.read(name=slot, key=key, format=format)
+                    loaded_dfs.append(df_loaded)
+
+                # Start with first DataFrame if available
+                if loaded_dfs:
+                    df = loaded_dfs[0]
+            else:
+                # Use input_data directly
+                if isinstance(input_data, dict) and "data" in input_data:
+                    df = input_data["data"]
+                else:
+                    df = input_data
+
+            if df is None:
+                raise ValueError("No input DataFrame available")
+
+            if not isinstance(df, pd.DataFrame):
+                # Try to convert to DataFrame
+                try:
+                    df = pd.DataFrame(df)
+                except Exception as e:
+                    raise ValueError(f"Could not convert input to DataFrame: {e}")
+
+            # Execute operations sequentially
+            for i, op in enumerate(operations):
+                df = await self._execute_operation(df, op, loaded_dfs if input_artifacts else [df])
+
+            # Store output artifact if specified
+            result = {"output": df}
+
+            if output_artifact:
+                artifact_store = ports.get("artifact_store")
+                if not artifact_store:
+                    raise ValueError("artifact_store port required when using output_artifact")
+
+                slot = output_artifact.get("slot")
+                key = output_artifact.get("key", f"{name}_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}")
+                format = output_artifact.get("format", "pickle")
+                compression = output_artifact.get("compression")
+                metadata = output_artifact.get("metadata")
+
+                if not slot:
+                    raise ValueError("output_artifact must specify 'slot' name")
+
+                # Write to artifact store
+                artifact_info = await artifact_store.write(
+                    name=slot,
+                    key=key,
+                    data=df,
+                    format=format,
+                    compression=compression,
+                    metadata=metadata,
+                )
+
+                result["artifact_info"] = asdict(artifact_info)
+                result["records"] = len(df)
+
+            return result
+
+        # Preserve function metadata
+        wrapped_fn.__name__ = f"pandas_transform_{name}"
+        wrapped_fn.__doc__ = f"Multi-operation pandas transform: {name}"
+
+        return wrapped_fn
+
+    async def _execute_operation(
+        self, df: pd.DataFrame, op: PandasOperation, input_dfs: list[pd.DataFrame]
+    ) -> pd.DataFrame:
+        """Execute a single pandas operation.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            Current DataFrame
+        op : PandasOperation
+            Operation to execute
+        input_dfs : list[pd.DataFrame]
+            Available input DataFrames
+
+        Returns
+        -------
+        pd.DataFrame
+            Transformed DataFrame
+        """
+        op_type = op.type or "transform"
+
+        if op_type == "transform":
+            return await self._execute_transform(df, op, input_dfs)
+
+        elif op_type == "map":
+            return await self._execute_map(df, op)
+
+        elif op_type == "filter":
+            return await self._execute_filter(df, op)
+
+        elif op_type == "assign":
+            return await self._execute_assign(df, op)
+
+        else:
+            raise ValueError(f"Unknown operation type: {op_type}")
+
+    async def _execute_transform(
+        self, df: pd.DataFrame, op: PandasOperation, input_dfs: list[pd.DataFrame]
+    ) -> pd.DataFrame:
+        """Execute a transform operation (calls a pandas method).
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            Current DataFrame
+        op : PandasOperation
+            Operation configuration
+        input_dfs : list[pd.DataFrame]
+            Available input DataFrames
+
+        Returns
+        -------
+        pd.DataFrame
+            Transformed DataFrame
+        """
+        if not op.method:
+            raise ValueError("Transform operation requires 'method' parameter")
+
+        # Resolve method
+        method = self._resolve_method(op.method)
+
+        # Prepare arguments (resolve template expressions)
+        args = []
+        if op.args:
+            for arg in op.args:
+                args.append(self._resolve_arg(arg, df, input_dfs))
+
+        # Prepare keyword arguments
+        kwargs = {}
+        if op.kwargs:
+            for k, v in op.kwargs.items():
+                kwargs[k] = self._resolve_arg(v, df, input_dfs)
+
+        # Execute method (handle both sync and async)
+        if asyncio.iscoroutinefunction(method):
+            result = await method(df, *args, **kwargs)
+        else:
+            result = method(df, *args, **kwargs)
+
+        return result
+
+    async def _execute_map(self, df: pd.DataFrame, op: PandasOperation) -> pd.DataFrame:
+        """Execute a map operation (column rename/mapping).
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            Current DataFrame
+        op : PandasOperation
+            Operation configuration with columns mapping
+
+        Returns
+        -------
+        pd.DataFrame
+            DataFrame with renamed columns
+        """
+        if not op.columns:
+            return df
+
+        return df.rename(columns=op.columns)
+
+    async def _execute_filter(self, df: pd.DataFrame, op: PandasOperation) -> pd.DataFrame:
+        """Execute a filter operation.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            Current DataFrame
+        op : PandasOperation
+            Operation configuration with filter condition
+
+        Returns
+        -------
+        pd.DataFrame
+            Filtered DataFrame
+        """
+        if not op.condition:
+            return df
+
+        # Evaluate condition
+        # Note: This is a simplified implementation - production code should validate
+        # the condition for security
+        condition_result = self._resolve_arg(op.condition, df, [df])
+
+        return df[condition_result]
+
+    async def _execute_assign(self, df: pd.DataFrame, op: PandasOperation) -> pd.DataFrame:
+        """Execute an assign operation (add new columns).
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            Current DataFrame
+        op : PandasOperation
+            Operation configuration
+
+        Returns
+        -------
+        pd.DataFrame
+            DataFrame with new columns
+        """
+        if not op.kwargs:
+            return df
+
+        # Prepare new column assignments
+        new_cols = {}
+        for col_name, col_expr in op.kwargs.items():
+            new_cols[col_name] = self._resolve_arg(col_expr, df, [df])
+
+        return df.assign(**new_cols)
+
+    def _resolve_method(self, method_path: str) -> Callable:
+        """Resolve a method from a path string or return callable directly.
+
+        Parameters
+        ----------
+        method_path : str
+            Path like "pandas.DataFrame.groupby" or "pandas.merge"
+
+        Returns
+        -------
+        Callable
+            The resolved method
+        """
+        # Already callable
+        if callable(method_path):
+            return method_path
+
+        # Parse module path
+        if "." not in method_path:
+            raise ValueError(f"Method path must contain '.', got: {method_path}")
+
+        try:
+            # Handle pandas class paths like pandas.DataFrame.sort_values
+            if method_path.startswith("pandas."):
+                parts = method_path.split(".")
+                if len(parts) >= 3 and parts[1] == "DataFrame":
+                    # It's a DataFrame method: pandas.DataFrame.method_name
+                    method_name = parts[2]
+                    return getattr(pd.DataFrame, method_name)
+                else:
+                    # It's a module-level function like pandas.merge
+                    module_path = ".".join(parts[:-1])
+                    attr_path = parts[-1]
+                    module = importlib.import_module(module_path)
+                    method = getattr(module, attr_path)
+
+                    if not callable(method):
+                        raise ValueError(f"'{method_path}' is not callable")
+
+                    return method
+            else:
+                # Standard module attribute resolution
+                module_path, attr_path = method_path.rsplit(".", 1)
+                module = importlib.import_module(module_path)
+                method = getattr(module, attr_path)
+
+                if not callable(method):
+                    raise ValueError(f"'{method_path}' is not callable")
+
+                return method
+        except Exception as e:
+            raise ValueError(f"Could not resolve method '{method_path}': {e}") from e
+
+    def _resolve_arg(self, arg: Any, df: pd.DataFrame, input_dfs: list[pd.DataFrame]) -> Any:
+        """Resolve an argument value (handles templates and expressions).
+
+        Parameters
+        ----------
+        arg : Any
+            Argument value or template expression
+        df : pd.DataFrame
+            Current DataFrame for context
+        input_dfs : list[pd.DataFrame]
+            All input DataFrames
+
+        Returns
+        -------
+        Any
+            Resolved argument value
+        """
+        # If it's a string template expression
+        if isinstance(arg, str) and "{{" in arg and "}}" in arg:
+            # Parse template expression
+            import re
+
+            pattern = r"\{\{\s*(.+?)\s*\}\}"
+            match = re.search(pattern, arg)
+
+            if match:
+                expr = match.group(1)
+
+                # Handle special variables
+                if expr == "df":
+                    return df
+                elif expr.startswith("input_artifacts["):
+                    # Extract index
+                    idx_match = re.search(r"input_artifacts\[(\d+)\]", expr)
+                    if idx_match:
+                        idx = int(idx_match.group(1))
+                        if 0 <= idx < len(input_dfs):
+                            return input_dfs[idx]
+                        else:
+                            raise IndexError(f"input_artifacts[{idx}] out of range")
+
+                # Try to evaluate as Python expression
+                try:
+                    # Safe evaluation - limited scope
+                    scope = {"df": df, "input_artifacts": input_dfs, "pd": pd}
+                    return eval(expr, {"__builtins__": {}}, scope)
+                except Exception:
+                    # Return as-is if evaluation fails
+                    return arg
+
+        # If it's a dict with lambda expression
+        if isinstance(arg, dict):
+            resolved = {}
+            for k, v in arg.items():
+                resolved[k] = self._resolve_arg(v, df, input_dfs)
+            return resolved
+
+        # Return as-is
+        return arg

hexdag_plugins/hexdag_etl/hexdag_etl/nodes/sql_extract_load.py

@@ -0,0 +1,112 @@
+"""SQL extraction and loading nodes for database operations."""
+
+from typing import Any
+
+from hexdag.core.domain.dag import NodeSpec
+from hexdag.core.registry import node
+from hexdag.core.registry.models import NodeSubtype
+
+from .base_node_factory import BaseNodeFactory
+
+
+@node(name="sql_extract", subtype=NodeSubtype.TOOL, namespace="etl")
+class SQLExtractNode(BaseNodeFactory):
+    """Extract data from SQL databases.
+
+    Placeholder implementation - to be completed with full SQLAlchemy integration.
+    """
+
+    def __call__(
+        self, name: str, query: str, database: str | None = None, deps: list[str] | None = None, **kwargs: Any
+    ) -> NodeSpec:
+        """Create SQL extract node.
+
+        Parameters
+        ----------
+        name : str
+            Node name
+        query : str
+            SQL query to execute
+        database : str, optional
+            Database connection reference
+        deps : list, optional
+            Dependencies
+        **kwargs : Any
+            Additional parameters
+
+        Returns
+        -------
+        NodeSpec
+            Node specification
+        """
+
+        async def wrapped_fn(input_data: dict, **ports: dict) -> dict:
+            """Placeholder implementation."""
+            return {"output": [], "metadata": {"query": query, "database": database, "status": "placeholder"}}
+
+        wrapped_fn.__name__ = f"sql_extract_{name}"
+
+        return self.create_node_with_mapping(
+            name=name,
+            wrapped_fn=wrapped_fn,
+            input_schema={"input_data": dict, "**ports": dict},
+            output_schema={"output": dict, "metadata": dict},
+            deps=deps or [],
+            **kwargs,
+        )
+
+
+@node(name="sql_load", subtype=NodeSubtype.TOOL, namespace="etl")
+class SQLLoadNode(BaseNodeFactory):
+    """Load data into SQL databases.
+
+    Placeholder implementation - to be completed with SQLAlchemy integration.
+    """
+
+    def __call__(
+        self,
+        name: str,
+        table: str,
+        mode: str = "append",
+        database: str | None = None,
+        deps: list[str] | None = None,
+        **kwargs: Any,
+    ) -> NodeSpec:
+        """Create SQL load node.
+
+        Parameters
+        ----------
+        name : str
+            Node name
+        table : str
+            Target table name
+        mode : str
+            Load mode: "append", "replace", "truncate_insert", "merge"
+        database : str, optional
+            Database connection reference
+        deps : list, optional
+            Dependencies
+        **kwargs : Any
+            Additional parameters
+
+        Returns
+        -------
+        NodeSpec
+            Node specification
+        """
+
+        async def wrapped_fn(input_data: dict, **ports: dict) -> dict:
+            """Placeholder implementation."""
+            row_count = len(input_data.get("output", [])) if isinstance(input_data, dict) else 0
+            return {"status": "loaded", "table": table, "rows": row_count}
+
+        wrapped_fn.__name__ = f"sql_load_{name}"
+
+        return self.create_node_with_mapping(
+            name=name,
+            wrapped_fn=wrapped_fn,
+            input_schema={"input_data": dict, "**ports": dict},
+            output_schema={"status": dict, "table": dict, "rows": dict},
+            deps=deps or [],
+            **kwargs,
+        )