hexdag 0.5.0.dev1__py3-none-any.whl

hexdag_plugins/hexdag_etl/hexdag_etl/nodes/file_io.py

@@ -0,0 +1,415 @@
+"""File I/O nodes for reading and writing data files.
+
+These nodes provide file-based input and output for ETL pipelines,
+supporting CSV, Parquet, JSON, and Excel formats.
+"""
+
+from pathlib import Path
+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 FileReaderOutput(BaseModel):
+    """Output model for FileReaderNode."""
+
+    data: Any  # DataFrame as dict for serialization
+    rows: int
+    columns: list[str]
+    file_path: str
+
+
+class FileWriterOutput(BaseModel):
+    """Output model for FileWriterNode."""
+
+    file_path: str
+    rows: int
+    format: str
+    success: bool
+
+
+@node(name="file_reader_node", subtype=NodeSubtype.FUNCTION, namespace="etl")
+class FileReaderNode(BaseNodeFactory):
+    """Node for reading data files into DataFrames.
+
+    Supports multiple file formats:
+    - CSV (.csv)
+    - Parquet (.parquet)
+    - JSON (.json, .jsonl)
+    - Excel (.xlsx, .xls)
+
+    Examples
+    --------
+    YAML pipeline::
+
+        - kind: etl:file_reader_node
+          metadata:
+            name: load_customers
+          spec:
+            file_path: data/customers.csv
+            format: csv
+            options:
+              sep: ","
+              encoding: utf-8
+          dependencies: []
+
+        - kind: etl:file_reader_node
+          metadata:
+            name: load_transactions
+          spec:
+            file_path: data/transactions.parquet
+            format: parquet
+          dependencies: []
+
+        - kind: etl:file_reader_node
+          metadata:
+            name: load_products
+          spec:
+            file_path: data/products.json
+            format: json
+            options:
+              orient: records
+          dependencies: []
+    """
+
+    def __call__(
+        self,
+        name: str,
+        file_path: str,
+        format: str | None = None,
+        options: dict[str, Any] | None = None,
+        deps: list[str] | None = None,
+        **kwargs: Any,
+    ) -> NodeSpec:
+        """Create a file reader node specification.
+
+        Parameters
+        ----------
+        name : str
+            Node name
+        file_path : str
+            Path to the input file (relative to workspace or absolute)
+        format : str, optional
+            File format: 'csv', 'parquet', 'json', 'excel'
+            Auto-detected from extension if not specified
+        options : dict, optional
+            Additional options passed to pandas read function
+        deps : list[str], optional
+            Dependency node names
+        **kwargs : Any
+            Additional node parameters
+
+        Returns
+        -------
+        NodeSpec
+            Node specification ready for execution
+        """
+        # Auto-detect format from file extension if not specified
+        if format is None:
+            format = self._detect_format(file_path)
+
+        # Create wrapped function
+        wrapped_fn = self._create_reader_function(name, file_path, format, options or {})
+
+        # Define schemas
+        input_schema = {"input_data": dict | None}
+        output_model = FileReaderOutput
+
+        input_model = self.create_pydantic_model(f"{name}Input", input_schema)
+
+        # Store parameters
+        node_params = {
+            "file_path": file_path,
+            "format": format,
+            "options": options,
+            **kwargs,
+        }
+
+        return NodeSpec(
+            name=name,
+            fn=wrapped_fn,
+            in_model=input_model,
+            out_model=output_model,
+            deps=frozenset(deps or []),
+            params=node_params,
+        )
+
+    def _detect_format(self, file_path: str) -> str:
+        """Detect file format from extension."""
+        path = Path(file_path)
+        ext = path.suffix.lower()
+
+        format_map = {
+            ".csv": "csv",
+            ".parquet": "parquet",
+            ".pq": "parquet",
+            ".json": "json",
+            ".jsonl": "jsonl",
+            ".xlsx": "excel",
+            ".xls": "excel",
+            ".feather": "feather",
+            ".pickle": "pickle",
+            ".pkl": "pickle",
+        }
+
+        if ext not in format_map:
+            raise ValueError(f"Unknown file format for extension '{ext}'. Supported: {list(format_map.keys())}")
+
+        return format_map[ext]
+
+    def _create_reader_function(
+        self,
+        name: str,
+        file_path: str,
+        format: str,
+        options: dict[str, Any],
+    ) -> Any:
+        """Create the file reading function."""
+
+        async def read_file(input_data: Any = None) -> dict[str, Any]:
+            """Read data file into DataFrame."""
+            # Resolve file path
+            path = Path(file_path)
+
+            if not path.exists():
+                raise FileNotFoundError(f"File not found: {path}")
+
+            # Read based on format
+            if format == "csv":
+                df = pd.read_csv(path, **options)
+            elif format == "parquet":
+                df = pd.read_parquet(path, **options)
+            elif format == "json":
+                df = pd.read_json(path, **options)
+            elif format == "jsonl":
+                df = pd.read_json(path, lines=True, **options)
+            elif format == "excel":
+                df = pd.read_excel(path, **options)
+            elif format == "feather":
+                df = pd.read_feather(path, **options)
+            elif format == "pickle":
+                df = pd.read_pickle(path, **options)
+            else:
+                raise ValueError(f"Unsupported format: {format}")
+
+            return {
+                "data": df,  # Keep as DataFrame for downstream nodes
+                "rows": len(df),
+                "columns": df.columns.tolist(),
+                "file_path": str(path.absolute()),
+            }
+
+        read_file.__name__ = f"file_reader_{name}"
+        read_file.__doc__ = f"Read file: {file_path}"
+
+        return read_file
+
+
+@node(name="file_writer_node", subtype=NodeSubtype.FUNCTION, namespace="etl")
+class FileWriterNode(BaseNodeFactory):
+    """Node for writing DataFrames to files.
+
+    Supports multiple file formats:
+    - CSV (.csv)
+    - Parquet (.parquet)
+    - JSON (.json, .jsonl)
+    - Excel (.xlsx)
+
+    Examples
+    --------
+    YAML pipeline::
+
+        - kind: etl:file_writer_node
+          metadata:
+            name: save_results
+          spec:
+            file_path: output/results.parquet
+            format: parquet
+            options:
+              compression: snappy
+          dependencies:
+            - transform_data
+
+        - kind: etl:file_writer_node
+          metadata:
+            name: export_csv
+          spec:
+            file_path: output/report.csv
+            format: csv
+            options:
+              index: false
+          dependencies:
+            - transform_data
+    """
+
+    def __call__(
+        self,
+        name: str,
+        file_path: str,
+        format: str | None = None,
+        input_key: str = "data",
+        options: dict[str, Any] | None = None,
+        create_dirs: bool = True,
+        deps: list[str] | None = None,
+        **kwargs: Any,
+    ) -> NodeSpec:
+        """Create a file writer node specification.
+
+        Parameters
+        ----------
+        name : str
+            Node name
+        file_path : str
+            Path for the output file
+        format : str, optional
+            File format: 'csv', 'parquet', 'json', 'excel'
+            Auto-detected from extension if not specified
+        input_key : str
+            Key in input data containing the DataFrame (default: 'data')
+        options : dict, optional
+            Additional options passed to pandas write function
+        create_dirs : bool
+            Create parent directories if they don't exist (default: True)
+        deps : list[str], optional
+            Dependency node names
+        **kwargs : Any
+            Additional node parameters
+
+        Returns
+        -------
+        NodeSpec
+            Node specification ready for execution
+        """
+        # Auto-detect format from file extension if not specified
+        if format is None:
+            format = self._detect_format(file_path)
+
+        # Create wrapped function
+        wrapped_fn = self._create_writer_function(name, file_path, format, input_key, options or {}, create_dirs)
+
+        # Define schemas
+        input_schema = {"input_data": dict}
+        output_model = FileWriterOutput
+
+        input_model = self.create_pydantic_model(f"{name}Input", input_schema)
+
+        # Store parameters
+        node_params = {
+            "file_path": file_path,
+            "format": format,
+            "input_key": input_key,
+            "options": options,
+            "create_dirs": create_dirs,
+            **kwargs,
+        }
+
+        return NodeSpec(
+            name=name,
+            fn=wrapped_fn,
+            in_model=input_model,
+            out_model=output_model,
+            deps=frozenset(deps or []),
+            params=node_params,
+        )
+
+    def _detect_format(self, file_path: str) -> str:
+        """Detect file format from extension."""
+        path = Path(file_path)
+        ext = path.suffix.lower()
+
+        format_map = {
+            ".csv": "csv",
+            ".parquet": "parquet",
+            ".pq": "parquet",
+            ".json": "json",
+            ".jsonl": "jsonl",
+            ".xlsx": "excel",
+            ".feather": "feather",
+            ".pickle": "pickle",
+            ".pkl": "pickle",
+        }
+
+        if ext not in format_map:
+            raise ValueError(f"Unknown file format for extension '{ext}'. Supported: {list(format_map.keys())}")
+
+        return format_map[ext]
+
+    def _create_writer_function(
+        self,
+        name: str,
+        file_path: str,
+        format: str,
+        input_key: str,
+        options: dict[str, Any],
+        create_dirs: bool,
+    ) -> Any:
+        """Create the file writing function."""
+
+        async def write_file(input_data: Any) -> dict[str, Any]:
+            """Write DataFrame to file."""
+            # Extract DataFrame from input
+            if isinstance(input_data, dict):
+                df = input_data.get(input_key)
+                if df is None:
+                    # Try to find a DataFrame in the input
+                    for key, value in input_data.items():
+                        if isinstance(value, pd.DataFrame):
+                            df = value
+                            break
+                        elif isinstance(value, dict) and "data" in value:
+                            df = value["data"]
+                            break
+            elif isinstance(input_data, pd.DataFrame):
+                df = input_data
+            else:
+                df = input_data
+
+            if df is None:
+                raise ValueError(f"No DataFrame found in input. Expected key: '{input_key}'")
+
+            if not isinstance(df, pd.DataFrame):
+                try:
+                    df = pd.DataFrame(df)
+                except Exception as e:
+                    raise ValueError(f"Could not convert input to DataFrame: {e}")
+
+            # Resolve file path and create directories
+            path = Path(file_path)
+
+            if create_dirs:
+                path.parent.mkdir(parents=True, exist_ok=True)
+
+            # Write based on format
+            if format == "csv":
+                df.to_csv(path, **options)
+            elif format == "parquet":
+                df.to_parquet(path, **options)
+            elif format == "json":
+                df.to_json(path, **options)
+            elif format == "jsonl":
+                df.to_json(path, orient="records", lines=True, **options)
+            elif format == "excel":
+                df.to_excel(path, **options)
+            elif format == "feather":
+                df.to_feather(path, **options)
+            elif format == "pickle":
+                df.to_pickle(path, **options)
+            else:
+                raise ValueError(f"Unsupported format: {format}")
+
+            return {
+                "file_path": str(path.absolute()),
+                "rows": len(df),
+                "format": format,
+                "success": True,
+            }
+
+        write_file.__name__ = f"file_writer_{name}"
+        write_file.__doc__ = f"Write file: {file_path}"
+
+        return write_file