diff --git a/.kokoro/system.sh b/.kokoro/system.sh index 91a35d1f22b9..2c75e048d265 100755 --- a/.kokoro/system.sh +++ b/.kokoro/system.sh @@ -43,7 +43,8 @@ run_package_test() { local PROJECT_ID local GOOGLE_APPLICATION_CREDENTIALS local NOX_FILE - local NOX_SESSION + # Inherit NOX_SESSION from environment to allow configs (like prerelease.cfg) to pass it in + local NOX_SESSION="${NOX_SESSION}" echo "------------------------------------------------------------" echo "Configuring environment for: ${package_name}" @@ -66,7 +67,8 @@ run_package_test() { PROJECT_ID=$(cat "${KOKORO_GFILE_DIR}/project-id.json") GOOGLE_APPLICATION_CREDENTIALS="${KOKORO_GFILE_DIR}/service-account.json" NOX_FILE="noxfile.py" - NOX_SESSION="system-3.12" + # Use inherited NOX_SESSION if set, otherwise fallback to system-3.12 + NOX_SESSION="${NOX_SESSION:-system-3.12}" ;; esac diff --git a/packages/bigframes/bigframes/core/blocks.py b/packages/bigframes/bigframes/core/blocks.py index 433d2f520a57..a23965dd1bef 100644 --- a/packages/bigframes/bigframes/core/blocks.py +++ b/packages/bigframes/bigframes/core/blocks.py @@ -2796,6 +2796,7 @@ def _is_monotonic( ) block = block.drop_columns([equal_monotonic_id, strict_monotonic_id]) + assert last_result_id is not None block, monotonic_result_id = block.apply_binary_op( last_result_id, last_notna_id, diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/aggregate_compiler.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/aggregate_compiler.py index 7d9510ce944d..94607bf04bcd 100644 --- a/packages/bigframes/bigframes/core/compile/ibis_compiler/aggregate_compiler.py +++ b/packages/bigframes/bigframes/core/compile/ibis_compiler/aggregate_compiler.py @@ -528,8 +528,9 @@ def _( column: ibis_types.Column, window=None, ) -> ibis_types.Value: + # Ibis FirstNonNullValue expects Value[Any, Columnar], Mypy struggles to see Column as compatible. return _apply_window_if_present( - ibis_ops.FirstNonNullValue(column).to_expr(), + ibis_ops.FirstNonNullValue(column).to_expr(), # type: ignore[arg-type] window, # type: ignore ) @@ -549,8 +550,9 @@ def _( column: ibis_types.Column, window=None, ) -> ibis_types.Value: + # Ibis LastNonNullValue expects Value[Any, Columnar], Mypy struggles to see Column as compatible. return _apply_window_if_present( - ibis_ops.LastNonNullValue(column).to_expr(), + ibis_ops.LastNonNullValue(column).to_expr(), # type: ignore[arg-type] window, # type: ignore ) @@ -803,8 +805,9 @@ def _to_ibis_boundary( ) -> Optional[ibis_expr_window.WindowBoundary]: if boundary is None: return None + # WindowBoundary expects Value[Any, Any], ibis_types.literal returns Scalar which Mypy doesn't see as compatible. return ibis_expr_window.WindowBoundary( - abs(boundary), + ibis_types.literal(boundary if boundary >= 0 else -boundary), # type: ignore[arg-type] preceding=boundary <= 0, # type:ignore ) diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/ibis_compiler.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/ibis_compiler.py index 1f29a253d550..d52a9e381b53 100644 --- a/packages/bigframes/bigframes/core/compile/ibis_compiler/ibis_compiler.py +++ b/packages/bigframes/bigframes/core/compile/ibis_compiler/ibis_compiler.py @@ -49,8 +49,8 @@ def compile_sql(request: configs.CompileRequest) -> configs.CompileResult: # Can only pullup slice if we are doing ORDER BY in outermost SELECT # Need to do this before replacing unsupported ops, as that will rewrite slice ops result_node = rewrites.pull_up_limits(result_node) - result_node = _replace_unsupported_ops(result_node) - result_node = result_node.bottom_up(rewrites.simplify_join) + result_node = cast(nodes.ResultNode, _replace_unsupported_ops(result_node)) + result_node = cast(nodes.ResultNode, result_node.bottom_up(rewrites.simplify_join)) # prune before pulling up order to avoid unnnecessary row_number() ops result_node = cast(nodes.ResultNode, rewrites.column_pruning(result_node)) result_node = rewrites.defer_order( diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/operations/geo_ops.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/operations/geo_ops.py index 32c368ff55fc..d52d982ceb2a 100644 --- a/packages/bigframes/bigframes/core/compile/ibis_compiler/operations/geo_ops.py +++ b/packages/bigframes/bigframes/core/compile/ibis_compiler/operations/geo_ops.py @@ -182,7 +182,7 @@ def st_buffer( @ibis_udf.scalar.builtin def st_distance( - a: ibis_dtypes.geography, b: ibis_dtypes.geography, use_spheroid: bool + a: ibis_dtypes.geography, b: ibis_dtypes.geography, use_spheroid: bool # type: ignore ) -> ibis_dtypes.float: # type: ignore """Convert string to geography.""" diff --git a/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py b/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py index 1331fff1f26c..7655ef62f3d3 100644 --- a/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py +++ b/packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py @@ -2168,9 +2168,12 @@ def obj_make_ref_json(objectref_json: ibis_dtypes.JSON) -> _OBJ_REF_IBIS_DTYPE: @ibis_udf.scalar.builtin(name="OBJ.GET_ACCESS_URL") -def obj_get_access_url( - obj_ref: _OBJ_REF_IBIS_DTYPE, mode: ibis_dtypes.String -) -> ibis_dtypes.JSON: # type: ignore +# Stub for BigQuery UDF, empty body is intentional. +# _OBJ_REF_IBIS_DTYPE is a variable holding a type, Mypy complains about it being used as type hint. +def obj_get_access_url( # type: ignore[empty-body] + obj_ref: _OBJ_REF_IBIS_DTYPE, # type: ignore[valid-type] + mode: ibis_dtypes.String +) -> ibis_dtypes.JSON: """Get access url (as ObjectRefRumtime JSON) from ObjectRef.""" diff --git a/packages/bigframes/bigframes/core/compile/sqlglot/compiler.py b/packages/bigframes/bigframes/core/compile/sqlglot/compiler.py index e343d8962d82..ba9e74a5e450 100644 --- a/packages/bigframes/bigframes/core/compile/sqlglot/compiler.py +++ b/packages/bigframes/bigframes/core/compile/sqlglot/compiler.py @@ -53,8 +53,8 @@ def compile_sql(request: configs.CompileRequest) -> configs.CompileResult: # Can only pullup slice if we are doing ORDER BY in outermost SELECT # Need to do this before replacing unsupported ops, as that will rewrite slice ops result_node = rewrite.pull_up_limits(result_node) - result_node = _replace_unsupported_ops(result_node) - result_node = result_node.bottom_up(rewrite.simplify_join) + result_node = typing.cast(nodes.ResultNode, _replace_unsupported_ops(result_node)) + result_node = typing.cast(nodes.ResultNode, result_node.bottom_up(rewrite.simplify_join)) # prune before pulling up order to avoid unnnecessary row_number() ops result_node = typing.cast(nodes.ResultNode, rewrite.column_pruning(result_node)) result_node = rewrite.defer_order( diff --git a/packages/bigframes/bigframes/core/expression_factoring.py b/packages/bigframes/bigframes/core/expression_factoring.py index b1bc5c99d457..43d518250238 100644 --- a/packages/bigframes/bigframes/core/expression_factoring.py +++ b/packages/bigframes/bigframes/core/expression_factoring.py @@ -243,7 +243,7 @@ def factor_aggregation(root: nodes.ColumnDef) -> FactoredAggregation: } root_scalar_expr = nodes.ColumnDef( - sub_expressions(root.expression, agg_outputs_dict), + sub_expressions(root.expression, cast(Mapping[expression.Expression, expression.Expression], agg_outputs_dict)), root.id, # type: ignore ) diff --git a/packages/bigframes/bigframes/core/local_data.py b/packages/bigframes/bigframes/core/local_data.py index 09111572f3c9..01d7e5570ca5 100644 --- a/packages/bigframes/bigframes/core/local_data.py +++ b/packages/bigframes/bigframes/core/local_data.py @@ -33,6 +33,7 @@ import bigframes.core.schema as schemata import bigframes.dtypes +from bigframes.core import identifiers from bigframes.core import pyarrow_utils @@ -155,7 +156,7 @@ def to_arrow( return schema, batches def is_nullable(self, column_id: identifiers.ColumnId) -> bool: - return self.data.column(column_id).null_count > 0 + return self.data.column(column_id.name).null_count > 0 def to_pyarrow_table( self, diff --git a/packages/bigframes/bigframes/core/nodes.py b/packages/bigframes/bigframes/core/nodes.py index 5297ceed9140..874ee3117f96 100644 --- a/packages/bigframes/bigframes/core/nodes.py +++ b/packages/bigframes/bigframes/core/nodes.py @@ -674,7 +674,7 @@ def fields(self) -> Sequence[Field]: Field( col_id, self.local_data_source.schema.get_type(source_id), - nullable=self.local_data_source.is_nullable(source_id), + nullable=self.local_data_source.is_nullable(identifiers.ColumnId(source_id)), ) for col_id, source_id in self.scan_list.items ) diff --git a/packages/bigframes/bigframes/core/rewrite/as_sql.py b/packages/bigframes/bigframes/core/rewrite/as_sql.py index cc4e05565203..eb823d1fed1d 100644 --- a/packages/bigframes/bigframes/core/rewrite/as_sql.py +++ b/packages/bigframes/bigframes/core/rewrite/as_sql.py @@ -291,8 +291,9 @@ def _extract_ctes_to_with_expr( root.top_down(lambda x: mapping.get(x, x)), cte_names, tuple( - cte_node.child.top_down(lambda x: mapping.get(x, x)) - for cte_node in topological_ctes # type: ignore + # Mypy loses context that cte_node is a CteNode with a child attribute, despite the isinstance filter above. + cte_node.child.top_down(lambda x: mapping.get(x, x)) # type: ignore[attr-defined] + for cte_node in topological_ctes ), ) diff --git a/packages/bigframes/bigframes/dataframe.py b/packages/bigframes/bigframes/dataframe.py index b89360c691d3..98340e5e377f 100644 --- a/packages/bigframes/bigframes/dataframe.py +++ b/packages/bigframes/bigframes/dataframe.py @@ -3926,12 +3926,13 @@ def round(self, decimals: Union[int, dict[Hashable, int]] = 0) -> DataFrame: bigframes.dtypes.BOOL_DTYPE }: if is_mapping: - if label in decimals: # type: ignore + decimals_dict = typing.cast(dict[typing.Hashable, int], decimals) + if label in decimals_dict: exprs.append( ops.round_op.as_expr( col_id, ex.const( - decimals[label], + decimals_dict[label], dtype=bigframes.dtypes.INT_DTYPE, # type: ignore ), ) @@ -4447,8 +4448,8 @@ def to_latex( ) -> str | None: return self.to_pandas(allow_large_results=allow_large_results).to_latex( buf, - columns=columns, - header=header, + columns=typing.cast(typing.Optional[list[str]], columns), + header=typing.cast(typing.Union[bool, list[str]], header), index=index, **kwargs, # type: ignore ) diff --git a/packages/bigframes/bigframes/series.py b/packages/bigframes/bigframes/series.py index fbcc949855c2..f0648117144a 100644 --- a/packages/bigframes/bigframes/series.py +++ b/packages/bigframes/bigframes/series.py @@ -2308,9 +2308,10 @@ def to_json( ) else: pd_series = self.to_pandas(allow_large_results=allow_large_results) + # Pandas Series.to_json only supports a subset of orients, but bigframes Series.to_json allows all of them. return pd_series.to_json( path_or_buf=path_or_buf, - orient=orient, + orient=orient, # type: ignore[arg-type] lines=lines, index=index, # type: ignore ) diff --git a/packages/bigframes/bigframes/session/iceberg.py b/packages/bigframes/bigframes/session/iceberg.py index 805d03aeeee3..0d2539f55545 100644 --- a/packages/bigframes/bigframes/session/iceberg.py +++ b/packages/bigframes/bigframes/session/iceberg.py @@ -98,9 +98,10 @@ def _extract_location_from_catalog_extension_data(data): class SchemaVisitor(pyiceberg.schema.SchemaVisitorPerPrimitiveType[bq.SchemaField]): - def schema( + # Override returns a tuple of fields instead of a single field, violating supertype signature but intentional for this visitor. + def schema( # type: ignore[override] self, schema: pyiceberg.schema.Schema, struct_result: bq.SchemaField - ) -> tuple[bq.SchemaField, ...]: # type: ignore + ) -> tuple[bq.SchemaField, ...]: return tuple(f for f in struct_result.fields) def struct( diff --git a/packages/bigframes/noxfile.py b/packages/bigframes/noxfile.py index cb57923288a6..b04435166410 100644 --- a/packages/bigframes/noxfile.py +++ b/packages/bigframes/noxfile.py @@ -20,6 +20,7 @@ import multiprocessing import os import pathlib +import re import shutil import time from typing import Dict, List @@ -61,6 +62,7 @@ "pytest-cov", "pytest-timeout", ] +UNIT_TEST_EXTERNAL_DEPENDENCIES: List[str] = [] UNIT_TEST_DEPENDENCIES: List[str] = [] UNIT_TEST_EXTRAS: List[str] = ["tests"] UNIT_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = { @@ -93,7 +95,7 @@ SYSTEM_TEST_EXTERNAL_DEPENDENCIES = [ "google-cloud-bigquery", ] -SYSTEM_TEST_EXTRAS: List[str] = ["tests"] +SYSTEM_TEST_EXTRAS: List[str] = [] SYSTEM_TEST_EXTRAS_BY_PYTHON: Dict[str, List[str]] = { # Make sure we leave some versions without "extras" so we know those # dependencies are actually optional. @@ -255,10 +257,14 @@ def run_unit(session, install_test_extra): @nox.session(python=ALL_PYTHON) -def unit(session): +@nox.parametrize("test_extra", [True, False]) +def unit(session, test_extra): if session.python in ("3.7", "3.8", "3.9"): session.skip("Python 3.9 and below are not supported") - run_unit(session, install_test_extra=True) + if test_extra: + run_unit(session, install_test_extra=test_extra) + else: + unit_noextras(session) @nox.session(python=ALL_PYTHON[-1]) @@ -358,17 +364,22 @@ def run_system( ) -@nox.session(python="3.12") -def system(session: nox.sessions.Session): + +@nox.session(python=SYSTEM_TEST_PYTHON_VERSIONS) +@nox.parametrize("test_extra", [True, False]) +def system(session: nox.sessions.Session, test_extra): """Run the system test suite.""" - # TODO(https://github.com/googleapis/google-cloud-python/issues/16489): Restore system test once this bug is fixed - run_system( - session=session, - prefix_name="system", - test_folder=os.path.join("tests", "system", "small"), - check_cov=True, - ) - # session.skip("Temporarily skip system test") + if session.python in ("3.7", "3.8", "3.9"): + session.skip("Python 3.9 and below are not supported") + if test_extra: + run_system( + session=session, + prefix_name="system", + test_folder=os.path.join("tests", "system", "small"), + check_cov=True, + ) + else: + system_noextras(session) @nox.session(python=DEFAULT_PYTHON_VERSION) @@ -921,69 +932,129 @@ def cleanup(session): @nox.session(python=DEFAULT_PYTHON_VERSION) -def core_deps_from_source(session): +@nox.parametrize( + "protobuf_implementation", + ["python", "upb"], +) +def core_deps_from_source(session, protobuf_implementation): """Run all tests with core dependencies installed from source rather than pulling the dependencies from PyPI. """ - # TODO(https://github.com/googleapis/google-cloud-python/issues/16014): - # Add core deps from source tests - session.skip("Core deps from source tests are not yet supported") + # Install all dependencies + session.install("-e", ".") -@nox.session(python=DEFAULT_PYTHON_VERSION) + # Install dependencies for the unit test environment + unit_deps_all = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_EXTERNAL_DEPENDENCIES + session.install(*unit_deps_all) + + # Install dependencies for the system test environment + system_deps_all = ( + SYSTEM_TEST_STANDARD_DEPENDENCIES + + SYSTEM_TEST_EXTERNAL_DEPENDENCIES + + SYSTEM_TEST_EXTRAS + ) + session.install(*system_deps_all) + + # Because we test minimum dependency versions on the minimum Python + # version, the first version we test with in the unit tests sessions has a + # constraints file containing all dependencies and extras. + with open( + CURRENT_DIRECTORY / "testing" / "constraints-3.10.txt", + encoding="utf-8", + ) as constraints_file: + constraints_text = constraints_file.read() + + # Ignore leading whitespace and comment lines. + # Fiona fails to build on GitHub CI because gdal-config is missing and no Python 3.14 wheels are available. + constraints_deps = [ + match.group(1) + for match in re.finditer( + r"^\s*(\S+)(?===\S+)", constraints_text, flags=re.MULTILINE + ) + if match.group(1) != "fiona" + ] + + # Install dependencies specified in `testing/constraints-X.txt`. + session.install(*constraints_deps) + + # TODO(https://github.com/googleapis/gapic-generator-python/issues/2358): `grpcio` and + # `grpcio-status` should be added to the list below so that they are installed from source, + # rather than PyPI. + # TODO(https://github.com/googleapis/gapic-generator-python/issues/2357): `protobuf` should be + # added to the list below so that it is installed from source, rather than PyPI + # Note: If a dependency is added to the `core_dependencies_from_source` list, + # the `prerel_deps` list in the `prerelease_deps` nox session should also be updated. + core_dependencies_from_source = [ + "googleapis-common-protos @ git+https://github.com/googleapis/google-cloud-python#egg=googleapis-common-protos&subdirectory=packages/googleapis-common-protos", + "google-api-core @ git+https://github.com/googleapis/google-cloud-python#egg=google-api-core&subdirectory=packages/google-api-core", + "google-auth @ git+https://github.com/googleapis/google-cloud-python#egg=google-auth&subdirectory=packages/google-auth", + "grpc-google-iam-v1 @ git+https://github.com/googleapis/google-cloud-python#egg=grpc-google-iam-v1&subdirectory=packages/grpc-google-iam-v1", + "proto-plus @ git+https://github.com/googleapis/google-cloud-python#egg=proto-plus&subdirectory=packages/proto-plus", + ] + + for dep in core_dependencies_from_source: + session.install(dep, "--no-deps", "--ignore-installed") + print(f"Installed {dep}") + + session.run( + "py.test", + "tests/unit", + env={ + "PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION": protobuf_implementation, + }, + ) + + +@nox.session(python=ALL_PYTHON[-1]) def prerelease_deps(session): """Run all tests with prerelease versions of dependencies installed.""" # TODO(https://github.com/googleapis/google-cloud-python/issues/16014): # Add prerelease deps tests - session.skip("prerelease deps tests are not yet supported") + unit_prerelease(session) + system_prerelease(session) -@nox.session(python=DEFAULT_PYTHON_VERSION) +# NOTE: this is based on mypy session that came directly from the bigframes split repo +# the split repo used 3.10, the monorepo uses 3.14 +@nox.session(python="3.14") def mypy(session): - """Run the type checker.""" - # TODO(https://github.com/googleapis/google-cloud-python/issues/16014): - # Add mypy tests previously used mypy session (below) failed to run in the monorepo - session.skip("mypy tests are not yet supported") - - -# @nox.session(python=ALL_PYTHON) -# def mypy(session): -# """Run type checks with mypy.""" -# # Editable mode is not compatible with mypy when there are multiple -# # package directories. See: -# # https://github.com/python/mypy/issues/10564#issuecomment-851687749 -# session.install(".") - -# # Just install the dependencies' type info directly, since "mypy --install-types" -# # might require an additional pass. -# deps = ( -# set( -# [ -# MYPY_VERSION, -# # TODO: update to latest pandas-stubs once we resolve bigframes issues. -# "pandas-stubs<=2.2.3.241126", -# "types-protobuf", -# "types-python-dateutil", -# "types-requests", -# "types-setuptools", -# "types-tabulate", -# "types-PyYAML", -# "polars", -# "anywidget", -# ] -# ) -# | set(SYSTEM_TEST_STANDARD_DEPENDENCIES) -# | set(UNIT_TEST_STANDARD_DEPENDENCIES) -# ) - -# session.install(*deps) -# shutil.rmtree(".mypy_cache", ignore_errors=True) -# session.run( -# "mypy", -# "bigframes", -# os.path.join("tests", "system"), -# os.path.join("tests", "unit"), -# "--check-untyped-defs", -# "--explicit-package-bases", -# '--exclude="^third_party"', -# ) + """Run type checks with mypy.""" + # Editable mode is not compatible with mypy when there are multiple + # package directories. See: + # https://github.com/python/mypy/issues/10564#issuecomment-851687749 + session.install(".") + + # Just install the dependencies' type info directly, since "mypy --install-types" + # might require an additional pass. + deps = ( + set( + [ + MYPY_VERSION, + # TODO: update to latest pandas-stubs once we resolve bigframes issues. + "pandas-stubs<=2.2.3.241126", + "types-protobuf", + "types-python-dateutil", + "types-requests", + "types-setuptools", + "types-tabulate", + "types-PyYAML", + "polars", + "anywidget", + ] + ) + | set(SYSTEM_TEST_STANDARD_DEPENDENCIES) + | set(UNIT_TEST_STANDARD_DEPENDENCIES) + ) + + session.install(*deps) + shutil.rmtree(".mypy_cache", ignore_errors=True) + session.run( + "mypy", + "bigframes", + os.path.join("tests", "system"), + os.path.join("tests", "unit"), + "--check-untyped-defs", + "--explicit-package-bases", + '--exclude="^third_party"', + ) diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py index 11cce647a56c..fd3aacc7e271 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/expressions/test_datetime_ops.py @@ -62,37 +62,37 @@ def test_datetime_to_integer_label(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[col_names] ops_map = { "fixed_freq": ops.DatetimeToIntegerLabelOp( - freq=pd.tseries.offsets.Day(), + freq=pd.tseries.offsets.Day(), # type: ignore[arg-type] origin="start", closed="left", # type: ignore ).as_expr("datetime_col", "timestamp_col"), "origin_epoch": ops.DatetimeToIntegerLabelOp( - freq=pd.tseries.offsets.Day(), + freq=pd.tseries.offsets.Day(), # type: ignore[arg-type] origin="epoch", closed="left", # type: ignore ).as_expr("datetime_col", "timestamp_col"), "origin_start_day": ops.DatetimeToIntegerLabelOp( - freq=pd.tseries.offsets.Day(), + freq=pd.tseries.offsets.Day(), # type: ignore[arg-type] origin="start_day", closed="left", # type: ignore ).as_expr("datetime_col", "timestamp_col"), "non_fixed_freq_weekly": ops.DatetimeToIntegerLabelOp( - freq=pd.tseries.offsets.Week(weekday=6), + freq=pd.tseries.offsets.Week(weekday=6), # type: ignore[arg-type] origin="start", closed="left", # type: ignore ).as_expr("datetime_col", "timestamp_col"), "non_fixed_freq_monthly": ops.DatetimeToIntegerLabelOp( - freq=pd.tseries.offsets.MonthEnd(), + freq=pd.tseries.offsets.MonthEnd(), # type: ignore[arg-type] origin="start", closed="left", # type: ignore ).as_expr("datetime_col", "timestamp_col"), "non_fixed_freq_quarterly": ops.DatetimeToIntegerLabelOp( - freq=pd.tseries.offsets.QuarterEnd(startingMonth=12), + freq=pd.tseries.offsets.QuarterEnd(startingMonth=12), # type: ignore[arg-type] origin="start", closed="left", # type: ignore ).as_expr("datetime_col", "timestamp_col"), "non_fixed_freq_yearly": ops.DatetimeToIntegerLabelOp( - freq=pd.tseries.offsets.YearEnd(), + freq=pd.tseries.offsets.YearEnd(), # type: ignore[arg-type] origin="start", closed="left", # type: ignore ).as_expr("datetime_col", "timestamp_col"), @@ -334,7 +334,7 @@ def test_integer_label_to_datetime_fixed(scalar_types_df: bpd.DataFrame, snapsho bf_df = scalar_types_df[col_names] ops_map = { "fixed_freq": ops.IntegerLabelToDatetimeOp( - freq=pd.tseries.offsets.Day(), + freq=pd.tseries.offsets.Day(), # type: ignore[arg-type] origin="start", label="left", # type: ignore ).as_expr("rowindex", "timestamp_col"), @@ -349,7 +349,7 @@ def test_integer_label_to_datetime_week(scalar_types_df: bpd.DataFrame, snapshot bf_df = scalar_types_df[col_names] ops_map = { "non_fixed_freq_weekly": ops.IntegerLabelToDatetimeOp( - freq=pd.tseries.offsets.Week(weekday=6), + freq=pd.tseries.offsets.Week(weekday=6), # type: ignore[arg-type] origin="start", label="left", # type: ignore ).as_expr("rowindex", "timestamp_col"), diff --git a/packages/bigframes/tests/unit/core/compile/sqlglot/tpch/conftest.py b/packages/bigframes/tests/unit/core/compile/sqlglot/tpch/conftest.py index 0fb034ac7091..8d38821eb9b7 100644 --- a/packages/bigframes/tests/unit/core/compile/sqlglot/tpch/conftest.py +++ b/packages/bigframes/tests/unit/core/compile/sqlglot/tpch/conftest.py @@ -157,7 +157,9 @@ def read_gbq_table_no_snapshot(*args, **kwargs): kwargs["enable_snapshot"] = False return original_read_gbq_table(*args, **kwargs) - session._loader.read_gbq_table = read_gbq_table_no_snapshot - session._executor = compiler_session.SQLCompilerExecutor() - return session + + with mock.patch.object( + session._loader, "read_gbq_table", new=read_gbq_table_no_snapshot + ): + yield session diff --git a/packages/bigframes/tests/unit/test_local_engine.py b/packages/bigframes/tests/unit/test_local_engine.py index 47e0360b9fef..fe5052771f2c 100644 --- a/packages/bigframes/tests/unit/test_local_engine.py +++ b/packages/bigframes/tests/unit/test_local_engine.py @@ -171,8 +171,11 @@ def test_polars_local_engine_agg(polars_session): pd_result = pd_df.agg(["sum", "count"]) # local engine appears to produce uint32 pandas.testing.assert_frame_equal( - bf_result, pd_result, check_dtype=False, check_index_type=False - ) # type: ignore + bf_result, # type: ignore[arg-type] + pd_result, + check_dtype=False, + check_index_type=False, + ) def test_polars_local_engine_groupby_sum(polars_session):