From 7d22540c6359564b5ce92b694a031bc8a681a5a5 Mon Sep 17 00:00:00 2001 From: Saurabh Maurya Date: Tue, 22 Sep 2026 13:51:58 +0000 Subject: [PATCH 1/3] feat(sql): add ai_generate, a Spark column function for Gemini Adds `ai_generate`, which calls a Gemini model once per row of a Spark DataFrame, and `ai_generate_udf`, which exposes the same function to Spark SQL. Requests are batched through Arrow, issued a few at a time per worker, and retried with full-jitter exponential backoff. A row that fails reports the failure in its `status` rather than failing the stage, so one bad row cannot lose a long query. The function always returns the same struct, whatever its arguments: STRUCT Keeping the return type fixed is what allows `endpoint`, `model_params` and `output_schema` to be ordinary runtime arguments: the query can be planned without knowing their values, so from SQL they may be passed by name, reordered, omitted, or even varied per row. `output_schema` is a Spark DDL string. It constrains the model through its response schema, and `result` holds the conforming JSON document, which the caller can project with `from_json` when typed columns are wanted. Choosing the Spark types at that point, rather than having this function infer them, is what keeps the return type fixed. `model_params` is a JSON string matching the Vertex AI `generateContent` request body, so any parameter the API accepts works without this library modelling it. JSON also decodes to plain floats, avoiding the Decimal values a Spark struct or VARIANT would have produced. Two behaviours are worth calling out, both verified against Spark: - `output_schema` can arrive as a column, so it has to be parsed on an executor, where `StructType.fromDDL` fails for want of a session. `_ddl.py` therefore parses the DDL directly into a response schema. Rather than assert its output by hand, its tests assert that it agrees with `fromDDL`, so the two parsers cannot drift apart. - The arguments are honoured per row and memoised by value. A literal reaches the executor as a repeated column and is indistinguishable from varying data, so reading one row as "the configuration" would silently apply different settings to different Arrow batches. Authentication uses Application Default Credentials; `project` and `location` are Python-only arguments and default to the global endpoint, where the current Gemini models are served. --- README.md | 82 ++ docs/api.rst | 4 + google/cloud/dataproc_ml/sql/__init__.py | 50 ++ google/cloud/dataproc_ml/sql/_ai_generate.py | 582 ++++++++++++++ google/cloud/dataproc_ml/sql/_client.py | 190 +++++ google/cloud/dataproc_ml/sql/_ddl.py | 346 ++++++++ google/cloud/dataproc_ml/sql/_model_params.py | 168 ++++ google/cloud/dataproc_ml/sql/_retry.py | 203 +++++ pyproject.toml | 1 + tests/integration/sql/__init__.py | 13 + tests/integration/sql/test_ai_generate.py | 329 ++++++++ tests/unit/sql/__init__.py | 13 + tests/unit/sql/test_ai_generate.py | 740 ++++++++++++++++++ tests/unit/sql/test_client.py | 67 ++ tests/unit/sql/test_ddl.py | 280 +++++++ tests/unit/sql/test_model_params.py | 182 +++++ tests/unit/sql/test_retry.py | 184 +++++ 17 files changed, 3434 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/_ddl.py create mode 100644 google/cloud/dataproc_ml/sql/_model_params.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_ddl.py create mode 100644 tests/unit/sql/test_model_params.py create mode 100644 tests/unit/sql/test_retry.py diff --git a/README.md b/README.md index 2a5cab8..279f2b1 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,88 @@ 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. It always returns the same struct, whatever the arguments: + +``` +STRUCT +``` + +`result` is the generated text, `full_response` is the whole model response as +JSON, and `status` is `SUCCESS`, `RATE_LIMITED`, `SAFETY_BLOCKED`, or a +description of the error. A row that fails reports it in `status` rather than +failing the query. + +```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|SUCCESS| +# |Tokyo| HND|SUCCESS| +# +-----+------+-------+ +``` + +Pass `output_schema` as a Spark DDL string to constrain the model to a shape. +`result` then holds a JSON document matching it, which you can project with +`from_json` when you want typed columns: + +```python +schema = "sentiment STRING, score DOUBLE" + +df.withColumn( + "review", + ai_generate(F.col("text"), output_schema=schema), +).select(F.from_json(F.col("review.result"), schema).alias("parsed")) +``` + +Model settings are passed as a JSON string matching the Vertex AI +`generateContent` request body, so anything the API supports works without +this library having to know about it: + +```python +ai_generate( + F.col("text"), + model_params='{"generation_config": {"temperature": 0.0, "max_output_tokens": 512}}', +) +``` + +The same function can be registered for use from Spark SQL, where the model +settings are ordinary arguments and may be passed by name in any order: + +```python +from google.cloud.dataproc_ml.sql import ai_generate_udf + +spark.udf.register("ai_generate", ai_generate_udf()) + +spark.sql(""" + SELECT ai_generate( + prompt => CONCAT('Airport code for ', city), + endpoint => 'gemini-3.6-flash', + output_schema => 'code STRING' + ).result + FROM cities +""").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..12b8b99 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -0,0 +1,50 @@ +# 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. + +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, where the model +settings are ordinary arguments and may be passed by name: + + >>> from google.cloud.dataproc_ml.sql import ai_generate_udf + >>> + >>> spark.udf.register("ai_generate", ai_generate_udf()) + >>> spark.sql(''' + ... SELECT ai_generate( + ... prompt => concat('Summarize: ', feedback), + ... output_schema => 'summary STRING, sentiment STRING' + ... ).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..711a01b --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -0,0 +1,582 @@ +# 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 functools +import json +import logging +import math +from typing import Any, Dict, 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 StringType, StructField, StructType + +from google.cloud.dataproc_ml.sql import _client +from google.cloud.dataproc_ml.sql import _ddl +from google.cloud.dataproc_ml.sql import _model_params +from google.cloud.dataproc_ml.sql import _retry + +logger = logging.getLogger(__name__) + +#: The model used when no endpoint is given. +DEFAULT_ENDPOINT = "gemini-3.6-flash" + +# Fixed request settings. These will be configurable in a follow-up PR. +_MAX_CONCURRENT_REQUESTS = 5 +_MAX_ATTEMPTS = 5 +_BASE_RETRY_DELAY_SECONDS = 1.0 +_MAX_RETRY_DELAY_SECONDS = 60.0 + +# How many distinct configurations to keep resolved. A query normally uses one, +# supplied as a literal; the cache only grows when the configuration is itself +# computed per row. +_CONFIG_CACHE_SIZE = 64 + +RESULT_FIELD = "result" +FULL_RESPONSE_FIELD = "full_response" +STATUS_FIELD = "status" + +#: The row succeeded. A null prompt also reports success, with a null result. +STATUS_SUCCESS = "SUCCESS" +#: The request was throttled and did not succeed within the retry budget. +STATUS_RATE_LIMITED = _retry.RATE_LIMITED_STATUS +#: The model refused to answer because the prompt or response was blocked. +STATUS_SAFETY_BLOCKED = "SAFETY_BLOCKED" + +#: The struct every call returns, whatever the arguments. Keeping the type +#: fixed is what lets ``output_schema`` be an ordinary runtime argument: the +#: query can be planned without knowing its value. +RETURN_TYPE = StructType( + [ + StructField(RESULT_FIELD, StringType()), + StructField(FULL_RESPONSE_FIELD, StringType()), + StructField(STATUS_FIELD, StringType()), + ] +) + +_RETURN_FIELDS = [field.name for field in RETURN_TYPE.fields] + +# Finish reasons that mean the content was withheld rather than generated. +_BLOCKED_FINISH_REASONS = frozenset( + {"SAFETY", "PROHIBITED_CONTENT", "BLOCKLIST", "SPII"} +) + +PromptArg = Union[Column, str, Sequence[Union[Column, str]]] +ModelParamsArg = Union[str, Dict[str, Any]] +OutputSchemaArg = Union[str, StructType] + + +def ai_generate( + prompt: PromptArg, + *, + endpoint: str = DEFAULT_ENDPOINT, + model_params: Optional[ModelParamsArg] = None, + output_schema: Optional[OutputSchemaArg] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _generate_content=None, +) -> Column: + """Generates text for every row using a Gemini model. + + 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. + + Requests are retried with exponential backoff and are issued a few at a + time per worker. + + .. 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, either a name such as + ``gemini-3.6-flash`` or a fully qualified + ``projects/.../publishers/google/models/...`` path. Defaults to + ``gemini-3.6-flash``. + model_params: Additional request settings, as a JSON string or an + equivalent dictionary, following the Vertex AI ``generateContent`` + request body without its ``contents`` field, for example + ``'{"generation_config": {"temperature": 0.2}}'``. + output_schema: The shape the model must answer in, as a Spark DDL + string such as ``"sentiment STRING, score INT"`` or as a + ``StructType``. The model is constrained to it, and ``result`` + holds the matching JSON document rather than free text. + 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``. ``result`` is the generated text, or a JSON document + when ``output_schema`` is set. ``full_response`` is the whole model + response as JSON. ``status`` is ``SUCCESS``, ``RATE_LIMITED``, + ``SAFETY_BLOCKED``, or a description of the error. + + Raises: + ValueError: If any argument is invalid. The arguments given here are + validated on the driver, so mistakes surface before a stage is + scheduled. + + 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", + ... ), + ... ) + """ + if not isinstance(endpoint, str) or not endpoint.strip(): + raise ValueError("endpoint must be a non-empty model name.") + + # Resolving the configuration here, rather than only on the executor, + # turns a malformed model_params or output_schema into an immediate error + # instead of a column of failed rows. + model_params_json = _to_model_params_json(model_params) + output_schema_ddl = _to_output_schema_ddl(output_schema) + _resolve_config(model_params_json, output_schema_ddl) + + generate = ai_generate_udf( + project=project, + location=location, + _generate_content=_generate_content, + ) + # An unset argument is passed as a typed NULL rather than left out, so + # that the executor always sees the same four columns. Without the cast an + # unset one would be a void column, which is not a string. + return generate( + _to_prompt_column(prompt), + sf.lit(endpoint).cast(StringType()), + sf.lit(model_params_json).cast(StringType()), + sf.lit(output_schema_ddl).cast(StringType()), + ) + + +def ai_generate_udf( + *, + project: Optional[str] = None, + location: Optional[str] = None, + _generate_content=None, +): + """Builds the ``AI.GENERATE`` user defined function for Spark SQL. + + >>> spark.udf.register("ai_generate", ai_generate_udf()) + >>> spark.sql(''' + ... SELECT ai_generate( + ... prompt => concat('Summarize: ', body), + ... endpoint => 'gemini-3.6-flash', + ... output_schema => 'summary STRING, sentiment STRING' + ... ).result + ... FROM articles + ... ''') + + Unlike :func:`ai_generate`, the model settings are not fixed when the + function is built. They are ordinary arguments, so a query can vary them, + and they may be omitted or passed by name in any order. + + Args: + project: The Google Cloud project billed for the calls. Inferred from + the environment when not set. + location: The Vertex AI 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 ``prompt`` and, optionally, + ``endpoint``, ``model_params`` and ``output_schema``. + """ + + @pandas_udf(RETURN_TYPE) + def _generate( + prompt: pd.Series, + endpoint: pd.Series = None, + model_params: pd.Series = None, + output_schema: pd.Series = None, + ) -> pd.DataFrame: + """Runs on the executor, once per Arrow batch. + + Spark omits an argument the caller did not supply, which arrives here + as ``None`` rather than as a column of nulls. That is the only way to + tell "not given" from "given as NULL", so it is what selects the + defaults. + """ + # Cached for the life of the worker process, so this is a lookup + # rather than a new connection pool per batch. + generate_content = _generate_content or _client.get_generate_content( + project, location + ) + size = len(prompt) + rows = _client.run_coroutine( + _generate_batch( + generate_content=generate_content, + prompts=prompt.tolist(), + endpoints=_argument_values(endpoint, size), + model_params=_argument_values(model_params, size), + output_schemas=_argument_values(output_schema, size), + ) + ) + return pd.DataFrame(rows, columns=_RETURN_FIELDS, index=prompt.index) + + return _generate + + +def _argument_values(argument: Optional[pd.Series], size: int) -> List[Any]: + """Returns one value per row for an argument that may be absent.""" + if argument is None: + return [None] * size + return argument.tolist() + + +def _to_model_params_json( + model_params: Optional[ModelParamsArg], +) -> Optional[str]: + """Renders the ``model_params`` argument as the JSON string the UDF takes. + + The user facing function accepts a dictionary as well, because building + one in Python is more natural than quoting JSON, but only one form has to + survive the trip to an executor. + """ + if model_params is None: + return None + if isinstance(model_params, str): + return model_params + if isinstance(model_params, dict): + try: + return json.dumps(model_params) + except (TypeError, ValueError) as e: + raise ValueError( + "model_params must be JSON serializable. Cause: " f"{e}" + ) from e + raise ValueError( + "model_params must be a JSON string or a dictionary, got " + f"{type(model_params).__name__}." + ) + + +def _to_output_schema_ddl( + output_schema: Optional[OutputSchemaArg], +) -> Optional[str]: + """Renders the ``output_schema`` argument as a DDL string. + + A ``StructType`` is accepted for convenience in Python and rendered with + Spark's own ``toDDL``, so that the executor only ever sees DDL and only + one parser has to agree with Spark. + """ + if output_schema is None: + return None + if isinstance(output_schema, str): + return output_schema + if isinstance(output_schema, StructType): + return output_schema.toDDL() + raise ValueError( + "output_schema must be a DDL string or a pyspark StructType, got " + f"{type(output_schema).__name__}." + ) + + +@functools.lru_cache(maxsize=_CONFIG_CACHE_SIZE) +def _resolve_config( + model_params_json: Optional[str], output_schema_ddl: Optional[str] +): + """Builds the generation config for one combination of arguments. + + Both arguments are ordinary columns, so in principle they can differ on + every row. In practice a query supplies literals, and the cache turns the + repeated work into a single parse per distinct configuration. The result + is only read, never mutated, so sharing one config across rows and threads + is safe. + + Args: + model_params_json: The ``model_params`` argument, or None. + output_schema_ddl: The ``output_schema`` argument, or None. + + Returns: + A ``google.genai.types.GenerateContentConfig``. + + Raises: + ValueError: If either argument is invalid. + """ + response_schema = ( + _ddl.parse_response_schema(output_schema_ddl) + if output_schema_ddl + else None + ) + flat_params = _model_params.normalize_model_params( + model_params_json, response_schema + ) + return _model_params.build_generation_config(flat_params) + + +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()) + 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__}." + ) + + +async def _generate_batch( + *, + generate_content, + prompts: List[Any], + endpoints: List[Any], + model_params: List[Any], + output_schemas: List[Any], +) -> List[Dict[str, Any]]: + """Generates a response for every prompt in an Arrow batch.""" + semaphore = asyncio.Semaphore(_MAX_CONCURRENT_REQUESTS) + + async def _one(index: int) -> Dict[str, Any]: + async with semaphore: + return await _generate_one( + generate_content=generate_content, + prompt=prompts[index], + endpoint=endpoints[index], + model_params=model_params[index], + output_schema=output_schemas[index], + ) + + return await asyncio.gather(*[_one(i) for i in range(len(prompts))]) + + +async def _generate_one( + *, + generate_content, + prompt: Any, + endpoint: Any, + model_params: Any, + output_schema: Any, +) -> 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(result=None, full_response=None, status=STATUS_SUCCESS) + + # ``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) + model = _argument_text(endpoint) or DEFAULT_ENDPOINT + + try: + config = _resolve_config( + _argument_text(model_params), _argument_text(output_schema) + ) + except ValueError as e: + # Only the driver-side entry point can reject these before the query + # runs. Called from SQL the arguments are columns, so a bad value has + # to be reported per row like any other failure. + return _row(result=None, full_response=None, status=str(e)) + + for attempt in range(_MAX_ATTEMPTS): + try: + response = await generate_content( + model=model, contents=text, config=config + ) + return _to_row(response) + 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", + model, + delay, + e, + ) + await asyncio.sleep(delay) + continue + return _row( + result=None, + full_response=None, + status=_retry.to_status(e), + ) + + # Unreachable: the loop either returns a row or retries. + raise AssertionError("The retry loop ended without producing a row.") + + +def _is_missing(value: Any) -> bool: + """Reports whether a value is null. + + 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 value. + """ + if value is None or value is pd.NA: + return True + return isinstance(value, float) and math.isnan(value) + + +def _argument_text(value: Any) -> Optional[str]: + """Normalizes one row's value of a string argument. + + An argument that was not supplied, was supplied as NULL, or was supplied + as an empty string all mean the same thing: use the default. + """ + if _is_missing(value): + return None + text = _to_text(value).strip() + return text or None + + +def _to_text(value: Any) -> str: + """Renders a value 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(value, str): + return value + if isinstance(value, float) and float(value).is_integer(): + return str(int(value)) + return str(value) + + +def _to_row(response) -> Dict[str, Any]: + """Converts a model response into one output row.""" + full_response = _dump_response(response) + + if _is_safety_blocked(response): + return _row( + result=None, + full_response=full_response, + status=STATUS_SAFETY_BLOCKED, + ) + + # When output_schema is set the model is constrained to it, so the text is + # already a JSON document conforming to the schema and is passed through + # unchanged rather than being parsed and rebuilt. + return _row( + result=_extract_text(response), + full_response=full_response, + status=STATUS_SUCCESS, + ) + + +def _row( + *, + result: Optional[str], + full_response: Optional[str], + status: str, +) -> Dict[str, Any]: + """Builds an output row.""" + return { + RESULT_FIELD: result, + FULL_RESPONSE_FIELD: full_response, + STATUS_FIELD: status, + } + + +def _is_safety_blocked(response) -> bool: + """Reports whether the model withheld the answer rather than generating it. + + A block is reported either on the prompt, before anything is generated, or + as the finish reason of the candidate. + """ + feedback = getattr(response, "prompt_feedback", None) + if getattr(feedback, "block_reason", None): + return True + + for candidate in getattr(response, "candidates", None) or []: + reason = getattr(candidate, "finish_reason", None) + if reason is None: + continue + # The client library returns an enum, but a plain string is accepted + # too so that a test or a future version needs no special handling. + name = getattr(reason, "name", None) or str(reason) + if name.upper() in _BLOCKED_FINISH_REASONS: + return True + return False + + +def _extract_text(response) -> Optional[str]: + """Returns the generated text. + + 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) -> str: + """Serializes the whole response into the ``full_response`` field. + + The spec fixes this field as a string so that the return type never + changes. On Spark 4 a caller who wants to query inside it can wrap it in + ``parse_json`` to get a VARIANT. + """ + try: + return response.model_dump_json(exclude_none=True) + except Exception: # pylint: disable=broad-except + logger.debug("Could not serialize the model response.", exc_info=True) + return "{}" 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/_ddl.py b/google/cloud/dataproc_ml/sql/_ddl.py new file mode 100644 index 0000000..731d8fe --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ddl.py @@ -0,0 +1,346 @@ +# 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. + +"""A dependency-free parser for the ``output_schema`` DDL string. + +``output_schema`` may be supplied as a runtime argument from Spark SQL, in +which case it reaches this library as a column value and has to be parsed on an +executor. Spark's own ``StructType.fromDDL`` cannot be used there: it calls +into the JVM through the active session, and a Python worker has none, so it +fails with ``SESSION_OR_CONTEXT_NOT_EXISTS``. + +This module therefore parses the DDL itself, and goes straight to the Gemini +response schema without building a Spark type on the way, since the response +schema is all the caller needs. The grammar accepted here is deliberately the +subset that the rest of this package supports, and the unit tests assert that +it agrees with ``StructType.fromDDL`` on every case they cover, so the two +parsers cannot drift apart unnoticed. +""" + +import re +from typing import Any, Dict, List, Optional, Tuple + +# How deeply an output schema may nest. A deeply nested schema is far more +# often a mistake than an intention, and constrained decoding grows harder for +# the model with every level. +MAX_NESTING_DEPTH = 5 + +# The scalar types this package supports, in every spelling Spark's own DDL +# parser accepts, mapped to their Gemini response schema type. Spark has no +# single-precision or width-limited integer in JSON terms, so the whole integer +# family collapses to INTEGER and both floating point types to NUMBER. +# +# Types with no faithful JSON counterpart are deliberately absent: DECIMAL +# would silently lose precision, and BINARY, DATE, TIMESTAMP and INTERVAL have +# no schema-level representation the model can be constrained to. They are +# rejected rather than approximated. +_SCALAR_TYPES = { + "STRING": "STRING", + "BOOLEAN": "BOOLEAN", + "BYTE": "INTEGER", + "TINYINT": "INTEGER", + "SHORT": "INTEGER", + "SMALLINT": "INTEGER", + "INT": "INTEGER", + "INTEGER": "INTEGER", + "LONG": "INTEGER", + "BIGINT": "INTEGER", + "FLOAT": "NUMBER", + "REAL": "NUMBER", + "DOUBLE": "NUMBER", +} + +_SUPPORTED_TYPE_NAMES = ( + "STRING, BOOLEAN, TINYINT, SMALLINT, INT, BIGINT, FLOAT, DOUBLE, ARRAY " + "and STRUCT" +) + +_TOKEN_PATTERN = re.compile( + r""" + (?P\s+) + | (?P`(?:[^`]|``)*`) + | (?P'(?:[^'\\]|\\.)*'|"(?:[^"\\]|\\.)*") + | (?P\d+) + | (?P[<>,:()]) + | (?P[A-Za-z_][A-Za-z_0-9]*) + """, + re.VERBOSE, +) + + +class _Token: + """One lexical token, carrying its offset for error messages.""" + + __slots__ = ("kind", "text", "position") + + def __init__(self, kind: str, text: str, position: int): + self.kind = kind + self.text = text + self.position = position + + @property + def upper(self) -> str: + return self.text.upper() + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"_Token({self.kind!r}, {self.text!r}, {self.position})" + + +def parse_response_schema(output_schema: str) -> Dict[str, Any]: + """Parses an ``output_schema`` string into a Gemini response schema. + + Args: + output_schema: A comma separated list of fields, such as + ``"sentiment STRING, score INT64"``. A schema wrapped in + ``STRUCT<...>`` is also accepted. + + Returns: + A schema dictionary suitable for the ``response_schema`` generation + setting, with the fields in the order they were declared. + + Raises: + ValueError: If the schema cannot be parsed, uses an unsupported type, + declares duplicate fields, or nests too deeply. + """ + if not isinstance(output_schema, str): + raise ValueError( + "output_schema must be a string, got " + f"{type(output_schema).__name__}." + ) + return _Parser(output_schema).parse() + + +def _tokenize(text: str) -> List[_Token]: + """Splits the schema into tokens, rejecting anything unrecognized.""" + tokens: List[_Token] = [] + position = 0 + while position < len(text): + match = _TOKEN_PATTERN.match(text, position) + if match is None: + raise _error( + text, position, f"unexpected character {text[position]!r}" + ) + position = match.end() + kind = match.lastgroup + if kind == "space": + continue + tokens.append(_Token(kind, match.group(), match.start())) + return tokens + + +def _error(source: str, position: int, detail: str) -> ValueError: + """Builds the error raised for any malformed schema.""" + return ValueError( + "Failed to parse the option `output_schema` at position " + f"{position}: {detail}. It must be a valid comma separated list of " + "field names and types, for example 'sentiment STRING, score INT64'. " + f"Got: {source!r}." + ) + + +class _Parser: + """A recursive descent parser over the supported DDL subset.""" + + def __init__(self, source: str): + self._source = source + self._tokens = _tokenize(source) + self._index = 0 + + # -- token helpers ---------------------------------------------------- + + def _peek(self) -> Optional[_Token]: + if self._index < len(self._tokens): + return self._tokens[self._index] + return None + + def _next(self) -> _Token: + token = self._peek() + if token is None: + raise self._fail("the schema ended unexpectedly") + self._index += 1 + return token + + def _at_punct(self, text: str) -> bool: + token = self._peek() + return ( + token is not None + and token.kind == "punct" + and token.text == text + ) + + def _at_word(self, word: str) -> bool: + token = self._peek() + return ( + token is not None + and token.kind == "word" + and token.upper == word + ) + + def _take_punct(self, text: str) -> None: + if not self._at_punct(text): + raise self._fail(f"expected {text!r}") + self._index += 1 + + def _fail(self, detail: str) -> ValueError: + token = self._peek() + position = ( + token.position if token is not None else len(self._source) + ) + found = f", found {token.text!r}" if token is not None else "" + return _error(self._source, position, f"{detail}{found}") + + # -- grammar ---------------------------------------------------------- + + def parse(self) -> Dict[str, Any]: + """Parses the whole schema and checks that nothing is left over.""" + # A schema may be written either as a bare field list or wrapped in + # STRUCT<...>, both of which Spark's own parser accepts. + if self._at_word("STRUCT"): + schema = self._parse_type(depth=1) + if self._peek() is not None: + raise self._fail("expected the end of the schema") + if schema.get("type") != "OBJECT": + raise self._fail("expected a struct") + return schema + + fields = self._parse_field_list(depth=1, terminator=None) + if self._peek() is not None: + raise self._fail("expected the end of the schema") + return _to_object(fields) + + def _parse_field_list( + self, depth: int, terminator: Optional[str] + ) -> List[Tuple[str, Dict[str, Any], bool]]: + """Parses one or more comma separated fields. + + Returns: + A list of ``(name, schema, required)`` triples in declaration + order. + """ + fields = [] + while True: + fields.append(self._parse_field(depth)) + if self._at_punct(","): + self._index += 1 + continue + break + + if terminator is not None: + self._take_punct(terminator) + + names = [name for name, _, _ in fields] + duplicates = sorted({n for n in names if names.count(n) > 1}) + if duplicates: + listed = ", ".join(duplicates) + raise ValueError( + f"Option `output_schema` has duplicate field names: {listed}." + ) + if not fields: + raise ValueError("output_schema must declare at least one field.") + return fields + + def _parse_field( + self, depth: int + ) -> Tuple[str, Dict[str, Any], bool]: + """Parses ``name [:] type [NOT NULL] [COMMENT '...']``.""" + name = self._parse_identifier() + # Inside STRUCT<> Spark allows, but does not require, a colon between + # the field name and its type. + if self._at_punct(":"): + self._index += 1 + schema = self._parse_type(depth) + + required = False + if self._at_word("NOT"): + self._index += 1 + if not self._at_word("NULL"): + raise self._fail("expected NULL after NOT") + self._index += 1 + required = True + + if self._at_word("COMMENT"): + self._index += 1 + token = self._peek() + if token is None or token.kind != "string": + raise self._fail("expected a quoted string after COMMENT") + self._index += 1 + + return name, schema, required + + def _parse_identifier(self) -> str: + """Parses a bare or backquoted field name.""" + token = self._next() + if token.kind == "word": + return token.text + if token.kind == "quoted": + # Spark escapes a backquote inside a quoted identifier by doubling + # it. + return token.text[1:-1].replace("``", "`") + raise _error( + self._source, + token.position, + f"expected a field name, found {token.text!r}", + ) + + def _parse_type(self, depth: int) -> Dict[str, Any]: + """Parses a scalar, ARRAY or STRUCT type.""" + if depth > MAX_NESTING_DEPTH: + raise ValueError( + "Option `output_schema` has too many nested levels. Maximum " + f"allowed is {MAX_NESTING_DEPTH}." + ) + + token = self._next() + if token.kind != "word": + raise _error( + self._source, + token.position, + f"expected a type name, found {token.text!r}", + ) + + name = token.upper + if name == "ARRAY": + self._take_punct("<") + element = self._parse_type(depth + 1) + self._take_punct(">") + return {"type": "ARRAY", "items": element} + + if name == "STRUCT": + self._take_punct("<") + fields = self._parse_field_list(depth + 1, terminator=">") + return _to_object(fields) + + if name in _SCALAR_TYPES: + return {"type": _SCALAR_TYPES[name]} + + raise _error( + self._source, + token.position, + f"unsupported field type {token.text!r}. Only " + f"{_SUPPORTED_TYPE_NAMES} are supported", + ) + + +def _to_object( + fields: List[Tuple[str, Dict[str, Any], bool]] +) -> Dict[str, Any]: + """Builds an OBJECT response schema from parsed fields.""" + schema: Dict[str, Any] = { + "type": "OBJECT", + "properties": {name: field for name, field, _ in fields}, + } + required = [name for name, _, is_required in fields if is_required] + if required: + schema["required"] = required + return schema 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..6863e1b --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_model_params.py @@ -0,0 +1,168 @@ +# 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 json +import re +from typing import Any, Dict, Optional, Union + +_CAMEL_BOUNDARY = re.compile(r"(? str: + """Converts a camelCase request body key to its snake_case equivalent.""" + return _CAMEL_BOUNDARY.sub("_", key).lower() + + +def _as_dict( + model_params: Optional[Union[str, Dict[str, Any]]] +) -> Dict[str, Any]: + """Decodes ``model_params`` into a dictionary. + + The string form is what the spec defines and the only form that can reach + an executor as a column value. A dictionary is also accepted because + building one is more natural from Python. + """ + if model_params is None: + return {} + + if isinstance(model_params, str): + text = model_params.strip() + if not text: + return {} + try: + decoded = json.loads(text) + except ValueError as e: + raise ValueError( + "The model_params is not valid JSON. It must be a JSON " + "object matching the generateContent request body, for " + 'example \'{"generation_config": {"temperature": 0.2}}\'. ' + f"Cause: {e}" + ) from e + if not isinstance(decoded, dict): + raise ValueError( + "The model_params must be a JSON object, got a " + f"{type(decoded).__name__}." + ) + return decoded + + if isinstance(model_params, dict): + return model_params + + raise ValueError( + "model_params must be a JSON string or a dictionary matching the " + f"generateContent request body, got {type(model_params).__name__}." + ) + + +def normalize_model_params( + model_params: Optional[Union[str, 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, as a JSON string or as + an already decoded dictionary, 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 are not valid JSON, contain + ``contents``, declare a response schema that conflicts with + ``output_schema``, or do not describe a valid request body. + """ + params = _as_dict(model_params) + + 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/_retry.py b/google/cloud/dataproc_ml/sql/_retry.py new file mode 100644 index 0000000..b6ca410 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_retry.py @@ -0,0 +1,203 @@ +# 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, 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 the request was throttled and outlived the retry budget. +RATE_LIMITED_STATUS = "RATE_LIMITED" + + +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(exception: BaseException) -> str: + """Converts a model error into the row's ``status`` value. + + Rate limiting gets its own name because it is the one failure a caller is + expected to react to programmatically, by slowing the query down or + retrying it later. Everything else is described rather than classified, + since the detail is what makes it actionable. + + Args: + exception: The exception raised while calling the model. + + Returns: + ``RATE_LIMITED``, or a human readable description. Never empty. + """ + if _canonical_code(exception) == "RESOURCE_EXHAUSTED": + return RATE_LIMITED_STATUS + return to_status_message(exception) + + +def to_status_message(exception: BaseException) -> str: + """Describes a model error in words. + + 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..ba6b334 --- /dev/null +++ b/tests/integration/sql/test_ai_generate.py @@ -0,0 +1,329 @@ +# 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 + +_AIRPORT_PROMPT = ( + "What is the airport code of the largest airport in {city}? " + "Answer with the three letter code only." +) + +# Keeping the model from thinking or improvising makes the assertions about +# content stable without making them weaker. +_DETERMINISTIC = json.dumps( + { + "generation_config": { + "temperature": 0.0, + "thinking_config": {"thinking_budget": 0}, + } + } +) + + +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"], "SUCCESS") + self.assertIsNotNone(row["result"]) + self.assertIn(row["expected"], row["result"].upper()) + + def test_the_return_type_is_always_the_same(self): + plain = self.cities.withColumn("g", ai_generate(self._prompt())) + structured = self.cities.withColumn( + "g", + ai_generate(self._prompt(), output_schema="code STRING"), + ) + + self.assertEqual( + plain.schema["g"].dataType.simpleString(), + "struct", + ) + self.assertEqual( + plain.schema["g"].dataType, structured.schema["g"].dataType + ) + + def test_full_response_carries_usage_metadata(self): + result = self.cities.limit(1).withColumn( + "g", ai_generate(self._prompt()) + ) + payload = json.loads(result.select("g.full_response").first()[0]) + + self.assertIn("candidates", payload) + self.assertIn("usage_metadata", payload) + + def test_full_response_can_be_queried_with_parse_json(self): + # The field is a string, so anything structural goes through + # parse_json. This is the documented way to reach inside it. + result = ( + self.cities.limit(1) + .withColumn("g", ai_generate(self._prompt())) + .selectExpr( + "variant_get(parse_json(g.full_response)," + " '$.usage_metadata.total_token_count', 'int') AS tokens" + ) + ) + self.assertGreater(result.first()["tokens"], 0) + + def test_model_params_are_applied(self): + rows = ( + self.cities.limit(2) + .withColumn( + "g", + ai_generate(self._prompt(), model_params=_DETERMINISTIC), + ) + .select("g.result", "g.status") + .collect() + ) + + for row in rows: + self.assertEqual(row["status"], "SUCCESS") + self.assertIsNotNone(row["result"]) + + def test_model_params_can_cap_the_output(self): + # A visible, checkable effect, rather than merely "it did not fail". + rows = ( + self.cities.limit(2) + .withColumn( + "g", + ai_generate( + [sf.lit("Write 500 words about "), sf.col("city")], + model_params=json.dumps( + { + "generation_config": { + "max_output_tokens": 20, + "thinking_config": {"thinking_budget": 0}, + } + } + ), + ), + ) + .select("g.result", "g.status") + .collect() + ) + + for row in rows: + self.assertEqual(row["status"], "SUCCESS") + self.assertLess(len(row["result"].split()), 60) + + def test_structured_output_is_json_in_the_result_field(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"], + ) + + rows = ( + reviews.withColumn( + "g", + ai_generate( + [ + sf.lit("Analyze this product review: "), + sf.col("review"), + ], + output_schema=( + "sentiment STRING, score INT, topics ARRAY" + ), + model_params=_DETERMINISTIC, + ), + ) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(len(rows), 2) + for row in rows: + with self.subTest(result=row["result"]): + self.assertEqual(row["status"], "SUCCESS") + # Constrained decoding guarantees this parses and conforms. + parsed = json.loads(row["result"]) + self.assertIsInstance(parsed["sentiment"], str) + self.assertIsInstance(parsed["score"], int) + self.assertIsInstance(parsed["topics"], list) + + def test_structured_output_can_be_projected_with_from_json(self): + # The caller chooses when to pay for parsing, and picks the Spark + # types, rather than the function deciding the return type for them. + schema = "sentiment STRING, score INT" + rows = ( + self.spark.createDataFrame( + [("Absolutely delightful, would buy again.",)], ["review"] + ) + .withColumn( + "g", + ai_generate( + [sf.lit("Analyze: "), sf.col("review")], + output_schema=schema, + model_params=_DETERMINISTIC, + ), + ) + .select(sf.from_json(sf.col("g.result"), schema).alias("parsed")) + .select("parsed.sentiment", "parsed.score") + .collect() + ) + + self.assertEqual(len(rows), 1) + self.assertIsInstance(rows[0]["sentiment"], str) + self.assertIsInstance(rows[0]["score"], int) + + def test_a_fully_qualified_endpoint_is_accepted(self): + from google.cloud.dataproc_ml.sql import _client + + project, _ = _client.resolve_project_and_location(None, None) + self.assertIsNotNone(project, "the test project must be discoverable") + + path = ( + f"projects/{project}/locations/global/publishers/google/models/" + "gemini-3.6-flash" + ) + rows = ( + self.cities.limit(1) + .withColumn("g", ai_generate(self._prompt(), endpoint=path)) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(rows[0]["status"], "SUCCESS") + self.assertIsNotNone(rows[0]["result"]) + + 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. + 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"]) + + def test_spark_sql_with_named_arguments(self): + # The call style the spec uses, end to end against the real model. + self.spark.udf.register("ai_generate", ai_generate_udf()) + self.cities.createOrReplaceTempView("cities") + + rows = self.spark.sql( + "SELECT city, expected, ai_generate(" + " prompt => CONCAT('Airport code of the largest airport in '," + " city, '? Three letters only.')," + " endpoint => 'gemini-3.6-flash'," + f" model_params => '{_DETERMINISTIC}'," + " output_schema => 'code STRING'" + ") AS g FROM cities" + ).collect() + + self.assertEqual(len(rows), 7) + for row in rows: + with self.subTest(city=row["city"]): + self.assertEqual(row["g"]["status"], "SUCCESS") + parsed = json.loads(row["g"]["result"]) + self.assertIn(row["expected"], parsed["code"].upper()) + + def test_spark_sql_argument_may_vary_per_row(self): + self.spark.udf.register("ai_generate", ai_generate_udf()) + self.spark.createDataFrame( + [ + ("France", "Answer with the capital city only."), + ("Japan", "Answer with the capital city only."), + ], + ["country", "instruction"], + ).createOrReplaceTempView("questions") + + rows = self.spark.sql( + "SELECT country, ai_generate(" + " prompt => CONCAT(instruction, ' Country: ', country)," + " output_schema => 'capital STRING'" + ").result AS r FROM questions ORDER BY country" + ).collect() + + capitals = [json.loads(row["r"])["capital"] for row in rows] + self.assertIn("Paris", capitals[0]) + self.assertIn("Tokyo", capitals[1]) + + +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..ea07cbf --- /dev/null +++ b/tests/unit/sql/test_ai_generate.py @@ -0,0 +1,740 @@ +# 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. + +"""Tests for ``ai_generate``.""" + +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 ( + ArrayType, + 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 blocked_response(reason=types.FinishReason.SAFETY): + """Builds a response whose content was withheld rather than generated.""" + return types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(role="model", parts=[]), + finish_reason=reason, + ) + ] + ) + + +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 + + +def echoes_the_request(): + """Returns a fake that reports the arguments it was called with.""" + + async def _generate_content(*, model, contents, config): + return response_with_text( + json.dumps( + { + "model": model, + "contents": contents, + "temperature": getattr(config, "temperature", None), + "response_schema": getattr( + config, "response_schema", None + ), + }, + default=str, + ) + ) + + return _generate_content + + +def run(coroutine): + """Runs one coroutine to completion.""" + return asyncio.run(coroutine) + + +def generate_one(**kwargs): + """Calls the per-row entry point with the arguments a batch would pass.""" + kwargs.setdefault("endpoint", None) + kwargs.setdefault("model_params", None) + kwargs.setdefault("output_schema", None) + return run(ai_generate_module._generate_one(**kwargs)) + + +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() + ) + + def setUp(self): + # The configuration cache is process wide, so it is cleared between + # tests to keep them independent. + ai_generate_module._resolve_config.cache_clear() + + +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): + """The return type is fixed, which is what makes the arguments dynamic.""" + + def test_fields(self): + self.assertEqual( + [field.name for field in ai_generate_module.RETURN_TYPE.fields], + ["result", "full_response", "status"], + ) + for field in ai_generate_module.RETURN_TYPE.fields: + with self.subTest(field=field.name): + self.assertIsInstance(field.dataType, StringType) + + def test_output_schema_does_not_change_the_return_type(self): + # The whole point of the fixed shape: a query can be planned without + # knowing what output_schema will be. + df = self.spark.createDataFrame([("hi",)], ["text"]) + plain = df.withColumn( + "g", ai_generate("text", _generate_content=answers_the_prompt()) + ) + structured = df.withColumn( + "g", + ai_generate( + "text", + output_schema="sentiment STRING, score INT", + _generate_content=answers_the_prompt(), + ), + ) + self.assertEqual( + plain.schema["g"].dataType, structured.schema["g"].dataType + ) + self.assertEqual( + plain.schema["g"].dataType.simpleString(), + "struct", + ) + + +class TestArgumentRendering(unittest.TestCase): + """Arguments are rendered to strings so they can travel as columns.""" + + def test_model_params_accepts_a_json_string(self): + self.assertEqual( + ai_generate_module._to_model_params_json('{"a": 1}'), '{"a": 1}' + ) + + def test_model_params_accepts_a_dictionary(self): + rendered = ai_generate_module._to_model_params_json({"a": 1}) + self.assertEqual(json.loads(rendered), {"a": 1}) + + def test_model_params_rejects_other_types(self): + for value in [42, ["a"], object()]: + with self.subTest(value=value): + with self.assertRaises(ValueError): + ai_generate_module._to_model_params_json(value) + + def test_model_params_rejects_values_that_are_not_json(self): + with self.assertRaises(ValueError): + ai_generate_module._to_model_params_json({"a": {1, 2}}) + + def test_output_schema_accepts_a_ddl_string(self): + self.assertEqual( + ai_generate_module._to_output_schema_ddl("a STRING"), "a STRING" + ) + + def test_output_schema_accepts_a_struct_type(self): + rendered = ai_generate_module._to_output_schema_ddl( + StructType( + [ + StructField("a", StringType()), + StructField("b", ArrayType(LongType())), + ] + ) + ) + # Whatever Spark renders must be readable by our own parser. + self.assertEqual( + ai_generate_module._ddl.parse_response_schema(rendered), + { + "type": "OBJECT", + "properties": { + "a": {"type": "STRING"}, + "b": {"type": "ARRAY", "items": {"type": "INTEGER"}}, + }, + }, + ) + + def test_output_schema_rejects_other_types(self): + for value in [42, ["a STRING"]]: + with self.subTest(value=value): + with self.assertRaises(ValueError): + ai_generate_module._to_output_schema_ddl(value) + + def test_argument_text_treats_missing_and_blank_alike(self): + for value in [None, "", " "]: + with self.subTest(value=value): + self.assertIsNone(ai_generate_module._argument_text(value)) + + def test_argument_text_trims(self): + self.assertEqual( + ai_generate_module._argument_text(" gemini-3.6-flash "), + "gemini-3.6-flash", + ) + + +class TestConfigResolution(unittest.TestCase): + """Turning the string arguments into a generation config.""" + + def setUp(self): + ai_generate_module._resolve_config.cache_clear() + + def test_output_schema_becomes_a_constrained_response_schema(self): + config = ai_generate_module._resolve_config( + None, "sentiment STRING, score INT" + ) + self.assertEqual(config.response_mime_type, "application/json") + self.assertEqual( + config.response_schema, + { + "type": "OBJECT", + "properties": { + "sentiment": {"type": "STRING"}, + "score": {"type": "INTEGER"}, + }, + }, + ) + + def test_model_params_are_unwrapped_into_the_config(self): + config = ai_generate_module._resolve_config( + '{"generation_config": {"temperature": 0.25,' + ' "max_output_tokens": 64}}', + None, + ) + self.assertEqual(config.temperature, 0.25) + self.assertEqual(config.max_output_tokens, 64) + + def test_json_numbers_stay_floats(self): + # A JSON string decodes to plain floats. The same value arriving as a + # Spark named_struct would be a Decimal, which is not serializable. + config = ai_generate_module._resolve_config( + '{"generation_config": {"temperature": 0.7}}', None + ) + self.assertIsInstance(config.temperature, float) + + def test_no_arguments_gives_an_empty_config(self): + config = ai_generate_module._resolve_config(None, None) + self.assertIsNone(config.response_schema) + self.assertIsNone(config.temperature) + + def test_the_same_arguments_are_resolved_once(self): + first = ai_generate_module._resolve_config(None, "a STRING") + second = ai_generate_module._resolve_config(None, "a STRING") + self.assertIs(first, second) + + def test_different_arguments_resolve_separately(self): + first = ai_generate_module._resolve_config(None, "a STRING") + second = ai_generate_module._resolve_config(None, "b INT") + self.assertIsNot(first, second) + + def test_invalid_arguments_are_rejected(self): + for model_params, output_schema in [ + ("{not json", None), + ('{"contents": []}', None), + (None, "a NOTATYPE"), + ('{"generation_config": {"temperature": "hot"}}', None), + ]: + with self.subTest( + model_params=model_params, output_schema=output_schema + ): + with self.assertRaises(ValueError): + ai_generate_module._resolve_config( + model_params, output_schema + ) + + +class TestGenerateOneRow(unittest.TestCase): + """The per-row path, including every value ``status`` can take.""" + + def setUp(self): + ai_generate_module._resolve_config.cache_clear() + + def test_a_successful_row(self): + row = generate_one( + generate_content=replies_with(response_with_text("hello")), + prompt="hi", + ) + self.assertEqual(row["result"], "hello") + self.assertEqual(row["status"], "SUCCESS") + self.assertIn("candidates", row["full_response"]) + + def test_full_response_is_a_json_string(self): + row = generate_one( + generate_content=replies_with(response_with_text("hello")), + prompt="hi", + ) + self.assertIsInstance(row["full_response"], str) + self.assertIsInstance(json.loads(row["full_response"]), dict) + + def test_a_null_prompt_is_not_sent_to_the_model(self): + def _fail(**_): + raise AssertionError("the model must not be called") + + for prompt in [None, float("nan")]: + with self.subTest(prompt=prompt): + row = generate_one(generate_content=_fail, prompt=prompt) + self.assertIsNone(row["result"]) + self.assertIsNone(row["full_response"]) + self.assertEqual(row["status"], "SUCCESS") + + def test_reasoning_parts_are_skipped(self): + row = generate_one( + generate_content=replies_with( + response_with_text("the answer", thought="thinking...") + ), + prompt="hi", + ) + self.assertEqual(row["result"], "the answer") + + def test_blocked_content_is_reported(self): + for reason in [ + types.FinishReason.SAFETY, + types.FinishReason.PROHIBITED_CONTENT, + ]: + with self.subTest(reason=reason): + row = generate_one( + generate_content=replies_with(blocked_response(reason)), + prompt="hi", + ) + self.assertIsNone(row["result"]) + self.assertEqual(row["status"], "SAFETY_BLOCKED") + # The reason is still recoverable from the full response. + self.assertIsNotNone(row["full_response"]) + + def test_an_ordinary_finish_reason_is_not_a_block(self): + row = generate_one( + generate_content=replies_with( + types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + role="model", parts=[types.Part(text="hi")] + ), + finish_reason=types.FinishReason.MAX_TOKENS, + ) + ] + ) + ), + prompt="hi", + ) + self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["result"], "hi") + + def test_rate_limiting_is_reported(self): + with mock.patch.object(ai_generate_module, "_MAX_ATTEMPTS", 1): + row = generate_one( + generate_content=replies_with( + api_error(429, "RESOURCE_EXHAUSTED", "slow down") + ), + prompt="hi", + ) + self.assertIsNone(row["result"]) + self.assertEqual(row["status"], "RATE_LIMITED") + + def test_a_request_error_is_described(self): + with mock.patch.object(ai_generate_module, "_MAX_ATTEMPTS", 1): + row = generate_one( + generate_content=replies_with( + api_error(404, "NOT_FOUND", "no such model") + ), + prompt="hi", + ) + self.assertIsNone(row["result"]) + self.assertIn("NOT_FOUND", row["status"]) + + def test_a_transient_failure_is_retried(self): + with mock.patch.object( + ai_generate_module, "_BASE_RETRY_DELAY_SECONDS", 0.0 + ): + row = generate_one( + generate_content=replies_with( + api_error(503, "UNAVAILABLE", "try again"), + response_with_text("recovered"), + ), + prompt="hi", + ) + self.assertEqual(row["result"], "recovered") + self.assertEqual(row["status"], "SUCCESS") + + def test_a_bug_in_this_library_is_not_hidden_in_status(self): + async def _raise(**_): + raise TypeError("this is a bug, not a backend failure") + + with self.assertRaises(TypeError): + generate_one(generate_content=_raise, prompt="hi") + + def test_an_invalid_argument_is_reported_per_row(self): + # Called from SQL the arguments are columns, so they cannot be + # rejected on the driver. + row = generate_one( + generate_content=replies_with(response_with_text("unused")), + prompt="hi", + output_schema="a NOTATYPE", + ) + self.assertIsNone(row["result"]) + self.assertIn("output_schema", row["status"]) + + def test_the_endpoint_argument_selects_the_model(self): + row = generate_one( + generate_content=echoes_the_request(), + prompt="hi", + endpoint="gemini-2.5-pro", + ) + self.assertEqual(json.loads(row["result"])["model"], "gemini-2.5-pro") + + def test_a_missing_endpoint_falls_back_to_the_default(self): + for endpoint in [None, "", " "]: + with self.subTest(endpoint=endpoint): + row = generate_one( + generate_content=echoes_the_request(), + prompt="hi", + endpoint=endpoint, + ) + self.assertEqual( + json.loads(row["result"])["model"], + ai_generate_module.DEFAULT_ENDPOINT, + ) + + def test_a_fully_qualified_endpoint_is_passed_through(self): + path = "projects/p/locations/l/publishers/google/models/m" + row = generate_one( + generate_content=echoes_the_request(), prompt="hi", endpoint=path + ) + self.assertEqual(json.loads(row["result"])["model"], path) + + def test_a_non_string_prompt_is_rendered(self): + # Arrow delivers a nullable BIGINT column as floats, so a whole number + # must not pick up a fraction just because another row was null. + row = generate_one( + generate_content=answers_the_prompt(), prompt=2.0 + ) + self.assertEqual(row["result"], "answer to 2") + + +class TestDriverValidation(SparkTestCase): + """The DataFrame entry point rejects bad arguments before the query runs.""" + + def test_rejects_an_empty_endpoint(self): + for endpoint in ["", " ", None, 42]: + with self.subTest(endpoint=endpoint): + with self.assertRaises(ValueError): + ai_generate("text", endpoint=endpoint) + + def test_rejects_an_invalid_output_schema(self): + with self.assertRaises(ValueError): + ai_generate("text", output_schema="a NOTATYPE") + + def test_rejects_invalid_model_params(self): + with self.assertRaises(ValueError): + ai_generate("text", model_params="{not json") + + def test_rejects_model_params_carrying_contents(self): + with self.assertRaises(ValueError): + ai_generate("text", model_params='{"contents": []}') + + def test_rejects_two_response_schemas(self): + with self.assertRaises(ValueError): + ai_generate( + "text", + model_params='{"generation_config":' + ' {"response_schema": {"type": "OBJECT"}}}', + output_schema="a STRING", + ) + + +class TestSparkExecution(SparkTestCase): + """End to end through a real Spark plan, with the backend faked out.""" + + def _documents(self): + return self.spark.createDataFrame( + [(1, "alpha"), (2, "beta"), (3, "gamma"), (4, None)], + "id INT, body STRING", + ) + + def test_every_row_gets_its_own_answer(self): + rows = ( + self._documents() + .withColumn( + "g", ai_generate("body", _generate_content=answers_the_prompt()) + ) + .select("id", "g.result", "g.status") + .orderBy("id") + .collect() + ) + + self.assertEqual( + [(r["id"], r["result"]) for r in rows], + [ + (1, "answer to alpha"), + (2, "answer to beta"), + (3, "answer to gamma"), + (4, None), + ], + ) + for row in rows: + self.assertEqual(row["status"], "SUCCESS") + + def test_a_failing_row_does_not_fail_the_query(self): + with mock.patch.object(ai_generate_module, "_MAX_ATTEMPTS", 1): + rows = ( + self._documents() + .withColumn( + "g", + ai_generate( + "body", + _generate_content=replies_with( + api_error(404, "NOT_FOUND", "no such model") + ), + ), + ) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(len(rows), 4) + for row in rows: + self.assertIsNone(row["result"]) + + def test_full_response_can_be_parsed_back_into_json(self): + # The spec fixes the field as a string; parse_json recovers a VARIANT + # for anyone who wants to query inside it. + result = ( + self._documents() + .limit(1) + .withColumn( + "g", ai_generate("body", _generate_content=answers_the_prompt()) + ) + .selectExpr("variant_get(parse_json(g.full_response)," + " '$.candidates[0].content.role', 'string') AS role") + ) + self.assertEqual(result.first()["role"], "model") + + def test_registered_for_spark_sql_with_named_arguments(self): + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self._documents().createOrReplaceTempView("documents") + + row = self.spark.sql( + "SELECT ai_generate(" + " prompt => body," + " endpoint => 'gemini-2.5-pro'," + " output_schema => 'sentiment STRING'" + ").result AS r FROM documents WHERE id = 1" + ).first() + + request = json.loads(row["r"]) + self.assertEqual(request["model"], "gemini-2.5-pro") + self.assertEqual(request["contents"], "alpha") + self.assertEqual( + request["response_schema"], + { + "type": "OBJECT", + "properties": {"sentiment": {"type": "STRING"}}, + }, + ) + + def test_named_arguments_may_be_reordered_and_omitted(self): + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self._documents().createOrReplaceTempView("documents") + + row = self.spark.sql( + "SELECT ai_generate(output_schema => 'a STRING', prompt => body)" + ".result AS r FROM documents WHERE id = 1" + ).first() + + request = json.loads(row["r"]) + self.assertEqual( + request["model"], ai_generate_module.DEFAULT_ENDPOINT + ) + self.assertIsNotNone(request["response_schema"]) + + def test_a_bad_argument_from_sql_is_reported_per_row(self): + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=answers_the_prompt()), + ) + self._documents().createOrReplaceTempView("documents") + + rows = self.spark.sql( + "SELECT ai_generate(prompt => body," + " output_schema => 'a NOTATYPE').status AS s FROM documents" + ).collect() + + self.assertEqual(len(rows), 4) + # The null-prompt row never reaches the configuration. + for row in rows[:3]: + self.assertIn("output_schema", row["s"]) + + def test_the_argument_may_vary_per_row(self): + # A column, not a literal: each row must get its own configuration. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self.spark.createDataFrame( + [("a", "gemini-2.5-pro"), ("b", "gemini-3.6-flash")], + "body STRING, model STRING", + ).createOrReplaceTempView("varying") + + rows = self.spark.sql( + "SELECT body, ai_generate(prompt => body, endpoint => model).result" + " AS r FROM varying ORDER BY body" + ).collect() + + self.assertEqual( + [json.loads(r["r"])["model"] for r in rows], + ["gemini-2.5-pro", "gemini-3.6-flash"], + ) + + def test_a_non_string_column_works_in_spark_sql(self): + # Spark SQL inserts no cast, so the worker has to render the value. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=answers_the_prompt()), + ) + self._documents().createOrReplaceTempView("documents") + + rows = self.spark.sql( + "SELECT ai_generate(id).result AS r FROM documents ORDER BY id" + ).collect() + + self.assertEqual( + [r["r"] for r in rows], + ["answer to 1", "answer to 2", "answer to 3", "answer to 4"], + ) + + +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_ddl.py b/tests/unit/sql/test_ddl.py new file mode 100644 index 0000000..32d8906 --- /dev/null +++ b/tests/unit/sql/test_ddl.py @@ -0,0 +1,280 @@ +# 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. + +"""Tests for the session-free ``output_schema`` DDL parser. + +The parser exists because Spark's own ``StructType.fromDDL`` needs a session +and so cannot run on an executor. The risk in owning a second parser is that it +quietly drifts from the first, so most of this file does not assert the +expected output by hand: it asserts that the two parsers agree. +""" + +import unittest + +from pyspark.sql import SparkSession +from pyspark.sql.types import ( + ArrayType, + BooleanType, + ByteType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + ShortType, + StringType, + StructType, +) + +from google.cloud.dataproc_ml.sql import _ddl + +# The reference conversion, expressed against Spark's own parsed types. Keeping +# it in the test rather than in the library is deliberate: it is the oracle the +# parser is checked against, not a second implementation to maintain. +_REFERENCE_SCALARS = { + StringType: "STRING", + BooleanType: "BOOLEAN", + ByteType: "INTEGER", + ShortType: "INTEGER", + IntegerType: "INTEGER", + LongType: "INTEGER", + FloatType: "NUMBER", + DoubleType: "NUMBER", +} + + +def _reference_schema(data_type: DataType): + """Converts a Spark type into the response schema the parser should emit.""" + if isinstance(data_type, StructType): + schema = { + "type": "OBJECT", + "properties": { + field.name: _reference_schema(field.dataType) + for field in data_type.fields + }, + } + required = [ + field.name for field in data_type.fields if not field.nullable + ] + if required: + schema["required"] = required + return schema + if isinstance(data_type, ArrayType): + return { + "type": "ARRAY", + "items": _reference_schema(data_type.elementType), + } + return {"type": _REFERENCE_SCALARS[type(data_type)]} + + +# Every schema here must parse identically under both parsers. +_AGREEING_SCHEMAS = [ + # The spec's own examples. + "sentiment STRING, urgency INT, categories ARRAY", + "summary STRING", + # The full set of supported scalars, in both spellings where Spark has two. + "a STRING, b BOOLEAN", + "a BYTE, b TINYINT, c SHORT, d SMALLINT", + "a INT, b INTEGER, c LONG, d BIGINT", + "a FLOAT, b REAL, c DOUBLE", + # Nesting. + "tags ARRAY", + "matrix ARRAY>", + "person STRUCT", + "person STRUCT", + "people ARRAY>>", + # Modifiers. + "name STRING NOT NULL, age INT", + "name STRING NOT NULL, age INT NOT NULL", + "a STRING COMMENT 'why', b INT", + "a STRING NOT NULL COMMENT 'why'", + "person STRUCT", + # Identifiers that need quoting, and fields named after types. + "`odd name` STRING", + "`select` STRING, `from` INT", + "int STRING, string INT", + # A schema wrapped in STRUCT<> rather than written as a bare field list. + "STRUCT>", + # Whitespace and casing. + " a string , b Int ", + "A STRING, b STRING", +] + + +class TestAgreesWithSpark(unittest.TestCase): + """Pins the parser to Spark's, so the two cannot drift apart.""" + + @classmethod + def setUpClass(cls): + # Only the oracle needs a session. The parser under test is called + # without one on an executor, which is why it exists. + cls.spark = SparkSession.builder.master("local[1]").getOrCreate() + + def test_matches_fromddl(self): + for schema in _AGREEING_SCHEMAS: + with self.subTest(schema=schema): + expected = _reference_schema(StructType.fromDDL(schema)) + self.assertEqual(_ddl.parse_response_schema(schema), expected) + + def test_field_order_is_preserved(self): + # JSON objects have no inherent order, but the properties are emitted + # in declaration order so that the prompt and the schema read the same + # way round. + schema = _ddl.parse_response_schema("zebra STRING, apple INT") + self.assertEqual(list(schema["properties"]), ["zebra", "apple"]) + + +class TestOutput(unittest.TestCase): + """Spot checks of the emitted schema, independent of Spark.""" + + def test_scalar_fields(self): + self.assertEqual( + _ddl.parse_response_schema("a STRING, b INT"), + { + "type": "OBJECT", + "properties": {"a": {"type": "STRING"}, "b": {"type": "INTEGER"}}, + }, + ) + + def test_not_null_becomes_required(self): + self.assertEqual( + _ddl.parse_response_schema("a STRING NOT NULL, b INT"), + { + "type": "OBJECT", + "properties": {"a": {"type": "STRING"}, "b": {"type": "INTEGER"}}, + "required": ["a"], + }, + ) + + def test_nullable_fields_are_not_required(self): + self.assertNotIn( + "required", _ddl.parse_response_schema("a STRING, b INT") + ) + + def test_array_of_struct(self): + self.assertEqual( + _ddl.parse_response_schema("xs ARRAY>"), + { + "type": "OBJECT", + "properties": { + "xs": { + "type": "ARRAY", + "items": { + "type": "OBJECT", + "properties": {"a": {"type": "STRING"}}, + }, + } + }, + }, + ) + + def test_backquoted_name_is_unquoted(self): + schema = _ddl.parse_response_schema("`odd name` STRING") + self.assertEqual(list(schema["properties"]), ["odd name"]) + + +class TestRejects(unittest.TestCase): + """Malformed or unsupported schemas must fail, and say why.""" + + @classmethod + def setUpClass(cls): + # Needed only by the cross-check that Spark rejects the same names. + cls.spark = SparkSession.builder.master("local[1]").getOrCreate() + + def test_rejects_malformed_schemas(self): + cases = [ + "", + " ", + "a", + "STRING", + "a STRING,", + ", a STRING", + "a STRUCT<>", + "a ARRAY<>", + "a STRUCT", + "VARCHAR(10)", + "CHAR(3)", + ]: + with self.subTest(type_name=type_name): + with self.assertRaises(ValueError) as context: + _ddl.parse_response_schema(f"a {type_name}") + self.assertIn("output_schema", str(context.exception)) + + def test_rejects_bigquery_type_names(self): + # The contract is Spark's DDL. Spark's own parser rejects these too, so + # accepting them here would be a silent divergence. + for type_name in ["INT64", "FLOAT64", "BOOL", "BYTES"]: + with self.subTest(type_name=type_name): + with self.assertRaises(ValueError): + _ddl.parse_response_schema(f"a {type_name}") + with self.assertRaises(Exception): + StructType.fromDDL(f"a {type_name}") + + def test_rejects_duplicate_field_names(self): + with self.assertRaises(ValueError) as context: + _ddl.parse_response_schema("a STRING, a INT") + self.assertIn("duplicate", str(context.exception).lower()) + + def test_rejects_duplicates_in_a_nested_struct(self): + with self.assertRaises(ValueError): + _ddl.parse_response_schema("s STRUCT") + + def test_rejects_schemas_nested_too_deeply(self): + too_deep = "a " + "ARRAY<" * 6 + "STRING" + ">" * 6 + with self.assertRaises(ValueError) as context: + _ddl.parse_response_schema(too_deep) + self.assertIn("nested", str(context.exception).lower()) + + def test_allows_schemas_at_the_nesting_limit(self): + at_limit = "a " + "ARRAY<" * 3 + "STRING" + ">" * 3 + self.assertIsNotNone(_ddl.parse_response_schema(at_limit)) + + def test_rejects_non_strings(self): + for value in [None, 42, ["a STRING"], {"a": "STRING"}]: + with self.subTest(value=value): + with self.assertRaises(ValueError): + _ddl.parse_response_schema(value) + + def test_error_names_the_option_and_shows_the_input(self): + with self.assertRaises(ValueError) as context: + _ddl.parse_response_schema("a NOTATYPE") + message = str(context.exception) + self.assertIn("output_schema", message) + self.assertIn("NOTATYPE", message) + + +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..f79dc5b --- /dev/null +++ b/tests/unit/sql/test_model_params.py @@ -0,0 +1,182 @@ +# 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 json +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_unsupported_types_are_rejected(self): + for value in [42, ["temperature"], object()]: + with self.subTest(value=value): + with self.assertRaises(ValueError): + _model_params.normalize_model_params(value) + + def test_non_dictionary_generation_config_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_model_params({"generation_config": 1}) + + +class TestJsonStringForm(unittest.TestCase): + """The spec defines model_params as a JSON string, and a SQL argument can + only ever arrive as one.""" + + def test_a_json_string_is_decoded(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_a_json_string_and_a_dictionary_agree(self): + params = {"generation_config": {"temperature": 0.2}} + self.assertEqual( + _model_params.normalize_model_params(json.dumps(params)), + _model_params.normalize_model_params(params), + ) + + def test_numbers_decode_to_plain_floats(self): + # Not Decimal, which is what a Spark named_struct or a VARIANT would + # produce, and which is not JSON serializable further down. + flat = _model_params.normalize_model_params( + '{"generation_config": {"temperature": 0.7}}' + ) + self.assertIsInstance(flat["temperature"], float) + + def test_an_empty_string_means_no_parameters(self): + for value in ["", " "]: + with self.subTest(value=value): + self.assertEqual( + _model_params.normalize_model_params(value), {} + ) + + def test_malformed_json_is_rejected(self): + for value in ["{not json", "temperature=0", "{'a': 1}", "{"]: + with self.subTest(value=value): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params(value) + self.assertIn("not valid JSON", str(ctx.exception)) + + def test_json_that_is_not_an_object_is_rejected(self): + for value in ["[1, 2]", '"hello"', "42", "null"]: + with self.subTest(value=value): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params(value) + self.assertIn("must be a JSON object", str(ctx.exception)) + + def test_contents_is_rejected_in_the_string_form_too(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params('{"contents": ["hello"]}') + self.assertIn("should not contain contents", str(ctx.exception)) + + +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_retry.py b/tests/unit/sql/test_retry.py new file mode 100644 index 0000000..27d60db --- /dev/null +++ b/tests/unit/sql/test_retry.py @@ -0,0 +1,184 @@ +# 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 TestStatus(unittest.TestCase): + """``status`` names rate limiting and describes everything else.""" + + def test_rate_limiting_gets_its_own_name(self): + for error in [ + _api_error(429, "RESOURCE_EXHAUSTED", "too many requests"), + _api_error(429, "", "too many requests"), + ]: + with self.subTest(error=error): + self.assertEqual(_retry.to_status(error), "RATE_LIMITED") + + def test_quota_exhaustion_is_rate_limiting(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(error), "RATE_LIMITED") + + def test_other_failures_are_described_rather_than_classified(self): + error = _api_error(404, "NOT_FOUND", "Publisher Model was not found") + self.assertEqual( + _retry.to_status(error), + "NOT_FOUND: Publisher Model was not found", + ) + + def test_transport_failures_are_described(self): + status = _retry.to_status(httpx.ConnectError("no route")) + self.assertNotEqual(status, "RATE_LIMITED") + self.assertIn("UNAVAILABLE", status) + + def test_status_is_never_empty(self): + self.assertTrue(_retry.to_status(_api_error(418, "", ""))) + + +class TestStatusMessages(unittest.TestCase): + """Failures that are not rate limiting are described in words.""" + + 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 91474b175f153653368eee13ef6d4b3fc6c372b8 Mon Sep 17 00:00:00 2001 From: Saurabh Maurya Date: Tue, 22 Sep 2026 17:06:28 +0000 Subject: [PATCH 2/3] feat(sql): file prompts, constant settings and fail-fast config errors Brings ai_generate in line with BigQuery's AI.GENERATE where the semantics are worth copying, and lets a prompt include a file. Result contract: - status is empty on success, so the dead-letter query is one predicate: WHERE g.status <> ''. A failure leads with its canonical error code (NOT_FOUND:, RESOURCE_EXHAUSTED:, SAFETY_BLOCKED:, ...) rather than this library maintaining a closed taxonomy of its own. - full_response is never null. It is "{}" when there is no response, so parse_json(full_response) is safe on every row. Mistakes in the call now fail the query instead of filling a column with failures, which is BigQuery's line and is what makes a typo cheap: - a malformed model_params or output_schema; - an endpoint that names no model, checked once per worker against the model catalogue before any generation; - a setting that differs from row to row. endpoint, model_params and output_schema configure the call, not the row, and BigQuery rejects a column for all three at plan time. Spark SQL cannot tell a literal from a column that happens to repeat, so constancy is checked per batch. Problems that really are per row, including a malformed file reference, are still reported in status. Multimodal prompts: - A prompt is a string or a struct whose fields are, in order, its parts. A string field is text and a field holding named_struct('uri', ...) is a file; a top-level named_struct('uri', ...) is itself a single file. A struct is SQL's only ordered container that can hold different types -- an array of mixed parts does not type-check. - Parts are read positionally. Field names are never consulted: an unnamed part is col1, col2..., but a bare column keeps its own name. - file() builds the same thing for the DataFrame API. - The media type is inferred from the file extension unless content_type is given. Requires Spark 4.0, now stated in the README and the docstrings: named arguments do not exist in any UDF form before it. --- README.md | 69 ++- google/cloud/dataproc_ml/sql/__init__.py | 16 +- google/cloud/dataproc_ml/sql/_ai_generate.py | 413 +++++++++++----- google/cloud/dataproc_ml/sql/_endpoint.py | 123 +++++ google/cloud/dataproc_ml/sql/_parts.py | 198 ++++++++ google/cloud/dataproc_ml/sql/_retry.py | 54 +-- tests/integration/sql/test_ai_generate.py | 316 ++++++++++-- tests/unit/sql/test_ai_generate.py | 481 +++++++++++++++---- tests/unit/sql/test_ddl.py | 10 +- tests/unit/sql/test_endpoint.py | 147 ++++++ tests/unit/sql/test_parts.py | 151 ++++++ tests/unit/sql/test_retry.py | 145 +++--- 12 files changed, 1757 insertions(+), 366 deletions(-) create mode 100644 google/cloud/dataproc_ml/sql/_endpoint.py create mode 100644 google/cloud/dataproc_ml/sql/_parts.py create mode 100644 tests/unit/sql/test_endpoint.py create mode 100644 tests/unit/sql/test_parts.py diff --git a/README.md b/README.md index 279f2b1..2f283e9 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,9 @@ Here are a couple of examples demonstrating how to use the handlers for distribu > Please review the [Vertex AI Generative > AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing). +> **Note:** The AI functions require Apache Spark 4.0 or later. Named arguments +> in Spark SQL, which the call style below relies on, do not exist before it. + `ai_generate` is a Spark column function that calls a Gemini model on every row. It always returns the same struct, whatever the arguments: @@ -37,10 +40,12 @@ row. It always returns the same struct, whatever the arguments: STRUCT ``` -`result` is the generated text, `full_response` is the whole model response as -JSON, and `status` is `SUCCESS`, `RATE_LIMITED`, `SAFETY_BLOCKED`, or a -description of the error. A row that fails reports it in `status` rather than -failing the query. +`result` is the generated text, and `full_response` is the whole model response +as JSON. `status` is empty when the row succeeded and otherwise describes the +failure, led by its canonical error code, so `WHERE status <> ''` collects the +rows that need attention. A row that fails reports it in `status` rather than +failing the query; a mistake in the *call*, such as a misspelled model, fails +the query immediately instead. ```python from pyspark.sql import SparkSession, functions as F @@ -56,12 +61,12 @@ result_df = df.withColumn( ).select("city", "generated.result", "generated.status") result_df.show() -# +-----+------+-------+ -# | city|result| status| -# +-----+------+-------+ -# |Paris| CDG|SUCCESS| -# |Tokyo| HND|SUCCESS| -# +-----+------+-------+ +# +-----+------+------+ +# | city|result|status| +# +-----+------+------+ +# |Paris| CDG| | +# |Tokyo| HND| | +# +-----+------+------+ ``` Pass `output_schema` as a Spark DDL string to constrain the model to a shape. @@ -88,6 +93,33 @@ ai_generate( ) ``` +`endpoint`, `model_params` and `output_schema` configure the call rather than +the row, so they have to be constant for the query. `prompt` is the argument +that varies. + +#### Prompts that include a file + +A prompt can mix text and files. Vertex AI reads the file itself, so nothing is +downloaded into Spark, and the media type is inferred from the file extension +unless you give it: + +```python +from google.cloud.dataproc_ml.sql import ai_generate, file + +invoices.withColumn( + "extracted", + ai_generate( + [F.lit("Extract the invoice details."), file(F.col("uri"))], + output_schema="invoice_number STRING, total DOUBLE", + ), +) +``` + +As with `prompt`, a plain string passed to `file` names a *column*; wrap a +literal path in `F.lit`. + +#### Spark SQL + The same function can be registered for use from Spark SQL, where the model settings are ordinary arguments and may be passed by name in any order: @@ -106,6 +138,23 @@ spark.sql(""" """).show() ``` +In SQL a prompt is either a string or a struct whose fields are, in order, the +parts of the prompt. A field holding `named_struct('uri', ...)` is a file, and +a top-level `named_struct('uri', ...)` is itself a single file: + +```sql +SELECT g.result +FROM ( + SELECT ai_generate( + prompt => struct('Extract the invoice details.', + named_struct('uri', uri)), + output_schema => 'invoice_number STRING, total DOUBLE' + ) AS g + FROM invoices +) +WHERE g.status = '' +``` + ### Generative AI (Gemini) Model Inference > **Note:** Using the `GenAiModelHandler` involves making API calls to diff --git a/google/cloud/dataproc_ml/sql/__init__.py b/google/cloud/dataproc_ml/sql/__init__.py index 12b8b99..d48b2f9 100644 --- a/google/cloud/dataproc_ml/sql/__init__.py +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -16,7 +16,7 @@ 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. +parsing. They require Apache Spark 4.0 or later. Example: >>> from pyspark.sql.functions import col @@ -26,6 +26,14 @@ ... "summary", ai_generate(col("feedback")).result ... ) +A prompt may include a file, which Vertex AI reads directly: + + >>> from google.cloud.dataproc_ml.sql import ai_generate, file + >>> + >>> described = df.withColumn( + ... "g", ai_generate([lit("Describe this image."), file(col("uri"))]) + ... ) + The same function can be registered for use from Spark SQL, where the model settings are ordinary arguments and may be passed by name: @@ -39,12 +47,18 @@ ... ).result ... FROM feedback_table ... ''') + +Every call returns ``STRUCT``. ``status`` is empty when the row succeeded, so the rows that +need attention are collected with ``WHERE g.status <> ''``. """ from ._ai_generate import ai_generate from ._ai_generate import ai_generate_udf +from ._ai_generate import file __all__ = ( "ai_generate", "ai_generate_udf", + "file", ) diff --git a/google/cloud/dataproc_ml/sql/_ai_generate.py b/google/cloud/dataproc_ml/sql/_ai_generate.py index 711a01b..6d118d8 100644 --- a/google/cloud/dataproc_ml/sql/_ai_generate.py +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -15,6 +15,7 @@ """``AI.GENERATE``: generative text and structured output over Spark columns.""" import asyncio +import dataclasses import functools import json import logging @@ -29,7 +30,9 @@ from google.cloud.dataproc_ml.sql import _client from google.cloud.dataproc_ml.sql import _ddl +from google.cloud.dataproc_ml.sql import _endpoint from google.cloud.dataproc_ml.sql import _model_params +from google.cloud.dataproc_ml.sql import _parts from google.cloud.dataproc_ml.sql import _retry logger = logging.getLogger(__name__) @@ -43,22 +46,30 @@ _BASE_RETRY_DELAY_SECONDS = 1.0 _MAX_RETRY_DELAY_SECONDS = 60.0 -# How many distinct configurations to keep resolved. A query normally uses one, -# supplied as a literal; the cache only grows when the configuration is itself -# computed per row. +# How many distinct configurations to keep resolved. A query uses one per +# distinct set of settings, and the settings are constant for a query. _CONFIG_CACHE_SIZE = 64 RESULT_FIELD = "result" FULL_RESPONSE_FIELD = "full_response" STATUS_FIELD = "status" -#: The row succeeded. A null prompt also reports success, with a null result. -STATUS_SUCCESS = "SUCCESS" -#: The request was throttled and did not succeed within the retry budget. -STATUS_RATE_LIMITED = _retry.RATE_LIMITED_STATUS -#: The model refused to answer because the prompt or response was blocked. +#: ``status`` when the row succeeded. An empty string rather than a word, so +#: that the dead-letter query is ``WHERE status <> ''`` and a caller never has +#: to know the vocabulary of success. A null prompt also reports success, with +#: a null result. +STATUS_SUCCESS = "" + +#: ``full_response`` when there is no model response to report. The field is +#: never null, so ``parse_json(full_response)`` is safe on every row. +EMPTY_RESPONSE = "{}" + +#: Leads the ``status`` of a row the model refused to answer. STATUS_SAFETY_BLOCKED = "SAFETY_BLOCKED" +#: Leads the ``status`` of a row whose prompt could not be understood. +STATUS_INVALID_ARGUMENT = "INVALID_ARGUMENT" + #: The struct every call returns, whatever the arguments. Keeping the type #: fixed is what lets ``output_schema`` be an ordinary runtime argument: the #: query can be planned without knowing its value. @@ -77,11 +88,55 @@ {"SAFETY", "PROHIBITED_CONTENT", "BLOCKLIST", "SPII"} ) -PromptArg = Union[Column, str, Sequence[Union[Column, str]]] +# The settings that describe the call rather than the row. Spark SQL cannot +# tell a literal from a column at plan time, so constancy is checked here. +_CONSTANT_ARGUMENTS = ("endpoint", "model_params", "output_schema") + + +@dataclasses.dataclass(frozen=True) +class FilePart: + """A file to include in a prompt. Build one with :func:`file`.""" + + uri: Column + content_type: Optional[Column] = None + + +PromptPart = Union[Column, str, FilePart] +PromptArg = Union[PromptPart, Sequence[PromptPart]] ModelParamsArg = Union[str, Dict[str, Any]] OutputSchemaArg = Union[str, StructType] +def file(uri: Union[Column, str], content_type=None) -> FilePart: + """References a file to be sent to the model as part of a prompt. + + The file is read by Vertex AI, not by Spark, so the service account + running the job needs read access to it. + + Following the convention of the rest of this API, a plain string names a + *column*; wrap a literal path in :func:`pyspark.sql.functions.lit`. + + Args: + uri: The file's URI, such as ``gs://bucket/invoice.pdf``, as a column + or the name of one. + content_type: The file's media type, for example + ``application/pdf``, as a column, the name of one, or a plain + string. Inferred from the file extension when not given. + + Returns: + A prompt part that :func:`ai_generate` turns into a file reference. + + Example: + >>> df.withColumn( + ... "g", + ... ai_generate( + ... [sf.lit("What is the invoice total?"), file(sf.col("uri"))] + ... ), + ... ) + """ + return FilePart(uri=uri, content_type=content_type) + + def ai_generate( prompt: PromptArg, *, @@ -96,7 +151,8 @@ def ai_generate( 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. + sequence of parts forms one prompt; text parts are concatenated in order, + and :func:`file` includes a file. Requests are retried with exponential backoff and are issued a few at a time per worker. @@ -106,13 +162,18 @@ def ai_generate( `Vertex AI Generative AI pricing `_ page. + .. note:: + Requires Apache Spark 4.0 or later. + Args: - prompt: The prompt column, the name of a prompt column, or a sequence - of parts to concatenate. + prompt: The prompt column, the name of a prompt column, a + :func:`file`, or a sequence of those forming one prompt. endpoint: The Gemini model to call, either a name such as ``gemini-3.6-flash`` or a fully qualified ``projects/.../publishers/google/models/...`` path. Defaults to - ``gemini-3.6-flash``. + ``gemini-3.6-flash``. The name is checked against the model + catalogue before any generation, so a typo fails the query rather + than filling a column with errors. model_params: Additional request settings, as a JSON string or an equivalent dictionary, following the Vertex AI ``generateContent`` request body without its ``contents`` field, for example @@ -131,9 +192,11 @@ def ai_generate( Returns: A column holding a ``STRUCT``. ``result`` is the generated text, or a JSON document - when ``output_schema`` is set. ``full_response`` is the whole model - response as JSON. ``status`` is ``SUCCESS``, ``RATE_LIMITED``, - ``SAFETY_BLOCKED``, or a description of the error. + when ``output_schema`` is set, and is null when the row failed. + ``full_response`` is the whole model response as JSON, or ``"{}"`` + when there is none. ``status`` is empty on success and otherwise + describes the failure, led by its canonical error code, so that + ``WHERE g.status <> ''`` collects the rows that need attention. Raises: ValueError: If any argument is invalid. The arguments given here are @@ -157,7 +220,7 @@ def ai_generate( # Resolving the configuration here, rather than only on the executor, # turns a malformed model_params or output_schema into an immediate error - # instead of a column of failed rows. + # instead of a failed stage. model_params_json = _to_model_params_json(model_params) output_schema_ddl = _to_output_schema_ddl(output_schema) _resolve_config(model_params_json, output_schema_ddl) @@ -197,8 +260,16 @@ def ai_generate_udf( ... ''') Unlike :func:`ai_generate`, the model settings are not fixed when the - function is built. They are ordinary arguments, so a query can vary them, - and they may be omitted or passed by name in any order. + function is built. They are ordinary arguments, so they may be omitted or + passed by name in any order — which requires Apache Spark 4.0 or later. + They must still be constant for the query: ``endpoint``, ``model_params`` + and ``output_schema`` describe the call, not the row, and a value that + varies from row to row fails the query. + + ``prompt`` is the one argument that varies. It is either a string or a + struct whose fields are, in order, the parts of the prompt: a string field + is text, and a field holding ``named_struct('uri', ...)`` is a file. A + top-level ``named_struct('uri', ...)`` is itself a single file. Args: project: The Google Cloud project billed for the calls. Inferred from @@ -215,7 +286,7 @@ def ai_generate_udf( @pandas_udf(RETURN_TYPE) def _generate( - prompt: pd.Series, + prompt: Union[pd.Series, pd.DataFrame], endpoint: pd.Series = None, model_params: pd.Series = None, output_schema: pd.Series = None, @@ -226,20 +297,33 @@ def _generate( as ``None`` rather than as a column of nulls. That is the only way to tell "not given" from "given as NULL", so it is what selects the defaults. + + A struct prompt arrives as a ``pandas.DataFrame`` with one column per + field; a string prompt arrives as a ``pandas.Series``. """ - # Cached for the life of the worker process, so this is a lookup - # rather than a new connection pool per batch. - generate_content = _generate_content or _client.get_generate_content( - project, location + model = _constant(endpoint, "endpoint") or DEFAULT_ENDPOINT + # A malformed setting is a mistake in the query, not a property of a + # row, so it is raised rather than written to every row's status. + config = _resolve_config( + _constant(model_params, "model_params"), + _constant(output_schema, "output_schema"), ) - size = len(prompt) + + if _generate_content is not None: + generate_content = _generate_content + else: + # Cached for the life of the worker process, so this is a lookup + # rather than a new connection pool per batch. + client = _client.get_client(project, location) + _endpoint.check(client, model, project, location) + generate_content = client.aio.models.generate_content + rows = _client.run_coroutine( _generate_batch( generate_content=generate_content, - prompts=prompt.tolist(), - endpoints=_argument_values(endpoint, size), - model_params=_argument_values(model_params, size), - output_schemas=_argument_values(output_schema, size), + prompts=_to_prompts(prompt), + model=model, + config=config, ) ) return pd.DataFrame(rows, columns=_RETURN_FIELDS, index=prompt.index) @@ -247,11 +331,57 @@ def _generate( return _generate -def _argument_values(argument: Optional[pd.Series], size: int) -> List[Any]: - """Returns one value per row for an argument that may be absent.""" - if argument is None: - return [None] * size - return argument.tolist() +def _constant(argument: Optional[pd.Series], name: str) -> Optional[str]: + """Returns the single value of a setting that must not vary by row. + + ``endpoint``, ``model_params`` and ``output_schema`` configure the call, + and BigQuery's ``AI.GENERATE`` likewise accepts only a literal for each. + Spark SQL cannot tell a literal from a column that happens to repeat, so + the rule is enforced here, per batch, where the values are visible. + + Args: + argument: The argument's column, or None if it was not supplied. + name: The argument's name, for the error message. + + Returns: + The value, or None if it was absent, null or blank. + + Raises: + ValueError: If the batch contains more than one distinct value. + """ + if argument is None or argument.empty: + return None + if argument.nunique(dropna=False) > 1: + seen = sorted(str(value) for value in argument.unique()[:3]) + raise ValueError( + f"{name} must be the same for every row: it configures the call, " + f"not the row, so it has to be a literal. Saw {seen}." + ) + return _argument_text(argument.iloc[0]) + + +def _to_prompts(prompt: Union[pd.Series, "pd.DataFrame"]) -> List[Any]: + """Turns one batch's prompt column into one prompt value per row. + + Args: + prompt: The prompt column, as a Series of strings or as a DataFrame + whose columns are the fields of a struct. + + Returns: + One entry per row: a string, or a list of parts for a struct prompt. + """ + if not isinstance(prompt, pd.DataFrame): + return prompt.tolist() + + if _parts.is_file_struct(list(prompt.columns)): + # ``named_struct('uri', path)`` is a single file, not a part list + # holding one string. See _parts.is_file_struct. + return [[fields] for fields in prompt.to_dict(orient="records")] + + # Field order is column order; the names are not meaningful. ``name=None`` + # yields plain tuples, which avoids the cost and the name mangling of + # building a namedtuple class per batch. + return [list(row) for row in prompt.itertuples(index=False, name=None)] def _to_model_params_json( @@ -307,11 +437,9 @@ def _resolve_config( ): """Builds the generation config for one combination of arguments. - Both arguments are ordinary columns, so in principle they can differ on - every row. In practice a query supplies literals, and the cache turns the - repeated work into a single parse per distinct configuration. The result - is only read, never mutated, so sharing one config across rows and threads - is safe. + The cache turns the per-batch parse into a single parse per distinct + configuration for the life of the worker. The result is only read, never + mutated, so sharing one config across rows and threads is safe. Args: model_params_json: The ``model_params`` argument, or None. @@ -335,91 +463,127 @@ def _resolve_config( def _to_prompt_column(prompt: PromptArg) -> Column: - """Builds the single string column that holds each row's prompt. + """Builds the single column that holds each row's prompt. + + A text-only prompt becomes one STRING column, concatenated if it was given + in parts. A prompt that includes a :func:`file` becomes a struct whose + fields are the parts in order, which is the only ordered container SQL has + that can hold values of different types. - Every path casts to STRING. A prompt column of another type would + Every text 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, FilePart): + return _to_file_column(prompt) + if isinstance(prompt, (Column, str)): + return _to_text_column(prompt) 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()) - return sf.concat(*[column.cast(StringType()) for column in columns]) + if len(parts) == 1: + return _to_prompt_column(parts[0]) + if not any(isinstance(part, FilePart) for part in parts): + return sf.concat(*[_to_text_column(part) for part in parts]) + # Aliasing the fields keeps them from accidentally spelling a file + # reference, and keeps the struct's own names out of the way: the + # worker reads parts by position, never by name. + return sf.struct( + *[ + _to_part_column(part).alias(f"_{position}") + for position, part in enumerate(parts, start=1) + ] + ) + raise ValueError( + "prompt must be a Column, a column name, a file(), or a sequence of " + f"those, got {type(prompt).__name__}." + ) + + +def _to_part_column(part: PromptPart) -> Column: + """Builds the column for one part of a multimodal prompt.""" + if isinstance(part, FilePart): + return _to_file_column(part) + return _to_text_column(part) + + +def _to_text_column(part: Union[Column, str]) -> Column: + """Builds the STRING column for one text part.""" + if isinstance(part, Column): + return part.cast(StringType()) + if isinstance(part, str): + return sf.col(part).cast(StringType()) raise ValueError( - "prompt must be a Column, a column name, or a sequence of those, got " - f"{type(prompt).__name__}." + "A text prompt part must be a Column or a column name, got " + f"{type(part).__name__}." ) +def _to_file_column(part: FilePart) -> Column: + """Builds the struct column that stands for one file reference. + + The field names are the ones the worker looks for, and are what makes a + top-level file prompt recognizable. + """ + fields = [_to_text_column(part.uri).alias(_parts.URI_FIELD)] + if part.content_type is not None: + content_type = part.content_type + if isinstance(content_type, str): + # Unlike a prompt, a media type is a constant far more often than + # it is a column, and 'application/pdf' cannot name a column. + content_type = sf.lit(content_type) + fields.append( + content_type.cast(StringType()).alias(_parts.CONTENT_TYPE_FIELD) + ) + return sf.struct(*fields) + + async def _generate_batch( *, generate_content, prompts: List[Any], - endpoints: List[Any], - model_params: List[Any], - output_schemas: List[Any], + model: str, + config, ) -> List[Dict[str, Any]]: """Generates a response for every prompt in an Arrow batch.""" semaphore = asyncio.Semaphore(_MAX_CONCURRENT_REQUESTS) - async def _one(index: int) -> Dict[str, Any]: + async def _one(prompt: Any) -> Dict[str, Any]: async with semaphore: return await _generate_one( generate_content=generate_content, - prompt=prompts[index], - endpoint=endpoints[index], - model_params=model_params[index], - output_schema=output_schemas[index], + prompt=prompt, + model=model, + config=config, ) - return await asyncio.gather(*[_one(i) for i in range(len(prompts))]) + return await asyncio.gather(*[_one(prompt) for prompt in prompts]) async def _generate_one( *, generate_content, prompt: Any, - endpoint: Any, - model_params: Any, - output_schema: Any, + model: str, + config, ) -> 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(result=None, full_response=None, status=STATUS_SUCCESS) - - # ``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) - model = _argument_text(endpoint) or DEFAULT_ENDPOINT - try: - config = _resolve_config( - _argument_text(model_params), _argument_text(output_schema) - ) + contents = _to_contents(prompt) except ValueError as e: - # Only the driver-side entry point can reject these before the query - # runs. Called from SQL the arguments are columns, so a bad value has - # to be reported per row like any other failure. - return _row(result=None, full_response=None, status=str(e)) + # A malformed part is a property of the row, unlike a malformed + # setting, so it is reported rather than raised. + return _row(result=None, status=f"{STATUS_INVALID_ARGUMENT}: {e}") + + if contents is None: + # A null prompt produces a null result without calling the model. + return _row(result=None, status=STATUS_SUCCESS) for attempt in range(_MAX_ATTEMPTS): try: response = await generate_content( - model=model, contents=text, config=config + model=model, contents=contents, config=config ) return _to_row(response) except Exception as e: # pylint: disable=broad-except @@ -443,16 +607,36 @@ async def _generate_one( ) await asyncio.sleep(delay) continue - return _row( - result=None, - full_response=None, - status=_retry.to_status(e), - ) + return _row(result=None, status=_retry.to_status(e)) # Unreachable: the loop either returns a row or retries. raise AssertionError("The retry loop ended without producing a row.") +def _to_contents(prompt: Any) -> Optional[Any]: + """Renders one row's prompt as the ``contents`` of a request. + + Args: + prompt: A string, a list of parts, or a null. + + Returns: + The request contents, or None if the row has no prompt at all. + + Raises: + ValueError: If a part is neither text nor a file reference. + """ + if isinstance(prompt, list): + return _parts.to_contents(prompt) + if _is_missing(prompt): + return None + # ``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. + return _to_text(prompt) + + def _is_missing(value: Any) -> bool: """Reports whether a value is null. @@ -466,10 +650,10 @@ def _is_missing(value: Any) -> bool: def _argument_text(value: Any) -> Optional[str]: - """Normalizes one row's value of a string argument. + """Normalizes the value of a string setting. - An argument that was not supplied, was supplied as NULL, or was supplied - as an empty string all mean the same thing: use the default. + A setting that was not supplied, was supplied as NULL, or was supplied as + an empty string all mean the same thing: use the default. """ if _is_missing(value): return None @@ -496,11 +680,13 @@ def _to_row(response) -> Dict[str, Any]: """Converts a model response into one output row.""" full_response = _dump_response(response) - if _is_safety_blocked(response): + blocked = _blocked_reason(response) + if blocked: return _row( result=None, full_response=full_response, - status=STATUS_SAFETY_BLOCKED, + status=f"{STATUS_SAFETY_BLOCKED}: the model withheld the " + f"response ({blocked}).", ) # When output_schema is set the model is constrained to it, so the text is @@ -516,10 +702,15 @@ def _to_row(response) -> Dict[str, Any]: def _row( *, result: Optional[str], - full_response: Optional[str], status: str, + full_response: str = EMPTY_RESPONSE, ) -> Dict[str, Any]: - """Builds an output row.""" + """Builds an output row. + + ``full_response`` defaults to an empty JSON document rather than to null, + so that every row can be fed to ``parse_json`` and so that a query cannot + tell a failure apart from a missing column. + """ return { RESULT_FIELD: result, FULL_RESPONSE_FIELD: full_response, @@ -527,26 +718,34 @@ def _row( } -def _is_safety_blocked(response) -> bool: - """Reports whether the model withheld the answer rather than generating it. +def _blocked_reason(response) -> Optional[str]: + """Returns why the model withheld the answer, or None if it did not. A block is reported either on the prompt, before anything is generated, or as the finish reason of the candidate. """ feedback = getattr(response, "prompt_feedback", None) - if getattr(feedback, "block_reason", None): - return True + block_reason = getattr(feedback, "block_reason", None) + if block_reason: + return _reason_name(block_reason) for candidate in getattr(response, "candidates", None) or []: reason = getattr(candidate, "finish_reason", None) if reason is None: continue - # The client library returns an enum, but a plain string is accepted - # too so that a test or a future version needs no special handling. - name = getattr(reason, "name", None) or str(reason) + name = _reason_name(reason) if name.upper() in _BLOCKED_FINISH_REASONS: - return True - return False + return name + return None + + +def _reason_name(reason) -> str: + """Names a reason the client library reports as an enum. + + A plain string is accepted too, so that a test or a future version of the + library needs no special handling. + """ + return getattr(reason, "name", None) or str(reason) def _extract_text(response) -> Optional[str]: @@ -579,4 +778,4 @@ def _dump_response(response) -> str: return response.model_dump_json(exclude_none=True) except Exception: # pylint: disable=broad-except logger.debug("Could not serialize the model response.", exc_info=True) - return "{}" + return EMPTY_RESPONSE diff --git a/google/cloud/dataproc_ml/sql/_endpoint.py b/google/cloud/dataproc_ml/sql/_endpoint.py new file mode 100644 index 0000000..8a73688 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_endpoint.py @@ -0,0 +1,123 @@ +# 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. + +"""Endpoint checking. + +A misspelled model name is a mistake in the query, not a property of a row, so +it should stop the query rather than fill a column with failures. The name is +checked once per worker process against the model catalogue, before any +generation is attempted. +""" + +import logging +import re +import threading +from typing import Dict, Optional, Tuple + +logger = logging.getLogger(__name__) + +# projects/

