diff --git a/README.md b/README.md index 2a5cab8..45cc5c8 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,166 @@ 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). + +> **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: + +``` +STRUCT +``` + +`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 +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.lit("Airport code for "), F.col("city")]), +).select("city", "generated.result", "generated.status") + +result_df.show() +# +-----+------+------+ +# | city|result|status| +# +-----+------+------+ +# |Paris| CDG| | +# |Tokyo| HND| | +# +-----+------+------+ +``` + +`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: + +```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}}', +) +``` + +`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. 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 + +spark.udf.register("ai_generate", ai_generate_udf()) + +spark.sql(""" + SELECT ai_generate( + 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: + +```sql +SELECT g.result +FROM ( + SELECT ai_generate( + struct('Extract the invoice details.', + named_struct('uri', uri)), + named_struct('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/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..67549c2 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -0,0 +1,65 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Declarative AI functions for Spark DataFrames and Spark SQL. + +These functions apply Google's Gemini models to a DataFrame column without any +of the usual plumbing: no client setup, no retry loops and no response +parsing. They work on Apache Spark 3.5 and later. + +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 + ... ) + +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 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( + ... concat('Summarize: ', feedback), + ... named_struct('output_schema', 'summary STRING') + ... ).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 new file mode 100644 index 0000000..8fce446 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -0,0 +1,940 @@ +# 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 dataclasses +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 _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__) + +#: 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 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" + +#: ``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. +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"} +) + +# 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") + +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: + """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, + *, + 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 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. + + .. note:: + Using this function incurs Vertex AI charges. See the + `Vertex AI Generative AI pricing + `_ page. + + .. note:: + 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 + :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``. 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 + ``'{"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, 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 + 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 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) + + generate = ai_generate_udf( + project=project, + location=location, + _generate_content=_generate_content, + ) + # 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()), + _to_options_column(model_params_json, output_schema_ddl), + ) + + +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( + ... concat('Summarize: ', body), + ... 'gemini-3.6-flash', + ... named_struct('output_schema', 'summary STRING') + ... ).result + ... FROM articles + ... ''') + + 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 + 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 + 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`` and ``options``. + """ + + @pandas_udf(RETURN_TYPE) + def _generate( + prompt: Union[pd.Series, pd.DataFrame], + endpoint: Union[pd.Series, pd.DataFrame] = None, + options: Union[pd.Series, pd.DataFrame] = 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. + + 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( + settings.get(MODEL_PARAMS_OPTION), + settings.get(OUTPUT_SCHEMA_OPTION), + ) + + 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=_to_prompts(prompt), + model=model, + config=config, + ) + ) + return pd.DataFrame(rows, columns=_RETURN_FIELDS, index=prompt.index) + + 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. + + ``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( + 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 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 _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] +): + """Builds the generation config for one combination of arguments. + + 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. + 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 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 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, 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.") + 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( + "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], + 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(prompt: Any) -> Dict[str, Any]: + async with semaphore: + return await _generate_one( + generate_content=generate_content, + prompt=prompt, + model=model, + config=config, + ) + + return await asyncio.gather(*[_one(prompt) for prompt in prompts]) + + +async def _generate_one( + *, + generate_content, + prompt: Any, + model: str, + config, +) -> Dict[str, Any]: + """Generates the response for one row, retrying transient failures.""" + try: + contents = _to_contents(prompt) + except ValueError as 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=contents, 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, 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. + + 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 the value of a string setting. + + 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 + 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) + + blocked = _blocked_reason(response) + if blocked: + return _row( + result=None, + full_response=full_response, + 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 + # 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], + status: str, + full_response: str = EMPTY_RESPONSE, +) -> Dict[str, Any]: + """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, + STATUS_FIELD: status, + } + + +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) + 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 + name = _reason_name(reason) + if name.upper() in _BLOCKED_FINISH_REASONS: + 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]: + """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 EMPTY_RESPONSE 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/_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/_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/_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 new file mode 100644 index 0000000..a48c3bd --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_retry.py @@ -0,0 +1,181 @@ +# 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 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", + } +) + +_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", +} + + + + +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: + """Describes a model error in words, for the row's ``status`` field. + + 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. + + Returns: + A human readable status message. Never empty. + """ + code = _canonical_code(exception) + message = getattr(exception, "message", None) or str(exception) + + if code in _DETAILED_CODES: + return f"{code}: {message}" + # 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( + 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..fcc3911 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", @@ -39,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/__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..e56253d --- /dev/null +++ b/tests/integration/sql/test_ai_generate.py @@ -0,0 +1,600 @@ +# 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. + +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 +import unittest + +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf + +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}? " + "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}, + } + } +) + +_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): + + @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) + + def _prompt(self): + before, after = _AIRPORT_PROMPT.split("{city}", maxsplit=1) + return [sf.lit(before), sf.col("city"), sf.lit(after)] + + def test_generates_an_answer_for_every_row(self): + rows = ( + self.cities.withColumn("g", ai_generate(self._prompt())) + .select("city", "expected", "g.result", "g.status") + .collect() + ) + + self.assertEqual(len(rows), 7) + for row in rows: + with self.subTest(city=row["city"]): + self.assertEqual(row["status"], "") + self.assertIsNotNone(row["result"]) + self.assertIn(row["expected"], row["result"].upper()) + + def test_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_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()) + ) + payload = json.loads(result.select("g.full_response").first()[0]) + + 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. + 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"], "") + 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"], "") + 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"], "") + # 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): + 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"], "") + self.assertIsNotNone(rows[0]["result"]) + + 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"), + ).collect() + + 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()) + 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_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(" + " 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() + + self.assertEqual(len(rows), 7) + 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_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 + # 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", "capital STRING"), + ("Japan", "city STRING"), + ], + ["country", "schema"], + ).repartition(1).createOrReplaceTempView("questions") + + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(" + " CONCAT('Capital of ', country, '?')," + " NULL," + " named_struct('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(" + " struct(" + " 'What brand name is on the label? Brand name only.'," + " named_struct('uri', uri)" + " )," + f" named_struct('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(" + " struct(" + " 'What is the invoice number? Digits only.'," + " named_struct('uri', uri, 'content_type', 'application/pdf')" + " )," + f" named_struct('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() + ) + + self.assertTrue(rows[0]["status"].startswith("INVALID_ARGUMENT:")) + self.assertIn("content_type", rows[0]["status"]) + + +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..ed077ad --- /dev/null +++ b/tests/unit/sql/test_ai_generate.py @@ -0,0 +1,1166 @@ +# 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, file +from tests.utils.spark_version import requires_spark_4 + +# 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 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) + + +def generate_one(**kwargs): + """Calls the per-row entry point with the arguments a batch would pass.""" + 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)) + + +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_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("!")] + ) + 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) + + 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.""" + + 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 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.""" + + 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"], "") + 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.assertEqual(row["full_response"], "{}") + self.assertEqual(row["status"], "") + + 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.assertTrue(row["status"].startswith("SAFETY_BLOCKED:")) + self.assertIn(reason.name, row["status"]) + # The reason is still recoverable from the full response. + self.assertIn("candidates", 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"], "") + self.assertEqual(row["result"], "hi") + + 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( + api_error(404, "NOT_FOUND", "no such model") + ), + prompt="hi", + ) + 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( + 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"], "") + + 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_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_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=answers_the_prompt(), + prompt=[{"uri": "gs://bucket/file-without-an-extension"}], + ) + self.assertIsNone(row["result"]) + self.assertTrue(row["status"].startswith("INVALID_ARGUMENT:")) + + def test_a_prompt_of_only_null_parts_is_a_null_prompt(self): + def _fail(**_): + raise AssertionError("the model must not be called") + + row = generate_one( + generate_content=_fail, prompt=[None, {"uri": None}] + ) + self.assertIsNone(row["result"]) + self.assertEqual(row["status"], "") + + def test_parts_are_sent_in_order(self): + row = generate_one( + 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"}, + ], + ) + + +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"], "") + + 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", "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) + + @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. + 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(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," + " 'gemini-2.5-pro'," + " named_struct('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_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()), + ) + self._documents().createOrReplaceTempView("documents") + + row = self.spark.sql( + "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. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=answers_the_prompt()), + ) + self._documents().createOrReplaceTempView("documents") + + with self.assertRaises(Exception) as caught: + self.spark.sql( + "SELECT ai_generate(body, NULL," + " named_struct('output_schema', 'a NOTATYPE')).status AS s" + " FROM documents" + ).collect() + self.assertIn("NOTATYPE", str(caught.exception)) + + 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()), + ) + self.spark.createDataFrame( + [("a", "gemini-2.5-pro"), ("b", "gemini-3.6-flash")], + "body STRING, model STRING", + ).repartition(1).createOrReplaceTempView("varying") + + with self.assertRaises(Exception) as caught: + self.spark.sql( + "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. + 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 ai_generate(body, model).result AS r FROM constant" + ).collect() + + self.assertEqual( + [json.loads(r["r"])["model"] for r in rows], + ["gemini-2.5-pro", "gemini-2.5-pro"], + ) + + 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"], + ) + + +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(" + " 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_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..e34ad9b --- /dev/null +++ b/tests/unit/sql/test_ddl.py @@ -0,0 +1,291 @@ +# 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 +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 +# 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() + + @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)) + 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_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_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_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 new file mode 100644 index 0000000..77f26b2 --- /dev/null +++ b/tests/unit/sql/test_retry.py @@ -0,0 +1,167 @@ +# 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`` 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(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: " + ) + ) + + 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), + "UNAVAILABLE: a retryable error occurred calling the remote " + "service.", + ) + self.assertNotIn("0x8badf00d", _retry.to_status(error)) + + def test_transport_failures_report_unavailable(self): + self.assertTrue( + _retry.to_status(httpx.ConnectError("no route")).startswith( + "UNAVAILABLE: " + ) + ) + + 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): + + 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() 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__})", +)