diff --git a/.github/workflows/release-python.yml b/.github/workflows/release-python.yml index 30f89988565e..b76216ed63b6 100644 --- a/.github/workflows/release-python.yml +++ b/.github/workflows/release-python.yml @@ -135,6 +135,7 @@ jobs: env: IS_LTS_CPU: ${{ matrix.package == 'polars-lts-cpu' }} IS_MACOS: ${{ matrix.os == 'macos-latest' }} + # IMPORTANT: All features enabled here should also be included in py-polars/polars/_cpu_check.py run: | if [[ "$IS_LTS_CPU" = true ]]; then FEATURES=+sse3,+ssse3,+sse4.1,+sse4.2,+popcnt diff --git a/py-polars/polars/_cpu_check.py b/py-polars/polars/_cpu_check.py index 3edb30df7694..77a343befe60 100644 --- a/py-polars/polars/_cpu_check.py +++ b/py-polars/polars/_cpu_check.py @@ -122,10 +122,6 @@ class CPUID_struct(ctypes.Structure): class CPUID: def __init__(self) -> None: - if _POLARS_ARCH != "x86-64": - msg = "CPUID is only available for x86" - raise SystemError(msg) - if _IS_WINDOWS: if _IS_64BIT: # VirtualAlloc seems to fail under some weird @@ -187,11 +183,7 @@ def __del__(self) -> None: self.win.VirtualFree(self.addr, 0, _MEM_RELEASE) -def read_cpu_flags() -> dict[str, bool]: - # Right now we only enable extra feature flags for x86. - if _POLARS_ARCH != "x86-64": - return {} - +def _read_cpu_flags() -> dict[str, bool]: # CPU flags from https://en.wikipedia.org/wiki/CPUID cpuid = CPUID() cpuid1 = cpuid(1, 0) @@ -223,12 +215,12 @@ def check_cpu_flags() -> None: return expected_cpu_flags = [f.lstrip("+") for f in _POLARS_FEATURE_FLAGS.split(",")] - supported_cpu_flags = read_cpu_flags() + supported_cpu_flags = _read_cpu_flags() missing_features = [] for f in expected_cpu_flags: if f not in supported_cpu_flags: - msg = f'unknown feature flag "{f}"' + msg = f"unknown feature flag: {f!r}" raise RuntimeError(msg) if not supported_cpu_flags[f]: diff --git a/py-polars/tests/unit/test_cpu_check.py b/py-polars/tests/unit/test_cpu_check.py new file mode 100644 index 000000000000..a131298e9329 --- /dev/null +++ b/py-polars/tests/unit/test_cpu_check.py @@ -0,0 +1,80 @@ +from unittest.mock import Mock + +import pytest + +from polars import _cpu_check +from polars._cpu_check import check_cpu_flags + + +@pytest.fixture() +def _feature_flags(monkeypatch: pytest.MonkeyPatch) -> None: + """Use the default set of feature flags.""" + feature_flags = "+sse3,+ssse3" + monkeypatch.setattr(_cpu_check, "_POLARS_FEATURE_FLAGS", feature_flags) + + +@pytest.mark.usefixtures("_feature_flags") +def test_check_cpu_flags( + monkeypatch: pytest.MonkeyPatch, recwarn: pytest.WarningsRecorder +) -> None: + cpu_flags = {"sse3": True, "ssse3": True} + mock_read_cpu_flags = Mock(return_value=cpu_flags) + monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags) + + check_cpu_flags() + + assert len(recwarn) == 0 + + +@pytest.mark.usefixtures("_feature_flags") +def test_check_cpu_flags_missing_features(monkeypatch: pytest.MonkeyPatch) -> None: + cpu_flags = {"sse3": True, "ssse3": False} + mock_read_cpu_flags = Mock(return_value=cpu_flags) + monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags) + + with pytest.warns(RuntimeWarning, match="Missing required CPU features") as w: + check_cpu_flags() + + assert "ssse3" in str(w[0].message) + + +def test_check_cpu_flags_unknown_flag( + monkeypatch: pytest.MonkeyPatch, +) -> None: + feature_flags = "+sse3,+ssse3,+HelloWorld!" + monkeypatch.setattr(_cpu_check, "_POLARS_FEATURE_FLAGS", feature_flags) + with pytest.raises(RuntimeError, match="unknown feature flag: 'HelloWorld!'"): + check_cpu_flags() + + +def test_check_cpu_flags_skipped_no_flags(monkeypatch: pytest.MonkeyPatch) -> None: + mock_read_cpu_flags = Mock() + monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags) + + check_cpu_flags() + + assert mock_read_cpu_flags.call_count == 0 + + +@pytest.mark.usefixtures("_feature_flags") +def test_check_cpu_flags_skipped_lts_cpu(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(_cpu_check, "_POLARS_LTS_CPU", True) + + mock_read_cpu_flags = Mock() + monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags) + + check_cpu_flags() + + assert mock_read_cpu_flags.call_count == 0 + + +@pytest.mark.usefixtures("_feature_flags") +def test_check_cpu_flags_skipped_env_var(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("POLARS_SKIP_CPU_CHECK", "1") + + mock_read_cpu_flags = Mock() + monkeypatch.setattr(_cpu_check, "_read_cpu_flags", mock_read_cpu_flags) + + check_cpu_flags() + + assert mock_read_cpu_flags.call_count == 0