Skip to content

Commit d765ce2

Browse files
committed
feat: experimental support for polars_map.Map datatype
1 parent 2b11dd1 commit d765ce2

16 files changed

Lines changed: 1023 additions & 73 deletions

File tree

‎pyproject.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,9 @@ dependencies = [
4747
]
4848

4949
[project.optional-dependencies]
50+
polars-map = [
51+
"polars-map>=0.1.1",
52+
]
5053
ibis = [
5154
"pyarrow>=18.0.0",
5255
"ibis-framework>=11.0.0",
@@ -151,6 +154,7 @@ dev = [
151154
"mermaid-py>=0.8.0",
152155
"mkdocs-metaxy",
153156
"pytest-postgresql>=8.0.0",
157+
"polars-map>=0.1.1",
154158
]
155159

156160
examples = [

‎src/metaxy/_utils.py‎

Lines changed: 67 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import warnings
12
from typing import Any, TypeAlias, cast, overload
23

34
import narwhals as nw
@@ -7,6 +8,51 @@
78
PolarsCompatibleFrame: TypeAlias = Frame | pl.DataFrame | pl.LazyFrame
89

910

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+
1056
def collect_to_polars(frame: PolarsCompatibleFrame) -> pl.DataFrame:
1157
"""Helper to convert a Narwhals frame into an eager Polars DataFrame.
1258
@@ -53,14 +99,30 @@ def switch_implementation_to_polars(frame: LazyFrameT) -> LazyFrameT: ...
5399
def switch_implementation_to_polars(frame: FrameT) -> FrameT:
54100
if frame.implementation == nw.Implementation.POLARS:
55101
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())
58112
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()
62114
else:
63115
raise ValueError(f"Unsupported frame type: {type(frame)}")
64116

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+
65127

66128
__all__ = ["collect_to_polars"]

‎src/metaxy/config/models/metaxy_config.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,11 @@ def _filter_incomplete_stores(cls, data: Any) -> Any:
189189

190190
hash_truncation_length: int = PydanticField(default=8, description="Truncate hash values to this length.", ge=8)
191191

192+
enable_polars_map: bool = PydanticField(
193+
default=False,
194+
description="Use `polars-map.Map` datatype for versioning columns instead of `pl.Struct`. Experimental.",
195+
)
196+
192197
auto_create_tables: bool = PydanticField(
193198
default=False,
194199
description="Auto-create tables when opening stores. It is not advised to enable this setting in production.",

‎src/metaxy/ext/metadata_stores/delta.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
from metaxy._decorators import public
1919
from metaxy._utils import collect_to_polars
20+
from metaxy.config import MetaxyConfig
2021
from metaxy.metadata_store.base import MetadataStore, MetadataStoreConfig
2122
from metaxy.metadata_store.types import AccessMode
2223
from metaxy.metadata_store.utils import is_local_path
@@ -26,6 +27,13 @@
2627
from metaxy.versioning.types import HashAlgorithm
2728

2829

30+
def _map_columns_from_delta_schema(schema: deltalake.Schema) -> list[str]:
31+
"""Return column names that have a native Arrow Map type in the Delta schema."""
32+
from deltalake.schema import MapType
33+
34+
return [field.name for field in schema.fields if isinstance(field.type, MapType)]
35+
36+
2937
@public
3038
class DeltaMetadataStoreConfig(MetadataStoreConfig):
3139
"""Configuration for DeltaMetadataStore.
@@ -268,6 +276,26 @@ def _write_feature(
268276
mode = write_opts.pop("mode", "append")
269277
storage_options = write_opts.pop("storage_options", None)
270278

279+
if MetaxyConfig.get().enable_polars_map:
280+
# Collect to Arrow, convert Struct *_by_field → native MapArray
281+
from metaxy.models.constants import METAXY_DATA_VERSION_BY_FIELD, METAXY_PROVENANCE_BY_FIELD
282+
from metaxy.versioning._arrow_map import convert_extension_maps_to_native, convert_structs_to_maps
283+
284+
df_native = self._cast_enum_to_string(collect_to_polars(df))
285+
df_native = convert_structs_to_maps(
286+
df_native, columns=[METAXY_PROVENANCE_BY_FIELD, METAXY_DATA_VERSION_BY_FIELD]
287+
)
288+
arrow_table = df_native.to_arrow()
289+
arrow_table = convert_extension_maps_to_native(arrow_table)
290+
deltalake.write_deltalake(
291+
table_uri,
292+
arrow_table,
293+
mode=mode,
294+
storage_options=storage_options,
295+
**(write_opts or {}),
296+
)
297+
return
298+
271299
# Check if we can use sink_delta (Polars >= 1.37, native Polars LazyFrame)
272300
can_sink = (
273301
df.implementation == nw.Implementation.POLARS
@@ -378,6 +406,14 @@ def _read_feature(
378406
storage_options=self.storage_options or None,
379407
)
380408

409+
if MetaxyConfig.get().enable_polars_map:
410+
delta_table = self._open_delta_table(feature_key, without_files=True)
411+
map_columns = _map_columns_from_delta_schema(delta_table.schema())
412+
if map_columns:
413+
from metaxy.versioning._arrow_map import convert_maps_to_polars_map
414+
415+
lf = convert_maps_to_polars_map(lf, columns=map_columns)
416+
381417
# Convert to Narwhals
382418
nw_lazy = nw.from_native(lf)
383419

‎src/metaxy/ext/metadata_stores/iceberg.py‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
if TYPE_CHECKING:
1111
from pyiceberg.catalog import Catalog
12+
from pyiceberg.schema import Schema
1213
from pyiceberg.table import Table
1314

1415
import narwhals as nw
@@ -19,6 +20,7 @@
1920

2021
from metaxy._decorators import public
2122
from metaxy._utils import collect_to_polars
23+
from metaxy.config import MetaxyConfig
2224
from metaxy.metadata_store.base import MetadataStore, MetadataStoreConfig
2325
from metaxy.metadata_store.types import AccessMode
2426
from metaxy.metadata_store.utils import is_local_path
@@ -31,6 +33,13 @@
3133
"""A ``(namespace, table_name)`` pair used by PyIceberg to locate a table within a catalog."""
3234

3335

36+
def _map_columns_from_iceberg_schema(schema: Schema) -> list[str]:
37+
"""Return column names that have a native Map type in the Iceberg schema."""
38+
from pyiceberg.types import MapType
39+
40+
return [field.name for field in schema.fields if isinstance(field.field_type, MapType)]
41+
42+
3443
def _strip_casts() -> Callable[[Any], Any]:
3544
"""Unwrap ``CAST(x AS type)`` → ``x`` to match PyIceberg's
3645
`row filter syntax <https://py.iceberg.apache.org/row-filter-syntax/>`_
@@ -257,7 +266,19 @@ def _write_feature(
257266
and Version(pl.__version__) >= Version("1.39.0")
258267
)
259268

260-
if can_sink:
269+
if MetaxyConfig.get().enable_polars_map:
270+
from metaxy.models.constants import METAXY_DATA_VERSION_BY_FIELD, METAXY_PROVENANCE_BY_FIELD
271+
from metaxy.versioning._arrow_map import convert_extension_maps_to_native, convert_structs_to_maps
272+
273+
df_polars = self._cast_enum_to_string(collect_to_polars(df))
274+
df_polars = convert_structs_to_maps(
275+
df_polars, columns=[METAXY_PROVENANCE_BY_FIELD, METAXY_DATA_VERSION_BY_FIELD]
276+
)
277+
arrow_table = df_polars.to_arrow()
278+
arrow_table = convert_extension_maps_to_native(arrow_table)
279+
iceberg_table = self._ensure_table(identifier, arrow_table.schema)
280+
iceberg_table.append(arrow_table, **kwargs)
281+
elif can_sink:
261282
lf_native = df.to_native()
262283
assert isinstance(lf_native, pl.LazyFrame)
263284
arrow_schema = pl.DataFrame(schema=self._cast_enum_to_string(lf_native).collect_schema()).to_arrow().schema
@@ -296,7 +317,17 @@ def _read_feature(
296317
if not self.catalog.table_exists(identifier):
297318
return None
298319

299-
nw_lazy = nw.from_native(pl.scan_iceberg(self.catalog.load_table(identifier)))
320+
iceberg_table = self.catalog.load_table(identifier)
321+
lf = pl.scan_iceberg(iceberg_table)
322+
323+
if MetaxyConfig.get().enable_polars_map:
324+
map_columns = _map_columns_from_iceberg_schema(iceberg_table.schema())
325+
if map_columns:
326+
from metaxy.versioning._arrow_map import convert_maps_to_polars_map
327+
328+
lf = convert_maps_to_polars_map(lf, columns=map_columns)
329+
330+
nw_lazy = nw.from_native(lf)
300331

301332
if filters:
302333
nw_lazy = nw_lazy.filter(*filters)

‎src/metaxy/metadata_store/base.py‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,20 @@
6161
pass
6262

6363

64+
def _is_map_column(df: Frame, col_name: str) -> bool:
65+
"""Check if a column in a narwhals frame is a Map type.
66+
67+
Returns False immediately if enable_polars_map is not set.
68+
Delegates to find_map_columns for backend-specific detection.
69+
"""
70+
if not MetaxyConfig.get().enable_polars_map:
71+
return False
72+
73+
from metaxy._utils import find_map_columns
74+
75+
return col_name in find_map_columns(df)
76+
77+
6478
# TypeVar for config types - used for typing from_config method
6579
MetadataStoreConfigT = TypeVar("MetadataStoreConfigT", bound="MetadataStoreConfig")
6680

@@ -1229,7 +1243,7 @@ def hash_struct_version_column(
12291243
df.implementation,
12301244
message=f"`{hash_column}` will be calculated in Polars.",
12311245
)
1232-
df = nw.from_native(df.lazy().collect().to_polars())
1246+
df = switch_implementation_to_polars(nw.from_native(df.lazy().collect()))
12331247

12341248
return cast(
12351249
Frame,
@@ -1647,10 +1661,12 @@ def _validate_schema(self, df: Frame) -> None:
16471661
if METAXY_PROVENANCE_BY_FIELD not in schema.names():
16481662
raise MetadataSchemaError(f"DataFrame must have '{METAXY_PROVENANCE_BY_FIELD}' column")
16491663

1650-
# Check that metaxy_provenance_by_field is a struct
1664+
# Check that metaxy_provenance_by_field is a struct or map
16511665
provenance_dtype = schema[METAXY_PROVENANCE_BY_FIELD]
1652-
if not isinstance(provenance_dtype, nw.Struct):
1653-
raise MetadataSchemaError(f"'{METAXY_PROVENANCE_BY_FIELD}' column must be a Struct, got {provenance_dtype}")
1666+
if not isinstance(provenance_dtype, nw.Struct) and not _is_map_column(df, METAXY_PROVENANCE_BY_FIELD):
1667+
raise MetadataSchemaError(
1668+
f"'{METAXY_PROVENANCE_BY_FIELD}' column must be a Struct or Map, got {provenance_dtype}"
1669+
)
16541670

16551671
# Note: metaxy_provenance is auto-computed if missing, so we don't validate it here
16561672

0 commit comments

Comments
 (0)