/locations//publishers//models/ +_QUALIFIED_PUBLISHER_PATH = re.compile( + r"^projects/[^/]+/locations/[^/]+/(publishers/[^/]+/models/.+)$" +) + +_LOCK = threading.Lock() +_CHECKED: Dict[Tuple[Optional[str], Optional[str], str], None] = {} + + +class UnknownEndpointError(ValueError): + """Raised when the endpoint does not name a model that can be called.""" + + +def normalize(endpoint: str) -> Optional[str]: + """Returns the name to look the endpoint up under, or None to skip. + + ``models.get`` resolves a bare name and a ``publishers/.../models/...`` + path, but not a fully qualified ``projects/.../locations/...`` one, which + it turns into a request to a URL that does not exist. A qualified path that + names a publisher model is therefore shortened to the part the catalogue + understands. Anything else qualified, such as a tuned model or an endpoint + in another project, is left unchecked rather than wrongly rejected. + + Args: + endpoint: The endpoint as the caller wrote it. + + Returns: + The name to look up, or None if the endpoint cannot be checked. + """ + endpoint = endpoint.strip() + match = _QUALIFIED_PUBLISHER_PATH.match(endpoint) + if match: + return match.group(1) + if endpoint.startswith("projects/"): + return None + return endpoint + + +def check( + client, endpoint: str, project: Optional[str], location: Optional[str] +) -> None: + """Verifies once per process that the endpoint names a real model. + + Args: + client: The ``google.genai.Client`` to ask. + endpoint: The endpoint as the caller wrote it. + project: The project the client is configured for, used as a cache key. + location: The location the client is configured for, likewise. + + Raises: + UnknownEndpointError: If the model does not exist or is not visible. + """ + key = (project, location, endpoint) + with _LOCK: + if key in _CHECKED: + return + + name = normalize(endpoint) + if name is None: + logger.debug("Not checking the fully qualified endpoint %s.", endpoint) + else: + try: + client.models.get(model=name) + except Exception as e: # pylint: disable=broad-except + if _is_not_found(e): + raise UnknownEndpointError( + f"The endpoint '{endpoint}' does not name a model " + "available to this project. Check the spelling, or pass a " + "fully qualified " + "'projects/.../locations/.../publishers/.../models/...' " + f"path. Cause: {e}" + ) from e + # A transient failure of the catalogue must not fail a query that + # would otherwise have worked; generation reports its own errors. + logger.warning( + "Could not check the endpoint '%s', continuing: %s", endpoint, e + ) + + with _LOCK: + _CHECKED[key] = None + + +def _is_not_found(exception: BaseException) -> bool: + """Reports whether an error means the model does not exist.""" + if getattr(exception, "code", None) == 404: + return True + status = getattr(exception, "status", None) + return isinstance(status, str) and status.upper() == "NOT_FOUND" + + +def clear_cache() -> None: + """Forgets which endpoints have been checked. Intended for tests.""" + with _LOCK: + _CHECKED.clear() diff --git a/google/cloud/dataproc_ml/sql/_parts.py b/google/cloud/dataproc_ml/sql/_parts.py new file mode 100644 index 0000000..398f6c2 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_parts.py @@ -0,0 +1,198 @@ +# 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. + +"""Multimodal prompt parts. + +A prompt is either text or an ordered sequence of parts. In Spark the only +ordered container that can hold different types is a struct, so a multimodal +prompt is a struct whose fields are, in order, the parts of the prompt. A field +holding a string is text; a field holding a struct with a ``uri`` is a file. + +By the time a struct column reaches a pandas UDF it has become a +``pandas.DataFrame`` with one column per field, and a nested struct field has +become a ``dict``. Parts are therefore read positionally and told apart by +their Python type. Field names are deliberately never consulted: an unnamed +part is called ``col1``, ``col2`` and so on, but a bare column keeps its own +name, so the names carry no reliable meaning. +""" + +import mimetypes +import posixpath +from typing import Any, Dict, List, Optional, Sequence + +#: The only field a file part must have. +URI_FIELD = "uri" +#: The optional field that overrides MIME detection. +CONTENT_TYPE_FIELD = "content_type" + +_FILE_FIELD_SETS = ( + frozenset({URI_FIELD}), + frozenset({URI_FIELD, CONTENT_TYPE_FIELD}), +) + +# Types that Python's mimetypes module does not know about but that Gemini +# accepts. Checked only after the standard table misses. +_EXTRA_MIME_TYPES = { + ".heic": "image/heic", + ".heif": "image/heif", + ".avif": "image/avif", + ".flac": "audio/flac", + ".m4a": "audio/mp4", + ".mkv": "video/x-matroska", +} + + +def is_file_struct(field_names: Sequence[str]) -> bool: + """Reports whether a struct is a lone file reference, not a part list. + + ``named_struct('uri', path)`` is the natural way to write a prompt that is + a single file, but structurally it is indistinguishable from a part list + holding one text part: both arrive as a one-column frame of strings. The + shorthand is resolved in favour of the file reading, because a part list of + exactly one string is what a plain string prompt is already for. + + Args: + field_names: The struct's field names, in order. + + Returns: + True if the struct should be read as one file reference. + """ + return frozenset(field_names) in _FILE_FIELD_SETS + + +def to_file_part(value: Dict[str, Any]): + """Converts one file struct into a Gemini part. + + Args: + value: The struct's fields, as delivered by Arrow. + + Returns: + A ``google.genai.types.Part`` referring to the file. + + Raises: + ValueError: If the struct has no usable ``uri``, carries unknown + fields, or names a file whose type cannot be determined. + """ + unknown = set(value) - {URI_FIELD, CONTENT_TYPE_FIELD} + if unknown: + raise ValueError( + "A file part accepts only 'uri' and 'content_type', got " + f"{sorted(unknown)}." + ) + + uri = value.get(URI_FIELD) + if not isinstance(uri, str) or not uri.strip(): + raise ValueError("A file part must have a non-empty 'uri' field.") + uri = uri.strip() + + content_type = value.get(CONTENT_TYPE_FIELD) + if _is_null(content_type): + content_type = None + elif not isinstance(content_type, str): + raise ValueError( + "A file part's 'content_type' must be a string, got " + f"{type(content_type).__name__}." + ) + mime_type = (content_type or "").strip() or guess_mime_type(uri) + + # Imported lazily so that importing this package does not pull in the + # Vertex AI dependency until an AI function is evaluated. + from google.genai import types # pylint: disable=import-outside-toplevel + + return types.Part.from_uri(file_uri=uri, mime_type=mime_type) + + +def guess_mime_type(uri: str) -> str: + """Infers a file's MIME type from its extension. + + Args: + uri: The file's URI, for example ``gs://bucket/invoice.pdf``. + + Returns: + The IANA MIME type. + + Raises: + ValueError: If the extension is missing or unrecognized, in which case + the caller has to say what the file is. + """ + # Only the final path segment matters, and taking it avoids a query string + # or a bucket name that happens to contain a dot being read as a suffix. + name = posixpath.basename(uri.split("?", 1)[0].rstrip("/")) + extension = posixpath.splitext(name)[1].lower() + + mime_type = _EXTRA_MIME_TYPES.get(extension) + if mime_type: + return mime_type + + mime_type, _ = mimetypes.guess_type(name) + if mime_type: + return mime_type + + raise ValueError( + f"Could not determine the media type of '{uri}' from its name. Give " + "it explicitly, for example " + "named_struct('uri', path, 'content_type', 'application/pdf')." + ) + + +def to_contents(parts: List[Any]) -> Optional[List[Any]]: + """Converts one row's parts into the ``contents`` of a request. + + Args: + parts: The struct's field values for this row, in field order. + + Returns: + The list of Gemini parts, or None if every part was null, which is + treated like a null prompt. + + Raises: + ValueError: If a part is neither text nor a file reference. + """ + # Imported lazily for the same reason as in to_file_part. + from google.genai import types # pylint: disable=import-outside-toplevel + + contents = [] + for position, part in enumerate(parts): + if _is_null(part): + # A null part contributes nothing, exactly as an empty string + # would. This keeps a row usable when only one of several optional + # parts is missing. + continue + if isinstance(part, str): + if part: + contents.append(types.Part.from_text(text=part)) + elif isinstance(part, dict): + if all(_is_null(field) for field in part.values()): + # A file reference whose every field is null is a missing + # part, not a malformed one: the row simply has no file. + continue + contents.append(to_file_part(part)) + else: + raise ValueError( + f"Prompt part {position + 1} must be text or a file struct " + f"with a 'uri' field, got {type(part).__name__}." + ) + + return contents or None + + +def _is_null(value: Any) -> bool: + """Reports whether an Arrow value is null. + + A struct field arrives as None, but a missing numeric arrives as NaN, and + a dict is never null itself. + """ + if value is None: + return True + return isinstance(value, float) and value != value diff --git a/google/cloud/dataproc_ml/sql/_retry.py b/google/cloud/dataproc_ml/sql/_retry.py index b6ca410..a48c3bd 100644 --- a/google/cloud/dataproc_ml/sql/_retry.py +++ b/google/cloud/dataproc_ml/sql/_retry.py @@ -30,20 +30,19 @@ # 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( +# Canonical error codes whose remote message is worth quoting verbatim: each +# describes a problem with the request itself, so the detail is actionable. +_DETAILED_CODES = frozenset( { "FAILED_PRECONDITION", "PERMISSION_DENIED", "NOT_FOUND", "INVALID_ARGUMENT", "OUT_OF_RANGE", + "RESOURCE_EXHAUSTED", } ) -_QUOTA_MESSAGE_PREFIX = "Quota exceeded for metric " - _HTTP_TO_CANONICAL_CODE = { 400: "INVALID_ARGUMENT", 401: "UNAUTHENTICATED", @@ -60,8 +59,7 @@ 504: "DEADLINE_EXCEEDED", } -#: Reported when the request was throttled and outlived the retry budget. -RATE_LIMITED_STATUS = "RATE_LIMITED" + def _api_error_types(): @@ -132,29 +130,13 @@ def _canonical_code(exception: BaseException) -> str: def to_status(exception: BaseException) -> str: - """Converts a model error into the row's ``status`` value. + """Describes a model error in words, for the row's ``status`` field. - Rate limiting gets its own name because it is the one failure a caller is - expected to react to programmatically, by slowing the query down or - retrying it later. Everything else is described rather than classified, - since the detail is what makes it actionable. - - Args: - exception: The exception raised while calling the model. - - Returns: - ``RATE_LIMITED``, or a human readable description. Never empty. - """ - if _canonical_code(exception) == "RESOURCE_EXHAUSTED": - return RATE_LIMITED_STATUS - return to_status_message(exception) - - -def to_status_message(exception: BaseException) -> str: - """Describes a model error in words. - - Request errors are surfaced verbatim, quota messages are passed through, - and everything else is summarized as a retryable error. + The canonical code leads the message so that a caller can classify a + failure by prefix — ``RESOURCE_EXHAUSTED`` for throttling, ``NOT_FOUND`` + for a missing file — without this library having to maintain a closed + taxonomy of its own. Success is reported as an empty string, so the + dead-letter filter is ``WHERE status != ''``. Args: exception: The exception raised while calling the model. @@ -165,16 +147,12 @@ def to_status_message(exception: BaseException) -> str: code = _canonical_code(exception) message = getattr(exception, "message", None) or str(exception) - if code in _PASSTHROUGH_CODES: + if code in _DETAILED_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." - ) + # Other backend failures are summarized rather than quoted: the remote + # message for these is usually an opaque internal identifier. + return f"{code}: a retryable error occurred calling the remote service." + def full_jitter_delay( diff --git a/tests/integration/sql/test_ai_generate.py b/tests/integration/sql/test_ai_generate.py index ba6b334..ba1576c 100644 --- a/tests/integration/sql/test_ai_generate.py +++ b/tests/integration/sql/test_ai_generate.py @@ -18,6 +18,10 @@ 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. + +The multimodal tests read fixtures from a public-to-the-project bucket; the +files are read by Vertex AI, not by Spark, so the caller's credentials are what +must be able to see them. """ import json @@ -26,7 +30,8 @@ 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 _client +from google.cloud.dataproc_ml.sql import ai_generate, ai_generate_udf, file _AIRPORT_PROMPT = ( "What is the airport code of the largest airport in {city}? " @@ -44,6 +49,15 @@ } ) +_FIXTURES = "gs://dataproc_ai_functions_test" +#: A product photo of a pet bed with the brand printed on a label. +_DOG_BED = f"{_FIXTURES}/images/playful-pup-dog-bed.png" +#: A close-up photograph of a white daisy. A JPEG, so that MIME detection is +#: exercised on more than one extension. +_FLOWER = f"{_FIXTURES}/images/100080576_f52e8ee070_n.jpg" +#: A one page invoice addressed to John Doe, number 001. +_INVOICE = f"{_FIXTURES}/pdfs/invoice.pdf" + class TestAiGenerate(unittest.TestCase): @@ -64,10 +78,6 @@ def setUpClass(cls): # 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)] @@ -82,7 +92,7 @@ def test_generates_an_answer_for_every_row(self): self.assertEqual(len(rows), 7) for row in rows: with self.subTest(city=row["city"]): - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["status"], "") self.assertIsNotNone(row["result"]) self.assertIn(row["expected"], row["result"].upper()) @@ -101,6 +111,14 @@ def test_the_return_type_is_always_the_same(self): plain.schema["g"].dataType, structured.schema["g"].dataType ) + def test_the_dead_letter_filter_finds_nothing_when_all_rows_succeed(self): + # The contract that makes the empty status worth having: one predicate + # separates the rows that need attention from the rest. + result = self.cities.withColumn("g", ai_generate(self._prompt())) + + self.assertEqual(result.where("g.status <> ''").count(), 0) + self.assertEqual(result.where("g.status = ''").count(), 7) + def test_full_response_carries_usage_metadata(self): result = self.cities.limit(1).withColumn( "g", ai_generate(self._prompt()) @@ -135,7 +153,7 @@ def test_model_params_are_applied(self): ) for row in rows: - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["status"], "") self.assertIsNotNone(row["result"]) def test_model_params_can_cap_the_output(self): @@ -161,7 +179,7 @@ def test_model_params_can_cap_the_output(self): ) for row in rows: - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["status"], "") self.assertLess(len(row["result"].split()), 60) def test_structured_output_is_json_in_the_result_field(self): @@ -194,7 +212,7 @@ def test_structured_output_is_json_in_the_result_field(self): self.assertEqual(len(rows), 2) for row in rows: with self.subTest(result=row["result"]): - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["status"], "") # Constrained decoding guarantees this parses and conforms. parsed = json.loads(row["result"]) self.assertIsInstance(parsed["sentiment"], str) @@ -227,8 +245,6 @@ def test_structured_output_can_be_projected_with_from_json(self): self.assertIsInstance(rows[0]["score"], int) def test_a_fully_qualified_endpoint_is_accepted(self): - from google.cloud.dataproc_ml.sql import _client - project, _ = _client.resolve_project_and_location(None, None) self.assertIsNotNone(project, "the test project must be discoverable") @@ -243,28 +259,20 @@ def test_a_fully_qualified_endpoint_is_accepted(self): .collect() ) - self.assertEqual(rows[0]["status"], "SUCCESS") + self.assertEqual(rows[0]["status"], "") self.assertIsNotNone(rows[0]["result"]) - 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. - rows = ( - self.cities.limit(2) - .withColumn( + def test_an_unknown_model_fails_the_query(self): + # A typo in the model name is a mistake in the query, not a property + # of a row: it must not cost a full scan producing a column of + # failures. BigQuery likewise rejects it before returning any row. + with self.assertRaises(Exception) as caught: + self.cities.withColumn( "g", - ai_generate( - self._prompt(), endpoint="gemini-does-not-exist" - ), - ) - .select("g.result", "g.status") - .collect() - ) + ai_generate(self._prompt(), endpoint="gemini-does-not-exist"), + ).collect() - self.assertEqual(len(rows), 2) - for row in rows: - self.assertIsNone(row["result"]) - self.assertIn("NOT_FOUND", row["status"]) + self.assertIn("gemini-does-not-exist", str(caught.exception)) def test_registered_for_spark_sql(self): self.spark.udf.register("ai_generate", ai_generate_udf()) @@ -299,30 +307,250 @@ def test_spark_sql_with_named_arguments(self): self.assertEqual(len(rows), 7) for row in rows: with self.subTest(city=row["city"]): - self.assertEqual(row["g"]["status"], "SUCCESS") + self.assertEqual(row["g"]["status"], "") parsed = json.loads(row["g"]["result"]) self.assertIn(row["expected"], parsed["code"].upper()) - def test_spark_sql_argument_may_vary_per_row(self): + def test_a_setting_that_varies_per_row_fails_the_query(self): + # endpoint, model_params and output_schema configure the call, not the + # row. BigQuery rejects a column for them at plan time; Spark SQL + # cannot tell a literal from a column, so the value is checked as it + # arrives and a second distinct one fails the query. self.spark.udf.register("ai_generate", ai_generate_udf()) self.spark.createDataFrame( [ - ("France", "Answer with the capital city only."), - ("Japan", "Answer with the capital city only."), + ("France", "capital STRING"), + ("Japan", "city STRING"), ], - ["country", "instruction"], - ).createOrReplaceTempView("questions") + ["country", "schema"], + ).repartition(1).createOrReplaceTempView("questions") - rows = self.spark.sql( - "SELECT country, ai_generate(" - " prompt => CONCAT(instruction, ' Country: ', country)," - " output_schema => 'capital STRING'" - ").result AS r FROM questions ORDER BY country" - ).collect() + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(" + " prompt => CONCAT('Capital of ', country, '?')," + " output_schema => schema" + ").result AS r FROM questions" + ).collect() + + self.assertIn("output_schema", str(caught.exception)) + + +class TestMultimodalPrompts(unittest.TestCase): + """Prompts that include a file, against the real model.""" + + @classmethod + def setUpClass(cls): + cls.spark = SparkSession.builder.getOrCreate() + cls.images = cls.spark.createDataFrame( + [(_DOG_BED,), (_FLOWER,)], ["uri"] + ) + + def test_text_and_an_image(self): + # The brand is printed on a label in the photo, so a correct answer + # can only come from having actually read the image. + rows = ( + self.spark.createDataFrame([(_DOG_BED,)], ["uri"]) + .withColumn( + "g", + ai_generate( + [ + sf.lit( + "What brand name is printed on the label of this " + "product? Answer with the brand name only." + ), + file(sf.col("uri")), + ], + model_params=_DETERMINISTIC, + ), + ) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(rows[0]["status"], "") + self.assertIn("playful pup", rows[0]["result"].lower()) + + def test_a_file_on_its_own_is_a_prompt(self): + # named_struct('uri', ...) with nothing else is one file, not a part + # list holding one string. + rows = ( + self.spark.createDataFrame([(_FLOWER,)], ["uri"]) + .withColumn( + "g", + ai_generate( + file(sf.col("uri")), + output_schema="subject STRING", + model_params=_DETERMINISTIC, + ), + ) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(rows[0]["status"], "") + subject = json.loads(rows[0]["result"])["subject"].lower() + self.assertTrue( + any(word in subject for word in ("flower", "daisy", "petal")), + subject, + ) + + def test_the_media_type_is_inferred_from_the_extension(self): + # One PNG and one JPEG in the same column, with no content_type given. + rows = ( + self.images.withColumn( + "g", + ai_generate( + [ + sf.lit( + "Is the main subject of this image a plant or a " + "manufactured object? Answer 'plant' or 'object'." + ), + file(sf.col("uri")), + ], + model_params=_DETERMINISTIC, + ), + ) + .select("uri", "g.result", "g.status") + .orderBy("uri") + .collect() + ) + + self.assertEqual(len(rows), 2) + for row in rows: + self.assertEqual(row["status"], "") + # Ordered by URI: the numeric flower filename sorts before "playful". + self.assertIn("plant", rows[0]["result"].lower()) + self.assertIn("object", rows[1]["result"].lower()) + + def test_an_explicit_content_type_is_used(self): + rows = ( + self.spark.createDataFrame([(_INVOICE, "application/pdf")], + ["uri", "kind"]) + .withColumn( + "g", + ai_generate( + [ + sf.lit("Who is this invoice addressed to?"), + file(sf.col("uri"), sf.col("kind")), + ], + model_params=_DETERMINISTIC, + ), + ) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(rows[0]["status"], "") + self.assertIn("john doe", rows[0]["result"].lower()) + + def test_a_pdf_with_structured_output(self): + rows = ( + self.spark.createDataFrame([(_INVOICE,)], ["uri"]) + .withColumn( + "g", + ai_generate( + [ + sf.lit("Extract the invoice details."), + file(sf.col("uri")), + ], + output_schema="invoice_number STRING, recipient STRING", + model_params=_DETERMINISTIC, + ), + ) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(rows[0]["status"], "") + parsed = json.loads(rows[0]["result"]) + self.assertIn("001", parsed["invoice_number"]) + self.assertIn("john doe", parsed["recipient"].lower()) + + def test_a_named_struct_is_a_file_in_spark_sql(self): + self.spark.udf.register("ai_generate", ai_generate_udf()) + self.spark.createDataFrame( + [(_DOG_BED,)], ["uri"] + ).createOrReplaceTempView("documents") + + row = self.spark.sql( + "SELECT ai_generate(" + " prompt => struct(" + " 'What brand name is on the label? Brand name only.'," + " named_struct('uri', uri)" + " )," + f" model_params => '{_DETERMINISTIC}'" + ") AS g FROM documents" + ).first() + + self.assertEqual(row["g"]["status"], "") + self.assertIn("playful pup", row["g"]["result"].lower()) + + def test_a_content_type_field_in_spark_sql(self): + self.spark.udf.register("ai_generate", ai_generate_udf()) + self.spark.createDataFrame( + [(_INVOICE,)], ["uri"] + ).createOrReplaceTempView("documents") + + row = self.spark.sql( + "SELECT ai_generate(" + " prompt => struct(" + " 'What is the invoice number? Digits only.'," + " named_struct('uri', uri, 'content_type', 'application/pdf')" + " )," + f" model_params => '{_DETERMINISTIC}'" + ") AS g FROM documents" + ).first() + + self.assertEqual(row["g"]["status"], "") + self.assertIn("001", row["g"]["result"]) + + def test_a_bad_uri_is_reported_per_row(self): + # Unlike a bad setting, a bad file is a property of one row: the good + # rows still produce answers. + rows = ( + self.spark.createDataFrame( + [(_FLOWER,), (f"{_FIXTURES}/images/no-such-image.png",)], + ["uri"], + ) + .withColumn( + "g", + ai_generate( + [sf.lit("Name the subject in one word."), + file(sf.col("uri"))], + model_params=_DETERMINISTIC, + ), + ) + .select("uri", "g.result", "g.status", "g.full_response") + .orderBy("uri") + .collect() + ) + + self.assertEqual(rows[0]["status"], "") + self.assertIsNotNone(rows[0]["result"]) + self.assertNotEqual(rows[1]["status"], "") + self.assertIsNone(rows[1]["result"]) + # Never null, so the column can always be given to parse_json. + self.assertEqual(rows[1]["full_response"], "{}") + + def test_a_uri_with_no_usable_extension_is_reported_per_row(self): + rows = ( + self.spark.createDataFrame( + [(f"{_FIXTURES}/images/mystery",)], ["uri"] + ) + .withColumn( + "g", + ai_generate( + [sf.lit("Describe this."), file(sf.col("uri"))], + model_params=_DETERMINISTIC, + ), + ) + .select("g.status") + .collect() + ) - capitals = [json.loads(row["r"])["capital"] for row in rows] - self.assertIn("Paris", capitals[0]) - self.assertIn("Tokyo", capitals[1]) + self.assertTrue(rows[0]["status"].startswith("INVALID_ARGUMENT:")) + self.assertIn("content_type", rows[0]["status"]) if __name__ == "__main__": diff --git a/tests/unit/sql/test_ai_generate.py b/tests/unit/sql/test_ai_generate.py index ea07cbf..dfe37d4 100644 --- a/tests/unit/sql/test_ai_generate.py +++ b/tests/unit/sql/test_ai_generate.py @@ -32,7 +32,7 @@ # 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 +from google.cloud.dataproc_ml.sql import ai_generate, ai_generate_udf, file # pylint: enable=ungrouped-imports @@ -126,6 +126,29 @@ async def _generate_content(*, model, contents, config): return _generate_content +def echoes_the_parts(): + """Returns a fake that reports the parts a multimodal prompt became.""" + + async def _generate_content(*, model, contents, config): + del model, config + if isinstance(contents, str): + return response_with_text(json.dumps([{"text": contents}])) + described = [] + for part in contents: + if part.text is not None: + described.append({"text": part.text}) + else: + described.append( + { + "uri": part.file_data.file_uri, + "mime_type": part.file_data.mime_type, + } + ) + return response_with_text(json.dumps(described)) + + return _generate_content + + def run(coroutine): """Runs one coroutine to completion.""" return asyncio.run(coroutine) @@ -133,9 +156,9 @@ def run(coroutine): def generate_one(**kwargs): """Calls the per-row entry point with the arguments a batch would pass.""" - kwargs.setdefault("endpoint", None) - kwargs.setdefault("model_params", None) - kwargs.setdefault("output_schema", None) + kwargs.setdefault("model", ai_generate_module.DEFAULT_ENDPOINT) + kwargs.setdefault("config", ai_generate_module._resolve_config(None, None)) + # pylint: disable-next=missing-kwoa return run(ai_generate_module._generate_one(**kwargs)) @@ -165,7 +188,7 @@ def test_a_string_names_a_column(self): column = ai_generate_module._to_prompt_column("text") self.assertEqual(df.select(column).first()[0], "hello") - def test_parts_are_concatenated_in_order(self): + def test_text_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("!")] @@ -197,6 +220,42 @@ def test_a_non_string_column_is_cast_to_string(self): self.assertEqual(value, "7") self.assertIsInstance(value, str) + def test_a_lone_file_becomes_a_uri_struct(self): + # The field names matter here: they are what tells the worker that the + # struct is one file rather than a list of parts. + df = self.spark.createDataFrame([("gs://b/x.pdf",)], ["uri"]) + column = ai_generate_module._to_prompt_column(file(sf.col("uri"))) + + field = df.select(column.alias("p")).schema["p"] + self.assertEqual( + [f.name for f in field.dataType.fields], ["uri"] + ) + self.assertEqual(df.select(column).first()[0]["uri"], "gs://b/x.pdf") + + def test_a_file_carries_its_content_type(self): + df = self.spark.createDataFrame([("gs://b/x",)], ["uri"]) + column = ai_generate_module._to_prompt_column( + file(sf.col("uri"), "application/pdf") + ) + + row = df.select(column).first()[0] + self.assertEqual(row["uri"], "gs://b/x") + self.assertEqual(row["content_type"], "application/pdf") + + def test_a_mixed_prompt_becomes_a_struct_of_parts(self): + df = self.spark.createDataFrame([("gs://b/x.pdf",)], ["uri"]) + column = ai_generate_module._to_prompt_column( + [sf.lit("Read this: "), file(sf.col("uri"))] + ) + + field = df.select(column.alias("p")).schema["p"] + # Positional names, so that a part can never be mistaken for a file + # and so that the worker is never tempted to read them. + self.assertEqual([f.name for f in field.dataType.fields], ["_1", "_2"]) + row = df.select(column).first()[0] + self.assertEqual(row["_1"], "Read this: ") + self.assertEqual(row["_2"]["uri"], "gs://b/x.pdf") + class TestReturnType(SparkTestCase): """The return type is fixed, which is what makes the arguments dynamic.""" @@ -300,6 +359,56 @@ def test_argument_text_trims(self): ) +class TestConstantArguments(unittest.TestCase): + """The settings configure the call, so they may not vary by row.""" + + def _series(self, *values): + import pandas as pd # pylint: disable=import-outside-toplevel + + return pd.Series(list(values), dtype="object") + + def test_an_absent_argument_has_no_value(self): + self.assertIsNone(ai_generate_module._constant(None, "endpoint")) + + def test_an_empty_batch_has_no_value(self): + self.assertIsNone( + ai_generate_module._constant(self._series(), "endpoint") + ) + + def test_a_repeated_value_is_the_value(self): + self.assertEqual( + ai_generate_module._constant( + self._series("gemini-3.6-flash", "gemini-3.6-flash"), + "endpoint", + ), + "gemini-3.6-flash", + ) + + def test_nulls_and_blanks_mean_the_default(self): + for values in [(None, None), ("", ""), (" ", " ")]: + with self.subTest(values=values): + self.assertIsNone( + ai_generate_module._constant( + self._series(*values), "endpoint" + ) + ) + + def test_a_varying_value_is_rejected(self): + for name in ai_generate_module._CONSTANT_ARGUMENTS: + with self.subTest(name=name): + with self.assertRaises(ValueError) as caught: + ai_generate_module._constant( + self._series("a", "b"), name + ) + self.assertIn(name, str(caught.exception)) + + def test_a_null_counts_as_a_distinct_value(self): + # Otherwise half a column of nulls would silently take the other + # half's setting. + with self.assertRaises(ValueError): + ai_generate_module._constant(self._series("a", None), "endpoint") + + class TestConfigResolution(unittest.TestCase): """Turning the string arguments into a generation config.""" @@ -382,7 +491,7 @@ def test_a_successful_row(self): prompt="hi", ) self.assertEqual(row["result"], "hello") - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["status"], "") self.assertIn("candidates", row["full_response"]) def test_full_response_is_a_json_string(self): @@ -401,8 +510,8 @@ def _fail(**_): with self.subTest(prompt=prompt): row = generate_one(generate_content=_fail, prompt=prompt) self.assertIsNone(row["result"]) - self.assertIsNone(row["full_response"]) - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["full_response"], "{}") + self.assertEqual(row["status"], "") def test_reasoning_parts_are_skipped(self): row = generate_one( @@ -424,9 +533,10 @@ def test_blocked_content_is_reported(self): prompt="hi", ) self.assertIsNone(row["result"]) - self.assertEqual(row["status"], "SAFETY_BLOCKED") + self.assertTrue(row["status"].startswith("SAFETY_BLOCKED:")) + self.assertIn(reason.name, row["status"]) # The reason is still recoverable from the full response. - self.assertIsNotNone(row["full_response"]) + self.assertIn("candidates", row["full_response"]) def test_an_ordinary_finish_reason_is_not_a_block(self): row = generate_one( @@ -444,21 +554,11 @@ def test_an_ordinary_finish_reason_is_not_a_block(self): ), prompt="hi", ) - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["status"], "") self.assertEqual(row["result"], "hi") - def test_rate_limiting_is_reported(self): - with mock.patch.object(ai_generate_module, "_MAX_ATTEMPTS", 1): - row = generate_one( - generate_content=replies_with( - api_error(429, "RESOURCE_EXHAUSTED", "slow down") - ), - prompt="hi", - ) - self.assertIsNone(row["result"]) - self.assertEqual(row["status"], "RATE_LIMITED") - - def test_a_request_error_is_described(self): + def test_a_failure_leaves_an_empty_json_document(self): + # Never null, so that parse_json(full_response) is safe on every row. with mock.patch.object(ai_generate_module, "_MAX_ATTEMPTS", 1): row = generate_one( generate_content=replies_with( @@ -466,8 +566,28 @@ def test_a_request_error_is_described(self): ), prompt="hi", ) - self.assertIsNone(row["result"]) - self.assertIn("NOT_FOUND", row["status"]) + self.assertEqual(row["full_response"], "{}") + self.assertEqual(json.loads(row["full_response"]), {}) + + def test_the_status_leads_with_the_canonical_code(self): + for code, status in [ + (429, "RESOURCE_EXHAUSTED"), + (404, "NOT_FOUND"), + (403, "PERMISSION_DENIED"), + ]: + with self.subTest(status=status): + with mock.patch.object( + ai_generate_module, "_MAX_ATTEMPTS", 1 + ): + row = generate_one( + generate_content=replies_with( + api_error(code, status, "the detail") + ), + prompt="hi", + ) + self.assertIsNone(row["result"]) + self.assertTrue(row["status"].startswith(f"{status}:")) + self.assertIn("the detail", row["status"]) def test_a_transient_failure_is_retried(self): with mock.patch.object( @@ -481,7 +601,7 @@ def test_a_transient_failure_is_retried(self): prompt="hi", ) self.assertEqual(row["result"], "recovered") - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["status"], "") def test_a_bug_in_this_library_is_not_hidden_in_status(self): async def _raise(**_): @@ -490,52 +610,43 @@ async def _raise(**_): with self.assertRaises(TypeError): generate_one(generate_content=_raise, prompt="hi") - def test_an_invalid_argument_is_reported_per_row(self): - # Called from SQL the arguments are columns, so they cannot be - # rejected on the driver. - row = generate_one( - generate_content=replies_with(response_with_text("unused")), - prompt="hi", - output_schema="a NOTATYPE", - ) - self.assertIsNone(row["result"]) - self.assertIn("output_schema", row["status"]) + def test_a_non_string_prompt_is_rendered(self): + # Arrow delivers a nullable BIGINT column as floats, so a whole number + # must not pick up a fraction just because another row was null. + row = generate_one(generate_content=answers_the_prompt(), prompt=2.0) + self.assertEqual(row["result"], "answer to 2") - def test_the_endpoint_argument_selects_the_model(self): + def test_a_malformed_part_is_reported_per_row(self): + # Unlike a malformed setting, a bad URI is a property of one row. row = generate_one( - generate_content=echoes_the_request(), - prompt="hi", - endpoint="gemini-2.5-pro", + generate_content=answers_the_prompt(), + prompt=[{"uri": "gs://bucket/file-without-an-extension"}], ) - self.assertEqual(json.loads(row["result"])["model"], "gemini-2.5-pro") + self.assertIsNone(row["result"]) + self.assertTrue(row["status"].startswith("INVALID_ARGUMENT:")) - def test_a_missing_endpoint_falls_back_to_the_default(self): - for endpoint in [None, "", " "]: - with self.subTest(endpoint=endpoint): - row = generate_one( - generate_content=echoes_the_request(), - prompt="hi", - endpoint=endpoint, - ) - self.assertEqual( - json.loads(row["result"])["model"], - ai_generate_module.DEFAULT_ENDPOINT, - ) + def test_a_prompt_of_only_null_parts_is_a_null_prompt(self): + def _fail(**_): + raise AssertionError("the model must not be called") - def test_a_fully_qualified_endpoint_is_passed_through(self): - path = "projects/p/locations/l/publishers/google/models/m" row = generate_one( - generate_content=echoes_the_request(), prompt="hi", endpoint=path + generate_content=_fail, prompt=[None, {"uri": None}] ) - self.assertEqual(json.loads(row["result"])["model"], path) + self.assertIsNone(row["result"]) + self.assertEqual(row["status"], "") - def test_a_non_string_prompt_is_rendered(self): - # Arrow delivers a nullable BIGINT column as floats, so a whole number - # must not pick up a fraction just because another row was null. + def test_parts_are_sent_in_order(self): row = generate_one( - generate_content=answers_the_prompt(), prompt=2.0 + generate_content=echoes_the_parts(), + prompt=["Read this: ", {"uri": "gs://b/x.pdf"}], + ) + self.assertEqual( + json.loads(row["result"]), + [ + {"text": "Read this: "}, + {"uri": "gs://b/x.pdf", "mime_type": "application/pdf"}, + ], ) - self.assertEqual(row["result"], "answer to 2") class TestDriverValidation(SparkTestCase): @@ -599,7 +710,7 @@ def test_every_row_gets_its_own_answer(self): ], ) for row in rows: - self.assertEqual(row["status"], "SUCCESS") + self.assertEqual(row["status"], "") def test_a_failing_row_does_not_fail_the_query(self): with mock.patch.object(ai_generate_module, "_MAX_ATTEMPTS", 1): @@ -614,13 +725,34 @@ def test_a_failing_row_does_not_fail_the_query(self): ), ), ) - .select("g.result", "g.status") + .select("g.result", "g.status", "g.full_response") .collect() ) self.assertEqual(len(rows), 4) for row in rows: self.assertIsNone(row["result"]) + self.assertEqual(row["full_response"], "{}") + + def test_the_dead_letter_filter_selects_the_failed_rows(self): + # The point of the empty status: one predicate finds what went wrong. + with mock.patch.object(ai_generate_module, "_MAX_ATTEMPTS", 1): + failed = ( + self._documents() + .withColumn( + "g", + ai_generate( + "body", + _generate_content=replies_with( + api_error(404, "NOT_FOUND", "no such model") + ), + ), + ) + .where("g.status <> ''") + .count() + ) + # The null-prompt row never calls the model, so it succeeds. + self.assertEqual(failed, 3) def test_full_response_can_be_parsed_back_into_json(self): # The spec fixes the field as a string; parse_json recovers a VARIANT @@ -675,30 +807,28 @@ def test_named_arguments_may_be_reordered_and_omitted(self): ).first() request = json.loads(row["r"]) - self.assertEqual( - request["model"], ai_generate_module.DEFAULT_ENDPOINT - ) + self.assertEqual(request["model"], ai_generate_module.DEFAULT_ENDPOINT) self.assertIsNotNone(request["response_schema"]) - def test_a_bad_argument_from_sql_is_reported_per_row(self): + def test_a_bad_setting_from_sql_fails_the_query(self): + # A mistake in the query is not a property of a row: it must not cost + # a full scan producing a column of failures. self.spark.udf.register( "ai_generate", ai_generate_udf(_generate_content=answers_the_prompt()), ) self._documents().createOrReplaceTempView("documents") - rows = self.spark.sql( - "SELECT ai_generate(prompt => body," - " output_schema => 'a NOTATYPE').status AS s FROM documents" - ).collect() + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(prompt => body," + " output_schema => 'a NOTATYPE').status AS s FROM documents" + ).collect() + self.assertIn("NOTATYPE", str(caught.exception)) - self.assertEqual(len(rows), 4) - # The null-prompt row never reaches the configuration. - for row in rows[:3]: - self.assertIn("output_schema", row["s"]) - - def test_the_argument_may_vary_per_row(self): - # A column, not a literal: each row must get its own configuration. + def test_a_setting_that_varies_per_row_fails_the_query(self): + # endpoint configures the call, not the row. BigQuery rejects a column + # outright; Spark SQL cannot, so the value is checked as it arrives. self.spark.udf.register( "ai_generate", ai_generate_udf(_generate_content=echoes_the_request()), @@ -706,16 +836,35 @@ def test_the_argument_may_vary_per_row(self): self.spark.createDataFrame( [("a", "gemini-2.5-pro"), ("b", "gemini-3.6-flash")], "body STRING, model STRING", - ).createOrReplaceTempView("varying") + ).repartition(1).createOrReplaceTempView("varying") + + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(prompt => body, endpoint => model).result" + " AS r FROM varying" + ).collect() + self.assertIn("endpoint", str(caught.exception)) + + def test_a_repeated_column_value_is_accepted(self): + # A literal and a column that happens to repeat are indistinguishable + # on the executor, so the constant one has to keep working. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self.spark.createDataFrame( + [("a", "gemini-2.5-pro"), ("b", "gemini-2.5-pro")], + "body STRING, model STRING", + ).createOrReplaceTempView("constant") rows = self.spark.sql( - "SELECT body, ai_generate(prompt => body, endpoint => model).result" - " AS r FROM varying ORDER BY body" + "SELECT ai_generate(prompt => body, endpoint => model).result AS r" + " FROM constant" ).collect() self.assertEqual( [json.loads(r["r"])["model"] for r in rows], - ["gemini-2.5-pro", "gemini-3.6-flash"], + ["gemini-2.5-pro", "gemini-2.5-pro"], ) def test_a_non_string_column_works_in_spark_sql(self): @@ -736,5 +885,171 @@ def test_a_non_string_column_works_in_spark_sql(self): ) +class TestMultimodalPrompts(SparkTestCase): + """Struct prompts, through a real Spark plan in both dialects.""" + + def _files(self): + return self.spark.createDataFrame( + [ + ("gs://bucket/invoice.pdf", "application/pdf"), + ("gs://bucket/photo.png", None), + ], + "uri STRING, content_type STRING", + ) + + def test_a_file_only_prompt(self): + rows = ( + self._files() + .withColumn( + "g", + ai_generate( + file(sf.col("uri")), + _generate_content=echoes_the_parts(), + ), + ) + .select("uri", "g.result") + .orderBy("uri") + .collect() + ) + + self.assertEqual( + [json.loads(row["result"]) for row in rows], + [ + [ + { + "uri": "gs://bucket/invoice.pdf", + "mime_type": "application/pdf", + } + ], + [{"uri": "gs://bucket/photo.png", "mime_type": "image/png"}], + ], + ) + + def test_text_and_a_file_keep_their_order(self): + rows = ( + self._files() + .limit(1) + .withColumn( + "g", + ai_generate( + [ + sf.lit("Total? "), + file(sf.col("uri")), + sf.lit(" Answer in digits."), + ], + _generate_content=echoes_the_parts(), + ), + ) + .select("g.result") + .collect() + ) + + self.assertEqual( + json.loads(rows[0]["result"]), + [ + {"text": "Total? "}, + { + "uri": "gs://bucket/invoice.pdf", + "mime_type": "application/pdf", + }, + {"text": " Answer in digits."}, + ], + ) + + def test_the_content_type_column_overrides_detection(self): + rows = ( + self._files() + .limit(1) + .withColumn( + "g", + ai_generate( + file(sf.col("uri"), sf.lit("text/plain")), + _generate_content=echoes_the_parts(), + ), + ) + .select("g.result") + .collect() + ) + + self.assertEqual( + json.loads(rows[0]["result"])[0]["mime_type"], "text/plain" + ) + + def test_a_named_struct_is_a_file_in_spark_sql(self): + self.spark.udf.register( + "ai_generate", ai_generate_udf(_generate_content=echoes_the_parts()) + ) + self._files().createOrReplaceTempView("files") + + rows = self.spark.sql( + "SELECT ai_generate(named_struct('uri', uri)).result AS r" + " FROM files ORDER BY uri" + ).collect() + + self.assertEqual( + [json.loads(row["r"])[0]["uri"] for row in rows], + ["gs://bucket/invoice.pdf", "gs://bucket/photo.png"], + ) + + def test_a_struct_of_parts_in_spark_sql(self): + self.spark.udf.register( + "ai_generate", ai_generate_udf(_generate_content=echoes_the_parts()) + ) + self._files().createOrReplaceTempView("files") + + rows = self.spark.sql( + "SELECT ai_generate(" + " prompt => struct('Summarize: ', named_struct('uri', uri))" + ").result AS r FROM files ORDER BY uri" + ).collect() + + self.assertEqual( + json.loads(rows[0]["r"]), + [ + {"text": "Summarize: "}, + { + "uri": "gs://bucket/invoice.pdf", + "mime_type": "application/pdf", + }, + ], + ) + + def test_an_unreadable_extension_is_reported_per_row(self): + self.spark.udf.register( + "ai_generate", ai_generate_udf(_generate_content=echoes_the_parts()) + ) + self.spark.createDataFrame( + [("gs://bucket/invoice.pdf",), ("gs://bucket/mystery",)], + "uri STRING", + ).createOrReplaceTempView("mixed") + + rows = self.spark.sql( + "SELECT uri, ai_generate(named_struct('uri', uri)) AS g" + " FROM mixed ORDER BY uri" + ).collect() + + # The good row still succeeds: only the row with the unusable name + # reports a problem. + self.assertEqual(rows[0]["g"]["status"], "") + self.assertTrue( + rows[1]["g"]["status"].startswith("INVALID_ARGUMENT:") + ) + + def test_an_unknown_field_is_rejected(self): + self.spark.udf.register( + "ai_generate", ai_generate_udf(_generate_content=echoes_the_parts()) + ) + self._files().createOrReplaceTempView("files") + + rows = self.spark.sql( + "SELECT ai_generate(struct(" + " named_struct('uri', uri, 'size', 1)" + ")).status AS s FROM files" + ).collect() + + for row in rows: + self.assertIn("uri", row["s"]) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/sql/test_ddl.py b/tests/unit/sql/test_ddl.py index 32d8906..b90fbcc 100644 --- a/tests/unit/sql/test_ddl.py +++ b/tests/unit/sql/test_ddl.py @@ -143,7 +143,10 @@ def test_scalar_fields(self): _ddl.parse_response_schema("a STRING, b INT"), { "type": "OBJECT", - "properties": {"a": {"type": "STRING"}, "b": {"type": "INTEGER"}}, + "properties": { + "a": {"type": "STRING"}, + "b": {"type": "INTEGER"}, + }, }, ) @@ -152,7 +155,10 @@ def test_not_null_becomes_required(self): _ddl.parse_response_schema("a STRING NOT NULL, b INT"), { "type": "OBJECT", - "properties": {"a": {"type": "STRING"}, "b": {"type": "INTEGER"}}, + "properties": { + "a": {"type": "STRING"}, + "b": {"type": "INTEGER"}, + }, "required": ["a"], }, ) diff --git a/tests/unit/sql/test_endpoint.py b/tests/unit/sql/test_endpoint.py new file mode 100644 index 0000000..39b5f02 --- /dev/null +++ b/tests/unit/sql/test_endpoint.py @@ -0,0 +1,147 @@ +# 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. + +"""Tests for the endpoint check.""" + +import unittest + +from google.genai import errors + +from google.cloud.dataproc_ml.sql import _endpoint + + +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 FakeModels: + """Records the names it was asked about, and optionally fails.""" + + def __init__(self, error=None): + self.error = error + self.asked = [] + + def get(self, *, model): + self.asked.append(model) + if self.error is not None: + raise self.error + return object() + + +class FakeClient: + + def __init__(self, error=None): + self.models = FakeModels(error) + + +class TestNormalize(unittest.TestCase): + """What ``models.get`` can actually be asked about.""" + + def test_a_bare_name_is_looked_up_as_is(self): + self.assertEqual( + _endpoint.normalize("gemini-3.6-flash"), "gemini-3.6-flash" + ) + + def test_a_name_is_trimmed(self): + self.assertEqual(_endpoint.normalize(" gemini-3.6-flash "), + "gemini-3.6-flash") + + def test_a_publisher_path_is_looked_up_as_is(self): + path = "publishers/google/models/gemini-3.6-flash" + self.assertEqual(_endpoint.normalize(path), path) + + def test_a_qualified_publisher_path_is_shortened(self): + # models.get turns a projects/... path into a URL that does not + # exist, so the part it understands is extracted. + self.assertEqual( + _endpoint.normalize( + "projects/p/locations/global/publishers/google/models/m" + ), + "publishers/google/models/m", + ) + + def test_any_other_qualified_resource_is_left_alone(self): + # A tuned model or a deployed endpoint is not in the catalogue, and + # must not be rejected merely because it cannot be checked. + for endpoint in [ + "projects/p/locations/us-central1/endpoints/123", + "projects/p/locations/us-central1/models/456", + ]: + with self.subTest(endpoint=endpoint): + self.assertIsNone(_endpoint.normalize(endpoint)) + + +class TestCheck(unittest.TestCase): + + def setUp(self): + _endpoint.clear_cache() + + def test_a_known_model_passes(self): + client = FakeClient() + _endpoint.check(client, "gemini-3.6-flash", "p", None) + self.assertEqual(client.models.asked, ["gemini-3.6-flash"]) + + def test_an_unknown_model_is_rejected(self): + client = FakeClient(api_error(404, "NOT_FOUND", "no such model")) + with self.assertRaises(_endpoint.UnknownEndpointError) as caught: + _endpoint.check(client, "gemini-nope", "p", None) + self.assertIn("gemini-nope", str(caught.exception)) + + def test_the_error_is_a_value_error(self): + # So that a caller catching ValueError, as the rest of this package + # does, still sees it. + client = FakeClient(api_error(404, "NOT_FOUND", "no such model")) + with self.assertRaises(ValueError): + _endpoint.check(client, "gemini-nope", "p", None) + + def test_a_catalogue_outage_does_not_fail_the_query(self): + # Generation reports its own errors; an unavailable catalogue must not + # break a query that would have worked. + client = FakeClient(api_error(503, "UNAVAILABLE", "try later")) + _endpoint.check(client, "gemini-3.6-flash", "p", None) + + def test_a_model_is_checked_once_per_process(self): + client = FakeClient() + for _ in range(3): + _endpoint.check(client, "gemini-3.6-flash", "p", None) + self.assertEqual(len(client.models.asked), 1) + + def test_each_model_is_checked_separately(self): + client = FakeClient() + _endpoint.check(client, "gemini-3.6-flash", "p", None) + _endpoint.check(client, "gemini-2.5-pro", "p", None) + self.assertEqual( + client.models.asked, ["gemini-3.6-flash", "gemini-2.5-pro"] + ) + + def test_the_cache_is_keyed_by_project_and_location(self): + client = FakeClient() + _endpoint.check(client, "gemini-3.6-flash", "p", None) + _endpoint.check(client, "gemini-3.6-flash", "other", None) + _endpoint.check(client, "gemini-3.6-flash", "p", "us-central1") + self.assertEqual(len(client.models.asked), 3) + + def test_an_uncheckable_endpoint_is_not_looked_up(self): + client = FakeClient() + _endpoint.check( + client, "projects/p/locations/l/endpoints/1", "p", None + ) + self.assertEqual(client.models.asked, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_parts.py b/tests/unit/sql/test_parts.py new file mode 100644 index 0000000..dc7978e --- /dev/null +++ b/tests/unit/sql/test_parts.py @@ -0,0 +1,151 @@ +# 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. + +"""Tests for the multimodal prompt parts.""" + +import unittest + +from google.cloud.dataproc_ml.sql import _parts + + +class TestFileStructRecognition(unittest.TestCase): + """Telling a lone file reference from a list of parts.""" + + def test_the_two_file_shapes(self): + for names in [["uri"], ["uri", "content_type"], + ["content_type", "uri"]]: + with self.subTest(names=names): + self.assertTrue(_parts.is_file_struct(names)) + + def test_anything_else_is_a_part_list(self): + for names in [ + [], + ["col1"], + ["uri", "col2"], + ["uri", "content_type", "size"], + ["url"], + ]: + with self.subTest(names=names): + self.assertFalse(_parts.is_file_struct(names)) + + +class TestMimeTypes(unittest.TestCase): + + def test_common_extensions(self): + for uri, expected in [ + ("gs://b/invoice.pdf", "application/pdf"), + ("gs://b/photo.JPG", "image/jpeg"), + ("gs://b/photo.png", "image/png"), + ("gs://b/clip.mp4", "video/mp4"), + ("gs://b/notes.txt", "text/plain"), + ]: + with self.subTest(uri=uri): + self.assertEqual(_parts.guess_mime_type(uri), expected) + + def test_extensions_python_does_not_know(self): + # Gemini accepts these, so a user must not have to spell them out. + for uri, expected in [ + ("gs://b/photo.heic", "image/heic"), + ("gs://b/song.flac", "audio/flac"), + ]: + with self.subTest(uri=uri): + self.assertEqual(_parts.guess_mime_type(uri), expected) + + def test_a_dot_in_the_bucket_is_not_an_extension(self): + with self.assertRaises(ValueError): + _parts.guess_mime_type("gs://my.bucket.name/object") + + def test_a_query_string_is_ignored(self): + self.assertEqual( + _parts.guess_mime_type("https://h/x.pdf?generation=17"), + "application/pdf", + ) + + def test_an_unknown_name_asks_for_the_type(self): + with self.assertRaises(ValueError) as caught: + _parts.guess_mime_type("gs://b/mystery") + self.assertIn("content_type", str(caught.exception)) + + +class TestFileParts(unittest.TestCase): + + def test_a_uri_alone(self): + part = _parts.to_file_part({"uri": "gs://b/x.pdf"}) + self.assertEqual(part.file_data.file_uri, "gs://b/x.pdf") + self.assertEqual(part.file_data.mime_type, "application/pdf") + + def test_the_content_type_wins_over_the_extension(self): + part = _parts.to_file_part( + {"uri": "gs://b/x.pdf", "content_type": "text/plain"} + ) + self.assertEqual(part.file_data.mime_type, "text/plain") + + def test_a_null_content_type_falls_back_to_the_extension(self): + part = _parts.to_file_part( + {"uri": "gs://b/x.pdf", "content_type": None} + ) + self.assertEqual(part.file_data.mime_type, "application/pdf") + + def test_the_uri_is_trimmed(self): + part = _parts.to_file_part({"uri": " gs://b/x.pdf "}) + self.assertEqual(part.file_data.file_uri, "gs://b/x.pdf") + + def test_a_missing_uri_is_rejected(self): + for value in [{}, {"uri": None}, {"uri": ""}, {"uri": " "}]: + with self.subTest(value=value): + with self.assertRaises(ValueError): + _parts.to_file_part(value) + + def test_unknown_fields_are_rejected(self): + with self.assertRaises(ValueError) as caught: + _parts.to_file_part({"uri": "gs://b/x.pdf", "size": 1}) + self.assertIn("size", str(caught.exception)) + + +class TestContents(unittest.TestCase): + """Turning one row's fields into the contents of a request.""" + + def test_text_and_files_keep_their_order(self): + contents = list( + _parts.to_contents(["before", {"uri": "gs://b/x.pdf"}, "after"]) + ) + self.assertEqual( + [ + part.text if part.text is not None else part.file_data.file_uri + for part in contents + ], + ["before", "gs://b/x.pdf", "after"], + ) + + def test_null_parts_are_skipped(self): + contents = list(_parts.to_contents([None, "text", float("nan")])) + self.assertEqual([part.text for part in contents], ["text"]) + + def test_an_empty_string_contributes_nothing(self): + self.assertIsNone(_parts.to_contents(["", ""])) + + def test_a_row_with_nothing_in_it_is_a_null_prompt(self): + for row in [[], [None], [{"uri": None}], [None, {"uri": None}]]: + with self.subTest(row=row): + self.assertIsNone(_parts.to_contents(row)) + + def test_a_part_of_the_wrong_type_is_rejected(self): + with self.assertRaises(ValueError) as caught: + _parts.to_contents(["ok", 42]) + # The position is one-based so that it reads like the query. + self.assertIn("part 2", str(caught.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_retry.py b/tests/unit/sql/test_retry.py index 27d60db..77f26b2 100644 --- a/tests/unit/sql/test_retry.py +++ b/tests/unit/sql/test_retry.py @@ -60,96 +60,79 @@ def test_backend_errors_are_reportable(self): class TestStatus(unittest.TestCase): - """``status`` names rate limiting and describes everything else.""" - - def test_rate_limiting_gets_its_own_name(self): - for error in [ - _api_error(429, "RESOURCE_EXHAUSTED", "too many requests"), - _api_error(429, "", "too many requests"), + """``status`` leads with the canonical code and describes the failure. + + Success is the empty string, so any non-empty status is a failure and the + dead-letter query never has to know a vocabulary. The code is a prefix + rather than the whole value so that a caller can still classify a failure + without this library maintaining a closed taxonomy of its own. + """ + + def test_the_canonical_code_leads(self): + for code, status in [ + (400, "INVALID_ARGUMENT"), + (403, "PERMISSION_DENIED"), + (404, "NOT_FOUND"), + (429, "RESOURCE_EXHAUSTED"), + (503, "UNAVAILABLE"), ]: - with self.subTest(error=error): - self.assertEqual(_retry.to_status(error), "RATE_LIMITED") - - def test_quota_exhaustion_is_rate_limiting(self): - message = ( - "Quota exceeded for metric aiplatform.googleapis.com/" - "generate_content_requests_per_minute_per_project" + with self.subTest(status=status): + error = _api_error(code, status, "the detail") + self.assertTrue( + _retry.to_status(error).startswith(f"{status}: ") + ) + + def test_the_code_is_derived_from_the_http_status_when_absent(self): + # The backend does not always fill in the canonical status. + self.assertTrue( + _retry.to_status(_api_error(429, "", "too many")).startswith( + "RESOURCE_EXHAUSTED: " + ) ) - error = _api_error(429, "RESOURCE_EXHAUSTED", message) - self.assertEqual(_retry.to_status(error), "RATE_LIMITED") - def test_other_failures_are_described_rather_than_classified(self): - error = _api_error(404, "NOT_FOUND", "Publisher Model was not found") + def test_actionable_failures_are_surfaced_verbatim(self): + for code, status, message in [ + ( + 400, + "INVALID_ARGUMENT", + "HTTP links are not supported for requests restricted by " + "VPCSC.", + ), + (404, "NOT_FOUND", "Publisher Model was not found"), + ( + 429, + "RESOURCE_EXHAUSTED", + "Quota exceeded for metric aiplatform.googleapis.com/" + "generate_content_requests_per_minute_per_project", + ), + ]: + with self.subTest(status=status): + error = _api_error(code, status, message) + self.assertEqual( + _retry.to_status(error), f"{status}: {message}" + ) + + def test_opaque_failures_are_summarised(self): + # The remote message for these is usually an internal identifier. + error = _api_error(503, "UNAVAILABLE", "0x8badf00d") self.assertEqual( _retry.to_status(error), - "NOT_FOUND: Publisher Model was not found", + "UNAVAILABLE: a retryable error occurred calling the remote " + "service.", ) + self.assertNotIn("0x8badf00d", _retry.to_status(error)) - def test_transport_failures_are_described(self): - status = _retry.to_status(httpx.ConnectError("no route")) - self.assertNotEqual(status, "RATE_LIMITED") - self.assertIn("UNAVAILABLE", status) - - def test_status_is_never_empty(self): - self.assertTrue(_retry.to_status(_api_error(418, "", ""))) - - -class TestStatusMessages(unittest.TestCase): - """Failures that are not rate limiting are described in words.""" - - 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_transport_failures_report_unavailable(self): + self.assertTrue( + _retry.to_status(httpx.ConnectError("no route")).startswith( + "UNAVAILABLE: " + ) ) - def test_status_is_never_empty(self): - error = _api_error(418, "", "") - self.assertTrue(_retry.to_status_message(error)) + def test_an_unmapped_failure_is_still_described(self): + status = _retry.to_status(_api_error(418, "", "")) + self.assertTrue(status) + self.assertTrue(status.startswith("UNKNOWN: ")) class TestFullJitterBackoff(unittest.TestCase): From 412e5520d48e33fd467475ec711823d60a182184 Mon Sep 17 00:00:00 2001 From: Saurabh Maurya Date: Wed, 23 Sep 2026 05:04:13 +0000 Subject: [PATCH 3/3] feat(sql): positional options struct, and support for Spark 3.5 The SQL call becomes ai_generate(prompt [, endpoint] [, options]), where options is a struct carrying model_params and output_schema. This drops the named-argument style taken from BigQuery, and with it the Spark 4.0 requirement: Dataproc images 2.2 and 2.3 both ship Spark 3.5.3, and only image 3.0 is on Spark 4. The previous commit said Spark 3.5 has no named arguments at all. That is wrong. 3.5 resolves them for built-in functions -- mask('AbCd', upperChar => 'X') returns 'XxXx', and an unknown name raises UNRECOGNIZED_PARAMETER_NAME. What 3.5 lacks is named-argument resolution for anything defined in Python. NamedArgumentExpression is Unevaluable and nothing rewrites it away ahead of a Python-defined function, so the query dies inside the stage that feeds the worker: [INTERNAL_ERROR] Cannot evaluate expression: prompt => input[0, string, true] at NamedArgumentExpression.eval(NamedArgumentExpression.scala:43) Measured for a plain udf, a scalar pandas_udf, and a Python UDTF -- the UDTF being the interesting case, since named arguments were added to Spark for table functions. Positional calls work for all three. So this is not a UDF-specific gap and there is no Python-side way around it; a Catalyst expression registered through SparkSessionExtensions would work but needs a Scala jar. None of these failures cost anything: the error is raised before the worker runs, so no generation is billed. The DataFrame API was never affected, since it passes its arguments positionally through Python. Shape of the new signature: - endpoint stays its own argument rather than folding into options, because it is the setting most likely to be edited. - Callers do not pad with nulls. ai_generate(prompt) and ai_generate(prompt, 'gemini-3.6-flash') are both valid: Spark declares no arity for a Python UDF, it applies whatever children the caller wrote, and the remaining parameters fall back to Python defaults. - The two optional arguments are told apart by type, so ai_generate(prompt, options) skips the endpoint. - An unknown field name in options is an error rather than being ignored, so a misspelt option cannot quietly do nothing. - This diverges from BigQuery's AI.GENERATE spelling. Portability of the call syntax was already partial, and running on the images customers have today is worth more. Also: - StructType.toDDL is 4.0-only, so output_schema is rendered here instead. - The pyspark floor drops to >=3.5.0,<5.0.0. - tests/utils/spark_version.py adds @requires_spark_4, which gates two tests: the fromDDL oracle for the DDL parser, and one reading full_response back with parse_json/variant_get. Both exercise Spark features the library does not itself depend on. Verified on pyspark 3.5.9 and 4.0.0: 160 and 162 unit tests, and 24 and 25 integration tests against a live Vertex endpoint. --- README.md | 51 ++++- google/cloud/dataproc_ml/sql/__init__.py | 11 +- google/cloud/dataproc_ml/sql/_ai_generate.py | 209 ++++++++++++++++--- pyproject.toml | 6 +- tests/integration/sql/test_ai_generate.py | 67 ++++-- tests/unit/sql/test_ai_generate.py | 139 ++++++++++-- tests/unit/sql/test_ddl.py | 5 + tests/utils/spark_version.py | 37 ++++ 8 files changed, 456 insertions(+), 69 deletions(-) create mode 100644 tests/utils/spark_version.py diff --git a/README.md b/README.md index 2f283e9..45cc5c8 100644 --- a/README.md +++ b/README.md @@ -30,8 +30,12 @@ Here are a couple of examples demonstrating how to use the handlers for distribu > Please review the [Vertex AI Generative > AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing). -> **Note:** The AI functions require Apache Spark 4.0 or later. Named arguments -> in Spark SQL, which the call style below relies on, do not exist before it. +> **Note:** The AI functions work on Apache Spark 3.5 and later. The only thing +> that needs Spark 4 is `parse_json` on `full_response`. Spark SQL settings +> travel in a struct rather than as `name => value` named arguments, because +> named arguments do not reach a Python UDF before Spark 4.0 — the query fails +> at run time with `[INTERNAL_ERROR] Cannot evaluate expression`. Named +> arguments do work on Spark 4 for anyone who prefers them. `ai_generate` is a Spark column function that calls a Gemini model on every row. It always returns the same struct, whatever the arguments: @@ -57,7 +61,7 @@ df = spark.createDataFrame([("Paris",), ("Tokyo",)], ["city"]) result_df = df.withColumn( "generated", - ai_generate(F.concat(F.lit("Airport code for "), F.col("city"))), + ai_generate([F.lit("Airport code for "), F.col("city")]), ).select("city", "generated.result", "generated.status") result_df.show() @@ -69,6 +73,17 @@ result_df.show() # +-----+------+------+ ``` +`full_response` is a JSON string, so it is readable on every supported +version. On Spark 4 you can also query inside it without leaving SQL: + +```python +# Spark 4 only: parse_json and variant_get were added in 4.0. +df.withColumn("generated", ai_generate(F.col("city"))).selectExpr( + "variant_get(parse_json(generated.full_response)," + " '$.usage_metadata.total_token_count', 'int') AS tokens" +) +``` + Pass `output_schema` as a Spark DDL string to constrain the model to a shape. `result` then holds a JSON document matching it, which you can project with `from_json` when you want typed columns: @@ -120,8 +135,10 @@ literal path in `F.lit`. #### Spark SQL -The same function can be registered for use from Spark SQL, where the model -settings are ordinary arguments and may be passed by name in any order: +The same function can be registered for use from Spark SQL. The call is +`ai_generate(prompt [, endpoint] [, options])`; trailing arguments may be left +out, and the `options` struct names each setting, so they can be written in any +order and omitted individually: ```python from google.cloud.dataproc_ml.sql import ai_generate_udf @@ -130,14 +147,26 @@ spark.udf.register("ai_generate", ai_generate_udf()) spark.sql(""" SELECT ai_generate( - prompt => CONCAT('Airport code for ', city), - endpoint => 'gemini-3.6-flash', - output_schema => 'code STRING' + CONCAT('Airport code for ', city), + 'gemini-3.6-flash', + named_struct('output_schema', 'code STRING') ).result FROM cities """).show() ``` +`ai_generate(body)` is a complete call. Because `endpoint` is a string and +`options` is a struct, the two are told apart by type, so there is never a +`NULL` to pad with: + +```sql +SELECT ai_generate(body, named_struct('output_schema', 'code STRING')).result +FROM cities +``` + +A name that is not a setting is an error rather than being ignored, so a typo +cannot quietly do nothing. + In SQL a prompt is either a string or a struct whose fields are, in order, the parts of the prompt. A field holding `named_struct('uri', ...)` is a file, and a top-level `named_struct('uri', ...)` is itself a single file: @@ -146,9 +175,9 @@ a top-level `named_struct('uri', ...)` is itself a single file: SELECT g.result FROM ( SELECT ai_generate( - prompt => struct('Extract the invoice details.', - named_struct('uri', uri)), - output_schema => 'invoice_number STRING, total DOUBLE' + struct('Extract the invoice details.', + named_struct('uri', uri)), + named_struct('output_schema', 'invoice_number STRING, total DOUBLE') ) AS g FROM invoices ) diff --git a/google/cloud/dataproc_ml/sql/__init__.py b/google/cloud/dataproc_ml/sql/__init__.py index d48b2f9..67549c2 100644 --- a/google/cloud/dataproc_ml/sql/__init__.py +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -16,7 +16,7 @@ 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 require Apache Spark 4.0 or later. +parsing. They work on Apache Spark 3.5 and later. Example: >>> from pyspark.sql.functions import col @@ -34,16 +34,17 @@ ... "g", ai_generate([lit("Describe this image."), file(col("uri"))]) ... ) -The same function can be registered for use from Spark SQL, where the model -settings are ordinary arguments and may be passed by name: +The same function can be registered for use from Spark SQL, where the call is +``ai_generate(prompt [, endpoint] [, options])`` and the settings travel in +the ``options`` struct: >>> from google.cloud.dataproc_ml.sql import ai_generate_udf >>> >>> spark.udf.register("ai_generate", ai_generate_udf()) >>> spark.sql(''' ... SELECT ai_generate( - ... prompt => concat('Summarize: ', feedback), - ... output_schema => 'summary STRING, sentiment STRING' + ... concat('Summarize: ', feedback), + ... named_struct('output_schema', 'summary STRING') ... ).result ... FROM feedback_table ... ''') diff --git a/google/cloud/dataproc_ml/sql/_ai_generate.py b/google/cloud/dataproc_ml/sql/_ai_generate.py index 6d118d8..8fce446 100644 --- a/google/cloud/dataproc_ml/sql/_ai_generate.py +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -92,6 +92,18 @@ # tell a literal from a column at plan time, so constancy is checked here. _CONSTANT_ARGUMENTS = ("endpoint", "model_params", "output_schema") +MODEL_PARAMS_OPTION = "model_params" +OUTPUT_SCHEMA_OPTION = "output_schema" + +#: The names allowed inside the ``options`` struct in Spark SQL. +#: +#: Settings travel in a struct rather than as named arguments because named +#: arguments do not reach a Python UDF before Spark 4.0: the analyzer leaves +#: the ``NamedArgumentExpression`` in place and it fails at run time with +#: ``[INTERNAL_ERROR] Cannot evaluate expression``. A struct carries the same +#: names positionally, so one spelling works on every supported version. +OPTION_NAMES = (MODEL_PARAMS_OPTION, OUTPUT_SCHEMA_OPTION) + @dataclasses.dataclass(frozen=True) class FilePart: @@ -163,7 +175,9 @@ def ai_generate( `_ page. .. note:: - Requires Apache Spark 4.0 or later. + Works on Apache Spark 3.5 and later. Passing a ``StructType`` for + ``output_schema`` is rendered here rather than by Spark, so it works + on 3.5 too; only ``parse_json(full_response)`` needs Spark 4. Args: prompt: The prompt column, the name of a prompt column, a @@ -230,14 +244,15 @@ def ai_generate( location=location, _generate_content=_generate_content, ) - # An unset argument is passed as a typed NULL rather than left out, so - # that the executor always sees the same four columns. Without the cast an - # unset one would be a void column, which is not a string. + # The settings are passed as a struct rather than as separate columns so + # that the executor sees exactly the signature a Spark SQL caller uses. + # Unset fields are typed NULLs rather than omitted, so the worker always + # receives the same three columns; without the cast an unset one would be + # a void column, which is not a string. return generate( _to_prompt_column(prompt), sf.lit(endpoint).cast(StringType()), - sf.lit(model_params_json).cast(StringType()), - sf.lit(output_schema_ddl).cast(StringType()), + _to_options_column(model_params_json, output_schema_ddl), ) @@ -252,19 +267,39 @@ def ai_generate_udf( >>> spark.udf.register("ai_generate", ai_generate_udf()) >>> spark.sql(''' ... SELECT ai_generate( - ... prompt => concat('Summarize: ', body), - ... endpoint => 'gemini-3.6-flash', - ... output_schema => 'summary STRING, sentiment STRING' + ... concat('Summarize: ', body), + ... 'gemini-3.6-flash', + ... named_struct('output_schema', 'summary STRING') ... ).result ... FROM articles ... ''') - Unlike :func:`ai_generate`, the model settings are not fixed when the - function is built. They are ordinary arguments, so they may be omitted or - passed by name in any order — which requires Apache Spark 4.0 or later. - They must still be constant for the query: ``endpoint``, ``model_params`` - and ``output_schema`` describe the call, not the row, and a value that - varies from row to row fails the query. + The signature is ``ai_generate(prompt [, endpoint] [, options])``. Every + argument is positional, and trailing ones may simply be left out: + ``ai_generate(body)`` is a complete call. + + Settings travel in the ``options`` struct, named by its field names, which + may appear in any order and be omitted individually. A name that is not a + setting is an error rather than being ignored, so a typo cannot silently + do nothing:: + + named_struct('output_schema', 'a STRING', 'model_params', '{"topP":1}') + + ``options`` may take the place of ``endpoint`` when only it is needed, so + there is never a NULL to pad with:: + + ai_generate(body, named_struct('output_schema', 'a STRING')) + + A struct is used instead of named arguments because ``name => value`` does + not reach a Python UDF before Spark 4.0. The analyzer leaves the + ``NamedArgumentExpression`` unresolved and the query fails at run time with + ``[INTERNAL_ERROR] Cannot evaluate expression``. Named arguments still work + on Spark 4.0 for callers who prefer them, since the parameters are named + ``prompt``, ``endpoint`` and ``options``. + + ``endpoint`` and the ``options`` fields must be constant for the query: + they describe the call, not the row, so a value that varies from row to + row fails the query. ``prompt`` is the one argument that varies. It is either a string or a struct whose fields are, in order, the parts of the prompt: a string field @@ -281,15 +316,14 @@ def ai_generate_udf( Returns: A pandas user defined function taking ``prompt`` and, optionally, - ``endpoint``, ``model_params`` and ``output_schema``. + ``endpoint`` and ``options``. """ @pandas_udf(RETURN_TYPE) def _generate( prompt: Union[pd.Series, pd.DataFrame], - endpoint: pd.Series = None, - model_params: pd.Series = None, - output_schema: pd.Series = None, + endpoint: Union[pd.Series, pd.DataFrame] = None, + options: Union[pd.Series, pd.DataFrame] = None, ) -> pd.DataFrame: """Runs on the executor, once per Arrow batch. @@ -301,12 +335,14 @@ def _generate( A struct prompt arrives as a ``pandas.DataFrame`` with one column per field; a string prompt arrives as a ``pandas.Series``. """ + endpoint, options = _split_settings(endpoint, options) model = _constant(endpoint, "endpoint") or DEFAULT_ENDPOINT # A malformed setting is a mistake in the query, not a property of a # row, so it is raised rather than written to every row's status. + settings = _to_options(options) config = _resolve_config( - _constant(model_params, "model_params"), - _constant(output_schema, "output_schema"), + settings.get(MODEL_PARAMS_OPTION), + settings.get(OUTPUT_SCHEMA_OPTION), ) if _generate_content is not None: @@ -331,6 +367,101 @@ def _generate( return _generate +def _split_settings(second, third): + """Sorts the two optional arguments into ``endpoint`` and ``options``. + + ``endpoint`` is always a string and ``options`` is always a struct, so the + two can be told apart by what arrived. That is what lets a caller who only + wants options skip ``endpoint`` without padding the call with a NULL:: + + ai_generate(body, named_struct('output_schema', 'a STRING')) + + Args: + second: Whatever was passed in the second position, or None. + third: Whatever was passed in the third position, or None. + + Returns: + A pair of (endpoint, options), either of which may be None. + + Raises: + ValueError: If a struct was passed in the second position and + something was passed in the third as well, which would mean two + options structs. + """ + if isinstance(second, pd.DataFrame): + if third is not None: + raise ValueError( + "ai_generate takes options once. The second argument is " + "already a struct, so the third cannot be given. The call is " + "ai_generate(prompt, endpoint, options), and endpoint may be " + "left out." + ) + return None, second + return second, third + + +def _to_options(options) -> Dict[str, Optional[str]]: + """Reads the settings out of the ``options`` struct. + + The field names are the setting names, so they may be written in any order + and any of them left out. An unrecognized name is rejected rather than + ignored: silently doing nothing in response to a typo is how a caller ends + up believing a setting took effect when it did not. + + Args: + options: The options column: a DataFrame for a struct, a Series if the + caller wrote NULL, or None if the argument was omitted. + + Returns: + The settings that were given, by name. + + Raises: + ValueError: If the argument is not a struct, or names something that + is not a setting. + """ + if options is None: + return {} + allowed = ", ".join(OPTION_NAMES) + if not isinstance(options, pd.DataFrame): + # A bare NULL in this position is how a caller spells "no options". + if options.isna().all(): + return {} + raise ValueError( + "The third argument to ai_generate must be a struct of settings, " + f"such as named_struct({OUTPUT_SCHEMA_OPTION!r}, 'a STRING'). " + f"Allowed names are {allowed}." + ) + + unknown = [name for name in options.columns if name not in OPTION_NAMES] + if unknown: + listed = ", ".join(sorted(unknown)) + raise ValueError( + f"Unknown option(s) {listed} passed to ai_generate. Allowed " + f"names are {allowed}. Note that named_struct() names its fields " + "positionally, so a field built without a name will not match." + ) + return { + name: _constant(options[name], f"options.{name}") + for name in options.columns + } + + +def _to_options_column( + model_params_json: Optional[str], output_schema_ddl: Optional[str] +) -> Column: + """Builds the options struct that :func:`ai_generate` passes to the UDF. + + Both fields are always present, as typed NULLs when unset, so that the + executor sees one shape no matter which settings were given. + """ + return sf.struct( + sf.lit(model_params_json).cast(StringType()).alias(MODEL_PARAMS_OPTION), + sf.lit(output_schema_ddl) + .cast(StringType()) + .alias(OUTPUT_SCHEMA_OPTION), + ) + + def _constant(argument: Optional[pd.Series], name: str) -> Optional[str]: """Returns the single value of a setting that must not vary by row. @@ -415,22 +546,50 @@ def _to_output_schema_ddl( ) -> Optional[str]: """Renders the ``output_schema`` argument as a DDL string. - A ``StructType`` is accepted for convenience in Python and rendered with - Spark's own ``toDDL``, so that the executor only ever sees DDL and only - one parser has to agree with Spark. + A ``StructType`` is accepted for convenience in Python and rendered here, + so that the executor only ever sees DDL and only one parser has to agree + with Spark. """ if output_schema is None: return None if isinstance(output_schema, str): return output_schema if isinstance(output_schema, StructType): - return output_schema.toDDL() + return _struct_to_ddl(output_schema) raise ValueError( "output_schema must be a DDL string or a pyspark StructType, got " f"{type(output_schema).__name__}." ) +def _struct_to_ddl(schema: StructType) -> str: + """Renders a ``StructType`` as the DDL the UDF takes. + + ``StructType.toDDL`` would do this, but it was only added in PySpark 4 and + this library supports 3.5. Rendering it here also means the string is the + same on every version, rather than whatever the local Spark produces. A + test asserts this agrees with ``toDDL`` where that exists. + """ + return ", ".join( + " ".join( + part + for part in ( + _quote_field_name(field.name), + field.dataType.simpleString(), + "" if field.nullable else "NOT NULL", + ) + if part + ) + for field in schema.fields + ) + + +def _quote_field_name(name: str) -> str: + """Backquotes a field name so that any name survives the round trip.""" + escaped = name.replace("`", "``") + return f"`{escaped}`" + + @functools.lru_cache(maxsize=_CONFIG_CACHE_SIZE) def _resolve_config( model_params_json: Optional[str], output_schema_ddl: Optional[str] diff --git a/pyproject.toml b/pyproject.toml index 9381ea2..fcc3911 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,8 +40,10 @@ test = [ "pytest", "torchvision>=0.21.0, <1.0.0", "pillow>=11.3.0, <12.0.0", - # For running tests, a full pyspark installation is needed. - "pyspark~=4.0.0", + # For running tests, a full pyspark installation is needed. The range + # spans the versions Dataproc ships: images 2.2 and 2.3 are Spark 3.5, + # image 3.0 is Spark 4. + "pyspark>=3.5.0, <5.0.0", ] dev = [ diff --git a/tests/integration/sql/test_ai_generate.py b/tests/integration/sql/test_ai_generate.py index ba1576c..e56253d 100644 --- a/tests/integration/sql/test_ai_generate.py +++ b/tests/integration/sql/test_ai_generate.py @@ -32,6 +32,7 @@ from google.cloud.dataproc_ml.sql import _client from google.cloud.dataproc_ml.sql import ai_generate, ai_generate_udf, file +from tests.utils.spark_version import requires_spark_4 _AIRPORT_PROMPT = ( "What is the airport code of the largest airport in {city}? " @@ -128,6 +129,7 @@ def test_full_response_carries_usage_metadata(self): self.assertIn("candidates", payload) self.assertIn("usage_metadata", payload) + @requires_spark_4 def test_full_response_can_be_queried_with_parse_json(self): # The field is a string, so anything structural goes through # parse_json. This is the documented way to reach inside it. @@ -289,18 +291,20 @@ def test_registered_for_spark_sql(self): for row in rows: self.assertIsNotNone(row["code"]) - def test_spark_sql_with_named_arguments(self): + def test_spark_sql_with_options(self): # The call style the spec uses, end to end against the real model. self.spark.udf.register("ai_generate", ai_generate_udf()) self.cities.createOrReplaceTempView("cities") rows = self.spark.sql( "SELECT city, expected, ai_generate(" - " prompt => CONCAT('Airport code of the largest airport in '," - " city, '? Three letters only.')," - " endpoint => 'gemini-3.6-flash'," - f" model_params => '{_DETERMINISTIC}'," - " output_schema => 'code STRING'" + " CONCAT('Airport code of the largest airport in '," + " city, '? Three letters only.')," + " 'gemini-3.6-flash'," + " named_struct(" + f" 'model_params', '{_DETERMINISTIC}'," + " 'output_schema', 'code STRING'" + " )" ") AS g FROM cities" ).collect() @@ -311,6 +315,44 @@ def test_spark_sql_with_named_arguments(self): parsed = json.loads(row["g"]["result"]) self.assertIn(row["expected"], parsed["code"].upper()) + def test_options_may_replace_endpoint(self): + # endpoint is a string and options is a struct, so they are told apart + # by type and a caller never has to write a NULL placeholder. + self.spark.udf.register("ai_generate", ai_generate_udf()) + self.cities.limit(2).createOrReplaceTempView("few_cities") + + rows = self.spark.sql( + "SELECT city, expected, ai_generate(" + " CONCAT('Airport code of the largest airport in '," + " city, '? Three letters only.')," + " named_struct(" + f" 'model_params', '{_DETERMINISTIC}'," + " 'output_schema', 'code STRING'" + " )" + ") AS g FROM few_cities" + ).collect() + + self.assertEqual(len(rows), 2) + for row in rows: + with self.subTest(city=row["city"]): + self.assertEqual(row["g"]["status"], "") + parsed = json.loads(row["g"]["result"]) + self.assertIn(row["expected"], parsed["code"].upper()) + + def test_an_unknown_option_fails_the_query(self): + # A typo must not look like a setting that quietly took effect. + self.spark.udf.register("ai_generate", ai_generate_udf()) + self.cities.createOrReplaceTempView("cities") + + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(city," + " named_struct('output_schemaa', 'code STRING')).result AS r" + " FROM cities" + ).collect() + + self.assertIn("output_schemaa", str(caught.exception)) + def test_a_setting_that_varies_per_row_fails_the_query(self): # endpoint, model_params and output_schema configure the call, not the # row. BigQuery rejects a column for them at plan time; Spark SQL @@ -328,8 +370,9 @@ def test_a_setting_that_varies_per_row_fails_the_query(self): with self.assertRaises(Exception) as caught: self.spark.sql( "SELECT ai_generate(" - " prompt => CONCAT('Capital of ', country, '?')," - " output_schema => schema" + " CONCAT('Capital of ', country, '?')," + " NULL," + " named_struct('output_schema', schema)" ").result AS r FROM questions" ).collect() @@ -475,11 +518,11 @@ def test_a_named_struct_is_a_file_in_spark_sql(self): row = self.spark.sql( "SELECT ai_generate(" - " prompt => struct(" + " struct(" " 'What brand name is on the label? Brand name only.'," " named_struct('uri', uri)" " )," - f" model_params => '{_DETERMINISTIC}'" + f" named_struct('model_params', '{_DETERMINISTIC}')" ") AS g FROM documents" ).first() @@ -494,11 +537,11 @@ def test_a_content_type_field_in_spark_sql(self): row = self.spark.sql( "SELECT ai_generate(" - " prompt => struct(" + " struct(" " 'What is the invoice number? Digits only.'," " named_struct('uri', uri, 'content_type', 'application/pdf')" " )," - f" model_params => '{_DETERMINISTIC}'" + f" named_struct('model_params', '{_DETERMINISTIC}')" ") AS g FROM documents" ).first() diff --git a/tests/unit/sql/test_ai_generate.py b/tests/unit/sql/test_ai_generate.py index dfe37d4..ed077ad 100644 --- a/tests/unit/sql/test_ai_generate.py +++ b/tests/unit/sql/test_ai_generate.py @@ -33,6 +33,7 @@ # 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, file +from tests.utils.spark_version import requires_spark_4 # pylint: enable=ungrouped-imports @@ -754,6 +755,7 @@ def test_the_dead_letter_filter_selects_the_failed_rows(self): # The null-prompt row never calls the model, so it succeeds. self.assertEqual(failed, 3) + @requires_spark_4 def test_full_response_can_be_parsed_back_into_json(self): # The spec fixes the field as a string; parse_json recovers a VARIANT # for anyone who wants to query inside it. @@ -768,7 +770,7 @@ def test_full_response_can_be_parsed_back_into_json(self): ) self.assertEqual(result.first()["role"], "model") - def test_registered_for_spark_sql_with_named_arguments(self): + def test_registered_for_spark_sql(self): self.spark.udf.register( "ai_generate", ai_generate_udf(_generate_content=echoes_the_request()), @@ -777,9 +779,9 @@ def test_registered_for_spark_sql_with_named_arguments(self): row = self.spark.sql( "SELECT ai_generate(" - " prompt => body," - " endpoint => 'gemini-2.5-pro'," - " output_schema => 'sentiment STRING'" + " body," + " 'gemini-2.5-pro'," + " named_struct('output_schema', 'sentiment STRING')" ").result AS r FROM documents WHERE id = 1" ).first() @@ -794,7 +796,9 @@ def test_registered_for_spark_sql_with_named_arguments(self): }, ) - def test_named_arguments_may_be_reordered_and_omitted(self): + def test_options_may_be_reordered(self): + # The struct's field names carry the meaning, so the order they are + # written in must not matter. self.spark.udf.register( "ai_generate", ai_generate_udf(_generate_content=echoes_the_request()), @@ -802,14 +806,103 @@ def test_named_arguments_may_be_reordered_and_omitted(self): self._documents().createOrReplaceTempView("documents") row = self.spark.sql( - "SELECT ai_generate(output_schema => 'a STRING', prompt => body)" - ".result AS r FROM documents WHERE id = 1" + "SELECT ai_generate(body, 'gemini-2.5-pro', named_struct(" + " 'output_schema', 'a STRING'," + " 'model_params', '{\"temperature\": 0.5}'" + ")).result AS r FROM documents WHERE id = 1" + ).first() + + request = json.loads(row["r"]) + self.assertIsNotNone(request["response_schema"]) + self.assertEqual(request["temperature"], 0.5) + + def test_options_may_replace_endpoint(self): + # endpoint is a string and options is a struct, so a caller who wants + # only options does not have to pad the call with a NULL. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self._documents().createOrReplaceTempView("documents") + + row = self.spark.sql( + "SELECT ai_generate(body, named_struct(" + " 'output_schema', 'a STRING'" + ")).result AS r FROM documents WHERE id = 1" ).first() request = json.loads(row["r"]) self.assertEqual(request["model"], ai_generate_module.DEFAULT_ENDPOINT) self.assertIsNotNone(request["response_schema"]) + def test_the_prompt_alone_is_a_complete_call(self): + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self._documents().createOrReplaceTempView("documents") + + row = self.spark.sql( + "SELECT ai_generate(body).result AS r FROM documents WHERE id = 1" + ).first() + + request = json.loads(row["r"]) + self.assertEqual(request["model"], ai_generate_module.DEFAULT_ENDPOINT) + self.assertIsNone(request["response_schema"]) + + def test_an_unknown_option_fails_the_query(self): + # Ignoring a name we do not know would let a typo look like a setting + # that took effect. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self._documents().createOrReplaceTempView("documents") + + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(body, named_struct(" + " 'output_schemaa', 'a STRING'" + ")).result AS r FROM documents" + ).collect() + message = str(caught.exception) + self.assertIn("output_schemaa", message) + self.assertIn("output_schema", message) + + def test_explicit_nulls_mean_the_defaults(self): + # Omitting an argument and writing NULL arrive differently on the + # executor, and must not behave differently. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self._documents().createOrReplaceTempView("documents") + + row = self.spark.sql( + "SELECT ai_generate(body, NULL, NULL).result AS r" + " FROM documents WHERE id = 1" + ).first() + + request = json.loads(row["r"]) + self.assertEqual(request["model"], ai_generate_module.DEFAULT_ENDPOINT) + self.assertIsNone(request["response_schema"]) + + def test_options_cannot_be_given_twice(self): + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self._documents().createOrReplaceTempView("documents") + + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(body," + " named_struct('output_schema', 'a STRING')," + " named_struct('model_params', '{}')" + ").result AS r FROM documents" + ).collect() + self.assertIn("options", str(caught.exception)) + def test_a_bad_setting_from_sql_fails_the_query(self): # A mistake in the query is not a property of a row: it must not cost # a full scan producing a column of failures. @@ -821,8 +914,9 @@ def test_a_bad_setting_from_sql_fails_the_query(self): with self.assertRaises(Exception) as caught: self.spark.sql( - "SELECT ai_generate(prompt => body," - " output_schema => 'a NOTATYPE').status AS s FROM documents" + "SELECT ai_generate(body, NULL," + " named_struct('output_schema', 'a NOTATYPE')).status AS s" + " FROM documents" ).collect() self.assertIn("NOTATYPE", str(caught.exception)) @@ -840,11 +934,29 @@ def test_a_setting_that_varies_per_row_fails_the_query(self): with self.assertRaises(Exception) as caught: self.spark.sql( - "SELECT ai_generate(prompt => body, endpoint => model).result" - " AS r FROM varying" + "SELECT ai_generate(body, model).result AS r FROM varying" ).collect() self.assertIn("endpoint", str(caught.exception)) + def test_an_option_that_varies_per_row_fails_the_query(self): + # The same rule has to hold for a setting inside the options struct. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=echoes_the_request()), + ) + self.spark.createDataFrame( + [("a", "x STRING"), ("b", "y STRING")], + "body STRING, schema STRING", + ).repartition(1).createOrReplaceTempView("varying_options") + + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(body, NULL," + " named_struct('output_schema', schema)).result AS r" + " FROM varying_options" + ).collect() + self.assertIn("output_schema", str(caught.exception)) + def test_a_repeated_column_value_is_accepted(self): # A literal and a column that happens to repeat are indistinguishable # on the executor, so the constant one has to keep working. @@ -858,8 +970,7 @@ def test_a_repeated_column_value_is_accepted(self): ).createOrReplaceTempView("constant") rows = self.spark.sql( - "SELECT ai_generate(prompt => body, endpoint => model).result AS r" - " FROM constant" + "SELECT ai_generate(body, model).result AS r FROM constant" ).collect() self.assertEqual( @@ -999,7 +1110,7 @@ def test_a_struct_of_parts_in_spark_sql(self): rows = self.spark.sql( "SELECT ai_generate(" - " prompt => struct('Summarize: ', named_struct('uri', uri))" + " struct('Summarize: ', named_struct('uri', uri))" ").result AS r FROM files ORDER BY uri" ).collect() diff --git a/tests/unit/sql/test_ddl.py b/tests/unit/sql/test_ddl.py index b90fbcc..e34ad9b 100644 --- a/tests/unit/sql/test_ddl.py +++ b/tests/unit/sql/test_ddl.py @@ -38,6 +38,7 @@ ) from google.cloud.dataproc_ml.sql import _ddl +from tests.utils.spark_version import requires_spark_4 # The reference conversion, expressed against Spark's own parsed types. Keeping # it in the test rather than in the library is deliberate: it is the oracle the @@ -121,7 +122,11 @@ def setUpClass(cls): # without one on an executor, which is why it exists. cls.spark = SparkSession.builder.master("local[1]").getOrCreate() + @requires_spark_4 def test_matches_fromddl(self): + # StructType.fromDDL is the oracle and was added in PySpark 4, so on + # 3.5 there is nothing to compare against. The parser itself is + # version independent and its other tests still run. for schema in _AGREEING_SCHEMAS: with self.subTest(schema=schema): expected = _reference_schema(StructType.fromDDL(schema)) diff --git a/tests/utils/spark_version.py b/tests/utils/spark_version.py new file mode 100644 index 0000000..915317b --- /dev/null +++ b/tests/utils/spark_version.py @@ -0,0 +1,37 @@ +# 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. + +"""Skips for tests that need a Spark feature the running version lacks. + +The library supports Spark 3.5 and later, but a few tests are written against +features that only exist in Spark 4: ``StructType.fromDDL``, used as an oracle +for the DDL parser, and the VARIANT functions, used to show that +``full_response`` can be queried. Neither is used by the library itself, so +they gate the test rather than the feature. +""" + +import unittest + +import pyspark + +_MAJOR = int(pyspark.__version__.split(".", maxsplit=1)[0]) + +#: True when the running PySpark is 4.0 or later. +IS_SPARK_4 = _MAJOR >= 4 + +#: Skips a test that needs an API added in PySpark 4. +requires_spark_4 = unittest.skipUnless( + IS_SPARK_4, + f"needs PySpark 4 (running {pyspark.__version__})", +)