33
44"""Utility functions for lance-ray."""
55
6- from typing import TYPE_CHECKING , Optional , Union
6+ from collections .abc import Mapping
7+ from typing import TYPE_CHECKING , Any , Optional , Union , cast
78
89import pyarrow as pa
910
1213
1314
1415def pd_to_arrow (
15- df : Union [pa .Table , "pd.DataFrame" , dict ], schema : Optional [pa .Schema ]
16+ df : Union [pa .Table , "pd.DataFrame" , Mapping [str , Any ]],
17+ schema : Optional [pa .Schema ],
1618) -> pa .Table :
1719 """Convert a pandas DataFrame to pyarrow Table."""
1820 from lance .dependencies import _PANDAS_AVAILABLE
1921 from lance .dependencies import pandas as pd
2022
21- if isinstance (df , dict ):
23+ if isinstance (df , Mapping ):
2224 return pa .Table .from_pydict (df , schema = schema )
2325 elif _PANDAS_AVAILABLE and isinstance (df , pd .DataFrame ):
2426 tbl = pa .Table .from_pandas (df , schema = schema )
@@ -31,4 +33,4 @@ def pd_to_arrow(
3133 return new_table
3234 elif isinstance (df , pa .Table ) and df .num_rows > 0 and schema is not None :
3335 return df .cast (schema )
34- return df
36+ return cast ( pa . Table , df )
0 commit comments