| 1589 | |
| 1590 | @functools.wraps(fn) |
| 1591 | def arrow_wrapper(*args, **kwargs): |
| 1592 | import numpy as np |
| 1593 | import pandas as pd |
| 1594 | import pyarrow as pa |
| 1595 | |
| 1596 | def to_arrow(val): |
| 1597 | """Convert a value to PyArrow Array if needed.""" |
| 1598 | if isinstance(val, (pa.Array, pa.ChunkedArray)): |
| 1599 | return val, False |
| 1600 | elif isinstance(val, pd.Series): |
| 1601 | return pa.Array.from_pandas(val), True |
| 1602 | elif isinstance(val, np.ndarray): |
| 1603 | return pa.array(val), False |
| 1604 | else: |
| 1605 | return val, False |
| 1606 | |
| 1607 | # Convert inputs to PyArrow and track pandas flags |
| 1608 | args_results = [to_arrow(arg) for arg in args] |
| 1609 | kwargs_results = {k: to_arrow(v) for k, v in kwargs.items()} |
| 1610 | |
| 1611 | converted_args = [v[0] for v in args_results] |
| 1612 | converted_kwargs = {k: v[0] for k, v in kwargs_results.items()} |
| 1613 | input_was_pandas = any(v[1] for v in args_results) or any( |
| 1614 | v[1] for v in kwargs_results.values() |
| 1615 | ) |
| 1616 | |
| 1617 | # Call function with converted inputs |
| 1618 | result = fn(*converted_args, **converted_kwargs) |
| 1619 | |
| 1620 | # Convert result back to pandas if input was pandas |
| 1621 | if input_was_pandas and isinstance(result, (pa.Array, pa.ChunkedArray)): |
| 1622 | result = result.to_pandas() |
| 1623 | |
| 1624 | return result |
| 1625 | |
| 1626 | return arrow_wrapper |
| 1627 | |