From 16b48897f3264f546db0fa3a4ca7b27f582dd110 Mon Sep 17 00:00:00 2001 From: Saurabh Maurya Date: Fri, 11 Sep 2026 15:08:52 +0000 Subject: [PATCH 1/2] feat(sql): add ai_generate, a Spark column function for Gemini Introduces google.cloud.dataproc_ml.sql.ai_generate, the Spark counterpart of BigQuery's AI.GENERATE. It calls a Gemini model once per row and returns a struct of result, full_response and status, so queries can move between the two engines with little change. The function is a pandas UDF, which means it works both as a DataFrame column expression and, through ai_generate_udf(), as a function registered with spark.udf.register for use from Spark SQL. The argument list is deliberately BigQuery's: prompt, endpoint, model_params and output_schema, with project and location standing in for connection_id. Concurrency and retry timing are implemented but kept as internal constants, as they are in BigQuery, rather than exposed as arguments that would be hard to withdraw later. Notable behaviour, all matching BigQuery: * A row that fails never fails the stage. Its error is reported in status and result is NULL. status is empty on success. * output_schema replaces result with one field per declared field, sorted by name. BigQuery type names (INT64, FLOAT64, BOOL) and their Spark spellings are both accepted. * full_response holds the complete model response, typed as VARIANT on Spark 4 and as a JSON string on older versions. Two things differ from BigQuery by necessity. There is no connection_id, because a BigQuery connection's service account cannot be impersonated from outside BigQuery; calls use Application Default Credentials instead. And the default model is gemini-3.6-flash, which is served only from the global endpoint, so the default location is global rather than a region. Tested with unit tests that exercise the full Spark path against an injected generate_content, and with integration tests that call Vertex AI for real. --- README.md | 55 ++ docs/api.rst | 4 + google/cloud/dataproc_ml/sql/__init__.py | 44 ++ google/cloud/dataproc_ml/sql/_ai_generate.py | 519 ++++++++++++++++++ google/cloud/dataproc_ml/sql/_client.py | 190 +++++++ google/cloud/dataproc_ml/sql/_model_params.py | 130 +++++ .../cloud/dataproc_ml/sql/_output_schema.py | 228 ++++++++ google/cloud/dataproc_ml/sql/_retry.py | 187 +++++++ pyproject.toml | 1 + tests/integration/sql/__init__.py | 13 + tests/integration/sql/test_ai_generate.py | 196 +++++++ tests/unit/sql/__init__.py | 13 + tests/unit/sql/test_ai_generate.py | 502 +++++++++++++++++ tests/unit/sql/test_client.py | 67 +++ tests/unit/sql/test_model_params.py | 127 +++++ tests/unit/sql/test_output_schema.py | 220 ++++++++ tests/unit/sql/test_retry.py | 149 +++++ 17 files changed, 2645 insertions(+) create mode 100644 google/cloud/dataproc_ml/sql/__init__.py create mode 100644 google/cloud/dataproc_ml/sql/_ai_generate.py create mode 100644 google/cloud/dataproc_ml/sql/_client.py create mode 100644 google/cloud/dataproc_ml/sql/_model_params.py create mode 100644 google/cloud/dataproc_ml/sql/_output_schema.py create mode 100644 google/cloud/dataproc_ml/sql/_retry.py create mode 100644 tests/integration/sql/__init__.py create mode 100644 tests/integration/sql/test_ai_generate.py create mode 100644 tests/unit/sql/__init__.py create mode 100644 tests/unit/sql/test_ai_generate.py create mode 100644 tests/unit/sql/test_client.py create mode 100644 tests/unit/sql/test_model_params.py create mode 100644 tests/unit/sql/test_output_schema.py create mode 100644 tests/unit/sql/test_retry.py diff --git a/README.md b/README.md index 2a5cab8..fdeb934 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,61 @@ pip install dataproc-ml Here are a couple of examples demonstrating how to use the handlers for distributed inference on a Spark DataFrame. +### AI Functions: `ai_generate` + +> **Note:** `ai_generate` makes API calls to Vertex AI, which will incur costs. +> Please review the [Vertex AI Generative +> AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing). + +`ai_generate` is a Spark column function that calls a Gemini model on every +row, mirroring the semantics of BigQuery's +[`AI.GENERATE`](https://cloud.google.com/bigquery/docs/reference/standard-sql/AI-functions#ai.generate). +It returns a `STRUCT` with a `result`, the raw `full_response`, and a `status` +that is empty on success. + +```python +from pyspark.sql import SparkSession, functions as F +from google.cloud.dataproc_ml.sql import ai_generate + +spark = SparkSession.builder.getOrCreate() + +df = spark.createDataFrame([("Paris",), ("Tokyo",)], ["city"]) + +result_df = df.withColumn( + "generated", + ai_generate(F.concat(F.lit("Airport code for "), F.col("city"))), +).select("city", "generated.result", "generated.status") + +result_df.show() +# +-----+------+------+ +# | city|result|status| +# +-----+------+------+ +# |Paris| CDG| | +# |Tokyo| HND| | +# +-----+------+------+ +``` + +Pass `output_schema` to get typed columns back instead of a single string: + +```python +df.withColumn( + "review", + ai_generate( + F.col("text"), + output_schema="sentiment STRING, score FLOAT64", + ), +) +``` + +The same function can be registered for use from Spark SQL: + +```python +from google.cloud.dataproc_ml.sql import ai_generate_udf + +spark.udf.register("ai_generate", ai_generate_udf()) +spark.sql("SELECT ai_generate('Airport code for Paris').result").show() +``` + ### Generative AI (Gemini) Model Inference > **Note:** Using the `GenAiModelHandler` involves making API calls to diff --git a/docs/api.rst b/docs/api.rst index 2155aa0..bdfbd43 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -5,5 +5,9 @@ This page contains the auto-generated API reference for the public modules in the ``dataproc-ml`` library. .. automodule:: google.cloud.dataproc_ml.inference + :members: + :show-inheritance: + +.. automodule:: google.cloud.dataproc_ml.sql :members: :show-inheritance: \ No newline at end of file diff --git a/google/cloud/dataproc_ml/sql/__init__.py b/google/cloud/dataproc_ml/sql/__init__.py new file mode 100644 index 0000000..98b7e6a --- /dev/null +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -0,0 +1,44 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Declarative AI functions for Spark DataFrames and Spark SQL. + +These functions apply Google's Gemini models to a DataFrame column without any +of the usual plumbing: no client setup, no retry loops and no response +parsing. They mirror the BigQuery ``AI.*`` functions, so a query written for +one engine reads the same on the other. + +Example: + >>> from pyspark.sql.functions import col + >>> from google.cloud.dataproc_ml.sql import ai_generate + >>> + >>> summarized = df.withColumn( + ... "summary", ai_generate(col("feedback")).result + ... ) + +The same function can be registered for use from Spark SQL: + + >>> from google.cloud.dataproc_ml.sql import ai_generate_udf + >>> + >>> spark.udf.register("ai_generate", ai_generate_udf()) + >>> spark.sql("SELECT ai_generate(feedback).result FROM feedback_table") +""" + +from ._ai_generate import ai_generate +from ._ai_generate import ai_generate_udf + +__all__ = ( + "ai_generate", + "ai_generate_udf", +) diff --git a/google/cloud/dataproc_ml/sql/_ai_generate.py b/google/cloud/dataproc_ml/sql/_ai_generate.py new file mode 100644 index 0000000..6a59139 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -0,0 +1,519 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``AI.GENERATE``: generative text and structured output over Spark columns.""" + +import asyncio +import json +import logging +import math +from typing import Any, Dict, Iterator, List, Optional, Sequence, Union + +import pandas as pd +from pyspark.sql import Column +from pyspark.sql import functions as sf +from pyspark.sql.pandas.functions import pandas_udf +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + StringType, + StructField, + StructType, +) + +from google.cloud.dataproc_ml.sql import _client +from google.cloud.dataproc_ml.sql import _model_params +from google.cloud.dataproc_ml.sql import _output_schema as _output_schema_util +from google.cloud.dataproc_ml.sql import _retry + +try: + # Spark 4 models semi-structured data as VARIANT, which is the closest + # equivalent of the JSON type that BigQuery's AI functions return. + from pyspark.sql.types import VariantType, VariantVal # pylint: disable=ungrouped-imports + + _VARIANT_SUPPORTED = True +except ImportError: # pragma: no cover - depends on the Spark version + _VARIANT_SUPPORTED = False + + +logger = logging.getLogger(__name__) + +#: The model used when no endpoint is given. +DEFAULT_ENDPOINT = "gemini-3.6-flash" + +# Fixed request settings. BigQuery does not expose these either; they are +# constants here so that the function behaves sensibly on a large cluster +# without asking the caller to tune anything. +_MAX_CONCURRENT_REQUESTS = 5 +_MAX_ATTEMPTS = 5 +_BASE_RETRY_DELAY_SECONDS = 1.0 +_MAX_RETRY_DELAY_SECONDS = 60.0 + +_RESULT_FIELD = "result" +_FULL_RESPONSE_FIELD = "full_response" +_STATUS_FIELD = "status" + +PromptArg = Union[Column, str, Sequence[Union[Column, str]]] + + +def ai_generate( + prompt: PromptArg, + *, + endpoint: str = DEFAULT_ENDPOINT, + model_params: Optional[Dict[str, Any]] = None, + output_schema: Optional[Union[str, StructType]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _generate_content=None, +) -> Column: + """Generates text for every row using a Gemini model. + + This is the Spark counterpart of BigQuery's ``AI.GENERATE`` function and + returns the same struct, so queries can be moved between the two engines + with little change. It takes the same arguments, except that the Vertex AI + call is authenticated with Application Default Credentials rather than a + BigQuery connection, so ``project`` and ``location`` take the place of + ``connection_id``. + + The prompt follows the usual PySpark convention: a string names a column, + and literal text must be wrapped in :func:`pyspark.sql.functions.lit`. A + sequence of parts is concatenated in order, like BigQuery's struct prompt. + + Requests are retried with exponential backoff and are issued a few at a + time per worker. Like BigQuery, this function does not expose those + settings. + + .. note:: + Using this function incurs Vertex AI charges. See the + `Vertex AI Generative AI pricing + `_ page. + + Args: + prompt: The prompt column, the name of a prompt column, or a sequence + of parts to concatenate. + endpoint: The Gemini model to call. Defaults to ``gemini-3.6-flash``. + model_params: Additional request settings, following the Vertex AI + ``generateContent`` request body without its ``contents`` field, + for example ``{"generation_config": {"temperature": 0.2}}``. + output_schema: An optional description of the fields the model should + return, for example ``"sentiment STRING, score INT64"``. When it + is set, the ``result`` field is replaced by one field per declared + field, sorted by name as BigQuery does. + project: The Google Cloud project billed for the calls. Inferred from + the environment when not set. + location: The Vertex AI location. Defaults to ``global``, which is + where the current Gemini models are served. Set it only when a + particular region is required, and check that the endpoint is + available there. + + Returns: + A column holding a struct. Without ``output_schema`` the fields are + ``result``, ``full_response`` and ``status``; with it, the declared + fields replace ``result``. ``full_response`` holds the complete model + response, typed as VARIANT on Spark 4 and as a JSON string on older + versions. ``status`` is empty when the row succeeded and holds the + error message otherwise. + + Example: + >>> df.withColumn("g", ai_generate(sf.col("review"))).select( + ... "g.result", "g.status" + ... ) + >>> df.withColumn( + ... "g", + ... ai_generate( + ... [sf.lit("Summarize in 10 words: "), sf.col("body")], + ... output_schema="summary STRING, sentiment STRING", + ... ), + ... ) + """ + generate = ai_generate_udf( + endpoint=endpoint, + model_params=model_params, + output_schema=output_schema, + project=project, + location=location, + _generate_content=_generate_content, + ) + return generate(_to_prompt_column(prompt)) + + +def ai_generate_udf( + *, + endpoint: str = DEFAULT_ENDPOINT, + model_params: Optional[Dict[str, Any]] = None, + output_schema: Optional[Union[str, StructType]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _generate_content=None, +): + """Builds a reusable ``AI.GENERATE`` user defined function. + + Use this when the same configuration is applied in several places, or to + make the function available to Spark SQL: + + >>> spark.udf.register("ai_generate", ai_generate_udf()) + >>> spark.sql("SELECT ai_generate(review).result FROM reviews") + + Args: + endpoint: See :func:`ai_generate`. + model_params: See :func:`ai_generate`. + output_schema: See :func:`ai_generate`. + project: See :func:`ai_generate`. + location: See :func:`ai_generate`. + _generate_content: Internal. Replaces the call to Vertex AI with an + awaitable of the same shape, so that tests can run the whole Spark + path without credentials or network access. + + Returns: + A pandas user defined function taking a single prompt column. + + Raises: + ValueError: If any argument is invalid. All validation happens here, on + the driver, so that mistakes surface before a stage is scheduled. + """ + if not isinstance(endpoint, str) or not endpoint.strip(): + raise ValueError("endpoint must be a non-empty model name.") + + parsed_schema = ( + _output_schema_util.parse_output_schema(output_schema) + if output_schema is not None + else None + ) + response_schema = ( + _output_schema_util.to_response_schema(parsed_schema) + if parsed_schema is not None + else None + ) + flat_params = _model_params.normalize_model_params( + model_params, response_schema + ) + return_type = _build_return_type(parsed_schema) + + def _predict( + batch_iterator: Iterator[pd.Series], + ) -> Iterator[pd.DataFrame]: + """Runs on the executor, once per task.""" + # Resolved once per task and shared with every other task running in + # this worker process. + generate_content = _generate_content or _client.get_generate_content( + project, location + ) + config = _model_params.build_generation_config(flat_params) + + for batch in batch_iterator: + rows = _client.run_coroutine( + _generate_batch( + generate_content=generate_content, + prompts=batch.tolist(), + endpoint=endpoint, + config=config, + schema=parsed_schema, + ) + ) + yield pd.DataFrame( + rows, + columns=[field.name for field in return_type.fields], + index=batch.index, + ) + + return pandas_udf(_predict, returnType=return_type) + + +def _to_prompt_column(prompt: PromptArg) -> Column: + """Builds the single string column that holds each row's prompt. + + Every path casts to STRING. A prompt column of another type would + otherwise reach the worker as a non-string value and be treated as a null + prompt, turning a whole column into nulls without any error. + """ + if isinstance(prompt, Column): + return prompt.cast(StringType()) + if isinstance(prompt, str): + return sf.col(prompt).cast(StringType()) + if isinstance(prompt, Sequence): + parts = list(prompt) + if not parts: + raise ValueError("prompt must contain at least one part.") + columns = [ + part if isinstance(part, Column) else sf.col(part) for part in parts + ] + if len(columns) == 1: + return columns[0].cast(StringType()) + # Concatenation matches BigQuery, where the parts of a struct prompt + # are combined in the order they are given. + return sf.concat(*[column.cast(StringType()) for column in columns]) + raise ValueError( + "prompt must be a Column, a column name, or a sequence of those, got " + f"{type(prompt).__name__}." + ) + + +def _build_return_type(parsed_schema: Optional[StructType]) -> StructType: + """Builds the struct returned by the function.""" + if parsed_schema is None: + fields = [StructField(_RESULT_FIELD, StringType())] + else: + fields = list(parsed_schema.fields) + for reserved in (_FULL_RESPONSE_FIELD, _STATUS_FIELD): + if any(field.name == reserved for field in fields): + raise ValueError( + f"output_schema must not declare a '{reserved}' field; it " + "is always added to the result." + ) + # Every field is nullable: a row that fails reports the failure in status + # and leaves the generated fields empty. + fields = [ + StructField(field.name, field.dataType, nullable=True) + for field in fields + ] + fields.append(StructField(_FULL_RESPONSE_FIELD, _full_response_data_type())) + fields.append(StructField(_STATUS_FIELD, StringType())) + return StructType(fields) + + +def _full_response_data_type() -> DataType: + """Returns the type of the ``full_response`` field. + + BigQuery returns the raw model response as JSON. Spark 4 has an equivalent + semi-structured type, VARIANT, which allows ``variant_get`` and + ``to_json`` to be used on the column. On older Spark versions, which have + no such type, the response is returned as a JSON string instead. + """ + return VariantType() if _VARIANT_SUPPORTED else StringType() + + +async def _generate_batch( + *, + generate_content, + prompts: List[Any], + endpoint: str, + config, + schema: Optional[StructType], +) -> List[Dict[str, Any]]: + """Generates a response for every prompt in an Arrow batch.""" + semaphore = asyncio.Semaphore(_MAX_CONCURRENT_REQUESTS) + + async def _one(prompt: Any) -> Dict[str, Any]: + async with semaphore: + return await _generate_one( + generate_content=generate_content, + prompt=prompt, + endpoint=endpoint, + config=config, + schema=schema, + ) + + return await asyncio.gather(*[_one(prompt) for prompt in prompts]) + + +async def _generate_one( + *, + generate_content, + prompt: Any, + endpoint: str, + config, + schema: Optional[StructType], +) -> Dict[str, Any]: + """Generates the response for one row, retrying transient failures.""" + if _is_missing(prompt): + # A null prompt produces a null result without calling the model. + return _row(schema, result=None, full_response=None, status="") + + for attempt in range(_MAX_ATTEMPTS): + try: + response = await generate_content( + model=endpoint, contents=prompt, config=config + ) + return _to_row(response, schema) + except Exception as e: # pylint: disable=broad-except + # Anything that is not a backend failure, such as a bug in this + # library, is left to fail the task instead of being hidden in the + # status field. + if not _retry.is_reportable(e): + raise + is_last_attempt = attempt == _MAX_ATTEMPTS - 1 + if _retry.is_retryable(e) and not is_last_attempt: + delay = _retry.full_jitter_delay( + attempt, + _BASE_RETRY_DELAY_SECONDS, + _MAX_RETRY_DELAY_SECONDS, + ) + logger.warning( + "Retrying request to %s in %.2fs after error: %s", + endpoint, + delay, + e, + ) + await asyncio.sleep(delay) + continue + return _row( + schema, + result=None, + full_response=None, + status=_retry.to_status_message(e), + ) + + # Unreachable: the loop either returns a row or retries. + raise AssertionError("The retry loop ended without producing a row.") + + +def _is_missing(prompt: Any) -> bool: + """Reports whether a prompt is null. + + Spark nulls arrive as ``None`` or, in a numeric batch, as NaN. + """ + if prompt is None: + return True + if isinstance(prompt, float) and math.isnan(prompt): + return True + return not isinstance(prompt, str) + + +def _to_row(response, schema: Optional[StructType]) -> Dict[str, Any]: + """Converts a model response into one output row.""" + full_response = _dump_response(response) + text = _extract_text(response) + + if schema is None: + return _row(schema, result=text, full_response=full_response, status="") + + if text is None: + # The call succeeded but produced no usable text, for example because + # the response was filtered. This is not an error. + return _row(schema, result=None, full_response=full_response, status="") + + try: + payload = json.loads(text) + except (TypeError, ValueError): + return _row( + schema, + result=None, + full_response=full_response, + status=_retry.PARSE_FAILURE_STATUS, + ) + + if not isinstance(payload, dict): + return _row( + schema, + result=None, + full_response=full_response, + status=_retry.PARSE_FAILURE_STATUS, + ) + + row = { + field.name: _coerce(payload.get(field.name), field.dataType) + for field in schema.fields + } + row[_FULL_RESPONSE_FIELD] = full_response + row[_STATUS_FIELD] = "" + return row + + +def _row( + schema: Optional[StructType], + *, + result: Optional[str], + full_response: Any, + status: str, +) -> Dict[str, Any]: + """Builds an output row, leaving the generated fields empty if needed.""" + if schema is None: + row: Dict[str, Any] = {_RESULT_FIELD: result} + else: + row = {field.name: None for field in schema.fields} + row[_FULL_RESPONSE_FIELD] = full_response + row[_STATUS_FIELD] = status + return row + + +def _extract_text(response) -> Optional[str]: + """Returns the generated text, following BigQuery's extraction rules. + + Only the first candidate is considered, reasoning parts are skipped, and + the first remaining text part wins. + """ + candidates = getattr(response, "candidates", None) or [] + if not candidates: + return None + content = getattr(candidates[0], "content", None) + for part in getattr(content, "parts", None) or []: + if getattr(part, "thought", None): + continue + text = getattr(part, "text", None) + if isinstance(text, str): + return text + return None + + +def _dump_response(response) -> Any: + """Serializes the whole response, as BigQuery's ``full_response`` does. + + Args: + response: The model response. + + Returns: + A Spark variant value, or a JSON string on Spark versions without the + variant type. + """ + try: + serialized = response.model_dump_json(exclude_none=True) + except Exception: # pylint: disable=broad-except + logger.debug("Could not serialize the model response.", exc_info=True) + serialized = "{}" + return _to_full_response(serialized) + + +def _to_full_response(serialized: str) -> Any: + """Converts serialized JSON into the value stored in ``full_response``.""" + if not _VARIANT_SUPPORTED: + return serialized + try: + return VariantVal.parseJson(serialized) + except Exception: # pylint: disable=broad-except + logger.debug("Could not encode the response as a variant.") + return VariantVal.parseJson("{}") + + +def _coerce(value: Any, data_type: DataType) -> Any: + """Converts a JSON value to the type declared in the output schema.""" + if value is None: + return None + try: + if isinstance(data_type, StringType): + return value if isinstance(value, str) else json.dumps(value) + if isinstance(data_type, BooleanType): + return bool(value) + if isinstance(data_type, (LongType, IntegerType)): + return int(value) + if isinstance(data_type, (DoubleType, FloatType)): + return float(value) + if isinstance(data_type, ArrayType): + if not isinstance(value, list): + return None + return [_coerce(item, data_type.elementType) for item in value] + if isinstance(data_type, StructType): + if not isinstance(value, dict): + return None + return { + field.name: _coerce(value.get(field.name), field.dataType) + for field in data_type.fields + } + except (TypeError, ValueError): + return None + return None diff --git a/google/cloud/dataproc_ml/sql/_client.py b/google/cloud/dataproc_ml/sql/_client.py new file mode 100644 index 0000000..7b1592b --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_client.py @@ -0,0 +1,190 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Process-level Vertex AI clients and event loop used by the AI functions. + +Spark imports a module only once per Python worker process, so the state held +here is shared by every task the worker runs. That is what lets one HTTPS +connection pool serve all the partitions a worker processes, instead of being +rebuilt for each one. +""" + +import asyncio +import logging +import os +import threading +from typing import Awaitable, Dict, Optional, Tuple, TypeVar + +logger = logging.getLogger(__name__) + +_T = TypeVar("_T") + +# The location used when neither the caller nor the environment supplies one. +# The global endpoint is used because the current Gemini models are served +# there and it spreads load across regions; a single region can be requested +# with the location argument when data residency requires it. +DEFAULT_LOCATION = "global" + +_LOOP_LOCK = threading.Lock() +_LOOP: Optional[asyncio.AbstractEventLoop] = None + +_CLIENT_LOCK = threading.Lock() +_CLIENTS: Dict[Tuple[Optional[str], Optional[str]], object] = {} + +# Counts how often a cached client was reused. Only read by tests, to confirm +# that clients really are shared across batches and tasks. +_client_cache_hits = 0 + + +def get_event_loop() -> asyncio.AbstractEventLoop: + """Returns the event loop shared by every task in this worker process. + + The asynchronous Vertex AI client owns an ``httpx.AsyncClient`` bound to + the loop that created it. Spark calls a pandas UDF once per Arrow batch, so + running :func:`asyncio.run` per batch would build and tear down a loop for + every batch and leave any cached client bound to a closed one. Instead a + single loop runs on a daemon thread for the life of the process. + + Returns: + The running event loop. + """ + global _LOOP # pylint: disable=global-statement + with _LOOP_LOCK: + if _LOOP is None or _LOOP.is_closed(): + _LOOP = asyncio.new_event_loop() + threading.Thread( + target=_LOOP.run_forever, + name="dataproc-ml-ai-functions", + daemon=True, + ).start() + return _LOOP + + +def run_coroutine(coro: Awaitable[_T]) -> _T: + """Runs a coroutine on the shared loop and waits for its result. + + Args: + coro: The coroutine to execute. + + Returns: + Whatever the coroutine returns. + """ + return asyncio.run_coroutine_threadsafe(coro, get_event_loop()).result() + + +def resolve_project_and_location( + project: Optional[str], location: Optional[str] +) -> Tuple[Optional[str], str]: + """Fills in the project from the environment and the location from the + default when they are unset. + + Workers on Dataproc, Dataproc Serverless and GKE resolve their identity + from Application Default Credentials, so the project normally does not + have to be passed explicitly. + + The location deliberately ignores ``GOOGLE_CLOUD_LOCATION`` and + ``GOOGLE_CLOUD_REGION``. A Dataproc cluster sets those to the region it + runs in, which is rarely a statement about where a model should be + called, and the current Gemini models are served only from the global + endpoint. Honoring them would route the default model to a region that + does not serve it, and because a failed row reports its error in + ``status`` rather than raising, that would surface as a column of NULLs. + A region can still be requested explicitly with the ``location`` + argument. + + Args: + project: An explicitly configured project, or None. + location: An explicitly configured location, or None. + + Returns: + A (project, location) tuple. The project may still be None when it + cannot be discovered, in which case the client reports the problem. + """ + project = project or os.environ.get("GOOGLE_CLOUD_PROJECT") + + if not project: + try: + # Imported lazily: only needed when the project is not configured. + import google.auth # pylint: disable=import-outside-toplevel + + _, project = google.auth.default() + except Exception as e: # pylint: disable=broad-except + logger.debug("Could not infer the project from ADC: %s", e) + + return project, location or DEFAULT_LOCATION + + +def get_client(project: Optional[str], location: Optional[str]): + """Returns a client for the given project and location, creating it once. + + The first call in a worker process builds the client; later calls in that + process, including calls from other Spark tasks, reuse it. + + Args: + project: The Google Cloud project, or None to infer it. + location: The Vertex AI location, or None to infer it. + + Returns: + A ``google.genai.Client`` configured for Vertex AI. + """ + global _client_cache_hits # pylint: disable=global-statement + + key = (project, location) + with _CLIENT_LOCK: + client = _CLIENTS.get(key) + if client is not None: + _client_cache_hits += 1 + return client + + resolved_project, resolved_location = resolve_project_and_location( + project, location + ) + logger.info( + "Creating a Vertex AI client for project=%s location=%s", + resolved_project, + resolved_location, + ) + + # Imported lazily so that importing this package does not pull in the + # Vertex AI dependency until an AI function is evaluated. + from google import genai # pylint: disable=import-outside-toplevel + + client = genai.Client( + vertexai=True, + project=resolved_project, + location=resolved_location, + ) + _CLIENTS[key] = client + return client + + +def get_generate_content(project: Optional[str], location: Optional[str]): + """Returns the asynchronous ``generate_content`` of the cached client. + + Args: + project: The Google Cloud project, or None to infer it. + location: The Vertex AI location, or None to infer it. + + Returns: + An awaitable callable accepting ``model``, ``contents`` and ``config``. + """ + return get_client(project, location).aio.models.generate_content + + +def clear_cache() -> None: + """Drops all cached clients. Intended for tests.""" + global _client_cache_hits # pylint: disable=global-statement + with _CLIENT_LOCK: + _CLIENTS.clear() + _client_cache_hits = 0 diff --git a/google/cloud/dataproc_ml/sql/_model_params.py b/google/cloud/dataproc_ml/sql/_model_params.py new file mode 100644 index 0000000..8ad8470 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_model_params.py @@ -0,0 +1,130 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Translation of the ``model_params`` option into a generation config. + +``model_params`` follows the Vertex AI ``generateContent`` request body, minus +the ``contents`` field which is supplied by the prompt. That body nests the +generation settings under ``generation_config`` whereas the client library +takes them as a flat configuration object, so this module flattens the request +body and validates it. +""" + +import re +from typing import Any, Dict, Optional + +_CAMEL_BOUNDARY = re.compile(r"(? str: + """Converts a camelCase request body key to its snake_case equivalent.""" + return _CAMEL_BOUNDARY.sub("_", key).lower() + + +def normalize_model_params( + model_params: Optional[Dict[str, Any]], + response_schema: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """Flattens ``model_params`` into generation config keyword arguments. + + Both camelCase and snake_case keys are accepted, matching the behavior of + the JSON request body. + + Args: + model_params: The user supplied request body, or None. + response_schema: A response schema derived from ``output_schema``, or + None when the function returns free form text. + + Returns: + A dictionary of keyword arguments for the generation config. + + Raises: + ValueError: If the parameters contain ``contents``, declare a response + schema that conflicts with ``output_schema``, or do not describe a + valid request body. + """ + params = model_params or {} + if not isinstance(params, dict): + raise ValueError( + "model_params must be a dictionary matching the generateContent " + f"request body, got {type(params).__name__}." + ) + + flat: Dict[str, Any] = {} + for key, value in params.items(): + snake_key = _to_snake_case(str(key)) + if snake_key == "contents": + raise ValueError(_CONTENTS_ERROR) + if snake_key in _GENERATION_CONFIG_KEYS: + if not isinstance(value, dict): + raise ValueError( + "model_params.generation_config must be a dictionary, got " + f"{type(value).__name__}." + ) + for nested_key, nested_value in value.items(): + flat[_to_snake_case(str(nested_key))] = nested_value + else: + flat[snake_key] = value + + if response_schema is not None: + if flat.get("response_schema") or flat.get("response_json_schema"): + raise ValueError(_BOTH_SCHEMAS_ERROR) + flat["response_mime_type"] = "application/json" + flat["response_schema"] = response_schema + + # Building the config here surfaces invalid parameters on the driver, + # before any executor work is scheduled. + build_generation_config(flat) + return flat + + +def build_generation_config(flat_params: Dict[str, Any]): + """Builds the client library generation config. + + Args: + flat_params: Keyword arguments produced by + :func:`normalize_model_params`. + + Returns: + A ``google.genai.types.GenerateContentConfig``. + + Raises: + ValueError: If the parameters do not describe a valid request body. + """ + # Imported lazily to keep the Vertex AI dependency off the import path of + # the package itself. + from google.genai import types # pylint: disable=import-outside-toplevel + + try: + return types.GenerateContentConfig(**flat_params) + except Exception as e: + raise ValueError( + "The model_params is invalid and does not conform to the " + "generateContent request body. See " + "https://cloud.google.com/vertex-ai/docs/reference/rest/v1/" + f"projects.locations.publishers.models/generateContent. Cause: {e}" + ) from e diff --git a/google/cloud/dataproc_ml/sql/_output_schema.py b/google/cloud/dataproc_ml/sql/_output_schema.py new file mode 100644 index 0000000..ba68a1f --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_output_schema.py @@ -0,0 +1,228 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Translation of the ``output_schema`` option into Spark and Gemini schemas. + +A single schema declaration drives two things: the Spark ``StructType`` that +the AI function returns, and the response schema sent to the model so that it +replies with conforming JSON. +""" + +import re +from typing import Any, Dict, Union + +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + StringType, + StructField, + StructType, +) + +# BigQuery spells some scalar types differently from Spark. Both spellings are +# accepted so that a schema can be copied straight from a BigQuery query. +_TYPE_ALIASES = { + "INT64": "BIGINT", + "FLOAT64": "DOUBLE", + "BOOL": "BOOLEAN", +} + +# Matches the field options syntax that BigQuery supports but this library does +# not parse yet, e.g. ``OPTIONS(description = '...')``. +_OPTIONS_PATTERN = re.compile(r"\bOPTIONS\s*\(", re.IGNORECASE) + +# Mirrors BigQuery's limit on how deeply an output schema may nest. +MAX_NESTING_DEPTH = 5 + +_SCALAR_RESPONSE_TYPES = { + StringType: "STRING", + BooleanType: "BOOLEAN", + LongType: "INTEGER", + IntegerType: "INTEGER", + DoubleType: "NUMBER", + FloatType: "NUMBER", +} + +_SUPPORTED_TYPE_NAMES = "STRING, INT64, FLOAT64, BOOL, ARRAY and STRUCT" + + +def parse_output_schema( + output_schema: Union[str, StructType], +) -> StructType: + """Parses the ``output_schema`` option into a Spark ``StructType``. + + Args: + output_schema: Either a comma separated list of fields, such as + ``"sentiment STRING, score INT64"``, or an already built + ``StructType``. + + Returns: + The validated struct, with fields sorted by name. + + Raises: + ValueError: If the schema cannot be parsed, uses an unsupported type, + or nests too deeply. + """ + if isinstance(output_schema, StructType): + struct = output_schema + elif isinstance(output_schema, str): + struct = _parse_ddl(output_schema) + else: + raise ValueError( + "output_schema must be a string or a pyspark StructType, got " + f"{type(output_schema).__name__}." + ) + + if not struct.fields: + raise ValueError("output_schema must declare at least one field.") + + _validate(struct, depth=1) + return _sorted_struct(struct) + + +def _parse_ddl(output_schema: str) -> StructType: + """Parses the schema string, translating BigQuery type names first.""" + if _OPTIONS_PATTERN.search(output_schema): + raise ValueError( + "The OPTIONS clause in output_schema is not supported yet. " + "Remove it and describe the fields in the prompt instead." + ) + + normalized = output_schema + for bigquery_name, spark_name in _TYPE_ALIASES.items(): + normalized = re.sub( + rf"\b{bigquery_name}\b", + spark_name, + normalized, + flags=re.IGNORECASE, + ) + + try: + parsed = StructType.fromDDL(normalized) + except Exception as e: + raise ValueError( + "Failed to parse the option `output_schema`. It must be a valid " + "comma separated list of field names and types, for example " + f"'sentiment STRING, score INT64'. Got: {output_schema!r}." + ) from e + + if not isinstance(parsed, StructType): + raise ValueError( + "Failed to parse the option `output_schema` into a struct: " + f"{output_schema!r}." + ) + return parsed + + +def _validate(data_type: DataType, depth: int) -> None: + """Recursively checks that only supported types are used.""" + if depth > MAX_NESTING_DEPTH: + raise ValueError( + "Option `output_schema` has too many nested levels. Maximum " + f"allowed is {MAX_NESTING_DEPTH}." + ) + + if isinstance(data_type, StructType): + names = [field.name for field in data_type.fields] + duplicates = {name for name in names if names.count(name) > 1} + if duplicates: + names = ", ".join(sorted(duplicates)) + raise ValueError( + f"Option `output_schema` has duplicate field names: {names}." + ) + for field in data_type.fields: + _validate(field.dataType, depth + 1) + elif isinstance(data_type, ArrayType): + _validate(data_type.elementType, depth + 1) + elif type(data_type) not in _SCALAR_RESPONSE_TYPES: + raise ValueError( + "Unsupported field type in output_schema: " + f"{data_type.simpleString()}." + f" Only {_SUPPORTED_TYPE_NAMES} are supported." + ) + + +def _sorted_struct(struct: StructType) -> StructType: + """Returns the struct with its fields, and nested fields, sorted by name. + + BigQuery returns the fields of a custom output schema in alphabetical + order rather than declaration order, and this library matches it so that + the same query produces the same column layout on both engines. + """ + return StructType( + [ + StructField( + field.name, + _sorted_type(field.dataType), + nullable=field.nullable, + ) + for field in sorted(struct.fields, key=lambda f: f.name) + ] + ) + + +def _sorted_type(data_type: DataType) -> DataType: + """Applies :func:`_sorted_struct` to any nested structs.""" + if isinstance(data_type, StructType): + return _sorted_struct(data_type) + if isinstance(data_type, ArrayType): + return ArrayType( + _sorted_type(data_type.elementType), + containsNull=data_type.containsNull, + ) + return data_type + + +def to_response_schema(struct: StructType) -> Dict[str, Any]: + """Converts a Spark struct into a Gemini response schema. + + Args: + struct: The parsed output schema. + + Returns: + A schema dictionary suitable for the ``response_schema`` generation + setting. + """ + return _to_response_schema(struct) + + +def _to_response_schema(data_type: DataType) -> Dict[str, Any]: + """Recursively converts a Spark type into a Gemini response schema.""" + if isinstance(data_type, StructType): + required = [ + field.name for field in data_type.fields if not field.nullable + ] + schema: Dict[str, Any] = { + "type": "OBJECT", + "properties": { + field.name: _to_response_schema(field.dataType) + for field in data_type.fields + }, + } + if required: + schema["required"] = required + return schema + + if isinstance(data_type, ArrayType): + return { + "type": "ARRAY", + "items": _to_response_schema(data_type.elementType), + } + + return {"type": _SCALAR_RESPONSE_TYPES[type(data_type)]} diff --git a/google/cloud/dataproc_ml/sql/_retry.py b/google/cloud/dataproc_ml/sql/_retry.py new file mode 100644 index 0000000..5f3c24d --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_retry.py @@ -0,0 +1,187 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Retry policy and error reporting for the AI functions. + +Transient backend failures are retried with full-jitter exponential backoff. +Errors that survive the retry budget are reported per row in the ``status`` +field of the result, using the same wording as the equivalent BigQuery AI +function, so that a single bad row never fails an entire Spark stage. +""" + +import logging +import random +from typing import Optional + +logger = logging.getLogger(__name__) + +# HTTP statuses that are worth retrying: rate limiting, transient backend +# failures and gateway timeouts. +RETRYABLE_HTTP_CODES = frozenset({408, 429, 500, 503, 504}) + +# Canonical error codes whose messages BigQuery passes through verbatim; these +# describe a problem with the request itself, so the detail is actionable. +_PASSTHROUGH_CODES = frozenset( + { + "FAILED_PRECONDITION", + "PERMISSION_DENIED", + "NOT_FOUND", + "INVALID_ARGUMENT", + "OUT_OF_RANGE", + } +) + +_QUOTA_MESSAGE_PREFIX = "Quota exceeded for metric " + +_HTTP_TO_CANONICAL_CODE = { + 400: "INVALID_ARGUMENT", + 401: "UNAUTHENTICATED", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 408: "DEADLINE_EXCEEDED", + 409: "ALREADY_EXISTS", + 413: "OUT_OF_RANGE", + 429: "RESOURCE_EXHAUSTED", + 499: "CANCELLED", + 500: "INTERNAL", + 501: "NOT_IMPLEMENTED", + 503: "UNAVAILABLE", + 504: "DEADLINE_EXCEEDED", +} + +# Reported when a structured response cannot be parsed as JSON. +PARSE_FAILURE_STATUS = ( + "Failed to parse the response - please check the full_response field." +) + + +def _api_error_types(): + """Returns the google-genai API error types, imported lazily.""" + from google.genai import errors # pylint: disable=import-outside-toplevel + + return errors.APIError + + +def _transport_error_types(): + """Returns the transport error types treated as retryable.""" + import httpx # pylint: disable=import-outside-toplevel + + return (httpx.TransportError, TimeoutError) + + +def is_reportable(exception: BaseException) -> bool: + """Reports whether an exception should become a row-level ``status``. + + Only failures that come from the model backend are reported per row. + Anything else, such as a bug in this library, is left to propagate so that + it is not silently hidden in a column of an otherwise successful query. + + Args: + exception: The exception raised while calling the model. + + Returns: + True if the exception should be converted into a status message. + """ + return isinstance( + exception, (_api_error_types(),) + _transport_error_types() + ) + + +def is_retryable(exception: BaseException) -> bool: + """Reports whether a failed request is worth retrying. + + Args: + exception: The exception raised while calling the model. + + Returns: + True for rate limiting, transient backend errors and transport errors. + """ + if isinstance(exception, _transport_error_types()): + return True + if isinstance(exception, _api_error_types()): + return _http_code(exception) in RETRYABLE_HTTP_CODES + return False + + +def _http_code(exception: BaseException) -> Optional[int]: + """Extracts the HTTP status code from an API error, if it has one.""" + code = getattr(exception, "code", None) + return code if isinstance(code, int) else None + + +def _canonical_code(exception: BaseException) -> str: + """Determines the canonical error code, e.g. ``RESOURCE_EXHAUSTED``.""" + status = getattr(exception, "status", None) + if isinstance(status, str) and status: + return status + code = _http_code(exception) + if code is not None and code in _HTTP_TO_CANONICAL_CODE: + return _HTTP_TO_CANONICAL_CODE[code] + if isinstance(exception, _transport_error_types()): + return "UNAVAILABLE" + return "UNKNOWN" + + +def to_status_message(exception: BaseException) -> str: + """Converts a model error into the row's ``status`` value. + + The wording matches BigQuery's AI functions: request errors are surfaced + verbatim, quota messages are passed through, and everything else is + summarized as a retryable error. + + Args: + exception: The exception raised while calling the model. + + Returns: + A human readable status message. Never empty. + """ + code = _canonical_code(exception) + message = getattr(exception, "message", None) or str(exception) + + if code in _PASSTHROUGH_CODES: + return f"{code}: {message}" + if code == "RESOURCE_EXHAUSTED" and message.startswith( + _QUOTA_MESSAGE_PREFIX + ): + return message + return ( + f"A retryable error occurred: {code} error from remote " + "service/endpoint." + ) + + +def full_jitter_delay( + attempt: int, + base_delay_seconds: float, + max_delay_seconds: float, + rng: Optional[random.Random] = None, +) -> float: + """Computes a full-jitter exponential backoff delay. + + The delay is drawn uniformly from ``[0, min(cap, base * 2 ** attempt)]``. + Randomising the whole interval, rather than only a fraction of it, stops + the tasks of a large cluster from retrying in lockstep after a shared + throttling event. + + Args: + attempt: The zero-based index of the attempt that just failed. + base_delay_seconds: The delay unit that is doubled on each attempt. + max_delay_seconds: The upper bound on the exponential term. + rng: Random source, overridable for tests. + + Returns: + The number of seconds to sleep before the next attempt. + """ + ceiling = min(max_delay_seconds, base_delay_seconds * (2**attempt)) + return (rng or random).uniform(0.0, max(0.0, ceiling)) diff --git a/pyproject.toml b/pyproject.toml index 4666af4..9381ea2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ requires-python = ">=3.12" dependencies = [ "google-cloud-aiplatform>=1.121.0, <2.0.0", "google-cloud-storage>=2.19.0, <3.0.0", + "google-genai>=1.38.0, <3.0.0", "pandas>=2.2.2, <3.0.0", "pyarrow>=18.1.0, <20.0.0", "tenacity>=8.5.0, <9.0.0", diff --git a/tests/integration/sql/__init__.py b/tests/integration/sql/__init__.py new file mode 100644 index 0000000..75edc9b --- /dev/null +++ b/tests/integration/sql/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/integration/sql/test_ai_generate.py b/tests/integration/sql/test_ai_generate.py new file mode 100644 index 0000000..d197c3b --- /dev/null +++ b/tests/integration/sql/test_ai_generate.py @@ -0,0 +1,196 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests that call Vertex AI for real. + +They rely on Application Default Credentials and, unless GOOGLE_CLOUD_PROJECT +is set, on the project discovered from the environment. The location is not +taken from the environment: these tests exercise the default global endpoint. +Running them incurs Vertex AI charges. +""" + +import json +import unittest + +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf + +from google.cloud.dataproc_ml.sql import ai_generate, ai_generate_udf +from google.cloud.dataproc_ml.sql import _ai_generate as ai_generate_module + +_AIRPORT_PROMPT = ( + "What is the airport code of the largest airport in {city}? " + "Answer with the three letter code only." +) + + +class TestAiGenerate(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.spark = SparkSession.builder.getOrCreate() + cls.cities = cls.spark.createDataFrame( + [ + ("Bengaluru", "BLR"), + ("London", "LHR"), + ("San Francisco", "SFO"), + ("Paris", "CDG"), + ("Tokyo", "HND"), + ("Sydney", "SYD"), + ("New York", "JFK"), + ], + ["city", "expected"], + # Several rows per batch, spread over more than one task. + ).repartition(3) + + @classmethod + def tearDownClass(cls): + cls.spark.stop() + + def _prompt(self): + before, after = _AIRPORT_PROMPT.split("{city}", maxsplit=1) + return [sf.lit(before), sf.col("city"), sf.lit(after)] + + def test_generates_an_answer_for_every_row(self): + rows = ( + self.cities.withColumn("g", ai_generate(self._prompt())) + .select("city", "expected", "g.result", "g.status") + .collect() + ) + + self.assertEqual(len(rows), 7) + for row in rows: + with self.subTest(city=row["city"]): + self.assertEqual(row["status"], "") + self.assertIsNotNone(row["result"]) + self.assertIn(row["expected"], row["result"].upper()) + + def test_full_response_carries_usage_metadata(self): + result = self.cities.limit(1).withColumn( + "g", ai_generate(self._prompt()) + ) + payload_column = ( + "to_json(g.full_response)" + if ai_generate_module._VARIANT_SUPPORTED + else "g.full_response" + ) + payload = json.loads( + result.selectExpr(f"{payload_column} AS payload").first()["payload"] + ) + + self.assertIn("candidates", payload) + self.assertIn("usage_metadata", payload) + + def test_model_params_are_applied(self): + result = ( + self.cities.limit(2) + .withColumn( + "g", + ai_generate( + self._prompt(), + model_params={ + "generation_config": { + "temperature": 0.0, + "thinking_config": {"thinking_budget": 0}, + } + }, + ), + ) + .select("g.result", "g.status") + .collect() + ) + + for row in result: + self.assertEqual(row["status"], "") + self.assertIsNotNone(row["result"]) + + def test_structured_output(self): + reviews = self.spark.createDataFrame( + [ + ("The battery lasts two days and the screen is gorgeous.",), + ("It broke after a week and support never replied.",), + ], + ["review"], + ) + + result = reviews.withColumn( + "g", + ai_generate( + [ + sf.lit("Analyze this product review: "), + sf.col("review"), + ], + output_schema=( + "sentiment STRING, score INT64, topics ARRAY" + ), + model_params={"generation_config": {"temperature": 0.0}}, + ), + ) + + # BigQuery returns custom schema fields sorted by name. + self.assertEqual( + [field.name for field in result.schema["g"].dataType.fields], + ["score", "sentiment", "topics", "full_response", "status"], + ) + + rows = result.select( + "g.sentiment", "g.score", "g.topics", "g.status" + ).collect() + self.assertEqual(len(rows), 2) + for row in rows: + with self.subTest(sentiment=row["sentiment"]): + self.assertEqual(row["status"], "") + self.assertIsInstance(row["sentiment"], str) + self.assertIsInstance(row["score"], int) + self.assertTrue(row["topics"]) + + def test_unknown_model_is_reported_per_row(self): + # A bad endpoint must not fail the query: every row reports the error + # in its status, exactly as BigQuery's AI.GENERATE does. + rows = ( + self.cities.limit(2) + .withColumn( + "g", + ai_generate( + self._prompt(), + endpoint="gemini-does-not-exist", + ), + ) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(len(rows), 2) + for row in rows: + self.assertIsNone(row["result"]) + self.assertIn("NOT_FOUND", row["status"]) + + def test_registered_for_spark_sql(self): + self.spark.udf.register("ai_generate", ai_generate_udf()) + self.cities.createOrReplaceTempView("cities") + + rows = self.spark.sql( + "SELECT city," + " ai_generate(CONCAT('Airport code of the largest airport in '," + " city, '? Three letters only.')).result AS code" + " FROM cities" + ).collect() + + self.assertEqual(len(rows), 7) + for row in rows: + self.assertIsNotNone(row["code"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/__init__.py b/tests/unit/sql/__init__.py new file mode 100644 index 0000000..75edc9b --- /dev/null +++ b/tests/unit/sql/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unit/sql/test_ai_generate.py b/tests/unit/sql/test_ai_generate.py new file mode 100644 index 0000000..bb940d6 --- /dev/null +++ b/tests/unit/sql/test_ai_generate.py @@ -0,0 +1,502 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import json +import unittest +from unittest import mock + +from google.genai import errors, types +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf +from pyspark.sql.types import LongType, StringType, StructField, StructType + +# pylint: disable=ungrouped-imports +from google.cloud.dataproc_ml.sql import _ai_generate as ai_generate_module +from google.cloud.dataproc_ml.sql import ai_generate, ai_generate_udf + +# pylint: enable=ungrouped-imports + + +def response_with_text(text, thought=None): + """Builds a model response containing a single answer part.""" + parts = [] + if thought is not None: + parts.append(types.Part(text=thought, thought=True)) + parts.append(types.Part(text=text)) + return types.GenerateContentResponse( + candidates=[ + types.Candidate(content=types.Content(role="model", parts=parts)) + ] + ) + + +def empty_response(): + """Builds a response with no usable text, as a filtered answer looks.""" + return types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(role="model", parts=[]), + finish_reason=types.FinishReason.SAFETY, + ) + ] + ) + + +def api_error(code, status, message): + """Builds an API error the way the client library reports one.""" + return errors.APIError( + code, {"error": {"code": code, "status": status, "message": message}} + ) + + +def replies_with(*outcomes): + """Returns a fake ``generate_content`` that replays the given outcomes. + + Each outcome is either a response to return or an exception to raise. The + returned coroutine function is defined here rather than at module level so + that it is serialized by value and can run inside a Spark worker. + + Args: + *outcomes: Responses or exceptions, consumed in order. The last one is + repeated once exhausted. + + Returns: + An awaitable callable with the same shape as the client method. + """ + remaining = list(outcomes) + + async def _generate_content(*, model, contents, config): + del model, contents, config + outcome = remaining.pop(0) if len(remaining) > 1 else remaining[0] + if isinstance(outcome, BaseException): + raise outcome + return outcome + + return _generate_content + + +def answers_the_prompt(): + """Returns a fake that echoes each prompt, so batching can be checked.""" + + async def _generate_content(*, model, contents, config): + del model, config + return response_with_text(f"answer to {contents}") + + return _generate_content + + +class SparkTestCase(unittest.TestCase): + """Base class holding a Spark session shared by every test module.""" + + @classmethod + def setUpClass(cls): + cls.spark = ( + SparkSession.builder.master("local[2]") + .appName("dataproc-ml-sql-tests") + .config("spark.sql.execution.arrow.maxRecordsPerBatch", "2") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + + +class TestPromptColumn(SparkTestCase): + + def test_a_string_names_a_column(self): + df = self.spark.createDataFrame([("hello",)], ["text"]) + column = ai_generate_module._to_prompt_column("text") + self.assertEqual(df.select(column).first()[0], "hello") + + def test_parts_are_concatenated_in_order(self): + df = self.spark.createDataFrame([("world",)], ["text"]) + column = ai_generate_module._to_prompt_column( + [sf.lit("hello "), sf.col("text"), sf.lit("!")] + ) + self.assertEqual(df.select(column).first()[0], "hello world!") + + def test_a_single_part_is_used_as_is(self): + df = self.spark.createDataFrame([("hello",)], ["text"]) + column = ai_generate_module._to_prompt_column([sf.col("text")]) + self.assertEqual(df.select(column).first()[0], "hello") + + def test_empty_prompt_is_rejected(self): + with self.assertRaises(ValueError): + ai_generate_module._to_prompt_column([]) + + def test_unsupported_prompt_is_rejected(self): + with self.assertRaises(ValueError): + ai_generate_module._to_prompt_column(42) + + def test_a_non_string_column_is_cast_to_string(self): + # Without the cast the worker would see an int, treat it as a null + # prompt and return nulls for the whole column without any error. + df = self.spark.createDataFrame([(7,)], ["n"]) + + for prompt in ("n", sf.col("n"), [sf.col("n")]): + with self.subTest(prompt=prompt): + column = ai_generate_module._to_prompt_column(prompt) + value = df.select(column).first()[0] + self.assertEqual(value, "7") + self.assertIsInstance(value, str) + + +class TestReturnType(SparkTestCase): + + def test_default_fields_match_bigquery(self): + struct = ai_generate_module._build_return_type(None) + self.assertEqual( + [field.name for field in struct.fields], + ["result", "full_response", "status"], + ) + self.assertIsInstance(struct["result"].dataType, StringType) + self.assertIsInstance(struct["status"].dataType, StringType) + + def test_output_schema_replaces_the_result_field(self): + struct = ai_generate_module._build_return_type( + StructType( + [ + StructField("sentiment", StringType()), + StructField("score", LongType()), + ] + ) + ) + self.assertEqual( + [field.name for field in struct.fields], + ["sentiment", "score", "full_response", "status"], + ) + + def test_reserved_field_names_are_rejected(self): + for reserved in ("full_response", "status"): + with self.subTest(field=reserved): + with self.assertRaises(ValueError): + ai_generate_module._build_return_type( + StructType([StructField(reserved, StringType())]) + ) + + def test_generated_fields_are_nullable(self): + # A failed row reports the failure in status and leaves the generated + # fields empty, so they must accept nulls. + struct = ai_generate_module._build_return_type( + StructType([StructField("a", StringType(), nullable=False)]) + ) + self.assertTrue(struct["a"].nullable) + + +class TestValidation(SparkTestCase): + """Configuration mistakes are reported on the driver, before any job.""" + + def test_empty_endpoint_is_rejected(self): + with self.assertRaises(ValueError): + ai_generate_udf(endpoint=" ") + + def test_invalid_model_params_are_rejected(self): + with self.assertRaises(ValueError): + ai_generate_udf(model_params={"contents": ["hi"]}) + + +class TestGenerateOneRow(SparkTestCase): + """The per row behavior, exercised without Spark for precision.""" + + def _run( + self, + generate_content, + prompt="a prompt", + *, + max_attempts=5, + **overrides, + ): + options = { + "generate_content": generate_content, + "prompt": prompt, + "endpoint": "gemini-3.6-flash", + "config": None, + "schema": None, + } + options.update(overrides) + # The retry settings are module constants. They are patched here so + # that the tests neither sleep nor depend on the shipped values. + with mock.patch.multiple( + ai_generate_module, + _MAX_ATTEMPTS=max_attempts, + _BASE_RETRY_DELAY_SECONDS=0.0, + _MAX_RETRY_DELAY_SECONDS=0.0, + ): + return asyncio.run(ai_generate_module._generate_one(**options)) + + def test_successful_row_has_empty_status(self): + row = self._run(replies_with(response_with_text("BLR"))) + self.assertEqual(row["result"], "BLR") + self.assertEqual(row["status"], "") + + def test_reasoning_parts_are_skipped(self): + row = self._run( + replies_with(response_with_text("BLR", thought="thinking...")) + ) + self.assertEqual(row["result"], "BLR") + + def test_response_without_text_produces_no_result(self): + # A filtered answer is not an error, so status stays empty. + row = self._run(replies_with(empty_response())) + self.assertIsNone(row["result"]) + self.assertEqual(row["status"], "") + + def test_null_prompt_skips_the_model(self): + called = [] + + async def _never(*, contents, **_kwargs): + called.append(contents) + raise AssertionError("the model must not be called") + + for prompt in (None, float("nan")): + with self.subTest(prompt=prompt): + row = self._run(_never, prompt=prompt) + self.assertIsNone(row["result"]) + self.assertEqual(row["status"], "") + self.assertEqual(called, []) + + def test_transient_errors_are_retried_then_succeed(self): + row = self._run( + replies_with( + api_error(429, "RESOURCE_EXHAUSTED", "slow down"), + api_error(503, "UNAVAILABLE", "backend down"), + response_with_text("BLR"), + ) + ) + self.assertEqual(row["result"], "BLR") + self.assertEqual(row["status"], "") + + def test_exhausted_retries_report_the_error(self): + row = self._run( + replies_with(api_error(429, "RESOURCE_EXHAUSTED", "slow down")), + max_attempts=3, + ) + self.assertIsNone(row["result"]) + self.assertEqual( + row["status"], + "A retryable error occurred: RESOURCE_EXHAUSTED error from remote " + "service/endpoint.", + ) + + def test_request_errors_are_not_retried(self): + attempts = [] + + async def _fail(*, model, contents, config): + del model, contents, config + attempts.append(1) + raise api_error(404, "NOT_FOUND", "Publisher Model was not found") + + row = self._run(_fail) + self.assertEqual(len(attempts), 1) + self.assertEqual( + row["status"], "NOT_FOUND: Publisher Model was not found" + ) + + def test_library_errors_are_not_hidden_in_status(self): + async def _bug(*, model, contents, config): + del model, contents, config + raise KeyError("a bug in the library") + + with self.assertRaises(KeyError): + self._run(_bug) + + def test_structured_output_is_mapped_to_fields(self): + schema = StructType( + [ + StructField("score", LongType()), + StructField("sentiment", StringType()), + ] + ) + row = self._run( + replies_with( + response_with_text( + json.dumps({"sentiment": "positive", "score": 5}) + ) + ), + schema=schema, + ) + self.assertEqual(row["sentiment"], "positive") + self.assertEqual(row["score"], 5) + self.assertEqual(row["status"], "") + self.assertNotIn("result", row) + + def test_unparsable_structured_output_is_reported(self): + schema = StructType([StructField("sentiment", StringType())]) + row = self._run( + replies_with(response_with_text("not json at all")), + schema=schema, + ) + self.assertIsNone(row["sentiment"]) + self.assertEqual( + row["status"], + "Failed to parse the response - please check the full_response " + "field.", + ) + self.assertIsNotNone(row["full_response"]) + + def test_missing_fields_become_null(self): + schema = StructType( + [ + StructField("sentiment", StringType()), + StructField("score", LongType()), + ] + ) + row = self._run( + replies_with( + response_with_text(json.dumps({"sentiment": "positive"})) + ), + schema=schema, + ) + self.assertEqual(row["sentiment"], "positive") + self.assertIsNone(row["score"]) + + +class TestSparkExecution(SparkTestCase): + """The function as users see it, running through real Spark tasks.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.df = cls.spark.createDataFrame( + [("Bengaluru",), ("London",), ("Tokyo",), ("Paris",), ("Oslo",)], + ["city"], + ).repartition(2) + + def test_default_schema_and_dot_access(self): + result = self.df.withColumn( + "g", + ai_generate(sf.col("city"), _generate_content=answers_the_prompt()), + ) + self.assertEqual( + [field.name for field in result.schema["g"].dataType.fields], + ["result", "full_response", "status"], + ) + + rows = result.select("city", "g.result", "g.status").collect() + self.assertEqual(len(rows), 5) + for row in rows: + city = row["city"] + self.assertEqual(row["result"], f"answer to {city}") + self.assertEqual(row["status"], "") + + def test_prompt_parts_reach_the_model(self): + result = self.df.withColumn( + "g", + ai_generate( + [sf.lit("Airport code for "), sf.col("city")], + _generate_content=answers_the_prompt(), + ), + ) + rows = result.select("city", "g.result").collect() + for row in rows: + city = row["city"] + self.assertEqual( + row["result"], f"answer to Airport code for {city}" + ) + + def test_full_response_holds_the_whole_response(self): + result = self.df.withColumn( + "g", + ai_generate(sf.col("city"), _generate_content=answers_the_prompt()), + ) + # The column is VARIANT on Spark 4 and a JSON string otherwise, so it + # is read back as text either way. + as_text = result.selectExpr( + "CAST(to_json(g.full_response) AS STRING) AS payload" + if ai_generate_module._VARIANT_SUPPORTED + else "g.full_response AS payload" + ).first()["payload"] + payload = json.loads(as_text) + self.assertIn("candidates", payload) + + def test_null_rows_produce_null_results(self): + df = self.spark.createDataFrame([("Oslo",), (None,)], ["city"]) + rows = ( + df.withColumn( + "g", + ai_generate( + sf.col("city"), _generate_content=answers_the_prompt() + ), + ) + .select("city", "g.result", "g.status") + .collect() + ) + by_city = {row["city"]: row for row in rows} + self.assertEqual(by_city["Oslo"]["result"], "answer to Oslo") + self.assertIsNone(by_city[None]["result"]) + self.assertEqual(by_city[None]["status"], "") + + def test_structured_output_columns(self): + async def _structured(*, model, contents, config): + del model, config + return response_with_text( + json.dumps({"city": contents, "letters": len(contents)}) + ) + + result = self.df.withColumn( + "g", + ai_generate( + sf.col("city"), + output_schema="letters INT64, city STRING", + _generate_content=_structured, + ), + ) + self.assertEqual( + [field.name for field in result.schema["g"].dataType.fields], + ["city", "letters", "full_response", "status"], + ) + + rows = result.select("g.city", "g.letters").collect() + for row in rows: + self.assertEqual(row["letters"], len(row["city"])) + + def test_registered_for_spark_sql(self): + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=answers_the_prompt()), + ) + self.df.createOrReplaceTempView("cities") + + rows = self.spark.sql( + "SELECT city, ai_generate(city).result AS answer," + " ai_generate(city).status AS status FROM cities" + ).collect() + + self.assertEqual(len(rows), 5) + for row in rows: + city = row["city"] + self.assertEqual(row["answer"], f"answer to {city}") + self.assertEqual(row["status"], "") + + def test_failed_rows_do_not_fail_the_query(self): + result = self.df.withColumn( + "g", + ai_generate( + sf.col("city"), + _generate_content=replies_with( + api_error(404, "NOT_FOUND", "Publisher Model not found") + ), + ), + ) + + rows = result.select("g.result", "g.status").collect() + self.assertEqual(len(rows), 5) + for row in rows: + self.assertIsNone(row["result"]) + self.assertEqual( + row["status"], "NOT_FOUND: Publisher Model not found" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_client.py b/tests/unit/sql/test_client.py new file mode 100644 index 0000000..b4bde32 --- /dev/null +++ b/tests/unit/sql/test_client.py @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +from unittest import mock + +from google.cloud.dataproc_ml.sql import _client + + +class TestResolveProjectAndLocation(unittest.TestCase): + + def test_explicit_values_win(self): + env = { + "GOOGLE_CLOUD_PROJECT": "from-env", + "GOOGLE_CLOUD_LOCATION": "europe-west4", + } + with mock.patch.dict("os.environ", env, clear=True): + project, location = _client.resolve_project_and_location( + "explicit-project", "asia-northeast1" + ) + + self.assertEqual(project, "explicit-project") + self.assertEqual(location, "asia-northeast1") + + def test_project_falls_back_to_the_environment(self): + with mock.patch.dict( + "os.environ", {"GOOGLE_CLOUD_PROJECT": "from-env"}, clear=True + ): + project, _ = _client.resolve_project_and_location(None, None) + + self.assertEqual(project, "from-env") + + def test_location_defaults_to_global(self): + with mock.patch.dict("os.environ", {}, clear=True): + _, location = _client.resolve_project_and_location("p", None) + + self.assertEqual(location, "global") + self.assertEqual(_client.DEFAULT_LOCATION, "global") + + def test_location_ignores_the_environment(self): + # A Dataproc cluster sets these to the region it runs in. Honoring + # them would send the default model to a region that does not serve + # it, and the failure would be reported per row rather than raised. + env = { + "GOOGLE_CLOUD_PROJECT": "p", + "GOOGLE_CLOUD_LOCATION": "europe-west4", + "GOOGLE_CLOUD_REGION": "us-central1", + } + with mock.patch.dict("os.environ", env, clear=True): + _, location = _client.resolve_project_and_location(None, None) + + self.assertEqual(location, "global") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_model_params.py b/tests/unit/sql/test_model_params.py new file mode 100644 index 0000000..27fdeaa --- /dev/null +++ b/tests/unit/sql/test_model_params.py @@ -0,0 +1,127 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from google.cloud.dataproc_ml.sql import _model_params + + +class TestNormalizeModelParams(unittest.TestCase): + + def test_no_params_produces_empty_config(self): + self.assertEqual(_model_params.normalize_model_params(None), {}) + + def test_generation_config_is_flattened(self): + flat = _model_params.normalize_model_params( + {"generation_config": {"temperature": 0.2, "top_p": 0.9}} + ) + self.assertEqual(flat, {"temperature": 0.2, "top_p": 0.9}) + + def test_camel_case_keys_are_accepted(self): + flat = _model_params.normalize_model_params( + {"generationConfig": {"maxOutputTokens": 128, "topK": 3}} + ) + self.assertEqual(flat, {"max_output_tokens": 128, "top_k": 3}) + + def test_nested_thinking_config_is_preserved(self): + flat = _model_params.normalize_model_params( + {"generation_config": {"thinking_config": {"thinking_budget": 0}}} + ) + self.assertEqual(flat, {"thinking_config": {"thinking_budget": 0}}) + + def test_top_level_request_fields_are_kept(self): + flat = _model_params.normalize_model_params( + { + "system_instruction": "Answer in one word.", + "generation_config": {"temperature": 0.0}, + } + ) + self.assertEqual( + flat, + { + "system_instruction": "Answer in one word.", + "temperature": 0.0, + }, + ) + + def test_contents_is_rejected(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params({"contents": ["hello"]}) + self.assertIn("should not contain contents", str(ctx.exception)) + + def test_invalid_parameters_are_rejected(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params( + {"generation_config": {"temperature": "warm"}} + ) + self.assertIn("model_params is invalid", str(ctx.exception)) + + def test_unknown_parameter_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_model_params({"not_a_real_field": 1}) + + def test_non_dictionary_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_model_params("temperature=0") + + def test_non_dictionary_generation_config_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_model_params({"generation_config": 1}) + + +class TestResponseSchemaHandling(unittest.TestCase): + + def test_output_schema_sets_json_mime_type(self): + schema = {"type": "OBJECT", "properties": {"a": {"type": "STRING"}}} + flat = _model_params.normalize_model_params(None, schema) + self.assertEqual(flat["response_mime_type"], "application/json") + self.assertEqual(flat["response_schema"], schema) + + def test_declaring_a_schema_twice_is_rejected(self): + schema = {"type": "OBJECT", "properties": {"a": {"type": "STRING"}}} + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params( + {"generation_config": {"response_schema": schema}}, schema + ) + self.assertIn( + "can only be set either in the model_params or the output_schema", + str(ctx.exception), + ) + + def test_response_schema_alone_is_allowed(self): + schema = {"type": "OBJECT", "properties": {"a": {"type": "STRING"}}} + flat = _model_params.normalize_model_params( + {"generation_config": {"response_schema": schema}} + ) + self.assertEqual(flat["response_schema"], schema) + + +class TestBuildGenerationConfig(unittest.TestCase): + + def test_values_reach_the_generation_config(self): + config = _model_params.build_generation_config( + {"temperature": 0.25, "max_output_tokens": 64} + ) + self.assertEqual(config.temperature, 0.25) + self.assertEqual(config.max_output_tokens, 64) + + def test_thinking_config_is_understood(self): + config = _model_params.build_generation_config( + {"thinking_config": {"thinking_budget": 0}} + ) + self.assertEqual(config.thinking_config.thinking_budget, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_output_schema.py b/tests/unit/sql/test_output_schema.py new file mode 100644 index 0000000..f6243c7 --- /dev/null +++ b/tests/unit/sql/test_output_schema.py @@ -0,0 +1,220 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from pyspark.sql import SparkSession +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DoubleType, + LongType, + StringType, + StructField, + StructType, +) + +from google.cloud.dataproc_ml.sql import _output_schema + + +class OutputSchemaTestCase(unittest.TestCase): + """Parsing a schema string uses Spark's parser, so a session is needed.""" + + @classmethod + def setUpClass(cls): + cls.spark = ( + SparkSession.builder.master("local[2]") + .appName("dataproc-ml-sql-tests") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + + def field_names(self, struct): + return [field.name for field in struct.fields] + + def field_type(self, struct, name): + return struct[name].dataType + + +class TestParsing(OutputSchemaTestCase): + + def test_bigquery_type_names_are_accepted(self): + struct = _output_schema.parse_output_schema( + "a STRING, b INT64, c FLOAT64, d BOOL" + ) + self.assertIsInstance(self.field_type(struct, "a"), StringType) + self.assertIsInstance(self.field_type(struct, "b"), LongType) + self.assertIsInstance(self.field_type(struct, "c"), DoubleType) + self.assertIsInstance(self.field_type(struct, "d"), BooleanType) + + def test_spark_type_names_are_accepted(self): + struct = _output_schema.parse_output_schema( + "a STRING, b BIGINT, c DOUBLE, d BOOLEAN" + ) + self.assertIsInstance(self.field_type(struct, "b"), LongType) + self.assertIsInstance(self.field_type(struct, "c"), DoubleType) + self.assertIsInstance(self.field_type(struct, "d"), BooleanType) + + def test_arrays_and_structs_are_supported(self): + struct = _output_schema.parse_output_schema( + "tags ARRAY, location STRUCT" + ) + self.assertIsInstance(self.field_type(struct, "tags"), ArrayType) + nested = self.field_type(struct, "location") + self.assertEqual(self.field_names(nested), ["city", "zip"]) + self.assertIsInstance(nested["zip"].dataType, LongType) + + def test_a_struct_type_can_be_passed_directly(self): + struct = _output_schema.parse_output_schema( + StructType( + [ + StructField("b", StringType()), + StructField("a", LongType()), + ] + ) + ) + self.assertEqual(self.field_names(struct), ["a", "b"]) + + def test_unparsable_schema_is_rejected(self): + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema("this is not a schema") + self.assertIn("Failed to parse", str(ctx.exception)) + + def test_empty_schema_is_rejected(self): + with self.assertRaises(ValueError): + _output_schema.parse_output_schema(StructType([])) + + def test_wrong_argument_type_is_rejected(self): + with self.assertRaises(ValueError): + _output_schema.parse_output_schema(42) + + +class TestFieldOrdering(OutputSchemaTestCase): + """BigQuery returns custom schema fields sorted by name.""" + + def test_fields_are_sorted_by_name(self): + struct = _output_schema.parse_output_schema( + "sentiment STRING, explanation STRING, score INT64" + ) + self.assertEqual( + self.field_names(struct), ["explanation", "score", "sentiment"] + ) + + def test_nested_struct_fields_are_sorted_too(self): + struct = _output_schema.parse_output_schema( + "location STRUCT" + ) + self.assertEqual( + self.field_names(self.field_type(struct, "location")), + ["city", "zip"], + ) + + def test_struct_inside_an_array_is_sorted(self): + struct = _output_schema.parse_output_schema( + "items ARRAY>" + ) + element = self.field_type(struct, "items").elementType + self.assertEqual(self.field_names(element), ["id", "name"]) + + +class TestValidation(OutputSchemaTestCase): + + def test_unsupported_types_are_rejected(self): + for declaration in ("d DATE", "t TIMESTAMP", "m MAP"): + with self.subTest(declaration=declaration): + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema(declaration) + self.assertIn("Unsupported field type", str(ctx.exception)) + + def test_options_clause_is_reported_as_unsupported(self): + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema( + "state STRING OPTIONS(description = 'the state')" + ) + self.assertIn("OPTIONS", str(ctx.exception)) + + def test_duplicate_field_names_are_rejected(self): + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema( + StructType( + [ + StructField("a", StringType()), + StructField("a", StringType()), + ] + ) + ) + self.assertIn("duplicate field names", str(ctx.exception)) + + def test_deep_nesting_is_rejected(self): + deep = "a STRUCT>>>" + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema(deep) + self.assertIn("too many nested levels", str(ctx.exception)) + + def test_moderate_nesting_is_allowed(self): + struct = _output_schema.parse_output_schema( + "a STRUCT>" + ) + self.assertEqual(self.field_names(struct), ["a"]) + + +class TestResponseSchema(OutputSchemaTestCase): + + def test_scalars_map_to_gemini_types(self): + struct = _output_schema.parse_output_schema( + "a STRING, b INT64, c FLOAT64, d BOOL" + ) + schema = _output_schema.to_response_schema(struct) + self.assertEqual(schema["type"], "OBJECT") + self.assertEqual( + { + name: value["type"] + for name, value in schema["properties"].items() + }, + {"a": "STRING", "b": "INTEGER", "c": "NUMBER", "d": "BOOLEAN"}, + ) + + def test_arrays_declare_their_item_type(self): + struct = _output_schema.parse_output_schema("tags ARRAY") + schema = _output_schema.to_response_schema(struct) + self.assertEqual( + schema["properties"]["tags"], + {"type": "ARRAY", "items": {"type": "STRING"}}, + ) + + def test_nested_structs_become_nested_objects(self): + struct = _output_schema.parse_output_schema( + "location STRUCT" + ) + schema = _output_schema.to_response_schema(struct) + self.assertEqual( + schema["properties"]["location"], + {"type": "OBJECT", "properties": {"city": {"type": "STRING"}}}, + ) + + def test_non_nullable_fields_are_required(self): + struct = _output_schema.parse_output_schema( + StructType( + [ + StructField("a", StringType(), nullable=False), + StructField("b", StringType(), nullable=True), + ] + ) + ) + schema = _output_schema.to_response_schema(struct) + self.assertEqual(schema["required"], ["a"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_retry.py b/tests/unit/sql/test_retry.py new file mode 100644 index 0000000..7f11877 --- /dev/null +++ b/tests/unit/sql/test_retry.py @@ -0,0 +1,149 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import random +import unittest + +import httpx +from google.genai import errors + +from google.cloud.dataproc_ml.sql import _retry + + +def _api_error(code, status, message): + """Builds an API error the way the client library reports one.""" + return errors.APIError( + code, {"error": {"code": code, "status": status, "message": message}} + ) + + +class TestErrorClassification(unittest.TestCase): + + def test_rate_limiting_and_transient_errors_are_retryable(self): + for code in (408, 429, 500, 503, 504): + with self.subTest(code=code): + error = _api_error(code, "UNAVAILABLE", "try again") + self.assertTrue(_retry.is_retryable(error)) + + def test_request_errors_are_not_retryable(self): + for code in (400, 401, 403, 404): + with self.subTest(code=code): + error = _api_error(code, "INVALID_ARGUMENT", "bad request") + self.assertFalse(_retry.is_retryable(error)) + + def test_transport_errors_are_retryable(self): + self.assertTrue(_retry.is_retryable(httpx.ConnectError("no route"))) + self.assertTrue(_retry.is_retryable(httpx.ReadTimeout("slow"))) + + def test_unrelated_exceptions_are_not_reportable(self): + # A bug in the library must fail the task rather than be hidden in the + # status column of an otherwise successful query. + self.assertFalse(_retry.is_reportable(KeyError("bug"))) + self.assertFalse(_retry.is_retryable(KeyError("bug"))) + + def test_backend_errors_are_reportable(self): + self.assertTrue( + _retry.is_reportable(_api_error(404, "NOT_FOUND", "no model")) + ) + self.assertTrue(_retry.is_reportable(httpx.ConnectError("no route"))) + + +class TestStatusMessages(unittest.TestCase): + """The wording matches BigQuery's AI functions.""" + + def test_request_errors_are_surfaced_verbatim(self): + error = _api_error( + 400, + "INVALID_ARGUMENT", + "HTTP links are not supported for requests restricted by VPCSC.", + ) + self.assertEqual( + _retry.to_status_message(error), + "INVALID_ARGUMENT: HTTP links are not supported for requests " + "restricted by VPCSC.", + ) + + def test_not_found_is_surfaced_verbatim(self): + error = _api_error(404, "NOT_FOUND", "Publisher Model was not found") + self.assertEqual( + _retry.to_status_message(error), + "NOT_FOUND: Publisher Model was not found", + ) + + def test_quota_messages_are_passed_through(self): + message = ( + "Quota exceeded for metric aiplatform.googleapis.com/" + "generate_content_requests_per_minute_per_project" + ) + error = _api_error(429, "RESOURCE_EXHAUSTED", message) + self.assertEqual(_retry.to_status_message(error), message) + + def test_other_errors_are_summarised_as_retryable(self): + error = _api_error(429, "RESOURCE_EXHAUSTED", "too many requests") + self.assertEqual( + _retry.to_status_message(error), + "A retryable error occurred: RESOURCE_EXHAUSTED error from remote " + "service/endpoint.", + ) + + def test_server_errors_are_summarised_as_retryable(self): + error = _api_error(503, "UNAVAILABLE", "backend down") + self.assertEqual( + _retry.to_status_message(error), + "A retryable error occurred: UNAVAILABLE error from remote " + "service/endpoint.", + ) + + def test_transport_errors_report_unavailable(self): + self.assertEqual( + _retry.to_status_message(httpx.ConnectError("no route")), + "A retryable error occurred: UNAVAILABLE error from remote " + "service/endpoint.", + ) + + def test_status_is_never_empty(self): + error = _api_error(418, "", "") + self.assertTrue(_retry.to_status_message(error)) + + +class TestFullJitterBackoff(unittest.TestCase): + + def test_delay_grows_exponentially_and_is_fully_randomized(self): + rng = random.Random(7) + for attempt in range(5): + with self.subTest(attempt=attempt): + ceiling = min(60.0, 1.0 * (2**attempt)) + samples = [ + _retry.full_jitter_delay(attempt, 1.0, 60.0, rng) + for _ in range(200) + ] + self.assertTrue(all(0.0 <= s <= ceiling for s in samples)) + # Full jitter draws from the whole interval, so samples should + # reach both ends rather than cluster near the ceiling. + self.assertLess(min(samples), ceiling * 0.2) + self.assertGreater(max(samples), ceiling * 0.8) + + def test_delay_is_capped(self): + rng = random.Random(1) + samples = [ + _retry.full_jitter_delay(20, 1.0, 30.0, rng) for _ in range(100) + ] + self.assertTrue(all(0.0 <= s <= 30.0 for s in samples)) + + def test_zero_base_delay_produces_no_wait(self): + self.assertEqual(_retry.full_jitter_delay(3, 0.0, 60.0), 0.0) + + +if __name__ == "__main__": + unittest.main() From a20e89c5bf5d3977c8e448c0f5855d104825c916 Mon Sep 17 00:00:00 2001 From: Saurabh Maurya Date: Fri, 11 Sep 2026 16:34:43 +0000 Subject: [PATCH 2/2] fix(sql): correct type and prompt handling raised in review Three fixes from the review bot on #8, each with a regression test: - Only rewrite a BigQuery type name where a type may appear. Matching the bare name also rewrote a field named after a type, so "int64 INT64" parsed as a field actually called BIGINT. The narrower match still accepts the NOT NULL and COMMENT modifiers that Spark allows after a type. - Treat only nulls as a missing prompt. Any other value is now rendered as text, so a non-string column no longer arrives as a whole column of nulls when the function is registered and called from Spark SQL, where no cast to STRING is inserted. A whole number keeps its integer spelling: Arrow delivers a nullable BIGINT column as floats, so otherwise a row's prompt would depend on whether an unrelated row in the same batch happened to be null. - Read a boolean field explicitly rather than through Python truthiness, which recorded the answer "false" as True. --- google/cloud/dataproc_ml/sql/_ai_generate.py | 65 ++++++++++++++++-- .../cloud/dataproc_ml/sql/_output_schema.py | 27 ++++++-- tests/unit/sql/test_ai_generate.py | 66 ++++++++++++++++++- tests/unit/sql/test_output_schema.py | 21 ++++++ 4 files changed, 164 insertions(+), 15 deletions(-) diff --git a/google/cloud/dataproc_ml/sql/_ai_generate.py b/google/cloud/dataproc_ml/sql/_ai_generate.py index 6a59139..202a799 100644 --- a/google/cloud/dataproc_ml/sql/_ai_generate.py +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -69,6 +69,12 @@ _FULL_RESPONSE_FIELD = "full_response" _STATUS_FIELD = "status" +# Spellings accepted when a BOOL field comes back as text rather than as a +# JSON boolean. Anything else is treated as an answer that does not fit the +# schema and becomes null. +_TRUE_WORDS = frozenset({"true", "t", "yes", "y", "1"}) +_FALSE_WORDS = frozenset({"false", "f", "no", "n", "0"}) + PromptArg = Union[Column, str, Sequence[Union[Column, str]]] @@ -335,10 +341,17 @@ async def _generate_one( # A null prompt produces a null result without calling the model. return _row(schema, result=None, full_response=None, status="") + # ``ai_generate`` casts its prompt column to STRING on the driver, but the + # function returned by ``ai_generate_udf`` can also be registered and + # called straight from Spark SQL, where no such cast is inserted. Render + # the value here so that a non-string column, say an id, still produces a + # prompt rather than being dropped. + text = _to_text(prompt) + for attempt in range(_MAX_ATTEMPTS): try: response = await generate_content( - model=endpoint, contents=prompt, config=config + model=endpoint, contents=text, config=config ) return _to_row(response, schema) except Exception as e: # pylint: disable=broad-except @@ -376,13 +389,28 @@ async def _generate_one( def _is_missing(prompt: Any) -> bool: """Reports whether a prompt is null. - Spark nulls arrive as ``None`` or, in a numeric batch, as NaN. + Spark nulls arrive as ``None``, as NaN in a numeric batch, or as + ``pandas.NA`` from a nullable column. Every other value, including a + number or a timestamp, is a real prompt that the caller renders as text. """ - if prompt is None: + if prompt is None or prompt is pd.NA: return True - if isinstance(prompt, float) and math.isnan(prompt): - return True - return not isinstance(prompt, str) + return isinstance(prompt, float) and math.isnan(prompt) + + +def _to_text(prompt: Any) -> str: + """Renders a prompt that Spark has not already cast to a string. + + A whole number is rendered without a fraction. Arrow hands a nullable + integer column to pandas as floats, so a BIGINT 2 arrives here as 2.0 as + soon as any row in the batch is null. Rendering that as "2.0" would let + one row's prompt depend on whether an unrelated row happened to be null. + """ + if isinstance(prompt, str): + return prompt + if isinstance(prompt, float) and float(prompt).is_integer(): + return str(int(prompt)) + return str(prompt) def _to_row(response, schema: Optional[StructType]) -> Dict[str, Any]: @@ -498,7 +526,7 @@ def _coerce(value: Any, data_type: DataType) -> Any: if isinstance(data_type, StringType): return value if isinstance(value, str) else json.dumps(value) if isinstance(data_type, BooleanType): - return bool(value) + return _to_boolean(value) if isinstance(data_type, (LongType, IntegerType)): return int(value) if isinstance(data_type, (DoubleType, FloatType)): @@ -517,3 +545,26 @@ def _coerce(value: Any, data_type: DataType) -> Any: except (TypeError, ValueError): return None return None + + +def _to_boolean(value: Any) -> Optional[bool]: + """Converts a JSON value to a boolean. + + ``bool()`` on its own is wrong here, because every non-empty string is + truthy: a model answering with the text "false" would be recorded as True. + A value that does not clearly denote a boolean becomes null, which is how + the rest of this module reports an answer that does not fit the requested + schema. + """ + if isinstance(value, bool): + return value + if isinstance(value, str): + text = value.strip().lower() + if text in _TRUE_WORDS: + return True + if text in _FALSE_WORDS: + return False + return None + if isinstance(value, (int, float)): + return bool(value) + return None diff --git a/google/cloud/dataproc_ml/sql/_output_schema.py b/google/cloud/dataproc_ml/sql/_output_schema.py index ba68a1f..d6c8f64 100644 --- a/google/cloud/dataproc_ml/sql/_output_schema.py +++ b/google/cloud/dataproc_ml/sql/_output_schema.py @@ -43,6 +43,24 @@ "BOOL": "BOOLEAN", } +# An alias is only rewritten where a type may appear, which is recognized by +# what surrounds it. Matching the bare name would also rewrite a *field* named +# after a type: "int64 INT64" would become "BIGINT BIGINT", quietly renaming +# the field to BIGINT instead of failing. A type is preceded by a space, a +# colon or "<", and is followed by the end of the field or by one of the +# modifiers that Spark allows after a type. +_ALIAS_PATTERNS = tuple( + ( + re.compile( + rf"(?<=[\s:<]){bigquery_name}" + r"(?=\s*(?:,|>|\)|$|NOT\s+NULL|COMMENT\b))", + re.IGNORECASE, + ), + spark_name, + ) + for bigquery_name, spark_name in _TYPE_ALIASES.items() +) + # Matches the field options syntax that BigQuery supports but this library does # not parse yet, e.g. ``OPTIONS(description = '...')``. _OPTIONS_PATTERN = re.compile(r"\bOPTIONS\s*\(", re.IGNORECASE) @@ -105,13 +123,8 @@ def _parse_ddl(output_schema: str) -> StructType: ) normalized = output_schema - for bigquery_name, spark_name in _TYPE_ALIASES.items(): - normalized = re.sub( - rf"\b{bigquery_name}\b", - spark_name, - normalized, - flags=re.IGNORECASE, - ) + for pattern, spark_name in _ALIAS_PATTERNS: + normalized = pattern.sub(spark_name, normalized) try: parsed = StructType.fromDDL(normalized) diff --git a/tests/unit/sql/test_ai_generate.py b/tests/unit/sql/test_ai_generate.py index bb940d6..52ad61c 100644 --- a/tests/unit/sql/test_ai_generate.py +++ b/tests/unit/sql/test_ai_generate.py @@ -20,7 +20,13 @@ from google.genai import errors, types from pyspark.sql import SparkSession from pyspark.sql import functions as sf -from pyspark.sql.types import LongType, StringType, StructField, StructType +from pyspark.sql.types import ( + BooleanType, + LongType, + StringType, + StructField, + StructType, +) # pylint: disable=ungrouped-imports from google.cloud.dataproc_ml.sql import _ai_generate as ai_generate_module @@ -265,6 +271,18 @@ async def _never(*, contents, **_kwargs): self.assertEqual(row["status"], "") self.assertEqual(called, []) + def test_a_non_string_prompt_is_rendered_as_text(self): + # Spark SQL can call the registered function on any column, and no + # cast is inserted there. Such a value must reach the model as text + # rather than be mistaken for a null. A whole number carried as a + # float, which is how Arrow delivers a nullable integer column, keeps + # its integer spelling. + cases = ((42, "42"), (2.0, "2"), (3.5, "3.5"), (True, "True")) + for prompt, expected in cases: + with self.subTest(prompt=prompt): + row = self._run(answers_the_prompt(), prompt=prompt) + self.assertEqual(row["result"], f"answer to {expected}") + def test_transient_errors_are_retried_then_succeed(self): row = self._run( replies_with( @@ -330,6 +348,33 @@ def test_structured_output_is_mapped_to_fields(self): self.assertEqual(row["status"], "") self.assertNotIn("result", row) + def test_a_boolean_field_reads_the_answer(self): + schema = StructType([StructField("spam", BooleanType())]) + # A tuple rather than a dict: True and 1 are the same dictionary key. + cases = ( + (True, True), + (False, False), + # A model that answers with text rather than a JSON boolean. + ("true", True), + ("false", False), + ("False", False), + ("no", False), + (1, True), + (0, False), + # Nothing sensible to read, so the field is null. + ("perhaps", None), + ([], None), + ) + for answer, expected in cases: + with self.subTest(answer=answer): + row = self._run( + replies_with( + response_with_text(json.dumps({"spam": answer})) + ), + schema=schema, + ) + self.assertIs(row["spam"], expected) + def test_unparsable_structured_output_is_reported(self): schema = StructType([StructField("sentiment", StringType())]) row = self._run( @@ -478,6 +523,25 @@ def test_registered_for_spark_sql(self): self.assertEqual(row["answer"], f"answer to {city}") self.assertEqual(row["status"], "") + def test_a_non_string_column_works_in_spark_sql(self): + # Spark SQL applies the registered function to the column as declared, + # with no cast inserted, so the worker has to render the value. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=answers_the_prompt()), + ) + numbers = self.spark.createDataFrame([(1,), (2,), (None,)], "n: bigint") + numbers.createOrReplaceTempView("numbers") + + rows = self.spark.sql( + "SELECT n, ai_generate(n).result AS answer FROM numbers" + ).collect() + + answers = {row["n"]: row["answer"] for row in rows} + self.assertEqual(answers[1], "answer to 1") + self.assertEqual(answers[2], "answer to 2") + self.assertIsNone(answers[None]) + def test_failed_rows_do_not_fail_the_query(self): result = self.df.withColumn( "g", diff --git a/tests/unit/sql/test_output_schema.py b/tests/unit/sql/test_output_schema.py index f6243c7..149fa90 100644 --- a/tests/unit/sql/test_output_schema.py +++ b/tests/unit/sql/test_output_schema.py @@ -66,6 +66,27 @@ def test_spark_type_names_are_accepted(self): self.assertIsInstance(self.field_type(struct, "c"), DoubleType) self.assertIsInstance(self.field_type(struct, "d"), BooleanType) + def test_a_field_named_after_a_type_keeps_its_name(self): + # Rewriting BigQuery type names must not touch the field names, or + # "int64 INT64" would silently declare a field called BIGINT. + struct = _output_schema.parse_output_schema( + "int64 INT64, bool BOOL, float64 FLOAT64" + ) + self.assertEqual(self.field_names(struct), ["bool", "float64", "int64"]) + self.assertIsInstance(self.field_type(struct, "int64"), LongType) + self.assertIsInstance(self.field_type(struct, "bool"), BooleanType) + self.assertIsInstance(self.field_type(struct, "float64"), DoubleType) + + def test_a_type_keeps_its_modifiers(self): + for schema in ( + "a INT64 NOT NULL", + "a INT64 COMMENT 'how many'", + "a INT64 NOT NULL, b BOOL", + ): + with self.subTest(schema=schema): + struct = _output_schema.parse_output_schema(schema) + self.assertIsInstance(self.field_type(struct, "a"), LongType) + def test_arrays_and_structs_are_supported(self): struct = _output_schema.parse_output_schema( "tags ARRAY, location STRUCT"