|
| 1 | +import warnings |
1 | 2 | from typing import Any, TypeAlias, cast, overload |
2 | 3 |
|
3 | 4 | import narwhals as nw |
|
7 | 8 | PolarsCompatibleFrame: TypeAlias = Frame | pl.DataFrame | pl.LazyFrame |
8 | 9 |
|
9 | 10 |
|
| 11 | +def _is_polars_map_dtype(dtype: pl.DataType) -> bool: |
| 12 | + """Check if a Polars dtype is a polars_map.Map extension type without importing polars_map.""" |
| 13 | + return isinstance(dtype, pl.datatypes.classes.BaseExtension) and getattr(dtype, "_name", None) == "polars_map.map" |
| 14 | + |
| 15 | + |
| 16 | +def find_map_columns(df: Frame) -> list[str]: |
| 17 | + """Return column names that have a Map type in the underlying frame. |
| 18 | +
|
| 19 | + Inspects the native backend (Polars, Arrow, Ibis) for Map-typed columns. |
| 20 | + """ |
| 21 | + if df.implementation == nw.Implementation.POLARS: |
| 22 | + from metaxy.config import MetaxyConfig |
| 23 | + |
| 24 | + if not MetaxyConfig.get().enable_polars_map: |
| 25 | + return [] |
| 26 | + |
| 27 | + native = df.to_native() |
| 28 | + schema = native.collect_schema() if isinstance(native, pl.LazyFrame) else native.schema # ty: ignore[union-attribute] |
| 29 | + return [name for name, dtype in schema.items() if _is_polars_map_dtype(dtype)] |
| 30 | + |
| 31 | + if df.implementation == nw.Implementation.PYARROW: |
| 32 | + import pyarrow as pa |
| 33 | + |
| 34 | + table = cast(pa.Table, df.to_native()) |
| 35 | + return [field.name for field in table.schema if pa.types.is_map(field.type)] |
| 36 | + |
| 37 | + if df.implementation == nw.Implementation.IBIS: |
| 38 | + import ibis |
| 39 | + import ibis.expr.datatypes as dt |
| 40 | + |
| 41 | + table = cast(ibis.Table, df.to_native()) |
| 42 | + return [name for name, dtype in table.schema().items() if isinstance(dtype, dt.Map)] |
| 43 | + |
| 44 | + warnings.warn( |
| 45 | + f"cannot identify Map columns for unsupported narwhals implementation {df.implementation}", |
| 46 | + stacklevel=2, |
| 47 | + ) |
| 48 | + return [] |
| 49 | + |
| 50 | + |
| 51 | +def has_polars_map_columns(df: pl.DataFrame | pl.LazyFrame) -> bool: |
| 52 | + """Check if a Polars DataFrame/LazyFrame has any polars_map.Map columns.""" |
| 53 | + return bool(find_map_columns(cast(Frame, nw.from_native(df)))) |
| 54 | + |
| 55 | + |
10 | 56 | def collect_to_polars(frame: PolarsCompatibleFrame) -> pl.DataFrame: |
11 | 57 | """Helper to convert a Narwhals frame into an eager Polars DataFrame. |
12 | 58 |
|
@@ -53,14 +99,30 @@ def switch_implementation_to_polars(frame: LazyFrameT) -> LazyFrameT: ... |
53 | 99 | def switch_implementation_to_polars(frame: FrameT) -> FrameT: |
54 | 100 | if frame.implementation == nw.Implementation.POLARS: |
55 | 101 | return frame |
56 | | - elif isinstance(frame, nw.DataFrame): |
57 | | - return nw.from_native(frame.to_polars()) |
| 102 | + |
| 103 | + # Detect map columns before conversion (they lose their type in Polars) |
| 104 | + from metaxy.config import MetaxyConfig |
| 105 | + |
| 106 | + map_columns: list[str] = [] |
| 107 | + if MetaxyConfig.get().enable_polars_map: |
| 108 | + map_columns = find_map_columns(frame) |
| 109 | + |
| 110 | + if isinstance(frame, nw.DataFrame): |
| 111 | + result = nw.from_native(frame.to_polars()) |
58 | 112 | elif isinstance(frame, nw.LazyFrame): |
59 | | - return nw.from_native( |
60 | | - frame.collect().to_polars(), |
61 | | - ).lazy() |
| 113 | + result = nw.from_native(frame.collect().to_polars()).lazy() |
62 | 114 | else: |
63 | 115 | raise ValueError(f"Unsupported frame type: {type(frame)}") |
64 | 116 |
|
| 117 | + # Reconstruct polars_map.Map columns that were lost during conversion |
| 118 | + if map_columns: |
| 119 | + from metaxy.versioning._arrow_map import convert_maps_to_polars_map |
| 120 | + |
| 121 | + native = result.to_native() |
| 122 | + native = convert_maps_to_polars_map(native, columns=map_columns) |
| 123 | + result = nw.from_native(native) |
| 124 | + |
| 125 | + return result |
| 126 | + |
65 | 127 |
|
66 | 128 | __all__ = ["collect_to_polars"] |
0 commit comments