Source code for kaxanuk.attribution_analysis.input_handlers.parquet_input

"""
Parquet input handler for market data files.

Read per-ticker Parquet files with efficient columnar loading, then delegate
to the shared processing pipeline for type casting, date parsing, and
date-range filtering.
"""

import logging
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

from kaxanuk.attribution_analysis.entities.configuration import Configuration
from kaxanuk.attribution_analysis.exceptions import DataLoadingError
from kaxanuk.attribution_analysis.input_handlers.input_handler_interface import InputHandlerInterface
from kaxanuk.attribution_analysis.input_handlers import input_handlers_helpers
from kaxanuk.attribution_analysis.market_data_pipeline.field_definitions import STANDARD_FIELDS_ORDER

logger = logging.getLogger(__name__)



class ParquetInput(InputHandlerInterface):
    """
    A specialized input handler for reading and processing Parquet files containing market data.

    This class implements the InputHandlerInterface to provide functionality for reading
    market data from Parquet files. It handles file existence checks, data loading,
    column mapping, type conversion, and date filtering operations. The class is designed
    to work with financial market data stored in Parquet format, providing efficient
    columnar data processing using Apache Arrow.

    The ParquetInput handler supports:
    - Efficient columnar data reading from Parquet files
    - Automatic type inference and conversion to decimal types for financial data
    - Date column parsing and validation
    - Column name mapping between input files and standardized field names
    - Date range filtering based on configuration parameters
    - Comprehensive error handling for file operations

    Attributes
    ----------
    input_dir : str
        The directory path where input files are located.
    """

    def __init__(
            self,
            *,
            input_dir: str,
    ) -> None:
        """
        Initialize the ParquetInput handler with directory configuration.

        Parameters
        ----------
        input_dir : str
            The directory path where Parquet files are located.

        Examples
        --------
        >>> handler = ParquetInput(
        ...     input_dir="/data/market_data/daily_prices",
        ... )
        >>> # Files will be read from: /data/market_data/daily_prices/{ticker}.parquet
        """
        super().__init__()
        self.input_dir = input_dir



    def file_exists(
            self,
            *,
            ticker: str,
            configuration: Configuration
    ) -> bool:
        """
        Check if a Parquet file exists for the specified ticker without loading the data.

        Parameters
        ----------
        ticker : str
            The ticker symbol for which to check file existence.
        configuration : Configuration
            The configuration object containing settings.

        Returns
        -------
        bool
            True if the Parquet file exists for the given ticker, False otherwise.
        """
        input_file_path = Path(self.input_dir) / f"{ticker}.parquet"
        return input_file_path.exists()




    def input_data(
            self,
            *,
            ticker: str,
            configuration: Configuration,
            equivalence_dict: dict[str, str]
    ) -> pa.Table:
        """
        Read and process market data from a Parquet file for the specified ticker.

        This method reads market data from a Parquet file and delegates all processing
        operations (type conversion, date parsing, decimal inference, validation,
        and filtering) to a common helper function. The Parquet-specific aspect is
        the efficient columnar reading with selective column loading.

        Parameters
        ----------
        ticker : str
            The ticker symbol for which to load market data.
        configuration : Configuration
            Configuration object containing:

            - start_date: Start date for data filtering
            - end_date: End date for data filtering
        equivalence_dict : dict[str, str]
            Dictionary mapping STANDARD field names to actual Parquet column
            names used in the files.

        Returns
        -------
        pa.Table
            A PyArrow table containing the processed market data.

        Raises
        ------
        DataLoadingError
            Raised in the following scenarios:

            - FileNotFoundError: When the Parquet file for the ticker doesn't exist
            - pa.lib.ArrowInvalid: When the Parquet format is invalid or corrupted
            - KeyError: When column mapping fails due to missing keys
            - Exception: For any other unexpected errors during processing
        """
        input_file_path = f"{self.input_dir}/{ticker}.parquet"

        try:
            # Get columns to read in standard order
            csv_columns_in_order = [
                equivalence_dict[std_name] for std_name in STANDARD_FIELDS_ORDER
            ]
            unique_csv_columns = list(dict.fromkeys(csv_columns_in_order))

            # Read Parquet with selective column loading
            table = pq.read_table(
                input_file_path,
                columns=unique_csv_columns
            )

            # Delegate all processing to common helper
            return input_handlers_helpers.process_table_columns(
                table=table,
                ticker=ticker,
                configuration=configuration,
                equivalence_dict=equivalence_dict,
            )

        except FileNotFoundError as error:
            raise DataLoadingError(
                f"Parquet file not found for ticker '{ticker}': {input_file_path}"
            ) from error

        except pa.lib.ArrowInvalid as error:
            raise DataLoadingError(
                f"Invalid Parquet format for ticker '{ticker}': {error}"
            ) from error

        except pa.lib.ArrowTypeError as error:
            raise DataLoadingError(
                f"Type conversion error in Parquet for ticker '{ticker}': {error}"
            ) from error

        except KeyError as error:
            raise DataLoadingError(
                f"Column mapping error for ticker '{ticker}': {error}. "
                f"Check your equivalence_dict and STANDARD_FIELDS_ORDER."
            ) from error

        except ValueError as error:
            raise DataLoadingError(
                f"Value error reading Parquet for ticker '{ticker}': {error}"
            ) from error

        except OSError as error:
            raise DataLoadingError(
                f"I/O error reading Parquet for ticker '{ticker}': {error}"
            ) from error