"""
Helper functions for input handlers.
This module provides utility functions used by various input handlers
(CSV, Parquet, Excel) for common data processing operations including:
- Column name mapping and transformation
- Data quality validation
- Date range filtering
- Decimal precision inference
- Date parsing with ISO format support
All functions are designed to work with PyArrow tables and arrays for
efficient columnar data processing.
"""
import datetime
import logging
from typing import Literal
import pyarrow as pa
import pyarrow.compute as pc
from kaxanuk.attribution_analysis.entities.configuration import Configuration
from kaxanuk.attribution_analysis.entities.backtest_results import DecimalPrecisionScale
from kaxanuk.attribution_analysis.exceptions import DataIntegrityError, DateParsingError
logger = logging.getLogger(__name__)
# Maximum precision for decimal types (PyArrow decimal256 limit)
MAX_PRECISION = 76
[docs]
def apply_column_mapping(
table: pa.Table,
equivalence_dict: dict[str, str]
) -> pa.Table:
"""
Apply column name mapping to a PyArrow table using an equivalence dictionary.
This function transforms column names from source format to standardized format
by applying the reverse mapping of the provided equivalence dictionary. If no
equivalence dictionary is provided, the original table is returned unchanged.
Parameters
----------
table : pa.Table
The PyArrow table whose column names need to be mapped.
equivalence_dict : dict[str, str]
Dictionary mapping standard column names to source column names.
The function will reverse this mapping internally.
Returns
-------
pa.Table
A new PyArrow table with mapped column names. Columns not found
in the mapping retain their original names.
Examples
--------
>>> import pyarrow as pa
>>> table = pa.table({'old_name': [1, 2, 3], 'other_col': [4, 5, 6]})
>>> equiv_dict = {'new_name': 'old_name'}
>>> result = apply_column_mapping(table, equiv_dict)
>>> result.column_names
['new_name', 'other_col']
"""
if not equivalence_dict:
return table
column_mapping = get_reverse_equivalence_dict(
equivalence_dict=equivalence_dict)
new_column_names = [
column_mapping.get(col_name, col_name)
for col_name in table.column_names
]
return pa.table(table.columns, names=new_column_names)
[docs]
def check_for_empty_values(
col: pa.ChunkedArray,
column_name: str,
ticker: str
) -> None:
"""
Validate a PyArrow column for data quality issues including empty and null values.
This function performs comprehensive data quality checks on a column, identifying
empty strings, null values, and whitespace-only entries. If any issues are found,
a detailed ValueError is raised with specific row indices and error descriptions.
Parameters
----------
col : pa.ChunkedArray
The PyArrow chunked array column to validate.
column_name : str
The name of the column being validated, used in error messages.
ticker : str
The ticker symbol associated with the data, used in error messages
for better context and debugging.
Returns
-------
None
This function doesn't return a value but raises an exception if issues
are found.
Raises
------
ValueError
If any data quality issues are detected, including:
- Empty strings ('')
- Null/NaN values
- Whitespace-only strings (e.g., ' ')
The exception message includes the ticker, column name, and specific
row indices where issues were found.
Examples
--------
>>> import pyarrow as pa
>>> col = pa.chunked_array([['valid', '', None, ' ']])
>>> check_for_empty_values(col, 'price', 'AAPL')
ValueError: Data quality issues in ticker 'AAPL', column 'price':
Empty strings at rows: [1]
Null values at rows: [2]
Whitespace-only at rows: [3]
"""
arr = col.combine_chunks()
null_mask = pc.is_null(arr)
null_indices = pc.indices_nonzero(null_mask).to_pylist()
valid_arr = pc.if_else(null_mask, pa.scalar("__NULL__"), arr)
empty_mask = pc.equal(valid_arr, "")
empty_indices = pc.indices_nonzero(empty_mask).to_pylist()
stripped = pc.utf8_trim_whitespace(valid_arr)
whitespace_mask = pc.and_(pc.equal(stripped, ""), pc.invert(empty_mask))
whitespace_mask = pc.and_(whitespace_mask, pc.invert(null_mask))
whitespace_indices = pc.indices_nonzero(whitespace_mask).to_pylist()
error_details = []
if empty_indices:
error_details.append(f"Empty strings at rows: {empty_indices}")
if null_indices:
error_details.append(f"Null values at rows: {null_indices}")
if whitespace_indices:
error_details.append(f"Whitespace-only at rows: {whitespace_indices}")
if error_details:
raise DataIntegrityError(
f"Data quality issues in ticker '{ticker}', column '{column_name}':\n"
+ "\n".join(error_details)
)
[docs]
def filter_by_date_range(
table: pa.Table,
start_date: datetime.date | Literal["auto"] | None,
end_date: datetime.date | Literal["auto"] | None
) -> pa.Table:
"""
Filter a PyArrow table to include only rows within the specified date range.
This function filters the input table based on a date column named 'date',
keeping only rows where the date falls within the specified range (inclusive).
If both start_date and end_date are None, the original table is returned unchanged.
Parameters
----------
table : pa.Table
The PyArrow table to filter. Must contain a column named 'date'.
start_date : datetime.date or {"auto"} or None
The earliest date to include in the filtered result. Can be None
or "auto" to indicate no lower bound. Should be compatible with
pa.date32() type.
end_date : datetime.date or {"auto"} or None
The latest date to include in the filtered result. Can be None
or "auto" to indicate no upper bound. Should be compatible with
pa.date32() type.
Returns
-------
pa.Table
A new PyArrow table containing only rows where the date column
falls within the specified range (inclusive of both bounds).
Notes
-----
- The table must have a column named 'date'
- Date comparisons are inclusive (>= start_date and <= end_date)
- If only one bound is specified, filtering is applied only for that bound
- Dates are converted to pa.date32() format for comparison
Examples
--------
>>> import pyarrow as pa
>>> from datetime import date
>>> table = pa.table({
... 'date': [date(2023, 1, 1), date(2023, 6, 15), date(2023, 12, 31)],
... 'value': [100, 200, 300]
... })
>>> filtered = filter_by_date_range(table, date(2023, 6, 1), date(2023, 12, 1))
>>> len(filtered)
1 # Only the June 15 row remains
"""
if start_date == "auto":
start_date = None
if end_date == "auto":
end_date = None
if start_date is None and end_date is None:
return table
from kaxanuk.attribution_analysis.market_data_pipeline.field_definitions import StandardField # noqa: PLC0415
date_col = table[StandardField.DATE.value]
filters = []
if start_date is not None:
filters.append(pc.greater_equal(date_col, pa.scalar(start_date, type=pa.date32())))
if end_date is not None:
filters.append(pc.less_equal(date_col, pa.scalar(end_date, type=pa.date32())))
combined_filter = filters[0] if len(filters) == 1 else pc.and_(*filters)
return pc.filter(table, combined_filter)
[docs]
def get_reverse_equivalence_dict(
equivalence_dict: dict[str, str]
) -> dict[str, str]:
"""
Create a reverse mapping dictionary by swapping keys and values.
This function takes an equivalence dictionary that maps standard column names
to source column names and returns a new dictionary with the mapping reversed,
allowing lookup from source names to standard names.
Parameters
----------
equivalence_dict : dict[str, str]
A dictionary mapping standard column names (keys) to source column
names (values).
Returns
-------
dict[str, str]
A new dictionary with reversed key-value pairs, mapping source
column names (keys) to standard column names (values).
Examples
--------
>>> equiv_dict = {'standard_name': 'source_name', 'price': 'Price_USD'}
>>> reversed_dict = get_reverse_equivalence_dict(equiv_dict)
>>> reversed_dict
{'source_name': 'standard_name', 'Price_USD': 'price'}
Notes
-----
This function assumes that all values in the input dictionary are unique.
If duplicate values exist, only the last key-value pair will be preserved
in the reversed dictionary.
"""
return {csv_col: standard_col for standard_col, csv_col in equivalence_dict.items()}
[docs]
def get_decimal_type(
precision: int,
scale: int,
precision_limit: int
) -> pa.DataType:
"""
Return the appropriate PyArrow decimal type for the given precision.
Select ``decimal128`` when *precision* fits within *precision_limit*,
otherwise fall back to ``decimal256``.
Parameters
----------
precision : int
Total number of significant digits.
scale : int
Number of digits after the decimal point.
precision_limit : int
Upper bound for ``decimal128``; values above this threshold use
``decimal256``.
Returns
-------
pa.DataType
Either ``pa.decimal128(precision, scale)`` or
``pa.decimal256(precision, scale)``.
"""
if precision > precision_limit:
return pa.decimal256(precision, scale)
return pa.decimal128(precision, scale)
[docs]
def infer_decimal_precision_and_scale(
col: pa.ChunkedArray
) -> DecimalPrecisionScale:
"""
Infer the optimal precision and scale for decimal representation of numeric data.
This function analyzes a PyArrow column containing numeric data (as strings) to
determine the appropriate precision (total number of digits) and scale (number
of decimal places) for decimal type conversion. It handles edge cases such as
maximum precision limits and ensures valid decimal type constraints.
Parameters
----------
col : pa.ChunkedArray
A PyArrow chunked array containing numeric data as strings
(e.g., "123.45", "67.890"). The function expects string
representation of numbers with optional decimal points.
Returns
-------
DecimalPrecisionScale
A result object containing:
- precision: Total number of significant digits (1 to MAX_PRECISION)
- scale: Number of digits after the decimal point (0 to precision)
Notes
-----
- MAX_PRECISION constant defines the maximum allowed precision (76)
- If inferred precision exceeds MAX_PRECISION, scale is reduced accordingly
- Minimum precision is 1, and scale cannot exceed precision
- Numbers without decimal points are treated as having scale 0
- The function uses PyArrow compute functions for efficient processing
Examples
--------
>>> import pyarrow as pa
>>> col = pa.chunked_array([["123.45", "67.890", "1000"]])
>>> result = infer_decimal_precision_and_scale(col)
>>> result.precision, result.scale
(6, 3) # 6 total digits, 3 decimal places (from "67.890")
>>> # Using result to create decimal type
>>> decimal_type = pa.decimal128(result.precision, result.scale)
Raises
------
pyarrow.ArrowInvalid
If PyArrow compute operations fail on invalid input data.
"""
dot_index = pc.find_substring(col, pattern=".")
total_len = pc.binary_length(col)
zero_i32 = pa.scalar(0, type=pa.int32())
integer_digits_arr = pc.if_else(pc.equal(dot_index, -1), total_len, dot_index)
decimal_digits_arr = pc.if_else(
pc.equal(dot_index, -1),
zero_i32,
pc.subtract(total_len, pc.add(dot_index, 1))
)
max_integers = pc.max(integer_digits_arr).as_py() or 0
max_decimals = pc.max(decimal_digits_arr).as_py() or 0
precision = max_integers + max_decimals
scale = max_decimals
if precision > MAX_PRECISION:
scale = max(0, scale - (precision - MAX_PRECISION))
precision = MAX_PRECISION
precision = max(precision, 1, scale)
return DecimalPrecisionScale(precision=precision, scale=scale)
[docs]
def parse_date_column(
date_col: pa.ChunkedArray
) -> pa.ChunkedArray:
"""
Parse a date column with strict ISO format validation and comprehensive error handling.
This function attempts to parse date strings using only ISO-compliant formats:
YYYY-MM-DD and YYYY/MM/DD. It employs multiple parsing strategies and provides
detailed error messages with sample data when parsing fails, making it easier
to diagnose data format issues.
Parameters
----------
date_col : pa.ChunkedArray
A PyArrow chunked array containing date strings to be parsed.
Expected formats are ISO-compliant date strings.
Returns
-------
pa.ChunkedArray
A PyArrow chunked array with dates converted to pa.date32()
type, suitable for date-based operations and filtering.
Raises
------
DateParsingError
If the date column cannot be parsed using any of the supported ISO
formats. The exception includes:
- Data type information
- Array length
- Sample values from the column
- List of supported formats
- Specific parsing error details
Notes
-----
- Only ISO date formats are accepted for data consistency:
- YYYY-MM-DD (ISO format with hyphens, e.g., "2023-12-25")
- YYYY/MM/DD (ISO format with slashes, e.g., "2023/12/25")
- The function logs debug information during parsing attempts
- Error messages include sample data to help identify format issues
- Parsed dates are returned as pa.date32() type for memory efficiency
Examples
--------
>>> import pyarrow as pa
>>> date_col = pa.chunked_array([["2023-01-15", "2023/06/30", "2023-12-25"]])
>>> parsed_dates = parse_date_column(date_col)
>>> parsed_dates.type
date32[day]
"""
date_array = date_col.combine_chunks()
parsing_strategies = [
('%Y-%m-%d', 'YYYY-MM-DD (ISO format)'),
('%Y/%m/%d', 'YYYY/MM/DD (ISO format with slashes)')
]
for date_format, description in parsing_strategies:
try:
logger.debug("Attempting to parse dates using format: %s", description)
parsed_timestamps = pc.strptime(date_array, format=date_format, unit='s')
return pc.cast(parsed_timestamps, pa.date32())
except (pa.ArrowInvalid, pa.ArrowTypeError, pa.ArrowNotImplementedError) as parse_error:
logger.debug("Failed to parse with %s: %s", description, parse_error)
# If none of the ISO formats work, raise error with useful information
first_value = date_array[0].as_py() if len(date_array) > 0 else "No data available"
data_type = str(date_array.type)
array_length = len(date_array)
try:
unique_sample = pc.unique(date_array)
sample_values = [
str(unique_sample[i].as_py())
for i in range(min(5, len(unique_sample)))
if unique_sample[i].is_valid
]
except (pa.ArrowInvalid, pa.ArrowTypeError, pa.ArrowNotImplementedError):
sample_values = [str(first_value)]
supported_formats = [desc for _, desc in parsing_strategies]
error_message = (
f"Unable to parse date column. Only ISO formats are accepted. "
f"Data type: {data_type}, Array length: {array_length}, "
f"Sample values: {sample_values}, Supported formats: {supported_formats}"
)
logger.error(error_message)
raise DateParsingError(error_message)
[docs]
def process_table_columns(
*,
table: pa.Table,
ticker: str,
configuration: Configuration,
equivalence_dict: dict[str, str],
) -> pa.Table:
"""
Process raw table columns into standardized, typed, and date-filtered output.
Consolidate the common processing pipeline shared by CSV, Parquet, and other
input handlers: parse date columns, validate and cast numeric columns to
decimal types, rename to standard field names, and filter by date range.
Parameters
----------
table : pa.Table
Raw input table with columns as strings.
ticker : str
Ticker symbol, used in validation error messages.
configuration : Configuration
Configuration containing ``start_date`` and ``end_date`` for date
range filtering.
equivalence_dict : dict[str, str]
Mapping from standard field names to the actual column names
present in *table*.
Returns
-------
pa.Table
Processed table with standard column names, typed columns
(dates as ``date32``, numerics as ``decimal256``), and rows filtered
to the configured date range.
Raises
------
DateParsingError
If the date column cannot be parsed.
ValueError
If data quality checks fail on any column.
"""
from kaxanuk.attribution_analysis.market_data_pipeline.field_definitions import ( # noqa: PLC0415
StandardField,
STANDARD_FIELDS_ORDER,
)
# Process each column
processed_columns: dict[str, pa.Array] = {}
for std_name in STANDARD_FIELDS_ORDER:
csv_col = equivalence_dict[std_name]
col = table[csv_col]
if std_name == StandardField.DATE.value:
processed_columns[csv_col] = parse_date_column(date_col=col)
else:
check_for_empty_values(col=col, column_name=csv_col, ticker=ticker)
decimal_spec = infer_decimal_precision_and_scale(col=col)
processed_columns[csv_col] = pc.cast(
col,
pa.decimal256(decimal_spec.precision, decimal_spec.scale)
)
# Build final table with standard names
new_columns = []
final_column_names = []
for std_name in STANDARD_FIELDS_ORDER:
csv_col = equivalence_dict[std_name]
new_columns.append(processed_columns[csv_col])
final_column_names.append(std_name)
processed_table = pa.table(new_columns, names=final_column_names)
# Filter by date range
return filter_by_date_range(
table=processed_table,
start_date=configuration.start_date,
end_date=configuration.end_date,
)
