From 8d3c242bcdc4b973ec7076ea73a6fc04960effb3 Mon Sep 17 00:00:00 2001 From: Niels Pardon Date: Tue, 13 Oct 2026 08:50:58 +0200 Subject: [PATCH] fix: update ruff to 0.16 and address its lint findings ruff 0.16 turns on a much larger default rule set, so bumping the pixi pin needs these fixes alongside it. Two of the findings were real bugs: - `substrait.narwhals.__all__` listed the classes instead of their names, so `from substrait.narwhals import *` raised a TypeError. - `test_wildcard_match` passed the builtin `bool()` (i.e. `False`) instead of the `boolean()` type builder, so it never exercised a boolean argument. TRY002/TRY004 are ignored rather than fixed: the public API raises plain `Exception`, and moving to typed exceptions is a separate, user-visible change. DTZ001 is ignored in tests, where naive datetimes model Substrait's timezone-less timestamps. --- .pre-commit-config.yaml | 2 +- examples/adbc_example.py | 37 ++- pixi.lock | 113 ++++----- pyproject.toml | 11 +- src/substrait/builders/extended_expression.py | 37 +-- src/substrait/builders/plan.py | 100 ++++---- src/substrait/builders/type.py | 2 +- src/substrait/dataframe/__init__.py | 2 +- src/substrait/dataframe/dtypes.py | 2 +- src/substrait/dataframe/expr.py | 146 ++++++------ .../dataframe/extension_relations.py | 8 +- src/substrait/dataframe/frame.py | 225 +++++++++--------- src/substrait/derivation_expression.py | 23 +- src/substrait/extension_registry/__init__.py | 12 +- src/substrait/extension_registry/collector.py | 7 +- .../extension_registry/exceptions.py | 4 - .../extension_registry/function_entry.py | 17 +- src/substrait/extension_registry/registry.py | 9 +- .../signature_checker_helpers.py | 4 +- src/substrait/narwhals/__init__.py | 2 +- src/substrait/narwhals/dataframe.py | 4 +- src/substrait/proto/__init__.py | 2 - src/substrait/simple_extension_utils.py | 6 +- src/substrait/sql/sql_to_substrait.py | 12 +- src/substrait/type_inference.py | 11 +- src/substrait/utils/__init__.py | 11 +- src/substrait/utils/display.py | 14 +- tests/builders/plan/test_lateral_join.py | 15 +- tests/builders/plan/test_schema_memo.py | 2 +- tests/dataframe/test_dtypes.py | 4 +- tests/dataframe/test_functions.py | 4 +- tests/dataframe/test_literals.py | 4 +- tests/extension_registry/test_collector.py | 7 +- .../test_registry_lookup.py | 4 +- tests/test_type_inference.py | 2 +- 35 files changed, 426 insertions(+), 439 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 038cd984..6943995a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ ci: repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.1.6 + rev: v0.16.9 hooks: - id: ruff args: [ --fix ] diff --git a/examples/adbc_example.py b/examples/adbc_example.py index a62e1fcb..95681796 100644 --- a/examples/adbc_example.py +++ b/examples/adbc_example.py @@ -35,22 +35,21 @@ def read_adbc_named_table(name: str, conn): return read_named_table(name, substrait_schema) -with adbc_driver_duckdb.dbapi.connect(":memory:") as conn: - with conn.cursor() as cur: - cur.adbc_ingest("AnswerToEverything", data) - - cur.executescript("INSTALL substrait;") - cur.executescript("LOAD substrait;") - - table = read_adbc_named_table("AnswerToEverything", conn) - table = filter( - table, - expression=scalar_function( - "extension:io.substrait:functions_comparison", - "gte", - expressions=[column("ints"), literal(3, i64())], - ), - ) - - cur.execute(table(registry).SerializeToString()) - print(cur.fetch_arrow_table()) +with adbc_driver_duckdb.dbapi.connect(":memory:") as conn, conn.cursor() as cur: + cur.adbc_ingest("AnswerToEverything", data) + + cur.executescript("INSTALL substrait;") + cur.executescript("LOAD substrait;") + + table = read_adbc_named_table("AnswerToEverything", conn) + table = filter( + table, + expression=scalar_function( + "extension:io.substrait:functions_comparison", + "gte", + expressions=[column("ints"), literal(3, i64())], + ), + ) + + cur.execute(table(registry).SerializeToString()) + print(cur.fetch_arrow_table()) diff --git a/pixi.lock b/pixi.lock index f45830b4..501f3de8 100644 --- a/pixi.lock +++ b/pixi.lock @@ -53,7 +53,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.14.10-h4196e79_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.16.10-hc0d6792_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda @@ -81,7 +81,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/openssl-3.6.0-h8e36d6e_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/python-3.13.11-h4c0d347_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/readline-8.3-hb682ff5_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.14.10-hc0dabaa_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.16.10-h42907a6_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/tk-8.6.13-noxft_h561c983_103.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/zstd-1.5.7-h85ac4a6_6.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda @@ -106,7 +106,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-64/openssl-3.6.0-h230baf5_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/python-3.13.11-h17c18a5_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/readline-8.3-h68b038d_0.conda - - conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.14.10-hb17bafe_0.conda + - conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.16.10-h6e8e395_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/tk-8.6.13-hf689a15_3.conda - pypi: https://files.pythonhosted.org/packages/52/cd/c1899c9a3bdb0eb7e94508967a0fc0ea20ceb925384ae0ba16c6e4d88d1b/substrait_extensions-0.103.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/89/11/719554af2bf75ecd0d83369c1796b6df585b9965de9925e1cfb8c21f42ff/substrait_protobuf-0.103.0-py3-none-any.whl @@ -126,7 +126,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.0-h5503f6c_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.13.11-hfc2f54d_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda - - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.14.10-hb0cad00_0.conda + - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.16.10-hc57303c_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_3.conda - pypi: https://files.pythonhosted.org/packages/52/cd/c1899c9a3bdb0eb7e94508967a0fc0ea20ceb925384ae0ba16c6e4d88d1b/substrait_extensions-0.103.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/89/11/719554af2bf75ecd0d83369c1796b6df585b9965de9925e1cfb8c21f42ff/substrait_protobuf-0.103.0-py3-none-any.whl @@ -144,7 +144,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/libzlib-1.3.1-h2466b09_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/openssl-3.6.0-h725018a_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.13.11-h09917c8_100_cp313.conda - - conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.14.10-h37e10c4_0.conda + - conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.16.10-h808fad9_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/tk-8.6.13-h2c6b04d_3.conda - conda: https://conda.anaconda.org/conda-forge/win-64/ucrt-10.0.26100.0-h57928b3_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/vc-14.3-h41ae7f8_34.conda @@ -179,7 +179,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.14.10-h4196e79_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.16.10-hc0d6792_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda @@ -221,7 +221,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/openssl-3.6.0-h8e36d6e_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/python-3.13.11-h4c0d347_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/readline-8.3-hb682ff5_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.14.10-hc0dabaa_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.16.10-h42907a6_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/tk-8.6.13-noxft_h561c983_103.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/zstd-1.5.7-h85ac4a6_6.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda @@ -260,7 +260,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-64/openssl-3.6.0-h230baf5_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/python-3.13.11-h17c18a5_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/readline-8.3-h68b038d_0.conda - - conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.14.10-hb17bafe_0.conda + - conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.16.10-h6e8e395_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/tk-8.6.13-hf689a15_3.conda - pypi: https://files.pythonhosted.org/packages/10/c7/79f57728f300079148ac4e4b988000dcbd1f58960040db89594424a58336/sqloxide-0.1.56.tar.gz - pypi: https://files.pythonhosted.org/packages/12/27/fb8d7338b4d551900fa3e580acbe7a0cf655d940e164cb5c00ec31961094/orderly_set-5.5.0-py3-none-any.whl @@ -294,7 +294,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.0-h5503f6c_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.13.11-hfc2f54d_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda - - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.14.10-hb0cad00_0.conda + - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.16.10-hc57303c_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_3.conda - pypi: https://files.pythonhosted.org/packages/12/27/fb8d7338b4d551900fa3e580acbe7a0cf655d940e164cb5c00ec31961094/orderly_set-5.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl @@ -326,7 +326,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/libzlib-1.3.1-h2466b09_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/openssl-3.6.0-h725018a_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.13.11-h09917c8_100_cp313.conda - - conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.14.10-h37e10c4_0.conda + - conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.16.10-h808fad9_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/tk-8.6.13-h2c6b04d_3.conda - conda: https://conda.anaconda.org/conda-forge/win-64/ucrt-10.0.26100.0-h57928b3_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/vc-14.3-h41ae7f8_34.conda @@ -376,7 +376,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.14.10-h4196e79_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.16.10-hc0d6792_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda @@ -407,7 +407,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/openssl-3.6.0-h8e36d6e_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/python-3.13.11-h4c0d347_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/readline-8.3-hb682ff5_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.14.10-hc0dabaa_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.16.10-h42907a6_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/tk-8.6.13-noxft_h561c983_103.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/zstd-1.5.7-h85ac4a6_6.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda @@ -435,7 +435,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-64/openssl-3.6.0-h230baf5_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/python-3.13.11-h17c18a5_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/readline-8.3-h68b038d_0.conda - - conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.14.10-hb17bafe_0.conda + - conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.16.10-h6e8e395_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/tk-8.6.13-hf689a15_3.conda - pypi: https://files.pythonhosted.org/packages/52/cd/c1899c9a3bdb0eb7e94508967a0fc0ea20ceb925384ae0ba16c6e4d88d1b/substrait_extensions-0.103.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/89/03/a851e84fcbb85214dc637b6378121ef9a0dd61b4c65264675d8a5c9b1ae7/antlr4_python3_runtime-4.13.2-py3-none-any.whl @@ -458,7 +458,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.0-h5503f6c_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.13.11-hfc2f54d_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda - - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.14.10-hb0cad00_0.conda + - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.16.10-hc57303c_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_3.conda - pypi: https://files.pythonhosted.org/packages/52/cd/c1899c9a3bdb0eb7e94508967a0fc0ea20ceb925384ae0ba16c6e4d88d1b/substrait_extensions-0.103.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/89/03/a851e84fcbb85214dc637b6378121ef9a0dd61b4c65264675d8a5c9b1ae7/antlr4_python3_runtime-4.13.2-py3-none-any.whl @@ -479,7 +479,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/libzlib-1.3.1-h2466b09_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/openssl-3.6.0-h725018a_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.13.11-h09917c8_100_cp313.conda - - conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.14.10-h37e10c4_0.conda + - conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.16.10-h808fad9_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/tk-8.6.13-h2c6b04d_3.conda - conda: https://conda.anaconda.org/conda-forge/win-64/ucrt-10.0.26100.0-h57928b3_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/vc-14.3-h41ae7f8_34.conda @@ -517,7 +517,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.0-h26f9b46_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.11-hc97d973_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.14.10-h4196e79_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.16.10-hc0d6792_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda @@ -548,7 +548,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/openssl-3.6.0-h8e36d6e_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/python-3.13.11-h4c0d347_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/readline-8.3-hb682ff5_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.14.10-hc0dabaa_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.16.10-h42907a6_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/tk-8.6.13-noxft_h561c983_103.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/zstd-1.5.7-h85ac4a6_6.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda @@ -576,7 +576,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-64/openssl-3.6.0-h230baf5_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/python-3.13.11-h17c18a5_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/readline-8.3-h68b038d_0.conda - - conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.14.10-hb17bafe_0.conda + - conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.16.10-h6e8e395_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-64/tk-8.6.13-hf689a15_3.conda - pypi: https://files.pythonhosted.org/packages/10/c7/79f57728f300079148ac4e4b988000dcbd1f58960040db89594424a58336/sqloxide-0.1.56.tar.gz - pypi: https://files.pythonhosted.org/packages/12/27/fb8d7338b4d551900fa3e580acbe7a0cf655d940e164cb5c00ec31961094/orderly_set-5.5.0-py3-none-any.whl @@ -599,7 +599,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.0-h5503f6c_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.13.11-hfc2f54d_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda - - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.14.10-hb0cad00_0.conda + - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.16.10-hc57303c_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_3.conda - pypi: https://files.pythonhosted.org/packages/12/27/fb8d7338b4d551900fa3e580acbe7a0cf655d940e164cb5c00ec31961094/orderly_set-5.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/52/cd/c1899c9a3bdb0eb7e94508967a0fc0ea20ceb925384ae0ba16c6e4d88d1b/substrait_extensions-0.103.0-py3-none-any.whl @@ -620,7 +620,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/libzlib-1.3.1-h2466b09_2.conda - conda: https://conda.anaconda.org/conda-forge/win-64/openssl-3.6.0-h725018a_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/python-3.13.11-h09917c8_100_cp313.conda - - conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.14.10-h37e10c4_0.conda + - conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.16.10-h808fad9_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/tk-8.6.13-h2c6b04d_3.conda - conda: https://conda.anaconda.org/conda-forge/win-64/ucrt-10.0.26100.0-h57928b3_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/vc-14.3-h41ae7f8_34.conda @@ -871,22 +871,23 @@ packages: purls: [] size: 345073 timestamp: 1765813471974 -- conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.14.10-h4196e79_0.conda +- conda: https://conda.anaconda.org/conda-forge/linux-64/ruff-0.16.10-hc0d6792_0.conda noarch: python - sha256: 997b45ce89554f677e4a30cd1d64565949be9d25c806727c3d844fee0d55d7d2 - md5: ddecdee0806589993d96a950ad51b927 + sha256: 1801965a5f5821b74a52076ba5c6435ac6d9111d83c7e63049b6813875c75734 + md5: 6cb114e481362b28ea3004c7619aecfb depends: - python - - libgcc >=14 + - libgcc >=15 - __glibc >=2.17,<3.0.a0 constrains: - __glibc >=2.17 license: MIT license_family: MIT purls: - - pkg:pypi/ruff?source=hash-mapping - size: 11472103 - timestamp: 1766094973645 + - pkg:pypi/ruff?source=compressed-mapping + run_exports: {} + size: 8923644 + timestamp: 1790964529417 - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_ha0e22de_103.conda sha256: 1544760538a40bcd8ace2b1d8ebe3eb5807ac268641f8acdc18c69c5ebfeaf64 md5: 86bc20552bf46075e3d92b67f089172d @@ -1124,21 +1125,22 @@ packages: purls: [] size: 357597 timestamp: 1765815673644 -- conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.14.10-hc0dabaa_0.conda +- conda: https://conda.anaconda.org/conda-forge/linux-aarch64/ruff-0.16.10-h42907a6_0.conda noarch: python - sha256: 4aa64574aba358480d03e86846e6a8444499e580b87b280f5d36ce1940e6e3f4 - md5: cd19c0de4d78564fc9e4758fa2280606 + sha256: 4176fe1afe09d2f92624c2495d976566c1629e321de45e34662a23809f72058b + md5: b61562a5b3e6e9f946322bb27ac6b115 depends: - python - - libgcc >=14 + - libgcc >=15 constrains: - __glibc >=2.17 license: MIT license_family: MIT purls: - - pkg:pypi/ruff?source=hash-mapping - size: 10995544 - timestamp: 1766094986744 + - pkg:pypi/ruff?source=compressed-mapping + run_exports: {} + size: 8625202 + timestamp: 1790964526260 - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/tk-8.6.13-noxft_h561c983_103.conda sha256: 154e73f6269f92ad5257aa2039278b083998fd19d371e150f307483fb93c07ae md5: 631db4799bc2bfe4daccf80bb3cbc433 @@ -1339,21 +1341,22 @@ packages: purls: [] size: 317819 timestamp: 1765813692798 -- conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.14.10-hb17bafe_0.conda +- conda: https://conda.anaconda.org/conda-forge/osx-64/ruff-0.16.10-h6e8e395_0.conda noarch: python - sha256: e1e48a3c7c9f097200eef38322a40c8f5d08fe68308d1e577e8270d353544ab5 - md5: cbfe1d44978259b92d7c9b221bfe59b8 + sha256: 7ecc285777343178db74fe2d36589c88a1f40931c1ad0f947a180172fdfea0d8 + md5: 59b3f584d902b8eab8959c1666d6ecba depends: - python - - __osx >=10.13 + - __osx >=11.0 constrains: - - __osx >=10.13 + - __osx >=11.0 license: MIT license_family: MIT purls: - - pkg:pypi/ruff?source=hash-mapping - size: 11405633 - timestamp: 1766095106770 + - pkg:pypi/ruff?source=compressed-mapping + run_exports: {} + size: 8816604 + timestamp: 1790964695397 - conda: https://conda.anaconda.org/conda-forge/osx-64/tk-8.6.13-hf689a15_3.conda sha256: 0d0b6cef83fec41bc0eb4f3b761c4621b7adfb14378051a8177bd9bb73d26779 md5: bd9f1de651dbd80b51281c694827f78f @@ -1495,10 +1498,10 @@ packages: purls: [] size: 313930 timestamp: 1765813902568 -- conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.14.10-hb0cad00_0.conda +- conda: https://conda.anaconda.org/conda-forge/osx-arm64/ruff-0.16.10-hc57303c_0.conda noarch: python - sha256: dcae2c2d8c3a07782c0ce900f44488e2c2a2acd3720850181165a4e0034219fc - md5: d506e661dcc8c2053b2661edb0d3c57c + sha256: 1689570fdffed6cebe9b4efc08c3e00d8d86fd7533e2fa1b4c9ff3c9c06fac75 + md5: d7b9a54d80315847badff420878967f3 depends: - python - __osx >=11.0 @@ -1507,9 +1510,10 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/ruff?source=hash-mapping - size: 10468468 - timestamp: 1766095111371 + - pkg:pypi/ruff?source=compressed-mapping + run_exports: {} + size: 8199552 + timestamp: 1790964593925 - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h892fb3f_3.conda sha256: ad0c67cb03c163a109820dc9ecf77faf6ec7150e942d1e8bb13e5d39dc058ab7 md5: a73d54a5abba6543cb2f0af1bfbd6851 @@ -1646,10 +1650,10 @@ packages: size: 16617922 timestamp: 1765019627175 python_site_packages_path: Lib/site-packages -- conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.14.10-h37e10c4_0.conda +- conda: https://conda.anaconda.org/conda-forge/win-64/ruff-0.16.10-h808fad9_0.conda noarch: python - sha256: 2ca5b7a6ab0a4cddec16f390e1c2d89b5cacd1780963745e463f1bb2e8ab4893 - md5: 987a8290e4bfd9e42e3c58be4481929c + sha256: e799af862ac154c34775e5f62de805ea605cc684bf926ce84e481dfdf0b795cc + md5: 15630e9d8178646a55e184cc40c6548e depends: - python - vc >=14.3,<15 @@ -1658,9 +1662,10 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/ruff?source=hash-mapping - size: 11908812 - timestamp: 1766095035171 + - pkg:pypi/ruff?source=compressed-mapping + run_exports: {} + size: 9291796 + timestamp: 1790964544005 - conda: https://conda.anaconda.org/conda-forge/win-64/tk-8.6.13-h2c6b04d_3.conda sha256: 4581f4ffb432fefa1ac4f85c5682cc27014bcd66e7beaa0ee330e927a7858790 md5: 7cb36e506a7dba4817970f8adb6396f9 diff --git a/pyproject.toml b/pyproject.toml index ca25e084..b07e13db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,15 @@ respect-gitignore = true # should target minimum supported version target-version = "py310" lint.extend-select = ["I"] +# The public API raises plain `Exception`; switching to typed exceptions is a separate, +# user-visible change. +lint.extend-ignore = ["TRY002", "TRY004"] + +[tool.ruff.lint.per-file-ignores] +# Naive datetimes are deliberate: they model Substrait's timezone-less timestamp literals. +"tests/**" = ["DTZ001"] +# Re-exports the generated protobuf modules wholesale. +"src/substrait/proto/__init__.py" = ["F401", "F403"] [tool.pixi.workspace] @@ -70,4 +79,4 @@ lint = "ruff check" lint_fix = "ruff check --fix" [tool.pixi.dependencies] -ruff = ">=0.14.10,<0.15" +ruff = ">=0.16.9,<0.17" diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py index 46464a37..223005af 100644 --- a/src/substrait/builders/extended_expression.py +++ b/src/substrait/builders/extended_expression.py @@ -3,9 +3,10 @@ import contextvars import itertools import uuid as uuid_module +from collections.abc import Callable, Iterable from datetime import date, datetime, time, timedelta, timezone from decimal import Decimal -from typing import Any, Callable, Iterable, Union +from typing import Any, Union import substrait.algebra_pb2 as stalg import substrait.extended_expression_pb2 as stee @@ -55,7 +56,7 @@ def fresh_rel_anchors(): UnboundExtendedExpression = Callable[ [stp.NamedStruct, ExtensionRegistry], stee.ExtendedExpression ] -ExtendedExpressionOrUnbound = Union[stee.ExtendedExpression, UnboundExtendedExpression] +ExtendedExpressionOrUnbound = stee.ExtendedExpression | UnboundExtendedExpression class EnumArg: @@ -83,7 +84,7 @@ def __repr__(self) -> str: def _alias_or_inferred( - alias: Union[Iterable[str], str, None], + alias: Iterable[str] | str | None, op: str, args: Iterable[str], ): @@ -456,7 +457,7 @@ def _make_literal(value: Any, type: stp.Type) -> stalg.Expression.Literal: def literal( - value: Any, type: stp.Type, alias: Union[Iterable[str], str, None] = None + value: Any, type: stp.Type, alias: Iterable[str] | str | None = None ) -> UnboundExtendedExpression: """Builds a resolver for ExtendedExpression containing a literal expression. @@ -480,7 +481,7 @@ def resolve( return build_scoped(resolve) -def outer_reference(field: Union[str, int], steps_out: int = 1): +def outer_reference(field: str | int, steps_out: int = 1): """A field reference to an enclosing query's column (a correlated reference). Only valid inside a subquery; ``steps_out`` counts query-nesting levels @@ -538,7 +539,7 @@ def __init__(self, rel_anchor: int, schema: stp.NamedStruct): self._rel_anchor = rel_anchor self._schema = schema - def column(self, field: Union[str, int]): + def column(self, field: str | int): """A correlated reference to the left row's ``field`` (name or index).""" rel_anchor = self._rel_anchor schema = self._schema @@ -570,7 +571,7 @@ def resolve( return build_scoped(resolve) -def column(field: Union[str, int], alias: Union[Iterable[str], str, None] = None): +def column(field: str | int, alias: Iterable[str] | str | None = None): """Builds a resolver for ExtendedExpression containing a FieldReference expression Accepts either an index or a field name of a desired field. @@ -624,8 +625,8 @@ def scalar_function( urn: str, function: str, expressions: Iterable[ExtendedExpressionOrUnbound], - alias: Union[Iterable[str], str, None] = None, - options: Union[dict, None] = None, + alias: Iterable[str] | str | None = None, + options: dict | None = None, ): """Builds a resolver for ExtendedExpression containing a ScalarFunction expression. @@ -678,14 +679,14 @@ def aggregate_function( urn: str, function: str, expressions: Iterable[ExtendedExpressionOrUnbound], - alias: Union[Iterable[str], str, None] = None, + alias: Iterable[str] | str | None = None, invocation: Union[ "stalg.AggregateFunction.AggregationInvocation.ValueType", None ] = None, sorts: Iterable[ tuple[ExtendedExpressionOrUnbound, "stalg.SortField.SortDirection.ValueType"] ] = (), - options: Union[dict, None] = None, + options: dict | None = None, ): """Builds a resolver for ExtendedExpression containing a AggregateFunction measure. @@ -751,8 +752,8 @@ def window_function( function: str, expressions: Iterable[ExtendedExpressionOrUnbound], partitions: Iterable[ExtendedExpressionOrUnbound] = [], - alias: Union[Iterable[str], str, None] = None, - options: Union[dict, None] = None, + alias: Iterable[str] | str | None = None, + options: dict | None = None, ): """Builds a resolver for ExtendedExpression containing a WindowFunction expression""" @@ -807,7 +808,7 @@ def resolve( def if_then( ifs: Iterable[tuple[ExtendedExpressionOrUnbound, ExtendedExpressionOrUnbound]], _else: ExtendedExpressionOrUnbound, - alias: Union[Iterable[str], str, None] = None, + alias: Iterable[str] | str | None = None, ): """Builds a resolver for ExtendedExpression containing an IfThen expression""" @@ -990,7 +991,7 @@ def resolve( def cast( input: ExtendedExpressionOrUnbound, type: stp.Type, - alias: Union[Iterable[str], str, None] = None, + alias: Iterable[str] | str | None = None, ): """Builds a resolver for ExtendedExpression containing a cast expression""" @@ -1078,7 +1079,7 @@ def _inner_rel(query, registry: ExtensionRegistry, base_schema) -> stalg.Rel: return rel -def scalar_subquery(query, alias: Union[str, None] = None): +def scalar_subquery(query, alias: str | None = None): """A scalar (one-row, one-column) subquery expression.""" def resolve(base_schema, registry): @@ -1091,7 +1092,7 @@ def resolve(base_schema, registry): return build_scoped(resolve) -def set_predicate(query, op, alias: Union[str, None] = None): +def set_predicate(query, op, alias: str | None = None): """An EXISTS / UNIQUE subquery predicate.""" def resolve(base_schema, registry): @@ -1106,7 +1107,7 @@ def resolve(base_schema, registry): return build_scoped(resolve) -def in_predicate(needles, query, alias: Union[str, None] = None): +def in_predicate(needles, query, alias: str | None = None): """A ``needles IN (subquery)`` predicate.""" def resolve(base_schema, registry): diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py index 115fdccb..60ed0e41 100644 --- a/src/substrait/builders/plan.py +++ b/src/substrait/builders/plan.py @@ -6,7 +6,7 @@ """ import re -from typing import Callable, Iterable, Optional, Union +from collections.abc import Callable, Iterable import substrait.algebra_pb2 as stalg import substrait.extended_expression_pb2 as stee @@ -43,7 +43,7 @@ UnboundPlan = Callable[[ExtensionRegistry], stp.Plan] -PlanOrUnbound = Union[stp.Plan, UnboundPlan] +PlanOrUnbound = stp.Plan | UnboundPlan def _create_default_version(): @@ -120,7 +120,7 @@ def _merge_input_subtrees(bound_inputs): carry no subtrees pass their root input through untouched, keeping output for the (overwhelmingly common) no-subtree case byte-identical. """ - combined: "list[stalg.Rel]" = [] + combined: list[stalg.Rel] = [] key_to_ordinal: dict = {} # serialized subtree bytes -> ordinal in `combined` rebased_root_inputs = [] for plan in bound_inputs: @@ -262,9 +262,9 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def read_named_table( - names: Union[str, Iterable[str]], + names: str | Iterable[str], named_struct: stt.NamedStruct, - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: if named_struct.struct.nullability is stt.Type.NULLABILITY_NULLABLE: raise Exception("NamedStruct must not contain a nullable struct") @@ -318,7 +318,7 @@ def _read_plan(named_struct: stt.NamedStruct, read_rel: stalg.ReadRel) -> stp.Pl def virtual_table( rows: Iterable[Iterable[ExtendedExpressionOrUnbound]], named_struct: stt.NamedStruct, - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """A ReadRel over inline rows (the VALUES clause). @@ -353,7 +353,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def local_files( named_struct: stt.NamedStruct, items: Iterable[stalg.ReadRel.LocalFiles.FileOrFiles], - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """A ReadRel over local/remote files; ``items`` are pre-built FileOrFiles.""" _require_schema(named_struct) @@ -373,7 +373,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def extension_table( named_struct: stt.NamedStruct, detail, - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """A ReadRel over a custom source; ``detail`` is a ``google.protobuf.Any``.""" _require_schema(named_struct) @@ -393,7 +393,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def project( plan: PlanOrUnbound, expressions: Iterable[ExtendedExpressionOrUnbound], - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """ Builds an UnboundPlan with ProjectRel as the root node. Expressions are appended to the parent relation fields to produce an output. @@ -443,7 +443,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def select( plan: PlanOrUnbound, expressions: Iterable[ExtendedExpressionOrUnbound], - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """ Builds an UnboundPlan with ProjectRel as the root node. Expressions make up the fields of an output relation. @@ -538,7 +538,7 @@ def _require_boolean_condition( def filter( plan: PlanOrUnbound, expression: ExtendedExpressionOrUnbound, - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_plan = _bind(plan, registry) @@ -567,12 +567,10 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def sort( plan: PlanOrUnbound, expressions: Iterable[ - Union[ - ExtendedExpressionOrUnbound, - tuple[ExtendedExpressionOrUnbound, stalg.SortField.SortDirection.ValueType], - ] + ExtendedExpressionOrUnbound + | tuple[ExtendedExpressionOrUnbound, stalg.SortField.SortDirection.ValueType] ], - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_plan = _bind(plan, registry) @@ -667,8 +665,8 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def fetch( plan: PlanOrUnbound, offset: ExtendedExpressionOrUnbound, - count: Optional[ExtendedExpressionOrUnbound], - extension: Optional[AdvancedExtension] = None, + count: ExtendedExpressionOrUnbound | None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_plan = _bind(plan, registry) @@ -707,8 +705,8 @@ def join( expression: ExtendedExpressionOrUnbound, type: stalg.JoinRel.JoinType, *, - post_join_filter: Optional[ExtendedExpressionOrUnbound] = None, - extension: Optional[AdvancedExtension] = None, + post_join_filter: ExtendedExpressionOrUnbound | None = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_left = _bind(left, registry) @@ -781,9 +779,9 @@ def lateral_join( right: Callable[[LateralInput], PlanOrUnbound], type: stalg.JoinRel.JoinType, *, - expression: Optional[ExtendedExpressionOrUnbound] = None, - post_join_filter: Optional[ExtendedExpressionOrUnbound] = None, - extension: Optional[AdvancedExtension] = None, + expression: ExtendedExpressionOrUnbound | None = None, + post_join_filter: ExtendedExpressionOrUnbound | None = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """A LateralJoinRel: the right (dependent) input is evaluated once per left row and may reference the current left row. @@ -883,7 +881,7 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def cross( left: PlanOrUnbound, right: PlanOrUnbound, - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_left = _bind(left, registry) @@ -920,9 +918,9 @@ def aggregate( grouping_expressions: Iterable[ExtendedExpressionOrUnbound], measures: Iterable[ExtendedExpressionOrUnbound], *, - grouping_sets: Optional[Iterable[Iterable[int]]] = None, - filters: Optional[Iterable[Optional[ExtendedExpressionOrUnbound]]] = None, - extension: Optional[AdvancedExtension] = None, + grouping_sets: Iterable[Iterable[int]] | None = None, + filters: Iterable[ExtendedExpressionOrUnbound | None] | None = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """Build an AggregateRel. @@ -1014,11 +1012,11 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def write_named_table( - table_names: Union[str, Iterable[str]], + table_names: str | Iterable[str], input: PlanOrUnbound, - create_mode: Union[stalg.WriteRel.CreateMode.ValueType, None] = None, - op: Union[stalg.WriteRel.WriteOp.ValueType, None] = None, - output_mode: Union[stalg.WriteRel.OutputMode.ValueType, None] = None, + create_mode: stalg.WriteRel.CreateMode.ValueType | None = None, + op: stalg.WriteRel.WriteOp.ValueType | None = None, + output_mode: stalg.WriteRel.OutputMode.ValueType | None = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_input = _bind(input, registry) @@ -1050,12 +1048,12 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def ddl( - names: Union[str, Iterable[str]], + names: str | Iterable[str], object_type: stalg.DdlRel.DdlObject.ValueType, op: stalg.DdlRel.DdlOp.ValueType, - table_schema: Optional[stt.NamedStruct] = None, - view_definition: Optional[PlanOrUnbound] = None, - extension: Optional[AdvancedExtension] = None, + table_schema: stt.NamedStruct | None = None, + view_definition: PlanOrUnbound | None = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """Build a DdlRel (CREATE / DROP of a TABLE or VIEW). @@ -1096,11 +1094,11 @@ def resolve(registry: ExtensionRegistry) -> stp.Plan: def update( - table_names: Union[str, Iterable[str]], + table_names: str | Iterable[str], table_schema: stt.NamedStruct, transformations: Iterable[tuple[int, ExtendedExpressionOrUnbound]], - condition: Optional[ExtendedExpressionOrUnbound] = None, - extension: Optional[AdvancedExtension] = None, + condition: ExtendedExpressionOrUnbound | None = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """Build an UpdateRel: set ``(column_index -> expression)`` where ``condition``.""" _names = [table_names] if isinstance(table_names, str) else list(table_names) @@ -1155,12 +1153,10 @@ def consistent_partition_window( window_functions: Iterable[ExtendedExpressionOrUnbound], partition_expressions: Iterable[ExtendedExpressionOrUnbound] = (), sorts: Iterable[ - Union[ - ExtendedExpressionOrUnbound, - tuple[ExtendedExpressionOrUnbound, stalg.SortField.SortDirection.ValueType], - ] + ExtendedExpressionOrUnbound + | tuple[ExtendedExpressionOrUnbound, stalg.SortField.SortDirection.ValueType] ] = (), - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_plan = _bind(plan, registry) @@ -1302,7 +1298,7 @@ def nested_loop_join( right: PlanOrUnbound, expression: ExtendedExpressionOrUnbound, type: stalg.NestedLoopJoinRel.JoinType.ValueType, - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """A NestedLoopJoinRel: join over the Cartesian product using ``expression``.""" @@ -1385,13 +1381,13 @@ def _physical_equi_join(rel_name, rel_cls): def builder( left: PlanOrUnbound, right: PlanOrUnbound, - left_keys: Iterable[Union[str, int]], - right_keys: Iterable[Union[str, int]], + left_keys: Iterable[str | int], + right_keys: Iterable[str | int], type, *, - post_join_filter: Optional[ExtendedExpressionOrUnbound] = None, - residual_expression: Optional[ExtendedExpressionOrUnbound] = None, - extension: Optional[AdvancedExtension] = None, + post_join_filter: ExtendedExpressionOrUnbound | None = None, + residual_expression: ExtendedExpressionOrUnbound | None = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: def resolve(registry: ExtensionRegistry) -> stp.Plan: bound_left = _bind(left, registry) @@ -1476,7 +1472,7 @@ def _detail_any(detail): return detail.to_any() if hasattr(detail, "to_any") else detail -def extension_leaf(detail, names: Optional[Iterable[str]] = None) -> UnboundPlan: +def extension_leaf(detail, names: Iterable[str] | None = None) -> UnboundPlan: """An ExtensionLeafRel (a custom source) from an ExtensionLeafDetail. Output names come from the detail's ``derive_schema`` (or ``names`` when a @@ -1593,9 +1589,9 @@ def top_n( tuple[ExtendedExpressionOrUnbound, stalg.SortField.SortDirection.ValueType] ], count: ExtendedExpressionOrUnbound, - offset: Optional[ExtendedExpressionOrUnbound] = None, + offset: ExtendedExpressionOrUnbound | None = None, with_ties: bool = False, - extension: Optional[AdvancedExtension] = None, + extension: AdvancedExtension | None = None, ) -> UnboundPlan: """A TopNRel: a fused sort + fetch (ORDER BY ... LIMIT). diff --git a/src/substrait/builders/type.py b/src/substrait/builders/type.py index 06b16573..75a56adf 100644 --- a/src/substrait/builders/type.py +++ b/src/substrait/builders/type.py @@ -1,4 +1,4 @@ -from typing import Iterable +from collections.abc import Iterable import substrait.type_pb2 as stt diff --git a/src/substrait/dataframe/__init__.py b/src/substrait/dataframe/__init__.py index e8f9fe25..55f00fd0 100644 --- a/src/substrait/dataframe/__init__.py +++ b/src/substrait/dataframe/__init__.py @@ -109,7 +109,7 @@ from substrait.dataframe.functions import f, functions_for from substrait.extension_registry import ExtensionRegistry -__all__ = [ +__all__ = [ # noqa: RUF022 -- grouped by topic # entry points "read_named_table", "from_records", diff --git a/src/substrait/dataframe/dtypes.py b/src/substrait/dataframe/dtypes.py index 867544f7..a9d99b37 100644 --- a/src/substrait/dataframe/dtypes.py +++ b/src/substrait/dataframe/dtypes.py @@ -24,7 +24,7 @@ class DataType: - __slots__ = ("_name", "_builder") + __slots__ = ("_builder", "_name") def __init__(self, name: str, builder): self._name = name diff --git a/src/substrait/dataframe/expr.py b/src/substrait/dataframe/expr.py index f39a55ba..688d9aae 100644 --- a/src/substrait/dataframe/expr.py +++ b/src/substrait/dataframe/expr.py @@ -22,7 +22,7 @@ from datetime import datetime as _datetime from datetime import time as _time from decimal import Decimal as _Decimal -from typing import Any, Union +from typing import Any import substrait.algebra_pb2 as stalg import substrait.extended_expression_pb2 as stee @@ -160,7 +160,7 @@ def infer_literal_type(value: Any) -> stp.Type: } -def _exact_unscaled(value: _Decimal, scale: int) -> Union[int, None]: +def _exact_unscaled(value: _Decimal, scale: int) -> int | None: """The unscaled integer of ``value`` at ``scale``, or ``None`` if representing it at ``scale`` would drop nonzero digits (i.e. the value needs a finer scale). @@ -266,14 +266,14 @@ def _resolve_over_urns( def _numeric_binary( - self_expr: "Expr", + self_expr: Expr, other: Any, - urns: Union[str, list], + urns: str | list, fn: str, *, swap: bool = False, decimal_literal: str = "natural", -) -> "Expr": +) -> Expr: """Build a binary comparison/arithmetic expression with literal coercion. A bare Python number is typed to match the *other* (column) operand at @@ -350,7 +350,7 @@ class _SubqueryReduction: """A subquery wrapped by :func:`any_` / :func:`all_`, consumed by a comparison operator to build a ``left ANY/ALL (subquery)`` expression.""" - __slots__ = ("reduction_op", "plan") + __slots__ = ("plan", "reduction_op") def __init__(self, reduction_op, plan): self.reduction_op = reduction_op @@ -410,17 +410,17 @@ def unbound(self) -> UnboundExtendedExpression: return self._unbound @staticmethod - def _coerce(value: Union["Expr", Any]) -> "Expr": + def _coerce(value: Expr | Any) -> Expr: if isinstance(value, Expr): return value return Expr(literal(value, infer_literal_type(value))) - def _scalar(self, urn: str, fn: str, *others: Any) -> "Expr": + def _scalar(self, urn: str, fn: str, *others: Any) -> Expr: args = [self._unbound] + [Expr._coerce(o)._unbound for o in others] return Expr(scalar_function(urn, fn, expressions=args)) # -- comparison ------------------------------------------------------- - def _compare(self, other: Any, fn: str) -> "Expr": + def _compare(self, other: Any, fn: str) -> Expr: # `col any_(df)/all_(df)` builds a SetComparison subquery instead. if isinstance(other, _SubqueryReduction): return Expr( @@ -432,22 +432,22 @@ def _compare(self, other: Any, fn: str) -> "Expr": self, other, FUNCTIONS_COMPARISON, fn, decimal_literal="peer" ) - def __lt__(self, other: Any) -> "Expr": + def __lt__(self, other: Any) -> Expr: return self._compare(other, "lt") - def __le__(self, other: Any) -> "Expr": + def __le__(self, other: Any) -> Expr: return self._compare(other, "lte") - def __gt__(self, other: Any) -> "Expr": + def __gt__(self, other: Any) -> Expr: return self._compare(other, "gt") - def __ge__(self, other: Any) -> "Expr": + def __ge__(self, other: Any) -> Expr: return self._compare(other, "gte") - def __eq__(self, other: Any) -> "Expr": # type: ignore[override] + def __eq__(self, other: object) -> Expr: # type: ignore[override] return self._compare(other, "equal") - def __ne__(self, other: Any) -> "Expr": # type: ignore[override] + def __ne__(self, other: object) -> Expr: # type: ignore[override] return self._compare(other, "not_equal") # Operator-overloaded ``__eq__`` means an Expr is not a normal value; like @@ -455,91 +455,91 @@ def __ne__(self, other: Any) -> "Expr": # type: ignore[override] __hash__ = None # type: ignore[assignment] # -- arithmetic ------------------------------------------------------- - def __add__(self, other: Any) -> "Expr": + def __add__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "add") - def __sub__(self, other: Any) -> "Expr": + def __sub__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "subtract") - def __mul__(self, other: Any) -> "Expr": + def __mul__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "multiply") - def __truediv__(self, other: Any) -> "Expr": + def __truediv__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "divide") - def __radd__(self, other: Any) -> "Expr": + def __radd__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "add", swap=True) - def __rsub__(self, other: Any) -> "Expr": + def __rsub__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "subtract", swap=True) - def __rmul__(self, other: Any) -> "Expr": + def __rmul__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "multiply", swap=True) - def __rtruediv__(self, other: Any) -> "Expr": + def __rtruediv__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "divide", swap=True) - def __mod__(self, other: Any) -> "Expr": + def __mod__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "modulus") - def __rmod__(self, other: Any) -> "Expr": + def __rmod__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "modulus", swap=True) - def __pow__(self, other: Any) -> "Expr": + def __pow__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "power") - def __rpow__(self, other: Any) -> "Expr": + def __rpow__(self, other: Any) -> Expr: return _numeric_binary(self, other, _ARITHMETIC_URNS, "power", swap=True) - def __neg__(self) -> "Expr": + def __neg__(self) -> Expr: return Expr( scalar_function(FUNCTIONS_ARITHMETIC, "negate", expressions=[self._unbound]) ) # -- boolean logic ---------------------------------------------------- - def __and__(self, other: Any) -> "Expr": + def __and__(self, other: Any) -> Expr: return self._scalar(FUNCTIONS_BOOLEAN, "and", other) - def __or__(self, other: Any) -> "Expr": + def __or__(self, other: Any) -> Expr: return self._scalar(FUNCTIONS_BOOLEAN, "or", other) - def __xor__(self, other: Any) -> "Expr": + def __xor__(self, other: Any) -> Expr: return self._scalar(FUNCTIONS_BOOLEAN, "xor", other) - def __invert__(self) -> "Expr": + def __invert__(self) -> Expr: return Expr( scalar_function(FUNCTIONS_BOOLEAN, "not", expressions=[self._unbound]) ) # -- helpers ---------------------------------------------------------- - def is_null(self) -> "Expr": + def is_null(self) -> Expr: return Expr( scalar_function( FUNCTIONS_COMPARISON, "is_null", expressions=[self._unbound] ) ) - def is_not_null(self) -> "Expr": + def is_not_null(self) -> Expr: return Expr( scalar_function( FUNCTIONS_COMPARISON, "is_not_null", expressions=[self._unbound] ) ) - def is_nan(self) -> "Expr": + def is_nan(self) -> Expr: return Expr( scalar_function(FUNCTIONS_COMPARISON, "is_nan", expressions=[self._unbound]) ) - def is_distinct_from(self, other: Any) -> "Expr": + def is_distinct_from(self, other: Any) -> Expr: """Null-safe inequality (``NULL`` distinct from a value / from ``NULL``).""" return self._scalar(FUNCTIONS_COMPARISON, "is_distinct_from", other) - def is_not_distinct_from(self, other: Any) -> "Expr": + def is_not_distinct_from(self, other: Any) -> Expr: """Null-safe equality (``NULL`` equals ``NULL``).""" return self._scalar(FUNCTIONS_COMPARISON, "is_not_distinct_from", other) - def between(self, low: Any, high: Any) -> "Expr": + def between(self, low: Any, high: Any) -> Expr: """Inclusive range test, ``low <= self <= high``. Like the ``f.*`` helpers, the bounds are not coerced to this column's @@ -547,7 +547,7 @@ def between(self, low: Any, high: Any) -> "Expr": """ return self._scalar(FUNCTIONS_COMPARISON, "between", low, high) - def is_in(self, options: Any) -> "Expr": + def is_in(self, options: Any) -> Expr: """True when this expression equals any value in ``options`` (SQL ``IN``). ``options`` is a collection of values or expressions, e.g. @@ -558,7 +558,7 @@ def is_in(self, options: Any) -> "Expr": bound = [Expr._coerce(o)._unbound for o in options] return Expr(singular_or_list(self._unbound, bound)) - def in_subquery(self, subquery: Any, alias: Union[str, None] = None) -> "Expr": + def in_subquery(self, subquery: Any, alias: str | None = None) -> Expr: """True when this expression is among a subquery's rows (``x IN (SELECT ...)``). ``subquery`` is a DataFrame producing a single output column. @@ -566,7 +566,7 @@ def in_subquery(self, subquery: Any, alias: Union[str, None] = None) -> "Expr": return Expr(_in_predicate([self._unbound], _plan_of(subquery), alias=alias)) # -- nested access ---------------------------------------------------- - def _append_segment(self, make_segment) -> "Expr": + def _append_segment(self, make_segment) -> Expr: """Append a nested ReferenceSegment child to the deepest segment.""" inner = self._unbound @@ -589,7 +589,7 @@ def resolve(base_schema, registry): return Expr(resolve) - def struct_field(self, index: int) -> "Expr": + def struct_field(self, index: int) -> Expr: """Access a nested struct field by position.""" return self._append_segment( lambda _s, _r: stalg.Expression.ReferenceSegment( @@ -597,7 +597,7 @@ def struct_field(self, index: int) -> "Expr": ) ) - def list_element(self, offset: int) -> "Expr": + def list_element(self, offset: int) -> Expr: """Access a list element by offset (also ``expr[offset]``).""" return self._append_segment( lambda _s, _r: stalg.Expression.ReferenceSegment( @@ -607,7 +607,7 @@ def list_element(self, offset: int) -> "Expr": ) ) - def __getitem__(self, offset: int) -> "Expr": + def __getitem__(self, offset: int) -> Expr: if not isinstance(offset, int) or isinstance(offset, bool): raise TypeError( "indexing selects a list element by integer offset; " @@ -615,7 +615,7 @@ def __getitem__(self, offset: int) -> "Expr": ) return self.list_element(offset) - def map_key(self, key: Any) -> "Expr": + def map_key(self, key: Any) -> Expr: """Access a map value by key.""" def make(base_schema, registry): @@ -632,7 +632,7 @@ def make(base_schema, registry): return self._append_segment(make) # -- higher-order list functions -------------------------------------- - def _higher_order(self, function: str, callback) -> "Expr": + def _higher_order(self, function: str, callback) -> Expr: """Apply a list higher-order function whose lambda ``callback`` receives an :class:`Expr` bound to the current list element.""" list_unbound = self._unbound @@ -698,7 +698,7 @@ def resolve_element(schema, _registry): return Expr(resolve) - def list_transform(self, fn) -> "Expr": + def list_transform(self, fn) -> Expr: """Map a function over each element of this list column (``transform``). ``fn`` receives an ``Expr`` for the current element, e.g. @@ -706,11 +706,11 @@ def list_transform(self, fn) -> "Expr": """ return self._higher_order("transform", fn) - def list_filter(self, fn) -> "Expr": + def list_filter(self, fn) -> Expr: """Keep list elements for which ``fn(element)`` is true (``filter``).""" return self._higher_order("filter", fn) - def switch(self, cases: dict, default: Any) -> "Expr": + def switch(self, cases: dict, default: Any) -> Expr: """Value-match CASE against literal keys:: col("code").switch({1: "one", 2: "two"}, default="other") @@ -724,7 +724,7 @@ def switch(self, cases: dict, default: Any) -> "Expr": ] return Expr(switch(self._unbound, ifs, Expr._coerce(default)._unbound)) - def distinct(self) -> "Expr": + def distinct(self) -> Expr: """Make this aggregate operate on distinct inputs (``COUNT(DISTINCT x)``). Only meaningful on an aggregate measure (e.g. ``f.count(col("x")).distinct()``). @@ -745,7 +745,7 @@ def resolve(base_schema, registry): def order_by( self, *keys: Any, descending: bool = False, nulls_last: bool = True - ) -> "Expr": + ) -> Expr: """Order the inputs to this aggregate (``string_agg(x ORDER BY ...)``). ``keys`` are column names or expressions; ``descending``/``nulls_last`` @@ -778,9 +778,9 @@ def over( *, descending: bool = False, nulls_last: bool = True, - rows: Union[tuple, None] = None, - range: Union[tuple, None] = None, - ) -> "Expr": + rows: tuple | None = None, + range: tuple | None = None, + ) -> Expr: """Turn a window function into a windowed expression (SQL ``OVER (...)``). ``partition_by`` / ``order_by`` are a column name/expression or a list of @@ -832,7 +832,7 @@ def resolve(base_schema, registry): return Expr(resolve) - def filter(self, predicate: Any) -> "Measure": + def filter(self, predicate: Any) -> Measure: """Restrict this aggregate to rows where ``predicate`` holds (SQL ``agg(x) FILTER (WHERE predicate)``). @@ -842,7 +842,7 @@ def filter(self, predicate: Any) -> "Measure": """ return Measure(self, Expr._coerce(predicate)) - def cast(self, type: Any) -> "Expr": + def cast(self, type: Any) -> Expr: """Cast this expression to ``type`` (a proto.Type or a type builder). The explicit escape hatch when automatic literal coercion is not enough, @@ -854,7 +854,7 @@ def cast(self, type: Any) -> "Expr": type = type() return Expr(cast(self._unbound, type)) - def alias(self, name: str) -> "Expr": + def alias(self, name: str) -> Expr: """Return a copy of this expression with its output name set to ``name``.""" return Expr(_alias(self._unbound, name)) @@ -875,13 +875,13 @@ def __init__(self, expr: Expr, predicate: Expr): self.expr = expr self.predicate = predicate - def alias(self, name: str) -> "Measure": + def alias(self, name: str) -> Measure: return Measure(self.expr.alias(name), self.predicate) - def distinct(self) -> "Measure": + def distinct(self) -> Measure: return Measure(self.expr.distinct(), self.predicate) - def order_by(self, *keys: Any, **kwargs: Any) -> "Measure": + def order_by(self, *keys: Any, **kwargs: Any) -> Measure: return Measure(self.expr.order_by(*keys, **kwargs), self.predicate) @@ -890,16 +890,16 @@ class When: __slots__ = ("_clauses", "_pending") - def __init__(self, clauses: list, pending: Union[Expr, None]): + def __init__(self, clauses: list, pending: Expr | None): self._clauses = clauses # list[(cond Expr, value Expr)] completed self._pending = pending # a condition Expr awaiting .then(), or None - def when(self, condition: Any) -> "When": + def when(self, condition: Any) -> When: if self._pending is not None: raise ValueError("call .then(...) before starting another .when(...)") return When(self._clauses, Expr._coerce(condition)) - def then(self, value: Any) -> "When": + def then(self, value: Any) -> When: if self._pending is None: raise ValueError(".then(...) must follow a .when(...)") return When(self._clauses + [(self._pending, Expr._coerce(value))], None) @@ -932,7 +932,7 @@ def coalesce(*exprs: Any) -> Expr: return Expr(scalar_function(FUNCTIONS_COMPARISON, "coalesce", expressions=args)) -def parameter(index: int, type: Any, alias: Union[str, None] = None) -> Expr: +def parameter(index: int, type: Any, alias: str | None = None) -> Expr: """A dynamic (runtime-bound) parameter of the given ``type``. ``index`` is the 0-based position bound via the plan's parameter bindings; @@ -943,7 +943,7 @@ def parameter(index: int, type: Any, alias: Union[str, None] = None) -> Expr: return Expr(_dynamic_parameter(index, type, alias)) -def current_timestamp(precision: int = 6, alias: Union[str, None] = None) -> Expr: +def current_timestamp(precision: int = 6, alias: str | None = None) -> Expr: """The query's execution timestamp (``precision_timestamp_tz``).""" return Expr( _execution_context_variable( @@ -956,7 +956,7 @@ def current_timestamp(precision: int = 6, alias: Union[str, None] = None) -> Exp ) -def current_date(alias: Union[str, None] = None) -> Expr: +def current_date(alias: str | None = None) -> Expr: """The query's execution date.""" return Expr( _execution_context_variable( @@ -967,7 +967,7 @@ def current_date(alias: Union[str, None] = None) -> Expr: ) -def current_timezone(alias: Union[str, None] = None) -> Expr: +def current_timezone(alias: str | None = None) -> Expr: """The query's execution timezone (a ``string``).""" return Expr( _execution_context_variable( @@ -978,18 +978,18 @@ def current_timezone(alias: Union[str, None] = None) -> Expr: ) -def scalar_subquery(subquery: Any, alias: Union[str, None] = None) -> Expr: +def scalar_subquery(subquery: Any, alias: str | None = None) -> Expr: """The single value of a one-row/one-column subquery (a DataFrame).""" return Expr(_scalar_subquery(_plan_of(subquery), alias=alias)) -def exists(subquery: Any, alias: Union[str, None] = None) -> Expr: +def exists(subquery: Any, alias: str | None = None) -> Expr: """``EXISTS (subquery)`` -- true when the subquery returns any row.""" op = stalg.Expression.Subquery.SetPredicate.PREDICATE_OP_EXISTS return Expr(_set_predicate(_plan_of(subquery), op, alias=alias)) -def unique(subquery: Any, alias: Union[str, None] = None) -> Expr: +def unique(subquery: Any, alias: str | None = None) -> Expr: """``UNIQUE (subquery)`` -- true when the subquery has no duplicate rows.""" op = stalg.Expression.Subquery.SetPredicate.PREDICATE_OP_UNIQUE return Expr(_set_predicate(_plan_of(subquery), op, alias=alias)) @@ -1007,7 +1007,7 @@ def all_(subquery: Any) -> _SubqueryReduction: return _SubqueryReduction(op, _plan_of(subquery)) -def col(name: Union[str, int]) -> Expr: +def col(name: str | int) -> Expr: """Reference an input column by name or index.""" return Expr(column(name)) @@ -1027,7 +1027,7 @@ def enum(value: str) -> EnumArg: return EnumArg(value) -def outer(name: Union[str, int], steps_out: int = 1) -> Expr: +def outer(name: str | int, steps_out: int = 1) -> Expr: """Reference a column from an enclosing query (a correlated reference). Only valid inside a subquery, e.g. a correlated ``exists``:: @@ -1041,7 +1041,7 @@ def outer(name: Union[str, int], steps_out: int = 1) -> Expr: return Expr(_outer_reference(name, steps_out)) -def lit(value: Any, type: Union[stp.Type, None] = None) -> Expr: +def lit(value: Any, type: stp.Type | None = None) -> Expr: """A literal expression. The Substrait type is inferred when omitted. Pass ``value=None`` to build a typed null; a ``type`` is required in that diff --git a/src/substrait/dataframe/extension_relations.py b/src/substrait/dataframe/extension_relations.py index a4fb889b..72579269 100644 --- a/src/substrait/dataframe/extension_relations.py +++ b/src/substrait/dataframe/extension_relations.py @@ -37,7 +37,7 @@ def to_any(self) -> Any: @classmethod @abc.abstractmethod - def from_any(cls, detail: Any) -> "ExtensionLeafDetail": + def from_any(cls, detail: Any) -> ExtensionLeafDetail: """Reconstruct a detail from its ``google.protobuf.Any``.""" @abc.abstractmethod @@ -56,7 +56,7 @@ def to_any(self) -> Any: @classmethod @abc.abstractmethod - def from_any(cls, detail: Any) -> "ExtensionSingleDetail": + def from_any(cls, detail: Any) -> ExtensionSingleDetail: """Reconstruct a detail from its ``google.protobuf.Any``.""" @abc.abstractmethod @@ -75,9 +75,9 @@ def to_any(self) -> Any: @classmethod @abc.abstractmethod - def from_any(cls, detail: Any) -> "ExtensionMultiDetail": + def from_any(cls, detail: Any) -> ExtensionMultiDetail: """Reconstruct a detail from its ``google.protobuf.Any``.""" @abc.abstractmethod - def derive_schema(self, inputs: "list[stt.Type.Struct]") -> stt.NamedStruct: + def derive_schema(self, inputs: list[stt.Type.Struct]) -> stt.NamedStruct: """The output schema given the input relations' types.""" diff --git a/src/substrait/dataframe/frame.py b/src/substrait/dataframe/frame.py index a0d52bf0..ed29b9ce 100644 --- a/src/substrait/dataframe/frame.py +++ b/src/substrait/dataframe/frame.py @@ -28,8 +28,9 @@ from __future__ import annotations +from collections.abc import Callable, Iterable from itertools import combinations -from typing import Any, Callable, Iterable, Optional, Union +from typing import Any import substrait.algebra_pb2 as stalg import substrait.plan_pb2 as stplan @@ -98,7 +99,7 @@ } -def _per_column(value: Any, n: int, name: str) -> "list[bool]": +def _per_column(value: Any, n: int, name: str) -> list[bool]: """Broadcast a bool, or validate a per-column list of bools, to length ``n``.""" if isinstance(value, (list, tuple)): if len(value) != n: @@ -110,8 +111,8 @@ def _per_column(value: Any, n: int, name: str) -> "list[bool]": def _normalize_grouping_sets( - keys: "tuple[Union[str, Expr], ...]", grouping_sets: Iterable[Iterable[Any]] -) -> "list[list[int]]": + keys: tuple[str | Expr, ...], grouping_sets: Iterable[Iterable[Any]] +) -> list[list[int]]: """Map each grouping set (of key names or positions) to index lists into ``keys``.""" name_to_index = {k: i for i, k in enumerate(keys) if isinstance(k, str)} result = [] @@ -132,14 +133,14 @@ def _normalize_grouping_sets( return result -def _split_measure(m: Union[Expr, Measure]): +def _split_measure(m: Expr | Measure): """Return ``(unbound_measure, unbound_filter_or_None)`` for an agg input.""" if isinstance(m, Measure): return _unbound(m.expr), _unbound(m.predicate) return _unbound(m), None -_default_registry: Optional[ExtensionRegistry] = None +_default_registry: ExtensionRegistry | None = None def default_registry() -> ExtensionRegistry: @@ -186,11 +187,11 @@ class LateralLeft: def __init__(self, handle: LateralInput): self._handle = handle - def col(self, name: Union[str, int]) -> Expr: + def col(self, name: str | int) -> Expr: """A correlated reference to the left row's column ``name`` (or index).""" return Expr(self._handle.column(name)) - def __getitem__(self, name: Union[str, int]) -> Expr: + def __getitem__(self, name: str | int) -> Expr: return self.col(name) @@ -202,11 +203,11 @@ class DataFrame: wraps this frame to satisfy the Narwhals backend protocol. """ - def __init__(self, plan, registry: Optional[ExtensionRegistry] = None): + def __init__(self, plan, registry: ExtensionRegistry | None = None): self._plan = plan self._registry = registry or default_registry() - def _next(self, plan) -> "DataFrame": + def _next(self, plan) -> DataFrame: return DataFrame(plan, self._registry) @property @@ -225,22 +226,20 @@ def f(self): self._functions_ns = cached return cached - def filter(self, predicate: Union[Expr, Any]) -> "DataFrame": + def filter(self, predicate: Expr | Any) -> DataFrame: return self._next(_plan.filter(self._plan, expression=_unbound(predicate))) - def select(self, *columns: Union[str, Expr]) -> "DataFrame": + def select(self, *columns: str | Expr) -> DataFrame: return self._next( _plan.select(self._plan, expressions=[_unbound(c) for c in columns]) ) - def with_columns( - self, *exprs: Union[str, Expr], **named: Union[Expr, Any] - ) -> "DataFrame": + def with_columns(self, *exprs: str | Expr, **named: Expr | Any) -> DataFrame: expressions = [_unbound(e) for e in exprs] expressions += [Expr._coerce(v).alias(k).unbound for k, v in named.items()] return self._next(_plan.project(self._plan, expressions=expressions)) - def rename(self, mapping: dict) -> "DataFrame": + def rename(self, mapping: dict) -> DataFrame: """Rename columns via a ``{old: new}`` mapping; others pass through. Implemented as a projection selecting every column, aliasing the renamed @@ -262,7 +261,7 @@ def resolve(registry: ExtensionRegistry): return self._next(build_scoped(resolve)) - def drop(self, *columns: str) -> "DataFrame": + def drop(self, *columns: str) -> DataFrame: """Drop the named columns, keeping the rest in their original order.""" drop_set = set(columns) inner = self._plan @@ -282,12 +281,12 @@ def resolve(registry: ExtensionRegistry): def unpivot( self, - on: Union[str, Iterable[str]], - index: Union[str, Iterable[str]] = (), + on: str | Iterable[str], + index: str | Iterable[str] = (), *, variable_name: str = "variable", value_name: str = "value", - ) -> "DataFrame": + ) -> DataFrame: """Unpivot ``on`` columns into ``variable``/``value`` rows (an ExpandRel). ``index`` columns are repeated on every output row. The ``on`` columns @@ -309,10 +308,10 @@ def unpivot( def sort( self, - *columns: Union[str, Expr], - descending: Union[bool, "list[bool]"] = False, - nulls_last: Union[bool, "list[bool]"] = True, - ) -> "DataFrame": + *columns: str | Expr, + descending: bool | list[bool] = False, + nulls_last: bool | list[bool] = True, + ) -> DataFrame: """Order rows by one or more columns. ``descending`` and ``nulls_last`` are each either a single bool applied @@ -329,7 +328,7 @@ def sort( ] return self._next(_plan.sort(self._plan, expressions=expressions)) - def limit(self, n: int, offset: int = 0) -> "DataFrame": + def limit(self, n: int, offset: int = 0) -> DataFrame: return self._next( _plan.fetch( self._plan, @@ -338,11 +337,11 @@ def limit(self, n: int, offset: int = 0) -> "DataFrame": ) ) - def head(self, n: int = 5) -> "DataFrame": + def head(self, n: int = 5) -> DataFrame: """The first ``n`` rows (alias for ``limit(n)``).""" return self.limit(n) - def offset(self, n: int) -> "DataFrame": + def offset(self, n: int) -> DataFrame: """Skip the first ``n`` rows, keeping all the rest.""" return self._next( _plan.fetch(self._plan, offset=lit(n, _type.i64()).unbound, count=None) @@ -351,13 +350,13 @@ def offset(self, n: int) -> "DataFrame": def top_n( self, n: int, - by: Union[str, Expr, Iterable[Union[str, Expr]]], + by: str | Expr | Iterable[str | Expr], *, - descending: Union[bool, "list[bool]"] = False, - nulls_last: Union[bool, "list[bool]"] = True, + descending: bool | list[bool] = False, + nulls_last: bool | list[bool] = True, offset: int = 0, with_ties: bool = False, - ) -> "DataFrame": + ) -> DataFrame: """The top ``n`` rows ordered by ``by`` (a fused sort + fetch, TopNRel). ``descending``/``nulls_last`` follow :meth:`sort`; ``with_ties`` keeps @@ -382,12 +381,12 @@ def top_n( def join( self, - other: "DataFrame", - on: Union[Expr, Any], + other: DataFrame, + on: Expr | Any, how: str = "inner", *, - post_filter: Union[Expr, Any, None] = None, - ) -> "DataFrame": + post_filter: Expr | Any | None = None, + ) -> DataFrame: """Join with another DataFrame. ``on`` is an expression evaluated against the concatenation of the left @@ -419,19 +418,19 @@ def join( ) ) - def cross_join(self, other: "DataFrame") -> "DataFrame": + def cross_join(self, other: DataFrame) -> DataFrame: """Cartesian product with ``other`` (every left row paired with every right row).""" return self._next(_plan.cross(self._plan, other._plan)) def lateral_join( self, - right: "Callable[[LateralLeft], DataFrame]", + right: Callable[[LateralLeft], DataFrame], how: str = "inner", *, - on: Union[Expr, Any, None] = None, - post_filter: Union[Expr, Any, None] = None, - ) -> "DataFrame": + on: Expr | Any | None = None, + post_filter: Expr | Any | None = None, + ) -> DataFrame: """Lateral join: evaluate the right frame once per row of this frame. ``right`` is a function of a :class:`LateralLeft` handle to this frame; @@ -472,8 +471,8 @@ def build_right(handle: LateralInput): ) def nested_loop_join( - self, other: "DataFrame", on: Union[Expr, Any], how: str = "inner" - ) -> "DataFrame": + self, other: DataFrame, on: Expr | Any, how: str = "inner" + ) -> DataFrame: """Physical nested-loop join: evaluate ``on`` over the Cartesian product. ``how`` accepts the same values as :meth:`join`. @@ -530,14 +529,14 @@ def _equi_join( def hash_join( self, - other: "DataFrame", - left_on: Union[str, int, Iterable[Union[str, int]]], - right_on: Union[str, int, Iterable[Union[str, int]], None] = None, + other: DataFrame, + left_on: str | int | Iterable[str | int], + right_on: str | int | Iterable[str | int] | None = None, how: str = "inner", *, - post_filter: Union[Expr, Any, None] = None, - residual: Union[Expr, Any, None] = None, - ) -> "DataFrame": + post_filter: Expr | Any | None = None, + residual: Expr | Any | None = None, + ) -> DataFrame: """Physical hash equi-join on key columns. ``left_on``/``right_on`` are column names/indices; ``right_on`` defaults @@ -560,14 +559,14 @@ def hash_join( def merge_join( self, - other: "DataFrame", - left_on: Union[str, int, Iterable[Union[str, int]]], - right_on: Union[str, int, Iterable[Union[str, int]], None] = None, + other: DataFrame, + left_on: str | int | Iterable[str | int], + right_on: str | int | Iterable[str | int] | None = None, how: str = "inner", *, - post_filter: Union[Expr, Any, None] = None, - residual: Union[Expr, Any, None] = None, - ) -> "DataFrame": + post_filter: Expr | Any | None = None, + residual: Expr | Any | None = None, + ) -> DataFrame: """Physical sort-merge equi-join on key columns (inputs assumed sorted). ``post_filter`` and ``residual`` behave as in :meth:`hash_join`. @@ -583,15 +582,15 @@ def merge_join( residual, ) - def repartition(self, n: int = 0) -> "DataFrame": + def repartition(self, n: int = 0) -> DataFrame: """Redistribute rows round-robin into ``n`` partitions (an ExchangeRel).""" return self._next(_plan.exchange(self._plan, partition_count=n)) - def broadcast(self) -> "DataFrame": + def broadcast(self) -> DataFrame: """Broadcast every row to all partitions (an ExchangeRel).""" return self._next(_plan.exchange(self._plan, broadcast=True)) - def extension(self, detail: Any) -> "DataFrame": + def extension(self, detail: Any) -> DataFrame: """Apply a custom single-input relation (ExtensionSingleRel). ``detail`` is an @@ -603,9 +602,7 @@ class via ``ExtensionRegistry.register_extension_relation`` for schema """ return self._next(_plan.extension_single(self._plan, detail)) - def extension_multi( - self, others: Iterable["DataFrame"], detail: Any - ) -> "DataFrame": + def extension_multi(self, others: Iterable[DataFrame], detail: Any) -> DataFrame: """A custom multi-input relation (ExtensionMultiRel) over this frame and ``others``; ``detail`` is an :class:`~substrait.dataframe.extension_relations.ExtensionMultiDetail`.""" @@ -615,11 +612,11 @@ def extension_multi( def hint( self, *, - row_count: Optional[float] = None, - record_size: Optional[float] = None, - alias: Optional[str] = None, - output_names: Optional[Iterable[str]] = None, - ) -> "DataFrame": + row_count: float | None = None, + record_size: float | None = None, + alias: str | None = None, + output_names: Iterable[str] | None = None, + ) -> DataFrame: """Attach non-semantic hints to the current relation (``RelCommon.Hint``). ``row_count`` / ``record_size`` are optimizer statistics; ``alias`` names @@ -655,7 +652,7 @@ def resolve(registry: ExtensionRegistry): return self._next(build_scoped(resolve)) - def cache(self) -> "DataFrame": + def cache(self) -> DataFrame: """Mark this DataFrame as a reusable common subplan (a CTE). The returned frame is emitted as a shared subtree (a leading ``rel`` entry @@ -670,7 +667,7 @@ def cache(self) -> "DataFrame": """ return self._next(_plan.reference(self._plan)) - def with_execution_behavior(self, variable_eval_mode: str) -> "DataFrame": + def with_execution_behavior(self, variable_eval_mode: str) -> DataFrame: """Set how often execution context variables are evaluated (plan-level). ``variable_eval_mode`` is ``"per_plan"`` (evaluate once for the whole @@ -688,7 +685,7 @@ def with_execution_behavior(self, variable_eval_mode: str) -> "DataFrame": ) from None return self._next(_plan.with_execution_behavior(self._plan, eval_mode)) - def union(self, *others: "DataFrame", distinct: bool = False) -> "DataFrame": + def union(self, *others: DataFrame, distinct: bool = False) -> DataFrame: """Concatenate rows of this DataFrame with ``others``. Defaults to keeping duplicates (``UNION ALL``); pass ``distinct=True`` @@ -701,7 +698,7 @@ def union(self, *others: "DataFrame", distinct: bool = False) -> "DataFrame": ) return self._set(others, op) - def intersect(self, *others: "DataFrame", distinct: bool = True) -> "DataFrame": + def intersect(self, *others: DataFrame, distinct: bool = True) -> DataFrame: """Rows present in this DataFrame and in every ``other`` (SQL ``INTERSECT``). Pass ``distinct=False`` for ``INTERSECT ALL``.""" op = ( @@ -711,7 +708,7 @@ def intersect(self, *others: "DataFrame", distinct: bool = True) -> "DataFrame": ) return self._set(others, op) - def except_(self, *others: "DataFrame", distinct: bool = True) -> "DataFrame": + def except_(self, *others: DataFrame, distinct: bool = True) -> DataFrame: """Rows in this DataFrame excluding any found in ``others`` (SQL ``EXCEPT``). Pass ``distinct=False`` for ``EXCEPT ALL``.""" op = ( @@ -721,9 +718,7 @@ def except_(self, *others: "DataFrame", distinct: bool = True) -> "DataFrame": ) return self._set(others, op) - def _set( - self, others: "tuple[DataFrame, ...]", op: "stalg.SetRel.SetOp" - ) -> "DataFrame": + def _set(self, others: tuple[DataFrame, ...], op: stalg.SetRel.SetOp) -> DataFrame: if not others: raise ValueError("a set operation needs at least one other DataFrame") inputs = [self._plan, *(o._plan for o in others)] @@ -731,9 +726,9 @@ def _set( def group_by( self, - *keys: Union[str, Expr], - grouping_sets: Optional[Iterable[Iterable[Any]]] = None, - ) -> "GroupBy": + *keys: str | Expr, + grouping_sets: Iterable[Iterable[Any]] | None = None, + ) -> GroupBy: """Begin an aggregation; follow with ``.agg(...)``. ``grouping_sets`` optionally supplies explicit GROUPING SETS as lists of @@ -747,13 +742,13 @@ def group_by( ) return GroupBy(self, keys, sets) - def rollup(self, *keys: Union[str, Expr]) -> "GroupBy": + def rollup(self, *keys: str | Expr) -> GroupBy: """Aggregate over the ROLLUP of ``keys``: the grouping sets ``(k0..kn), (k0..kn-1), ..., (k0), ()``.""" sets = [list(range(i)) for i in range(len(keys), -1, -1)] return GroupBy(self, keys, sets) - def cube(self, *keys: Union[str, Expr]) -> "GroupBy": + def cube(self, *keys: str | Expr) -> GroupBy: """Aggregate over the CUBE of ``keys``: every subset of the keys.""" n = len(keys) sets = [ @@ -763,17 +758,17 @@ def cube(self, *keys: Union[str, Expr]) -> "GroupBy": def aggregate( self, - group_by: Union[str, Expr, Iterable[Union[str, Expr]]] = (), - *measures: Union[Expr, Measure], - ) -> "DataFrame": + group_by: str | Expr | Iterable[str | Expr] = (), + *measures: Expr | Measure, + ) -> DataFrame: """One-shot aggregation. See also the fluent ``group_by().agg()``.""" if isinstance(group_by, (str, Expr)): group_by = [group_by] return GroupBy(self, tuple(group_by), None).agg(*measures) def write_named_table( - self, name: Union[str, Iterable[str]], *, mode: str = "error", op: str = "ctas" - ) -> "DataFrame": + self, name: str | Iterable[str], *, mode: str = "error", op: str = "ctas" + ) -> DataFrame: """Write these rows to a named table (a ``WriteRel`` sink). ``op`` is the write operation: ``ctas`` (create-table-as-select, @@ -799,7 +794,7 @@ def write_named_table( ) ) - def _finalize(self, registry: Optional[ExtensionRegistry]) -> stplan.Plan: + def _finalize(self, registry: ExtensionRegistry | None) -> stplan.Plan: """Build the plan and normalize it for output. The DataFrame layer emits correlated (outer) references in the id-based form (``rel_reference``): a lateral join's builder assigns them directly, and any offset-based @@ -815,7 +810,7 @@ def to_plan(self) -> stplan.Plan: return self._finalize(self._registry) # Kept for parity with the substrait.narwhals (Narwhals) wrapper's API. - def to_substrait(self, registry: Optional[ExtensionRegistry] = None) -> stplan.Plan: + def to_substrait(self, registry: ExtensionRegistry | None = None) -> stplan.Plan: return self._finalize(registry or self._registry) @@ -825,14 +820,14 @@ class GroupBy: def __init__( self, df: DataFrame, - keys: Iterable[Union[str, Expr]], - grouping_sets: Optional["list[list[int]]"] = None, + keys: Iterable[str | Expr], + grouping_sets: list[list[int]] | None = None, ): self._df = df self._keys = list(keys) self._grouping_sets = grouping_sets - def agg(self, *measures: Union[Expr, Measure]) -> DataFrame: + def agg(self, *measures: Expr | Measure) -> DataFrame: unbound, filters = ( zip(*(_split_measure(m) for m in measures)) if measures else ((), ()) ) @@ -847,9 +842,7 @@ def agg(self, *measures: Union[Expr, Measure]) -> DataFrame: ) -def extension_leaf( - detail: Any, registry: Optional[ExtensionRegistry] = None -) -> DataFrame: +def extension_leaf(detail: Any, registry: ExtensionRegistry | None = None) -> DataFrame: """Start a DataFrame from a custom leaf relation (ExtensionLeafRel). ``detail`` is an @@ -860,9 +853,9 @@ def extension_leaf( def read_named_table( - name: Union[str, Iterable[str]], + name: str | Iterable[str], schema: Any, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """Start a DataFrame from a named table and its schema. @@ -876,7 +869,7 @@ def read_named_table( def from_records( data: Iterable[Any], schema: Any, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """Start a DataFrame from inline rows (a ``VirtualTable`` / VALUES clause). @@ -901,9 +894,9 @@ def from_records( def _read_local_files( - paths: Union[str, Iterable[str]], + paths: str | Iterable[str], schema: Any, - registry: Optional[ExtensionRegistry], + registry: ExtensionRegistry | None, **file_format: Any, ) -> DataFrame: ns = _to_named_struct(schema) @@ -914,9 +907,9 @@ def _read_local_files( def read_parquet( - paths: Union[str, Iterable[str]], + paths: str | Iterable[str], schema: Any, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """Read one or more Parquet files into a DataFrame.""" opts = stalg.ReadRel.LocalFiles.FileOrFiles.ParquetReadOptions() @@ -924,9 +917,9 @@ def read_parquet( def read_orc( - paths: Union[str, Iterable[str]], + paths: str | Iterable[str], schema: Any, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """Read one or more ORC files into a DataFrame.""" opts = stalg.ReadRel.LocalFiles.FileOrFiles.OrcReadOptions() @@ -934,9 +927,9 @@ def read_orc( def read_arrow( - paths: Union[str, Iterable[str]], + paths: str | Iterable[str], schema: Any, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """Read one or more Arrow IPC files into a DataFrame.""" opts = stalg.ReadRel.LocalFiles.FileOrFiles.ArrowReadOptions() @@ -944,12 +937,12 @@ def read_arrow( def read_csv( - paths: Union[str, Iterable[str]], + paths: str | Iterable[str], schema: Any, *, delimiter: str = ",", header_lines_to_skip: int = 1, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """Read one or more delimiter-separated text files (CSV/TSV) into a DataFrame.""" opts = stalg.ReadRel.LocalFiles.FileOrFiles.DelimiterSeparatedTextReadOptions( @@ -961,18 +954,18 @@ def read_csv( def read_extension_table( schema: Any, detail: Any, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """Start a DataFrame from a custom source; ``detail`` is a ``google.protobuf.Any``.""" return DataFrame(_plan.extension_table(_to_named_struct(schema), detail), registry) def create_table( - name: Union[str, Iterable[str]], + name: str | Iterable[str], schema: Any, *, replace: bool = False, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """A ``CREATE TABLE`` DDL statement (``CREATE OR REPLACE`` when ``replace``).""" op = ( @@ -990,11 +983,11 @@ def create_table( def create_view( - name: Union[str, Iterable[str]], + name: str | Iterable[str], query: DataFrame, *, replace: bool = False, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """A ``CREATE VIEW`` DDL statement backed by ``query`` (a DataFrame).""" op = ( @@ -1007,10 +1000,10 @@ def create_view( def drop_table( - name: Union[str, Iterable[str]], + name: str | Iterable[str], *, if_exists: bool = False, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """A ``DROP TABLE`` DDL statement (``DROP TABLE IF EXISTS`` when ``if_exists``).""" op = stalg.DdlRel.DDL_OP_DROP_IF_EXIST if if_exists else stalg.DdlRel.DDL_OP_DROP @@ -1018,10 +1011,10 @@ def drop_table( def drop_view( - name: Union[str, Iterable[str]], + name: str | Iterable[str], *, if_exists: bool = False, - registry: Optional[ExtensionRegistry] = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """A ``DROP VIEW`` DDL statement (``DROP VIEW IF EXISTS`` when ``if_exists``).""" op = stalg.DdlRel.DDL_OP_DROP_IF_EXIST if if_exists else stalg.DdlRel.DDL_OP_DROP @@ -1029,12 +1022,12 @@ def drop_view( def update_table( - name: Union[str, Iterable[str]], + name: str | Iterable[str], schema: Any, assignments: dict, *, - where: Union[Expr, Any, None] = None, - registry: Optional[ExtensionRegistry] = None, + where: Expr | Any | None = None, + registry: ExtensionRegistry | None = None, ) -> DataFrame: """An ``UPDATE`` statement: ``assignments`` maps a column (name or index) to a new-value expression, applied where ``where`` holds (all rows if omitted).""" diff --git a/src/substrait/derivation_expression.py b/src/substrait/derivation_expression.py index be9d9690..7d01e6ee 100644 --- a/src/substrait/derivation_expression.py +++ b/src/substrait/derivation_expression.py @@ -1,5 +1,4 @@ import operator -from typing import Optional from antlr4 import CommonTokenStream, InputStream from antlr4.error.ErrorListener import ErrorListener @@ -77,9 +76,13 @@ def _evaluate(x, values: dict): ) elif isinstance(x, SubstraitTypeParser.LiteralNumberContext): return int(x.Number().symbol.text) - elif isinstance(x, SubstraitTypeParser.ParameterNameContext): - return values[x.Identifier().symbol.text] - elif isinstance(x, SubstraitTypeParser.NumericParameterNameContext): + elif isinstance( + x, + ( + SubstraitTypeParser.ParameterNameContext, + SubstraitTypeParser.NumericParameterNameContext, + ), + ): return values[x.Identifier().symbol.text] elif isinstance(x, SubstraitTypeParser.ParenExpressionContext): return _evaluate(x.expr(), values) @@ -201,9 +204,7 @@ def _evaluate(x, values: dict): ) ) elif isinstance(parametrized_type, SubstraitTypeParser.StructContext): - types = list( - map(lambda x: _evaluate(x, values), parametrized_type.expr()) - ) + types = [_evaluate(x, values) for x in parametrized_type.expr()] return Type( struct=Type.Struct( types=types, @@ -228,11 +229,9 @@ def _evaluate(x, values: dict): ) ) elif isinstance(parametrized_type, SubstraitTypeParser.NStructContext): - names = list(map(lambda k: k.getText(), parametrized_type.Identifier())) + names = [k.getText() for k in parametrized_type.Identifier()] struct = Type.Struct( - types=list( - map(lambda k: _evaluate(k, values), parametrized_type.expr()) - ), + types=[_evaluate(k, values) for k in parametrized_type.expr()], nullability=nullability, ) return NamedStruct( @@ -286,5 +285,5 @@ def _parse(x: str): return parser.expr() -def evaluate(x: str, values: Optional[dict] = None): +def evaluate(x: str, values: dict | None = None): return _evaluate(_parse(x), values) diff --git a/src/substrait/extension_registry/__init__.py b/src/substrait/extension_registry/__init__.py index 25c416e7..2ea57067 100644 --- a/src/substrait/extension_registry/__init__.py +++ b/src/substrait/extension_registry/__init__.py @@ -22,16 +22,16 @@ "ExtensionCollector", "ExtensionRegistry", "FunctionEntry", + "FunctionType", + "UnhandledParameterizedTypeError", + "UnrecognizedSubstraitTypeError", + "_bind_type_parameter", + "_check_integer_constraint", "build_scope", "build_scoped", + "covers", "current_collector", "function_reference", - "FunctionType", "normalize_substrait_type_names", - "_check_integer_constraint", "types_equal", - "_bind_type_parameter", - "covers", - "UnrecognizedSubstraitTypeError", - "UnhandledParameterizedTypeError", ] diff --git a/src/substrait/extension_registry/collector.py b/src/substrait/extension_registry/collector.py index 640e3108..e3128c52 100644 --- a/src/substrait/extension_registry/collector.py +++ b/src/substrait/extension_registry/collector.py @@ -22,7 +22,6 @@ import contextvars import functools import itertools -from typing import Optional, Union import substrait.extended_expression_pb2 as stee import substrait.extensions.extensions_pb2 as ste @@ -34,9 +33,9 @@ # The name is the compound form carried by SimpleExtensionDeclaration (e.g. # "add:i64_i64"), which is what makes an identity resolvable without the catalog. # The URN is None for a declaration that names no resolvable extension URN. -FunctionIdentity = tuple[Optional[str], str] +FunctionIdentity = tuple[str | None, str] -ExtensionCarrier = Union[stplan.Plan, stee.ExtendedExpression] +ExtensionCarrier = stplan.Plan | stee.ExtendedExpression def _identity_text(identity: FunctionIdentity) -> str: @@ -189,7 +188,7 @@ def anchor_for_urn(urn: str) -> int: _collector: contextvars.ContextVar = contextvars.ContextVar("_collector", default=None) -def current_collector() -> Optional[ExtensionCollector]: +def current_collector() -> ExtensionCollector | None: """The collector for the build in progress, or None outside a build.""" return _collector.get() diff --git a/src/substrait/extension_registry/exceptions.py b/src/substrait/extension_registry/exceptions.py index 37b05be4..0bdce244 100644 --- a/src/substrait/extension_registry/exceptions.py +++ b/src/substrait/extension_registry/exceptions.py @@ -3,10 +3,6 @@ class UnrecognizedSubstraitTypeError(Exception): """Raised when an unrecognized Substrait type is encountered.""" - pass - class UnhandledParameterizedTypeError(Exception): """Raised when an unhandled ANTLR parameterized type context is encountered.""" - - pass diff --git a/src/substrait/extension_registry/function_entry.py b/src/substrait/extension_registry/function_entry.py index d14d4bde..89539db5 100644 --- a/src/substrait/extension_registry/function_entry.py +++ b/src/substrait/extension_registry/function_entry.py @@ -1,7 +1,6 @@ """Function entry class for extension registry.""" from enum import Enum -from typing import Optional, Union from substrait.type_pb2 import Type from substrait_extensions.extensions import simple_extensions as se @@ -22,7 +21,7 @@ def __init__( self, urn: str, name: str, - impl: Union[se.Impl, se.Impl1, se.Impl2], + impl: se.Impl | se.Impl1 | se.Impl2, function_type: FunctionType = FunctionType.SCALAR, ) -> None: self.name = name @@ -48,7 +47,7 @@ def __init__( def __repr__(self) -> str: return f"{self.name}:{'_'.join(self.normalized_inputs)}" - def satisfies_signature(self, signature: tuple | list) -> Optional[Type]: + def satisfies_signature(self, signature: tuple | list) -> Type | None: """Match ``signature`` against this overload, returning its output type. ``signature`` interleaves value-operand ``Type``\\ s with enumeration @@ -86,18 +85,16 @@ def satisfies_signature(self, signature: tuple | list) -> Optional[Type]: return None try: output_type = evaluate(self.impl.return_, parameters) - except Exception: + except Exception: # noqa: BLE001 return None # return type cannot be derived for these arguments if self.nullability == se.NullabilityHandling.MIRROR and isinstance( output_type, Type ): sig_contains_nullable = any( - [ - p.__getattribute__(p.WhichOneof("kind")).nullability - == Type.NULLABILITY_NULLABLE - for p in signature - if isinstance(p, Type) - ] + p.__getattribute__(p.WhichOneof("kind")).nullability + == Type.NULLABILITY_NULLABLE + for p in signature + if isinstance(p, Type) ) kind = output_type.WhichOneof("kind") if kind is not None: diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py index c7e9e117..c2e37893 100644 --- a/src/substrait/extension_registry/registry.py +++ b/src/substrait/extension_registry/registry.py @@ -4,7 +4,6 @@ from collections import defaultdict from importlib.resources import files as importlib_files from pathlib import Path -from typing import Optional, Union try: # Python 3.11+ from importlib.resources.abc import Traversable @@ -57,7 +56,7 @@ def __init__(self, load_default_extensions=True) -> None: def register_extension_yaml( self, - fname: Union[str, Path, Traversable], + fname: str | Path | Traversable, ) -> None: """Register extensions from a YAML file. Args: @@ -164,7 +163,7 @@ def lookup_function( urn: str, function_name: str, signature: tuple[Type] | list[Type], - ) -> Optional[tuple[FunctionEntry, Type]]: + ) -> tuple[FunctionEntry, Type] | None: """Look up a function within a specific URN.""" matches = self._find_matching_functions(function_name, signature, [urn]) return matches[0] if matches else None @@ -190,8 +189,8 @@ def find_function( self, function_name: str, signature: tuple[Type] | list[Type], - urns: Optional[list[str]] = None, - ) -> Optional[tuple[FunctionEntry, Type]]: + urns: list[str] | None = None, + ) -> tuple[FunctionEntry, Type] | None: """Find the best-matching function for ``function_name`` across ``urns``. Searches ``urns`` in order (every registered URN when ``None``) and returns diff --git a/src/substrait/extension_registry/signature_checker_helpers.py b/src/substrait/extension_registry/signature_checker_helpers.py index 14f3143d..1ddf78c6 100644 --- a/src/substrait/extension_registry/signature_checker_helpers.py +++ b/src/substrait/extension_registry/signature_checker_helpers.py @@ -1,7 +1,5 @@ """Helper functions for extension registry.""" -from typing import Dict - from substrait.type_pb2 import Type from substrait_antlr.substrait_type.SubstraitTypeParser import SubstraitTypeParser @@ -10,7 +8,7 @@ from .exceptions import UnhandledParameterizedTypeError, UnrecognizedSubstraitTypeError # Type aliases -TypeParameterMapping = Dict[str, object] +TypeParameterMapping = dict[str, object] # mapping from argument types to shortened signature names: https://substrait.io/extensions/#function-signature-compound-names diff --git a/src/substrait/narwhals/__init__.py b/src/substrait/narwhals/__init__.py index 4e35d951..1b49e3c3 100644 --- a/src/substrait/narwhals/__init__.py +++ b/src/substrait/narwhals/__init__.py @@ -3,7 +3,7 @@ from substrait.narwhals.dataframe import DataFrame from substrait.narwhals.expression import Expression -__all__ = [DataFrame, Expression] +__all__ = ["DataFrame", "Expression"] def col(name: str) -> Expression: diff --git a/src/substrait/narwhals/dataframe.py b/src/substrait/narwhals/dataframe.py index 2a45008d..bf2c62c0 100644 --- a/src/substrait/narwhals/dataframe.py +++ b/src/substrait/narwhals/dataframe.py @@ -14,7 +14,7 @@ :mod:`substrait.dataframe.frame`. """ -from typing import Iterable, Union +from collections.abc import Iterable import substrait.narwhals from substrait.builders.plan import select @@ -47,7 +47,7 @@ def __narwhals_namespace__(self): return substrait.narwhals def select( - self, *exprs: Union[Expression, Iterable[Expression]], **named_exprs: Expression + self, *exprs: Expression | Iterable[Expression], **named_exprs: Expression ) -> "DataFrame": expressions = [e.expr for e in exprs] + [ expr.alias(alias).expr for alias, expr in named_exprs.items() diff --git a/src/substrait/proto/__init__.py b/src/substrait/proto/__init__.py index 184dc46b..d36c05c7 100644 --- a/src/substrait/proto/__init__.py +++ b/src/substrait/proto/__init__.py @@ -1,5 +1,3 @@ -# ruff: noqa: F401 F403 - import substrait.algebra_pb2 as algebra import substrait.extended_expression_pb2 as extended_expression import substrait.extensions.extensions_pb2 as extensions diff --git a/src/substrait/simple_extension_utils.py b/src/substrait/simple_extension_utils.py index 738bede7..ad22f5f9 100644 --- a/src/substrait/simple_extension_utils.py +++ b/src/substrait/simple_extension_utils.py @@ -1,9 +1,7 @@ -from typing import Optional, Union - from substrait_extensions.extensions import simple_extensions as se -def build_deprecation_status(d: Optional[dict]) -> Optional[se.DeprecationStatus]: +def build_deprecation_status(d: dict | None) -> se.DeprecationStatus | None: if d is None: return None return se.DeprecationStatus( @@ -13,7 +11,7 @@ def build_deprecation_status(d: Optional[dict]) -> Optional[se.DeprecationStatus ) -def build_arg(d: dict) -> Union[se.ValueArg, se.TypeArg, se.EnumerationArg]: +def build_arg(d: dict) -> se.ValueArg | se.TypeArg | se.EnumerationArg: if "value" in d: return se.ValueArg( value=d["value"], diff --git a/src/substrait/sql/sql_to_substrait.py b/src/substrait/sql/sql_to_substrait.py index 97e317b8..335c9152 100644 --- a/src/substrait/sql/sql_to_substrait.py +++ b/src/substrait/sql/sql_to_substrait.py @@ -1,6 +1,6 @@ import random import string -from typing import Callable +from collections.abc import Callable from deepdiff import DeepDiff from sqloxide import parse_sql @@ -58,10 +58,10 @@ def translate_expression( registry: ExtensionRegistry, measures: list[UnboundExtendedExpression], groupings: list[dict], - alias: str = None, + alias: str | None = None, ) -> UnboundExtendedExpression: assert len(ast) == 1 - op = list(ast.keys())[0] + op = next(iter(ast)) if groupings: # This means we are parsing a projection after a grouping @@ -173,7 +173,7 @@ def translate_expression( def translate(ast: dict, schema_resolver: SchemaResolver, registry: ExtensionRegistry): assert len(ast) == 1 - op = list(ast.keys())[0] + op = next(iter(ast)) ast = ast[op] if op == "Query": @@ -239,7 +239,7 @@ def translate(ast: dict, schema_resolver: SchemaResolver, registry: ExtensionReg registry=registry, ) - join_type = list(_join["join_operator"].keys())[0] + join_type = next(iter(_join["join_operator"])) expression = translate_expression( _join["join_operator"][join_type]["On"], @@ -253,7 +253,7 @@ def translate(ast: dict, schema_resolver: SchemaResolver, registry: ExtensionReg relation, right, expression, join_type_mapping[join_type] )(registry) - if "selection" in ast and ast["selection"]: + if ast.get("selection"): where_expression = translate_expression( ast["selection"], schema_resolver=schema_resolver, diff --git a/src/substrait/type_inference.py b/src/substrait/type_inference.py index f58f6492..b5a071fa 100644 --- a/src/substrait/type_inference.py +++ b/src/substrait/type_inference.py @@ -1,6 +1,5 @@ import contextlib import contextvars -from typing import Optional import substrait.algebra_pb2 as stalg import substrait.extended_expression_pb2 as stee @@ -23,7 +22,7 @@ class _SubtreeScope: callers that pass a plain list of subtrees still work unchanged. """ - __slots__ = ("_rels", "_schemas", "_resolving") + __slots__ = ("_rels", "_resolving", "_schemas") def __init__(self, rels): self._rels = list(rels) @@ -82,14 +81,14 @@ class _AnchorScope: largest cost of assembling a long pipeline. """ - __slots__ = ("_rels", "_index", "_subtrees", "_schemas", "_resolving", "_parent") + __slots__ = ("_index", "_parent", "_rels", "_resolving", "_schemas", "_subtrees") def __init__(self, index, subtrees, *, parent=None): # ``index`` is always a factory, never the index itself: a caller with an # already-built dict passes ``dict`` or a lambda, which keeps this from having # to tell the two apart -- a test that got it wrong would report the mistake as # an uncallable-object TypeError from deep inside a later lookup. - self._rels: Optional[dict] = None + self._rels: dict | None = None self._index = index self._subtrees = subtrees self._schemas: dict = {} @@ -227,7 +226,7 @@ class _SchemaMemo: results wholesale would not have that property. """ - __slots__ = ("_structs", "_pending", "_resolving") + __slots__ = ("_pending", "_resolving", "_structs") def __init__(self) -> None: # id(Rel) -> (Rel, struct) resolved, and id(Rel) -> (Rel, Plan) not yet. The @@ -306,7 +305,7 @@ def _release_boundaries_of( if unresolved is not None: unswept.append((unresolved[1], False)) - def struct_of(self, rel: stalg.Rel, registry) -> Optional[stt.Type.Struct]: + def struct_of(self, rel: stalg.Rel, registry) -> stt.Type.Struct | None: """``rel``'s remembered output struct, or None if nothing was recorded.""" key = id(rel) known = self._structs.get(key) diff --git a/src/substrait/utils/__init__.py b/src/substrait/utils/__init__.py index 2ed7ce88..c6a0e9a8 100644 --- a/src/substrait/utils/__init__.py +++ b/src/substrait/utils/__init__.py @@ -2,7 +2,7 @@ Utility and debugging functions for Substrait. """ -from typing import Iterable +from collections.abc import Iterable import substrait.algebra_pb2 as stalg import substrait.extensions.extensions_pb2 as ste @@ -650,9 +650,12 @@ def convert_rel(rel, scope): for name, expr in _iter_named_direct_expressions(node): if single_input is not None: binding = single_input - elif reshaped_join and name in _JOIN_COMBINED_SCOPED_FIELDS: - binding = None - elif projected_read and name in _READ_FILTER_FIELDS: + elif ( + reshaped_join + and name in _JOIN_COMBINED_SCOPED_FIELDS + or projected_read + and name in _READ_FILTER_FIELDS + ): binding = None else: binding = rel diff --git a/src/substrait/utils/display.py b/src/substrait/utils/display.py index f42b590d..535ec270 100644 --- a/src/substrait/utils/display.py +++ b/src/substrait/utils/display.py @@ -56,7 +56,10 @@ class PlanPrinter: """Concise pretty printer for Substrait plans and expressions""" def __init__( - self, indent_size: int = 2, show_metadata: bool = False, use_colors: bool = None + self, + indent_size: int = 2, + show_metadata: bool = False, + use_colors: bool | None = None, ): self.indent_size = indent_size self.show_metadata = show_metadata @@ -87,10 +90,7 @@ def _detect_color_support(self) -> bool: return False # Check if we're in Windows Command Prompt (which has limited color support) - if os.name == "nt" and "ANSICON" not in os.environ: - return False - - return True + return not (os.name == "nt" and "ANSICON" not in os.environ) def _color(self, text: str, color: str) -> str: """Apply color to text if colors are enabled""" @@ -395,7 +395,7 @@ def _stream_extension_single_rel( stream.write( f"{self._get_indent_with_arrow(depth + 2)}\n" ) - except Exception: + except Exception: # noqa: BLE001 stream.write( f"{self._get_indent_with_arrow(depth + 2)}\n" ) @@ -434,7 +434,7 @@ def _stream_extension_multi_rel( stream.write( f"{self._get_indent_with_arrow(depth + 2)}\n" ) - except Exception: + except Exception: # noqa: BLE001 stream.write( f"{self._get_indent_with_arrow(depth + 2)}\n" ) diff --git a/tests/builders/plan/test_lateral_join.py b/tests/builders/plan/test_lateral_join.py index 19a7ed5b..115421c1 100644 --- a/tests/builders/plan/test_lateral_join.py +++ b/tests/builders/plan/test_lateral_join.py @@ -185,11 +185,10 @@ def test_lateral_join_post_join_filter_on_dropped_side_raises(): # A left-semi lateral join drops the right side from its output, so a # post_join_filter on a right column cannot resolve -- it fails fast rather # than emitting a dangling reference against the combined inputs. - with pytest.raises(ValueError, match="not in list"): - with fresh_rel_anchors(): - lateral_join( - _left(), - _uncorrelated_right, - type=stalg.JoinRel.JOIN_TYPE_LEFT_SEMI, - post_join_filter=column("w"), # right-only, absent from output - )(registry) + with pytest.raises(ValueError, match="not in list"), fresh_rel_anchors(): + lateral_join( + _left(), + _uncorrelated_right, + type=stalg.JoinRel.JOIN_TYPE_LEFT_SEMI, + post_join_filter=column("w"), # right-only, absent from output + )(registry) diff --git a/tests/builders/plan/test_schema_memo.py b/tests/builders/plan/test_schema_memo.py index acd9f5df..748c4117 100644 --- a/tests/builders/plan/test_schema_memo.py +++ b/tests/builders/plan/test_schema_memo.py @@ -17,7 +17,7 @@ import substrait.plan_pb2 as stp import substrait.type_pb2 as stt -import substrait.type_inference as type_inference +from substrait import type_inference from substrait.builders.extended_expression import column, literal from substrait.builders.plan import ( cross, diff --git a/tests/dataframe/test_dtypes.py b/tests/dataframe/test_dtypes.py index a49db0c2..b4d9f015 100644 --- a/tests/dataframe/test_dtypes.py +++ b/tests/dataframe/test_dtypes.py @@ -237,7 +237,7 @@ def test_two_column_mismatch_still_raises_without_cast(): struct=struct(types=[i32(nullable=False), i64(nullable=False)], nullable=False), ) # Coercion only applies to literals, not between two columns. - with pytest.raises(Exception): + with pytest.raises(Exception, match="No matching overload"): _project_expr(ns, col("a") + col("b")) @@ -339,7 +339,7 @@ def test_decimal_compared_to_decimal_literal_coerces_to_column_type(): def test_decimal_plus_float_literal_still_raises(): # A float against a decimal peer stays fp and does not resolve; pass a Decimal. - with pytest.raises(Exception): + with pytest.raises(Exception, match="No matching overload"): _project_expr(_dec_ns(DEC), col("a") + 1.5) diff --git a/tests/dataframe/test_functions.py b/tests/dataframe/test_functions.py index 5b37553d..517a50a1 100644 --- a/tests/dataframe/test_functions.py +++ b/tests/dataframe/test_functions.py @@ -46,12 +46,12 @@ def test_keyword_names_are_suffixed_and_raw_reachable(): def test_unknown_function_raises_attributeerror(): with pytest.raises(AttributeError, match="no Substrait function"): - sub.f.definitely_not_a_function + _ = sub.f.definitely_not_a_function def test_dunder_access_does_not_trigger_build(): with pytest.raises(AttributeError): - sub.f.__wrapped__ + _ = sub.f.__wrapped__ # --------------------------------------------------------------------------- diff --git a/tests/dataframe/test_literals.py b/tests/dataframe/test_literals.py index 946803d8..a1336b9e 100644 --- a/tests/dataframe/test_literals.py +++ b/tests/dataframe/test_literals.py @@ -73,7 +73,7 @@ def test_every_concrete_type_kind_can_build_a_literal(): "fixed_char": ("x", t.fixed_char(1)), "varchar": ("x", t.var_char(1)), "fixed_binary": (b"x", t.fixed_binary(1)), - "decimal": (Decimal("1"), t.decimal(0, 10)), + "decimal": (Decimal(1), t.decimal(0, 10)), "precision_time": (0, t.precision_time(6)), "precision_timestamp": (0, t.precision_timestamp(6)), "precision_timestamp_tz": (0, t.precision_timestamp_tz(6)), @@ -230,7 +230,7 @@ def test_lit_decimal_infers_scale_and_precision(): (Decimal("5E2"), 0, 3), # unscaled 500 (Decimal("2E1"), 0, 2), # unscaled 20 (Decimal("123.45"), 2, 5), # unscaled 12345 - (Decimal("100"), 0, 3), + (Decimal(100), 0, 3), ], ) def test_lit_decimal_precision_covers_unscaled_value(value, scale, precision): diff --git a/tests/extension_registry/test_collector.py b/tests/extension_registry/test_collector.py index 8326d78c..8f1ad4b6 100644 --- a/tests/extension_registry/test_collector.py +++ b/tests/extension_registry/test_collector.py @@ -284,10 +284,9 @@ def test_outside_a_build_scope_is_an_error(self): function_reference(ARITHMETIC, "add:i64_i64") def test_nested_scope_defers_to_the_owner(self): - with build_scope() as (outer, outer_owns): - with build_scope() as (inner, inner_owns): - assert outer_owns and not inner_owns - assert inner is outer + with build_scope() as (outer, outer_owns), build_scope() as (inner, inner_owns): + assert outer_owns and not inner_owns + assert inner is outer class TestPlanLocalAnchors: diff --git a/tests/extension_registry/test_registry_lookup.py b/tests/extension_registry/test_registry_lookup.py index 32d42a4a..1e8aeba4 100644 --- a/tests/extension_registry/test_registry_lookup.py +++ b/tests/extension_registry/test_registry_lookup.py @@ -1,6 +1,6 @@ from substrait.type_pb2 import Type -from substrait.builders.type import decimal, i8, i16 +from substrait.builders.type import boolean, decimal, i8, i16 def test_non_existing_urn(registry): @@ -48,7 +48,7 @@ def test_wildcard_match(registry): assert registry.lookup_function( urn="extension:test:functions", function_name="add", - signature=[i8(nullable=False), i8(nullable=False), bool()], + signature=[i8(nullable=False), i8(nullable=False), boolean(nullable=False)], )[1] == Type(i16=Type.I16(nullability=Type.NULLABILITY_REQUIRED)) diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py index 25730448..54b98951 100644 --- a/tests/test_type_inference.py +++ b/tests/test_type_inference.py @@ -1569,7 +1569,7 @@ def test_anchor_index_is_built_only_when_a_rel_reference_needs_it(monkeypatch): # call. Plans carrying an id-based OuterReference are the exception, and the # builders re-infer their input's schema at every level, so the index is built on # demand rather than for every inference. - import substrait.type_inference as type_inference + from substrait import type_inference iter_plan_rels = type_inference.iter_plan_rels indexed = []