Skip to content

Commit

Permalink
refactor(python): change Flavor to CompatLevel
Browse files Browse the repository at this point in the history
  • Loading branch information
ruihe774 committed Jul 6, 2024
1 parent 06b57e6 commit 0401852
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 63 deletions.
64 changes: 35 additions & 29 deletions py-polars/polars/dataframe/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@
TooManyRowsReturnedError,
)
from polars.functions import col, lit
from polars.interchange.protocol import Flavor
from polars.interchange.protocol import CompatLevel
from polars.schema import Schema
from polars.selectors import _expand_selector_dicts, _expand_selectors

Expand Down Expand Up @@ -1385,7 +1385,8 @@ def item(self, row: int | None = None, column: int | str | None = None) -> Any:
)
return s.get_index_signed(row)

def to_arrow(self, *, future: Flavor | None = None) -> pa.Table:
@deprecate_renamed_parameter("future", "compat_level", version="1.0.1")
def to_arrow(self, *, compat_level: CompatLevel | None = None) -> pa.Table:
"""
Collect the underlying arrow arrays in an Arrow Table.
Expand All @@ -1396,8 +1397,9 @@ def to_arrow(self, *, future: Flavor | None = None) -> pa.Table:
Parameters
----------
future
Use a specific version of Polars' internal data structures.
compat_level
Use a specific compatibility level
when exporting Polars' internal data structures.
Examples
--------
Expand All @@ -1415,12 +1417,12 @@ def to_arrow(self, *, future: Flavor | None = None) -> pa.Table:
if not self.width: # 0x0 dataframe, cannot infer schema from batches
return pa.table({})

if future is None:
future = False # type: ignore[assignment]
elif isinstance(future, Flavor):
future = future._version # type: ignore[attr-defined]
if compat_level is None:
compat_level = False # type: ignore[assignment]
elif isinstance(compat_level, CompatLevel):
compat_level = compat_level._version # type: ignore[attr-defined]

record_batches = self._df.to_arrow(future)
record_batches = self._df.to_arrow(compat_level)
return pa.Table.from_batches(record_batches)

@overload
Expand Down Expand Up @@ -3293,7 +3295,7 @@ def write_ipc(
file: None,
*,
compression: IpcCompression = "uncompressed",
future: Flavor | None = None,
compat_level: CompatLevel | None = None,
) -> BytesIO: ...

@overload
Expand All @@ -3302,15 +3304,16 @@ def write_ipc(
file: str | Path | IO[bytes],
*,
compression: IpcCompression = "uncompressed",
future: Flavor | None = None,
compat_level: CompatLevel | None = None,
) -> None: ...

@deprecate_renamed_parameter("future", "compat_level", version="1.0.1")
def write_ipc(
self,
file: str | Path | IO[bytes] | None,
*,
compression: IpcCompression = "uncompressed",
future: Flavor | None = None,
compat_level: CompatLevel | None = None,
) -> BytesIO | None:
"""
Write to Arrow IPC binary stream or Feather file.
Expand All @@ -3324,8 +3327,9 @@ def write_ipc(
written. If set to `None`, the output is returned as a BytesIO object.
compression : {'uncompressed', 'lz4', 'zstd'}
Compression method. Defaults to "uncompressed".
future
Use a specific version of Polars' internal data structures.
compat_level
Use a specific compatibility level
when exporting Polars' internal data structures.
Examples
--------
Expand All @@ -3347,15 +3351,15 @@ def write_ipc(
elif isinstance(file, (str, Path)):
file = normalize_filepath(file)

if future is None:
future = True # type: ignore[assignment]
elif isinstance(future, Flavor):
future = future._version # type: ignore[attr-defined]
if compat_level is None:
compat_level = True # type: ignore[assignment]
elif isinstance(compat_level, CompatLevel):
compat_level = compat_level._version # type: ignore[attr-defined]

if compression is None:
compression = "uncompressed"

self._df.write_ipc(file, compression, future)
self._df.write_ipc(file, compression, compat_level)
return file if return_bytes else None # type: ignore[return-value]

@overload
Expand All @@ -3364,7 +3368,7 @@ def write_ipc_stream(
file: None,
*,
compression: IpcCompression = "uncompressed",
future: Flavor | None = None,
compat_level: CompatLevel | None = None,
) -> BytesIO: ...

@overload
Expand All @@ -3373,15 +3377,16 @@ def write_ipc_stream(
file: str | Path | IO[bytes],
*,
compression: IpcCompression = "uncompressed",
future: Flavor | None = None,
compat_level: CompatLevel | None = None,
) -> None: ...

@deprecate_renamed_parameter("future", "compat_level", version="1.0.1")
def write_ipc_stream(
self,
file: str | Path | IO[bytes] | None,
*,
compression: IpcCompression = "uncompressed",
future: Flavor | None = None,
compat_level: CompatLevel | None = None,
) -> BytesIO | None:
"""
Write to Arrow IPC record batch stream.
Expand All @@ -3395,8 +3400,9 @@ def write_ipc_stream(
be written. If set to `None`, the output is returned as a BytesIO object.
compression : {'uncompressed', 'lz4', 'zstd'}
Compression method. Defaults to "uncompressed".
future
Use a specific version of Polars' internal data structures.
compat_level
Use a specific compatibility level
when exporting Polars' internal data structures.
Examples
--------
Expand All @@ -3418,15 +3424,15 @@ def write_ipc_stream(
elif isinstance(file, (str, Path)):
file = normalize_filepath(file)

if future is None:
future = True # type: ignore[assignment]
elif isinstance(future, Flavor):
future = future._version # type: ignore[attr-defined]
if compat_level is None:
compat_level = True # type: ignore[assignment]
elif isinstance(compat_level, CompatLevel):
compat_level = compat_level._version # type: ignore[attr-defined]

if compression is None:
compression = "uncompressed"

self._df.write_ipc_stream(file, compression, future)
self._df.write_ipc_stream(file, compression, compat_level)
return file if return_bytes else None # type: ignore[return-value]

def write_parquet(
Expand Down
40 changes: 21 additions & 19 deletions py-polars/polars/interchange/protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -259,43 +259,45 @@ class CopyNotAllowedError(RuntimeError):
"""Exception raised when a copy is required, but `allow_copy` is set to `False`."""


class Flavor:
"""Data structure flavor."""
class CompatLevel:
"""Data structure compatibility level."""

def __init__(self) -> None:
msg = "it is not allowed to create a Flavor object"
msg = "it is not allowed to create a CompatLevel object"
raise TypeError(msg)

@staticmethod
def _with_version(version: int) -> Flavor:
flavor = Flavor.__new__(Flavor)
flavor._version = version # type: ignore[attr-defined]
return flavor
def _with_version(version: int) -> CompatLevel:
compat_level = CompatLevel.__new__(CompatLevel)
compat_level._version = version # type: ignore[attr-defined]
return compat_level

@staticmethod
def _highest() -> Flavor:
return Flavor._future1 # type: ignore[attr-defined]
def _newest() -> CompatLevel:
return CompatLevel._future1 # type: ignore[attr-defined]

@staticmethod
def highest() -> Flavor:
def newest() -> CompatLevel:
"""
Get the highest supported flavor.
Get the highest supported compatibility level.
.. warning::
Highest flavor is considered **unstable**. It may be changed
Highest compatibility level is considered **unstable**. It may be changed
at any point without it being considered a breaking change.
"""
issue_unstable_warning("Using the highest flavor is considered unstable.")
return Flavor._highest()
issue_unstable_warning(
"Using the highest compatibility level is considered unstable."
)
return CompatLevel._newest()

@staticmethod
def compatible() -> Flavor:
"""Get the flavor that is compatible with older arrow implementation."""
return Flavor._compatible # type: ignore[attr-defined]
def oldest() -> CompatLevel:
"""Get the most compatible level."""
return CompatLevel._compatible # type: ignore[attr-defined]

def __repr__(self) -> str:
return f"<{self.__class__.__module__}.{self.__class__.__qualname__}: {self._version}>" # type: ignore[attr-defined]


Flavor._compatible = Flavor._with_version(0) # type: ignore[attr-defined]
Flavor._future1 = Flavor._with_version(1) # type: ignore[attr-defined]
CompatLevel._compatible = CompatLevel._with_version(0) # type: ignore[attr-defined]
CompatLevel._future1 = CompatLevel._with_version(1) # type: ignore[attr-defined]
24 changes: 13 additions & 11 deletions py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@
from polars.dependencies import pandas as pd
from polars.dependencies import pyarrow as pa
from polars.exceptions import ComputeError, ModuleUpgradeRequiredError, ShapeError
from polars.interchange.protocol import Flavor
from polars.interchange.protocol import CompatLevel
from polars.series.array import ArrayNameSpace
from polars.series.binary import BinaryNameSpace
from polars.series.categorical import CatNameSpace
Expand Down Expand Up @@ -503,14 +503,14 @@ def _from_buffers(
validity = validity._s
return cls._from_pyseries(PySeries._from_buffers(dtype, data, validity))

def _highest_flavor(self) -> int:
def _newest_compat_level(self) -> int:
"""
Get the highest supported flavor version.
This is only used by pyo3-polars,
and it is simpler not to make it a static method.
"""
return Flavor._highest()._version # type: ignore[attr-defined]
return CompatLevel._newest()._version # type: ignore[attr-defined]

@property
def dtype(self) -> DataType:
Expand Down Expand Up @@ -4352,16 +4352,18 @@ def to_torch(self) -> torch.Tensor:
# tensor.rename(self.name)
return tensor

def to_arrow(self, *, future: Flavor | None = None) -> pa.Array:
@deprecate_renamed_parameter("future", "compat_level", version="1.0.1")
def to_arrow(self, *, compat_level: CompatLevel | None = None) -> pa.Array:
"""
Return the underlying Arrow array.
If the Series contains only a single chunk this operation is zero copy.
Parameters
----------
future
Use a specific version of Polars' internal data structures.
compat_level
Use a specific compatibility level
when exporting Polars' internal data structures.
Examples
--------
Expand All @@ -4375,11 +4377,11 @@ def to_arrow(self, *, future: Flavor | None = None) -> pa.Array:
3
]
"""
if future is None:
future = False # type: ignore[assignment]
elif isinstance(future, Flavor):
future = future._version # type: ignore[attr-defined]
return self._s.to_arrow(future)
if compat_level is None:
compat_level = False # type: ignore[assignment]
elif isinstance(compat_level, CompatLevel):
compat_level = compat_level._version # type: ignore[attr-defined]
return self._s.to_arrow(compat_level)

def to_pandas(
self, *, use_pyarrow_extension_array: bool = False, **kwargs: Any
Expand Down
8 changes: 4 additions & 4 deletions py-polars/tests/unit/interop/test_interop.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@

import polars as pl
from polars.exceptions import ComputeError, UnstableWarning
from polars.interchange.protocol import Flavor
from polars.interchange.protocol import CompatLevel
from polars.testing import assert_frame_equal, assert_series_equal


Expand Down Expand Up @@ -708,7 +708,7 @@ def test_from_numpy_different_resolution_invalid() -> None:
def test_highest_flavor(monkeypatch: pytest.MonkeyPatch) -> None:
# change these if flavor version bumped
monkeypatch.setenv("POLARS_WARN_UNSTABLE", "1")
assert Flavor.compatible()._version == 0 # type: ignore[attr-defined]
assert CompatLevel.oldest()._version == 0 # type: ignore[attr-defined]
with pytest.warns(UnstableWarning):
assert Flavor.highest()._version == 1 # type: ignore[attr-defined]
assert pl.Series([1])._highest_flavor() == 1
assert CompatLevel.newest()._version == 1 # type: ignore[attr-defined]
assert pl.Series([1])._newest_compat_level() == 1

0 comments on commit 0401852

Please sign in to comment.