feat(sql): add ai_generate, a Spark column function for Gemini - #8
saurabh-net wants to merge 2 commits into
Conversation
Introduces google.cloud.dataproc_ml.sql.ai_generate, the Spark
counterpart of BigQuery's AI.GENERATE. It calls a Gemini model once per
row and returns a struct of result, full_response and status, so queries
can move between the two engines with little change.
The function is a pandas UDF, which means it works both as a DataFrame
column expression and, through ai_generate_udf(), as a function
registered with spark.udf.register for use from Spark SQL.
The argument list is deliberately BigQuery's: prompt, endpoint,
model_params and output_schema, with project and location standing in
for connection_id. Concurrency and retry timing are implemented but kept
as internal constants, as they are in BigQuery, rather than exposed as
arguments that would be hard to withdraw later.
Notable behaviour, all matching BigQuery:
* A row that fails never fails the stage. Its error is reported in
status and result is NULL. status is empty on success.
* output_schema replaces result with one field per declared field,
sorted by name. BigQuery type names (INT64, FLOAT64, BOOL) and their
Spark spellings are both accepted.
* full_response holds the complete model response, typed as VARIANT on
Spark 4 and as a JSON string on older versions.
Two things differ from BigQuery by necessity. There is no connection_id,
because a BigQuery connection's service account cannot be impersonated
from outside BigQuery; calls use Application Default Credentials
instead. And the default model is gemini-3.6-flash, which is served only
from the global endpoint, so the default location is global rather than
a region.
Tested with unit tests that exercise the full Spark path against an
injected generate_content, and with integration tests that call Vertex
AI for real.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces declarative AI functions (ai_generate and ai_generate_udf) to apply Gemini models over Spark DataFrame columns, mirroring BigQuery's AI.GENERATE semantics. The code review feedback is highly constructive and identifies several critical bugs: a boolean coercion flaw where string representations like 'false' evaluate to True, an issue where non-string columns are incorrectly treated as missing and silently nulled, and a regex replacement bug that could corrupt column names matching BigQuery type names. Addressing these issues will ensure robust type handling and schema parsing.
| if isinstance(data_type, BooleanType): | ||
| return bool(value) |
There was a problem hiding this comment.
In Python, bool("false") evaluates to True because any non-empty string is truthy. If the model returns a string representation of a boolean (e.g., "false" or "False"), or if the JSON parser parses it as a string, it will be incorrectly coerced to True. Use a safer check for string-to-boolean coercion.
| if isinstance(data_type, BooleanType): | |
| return bool(value) | |
| if isinstance(data_type, BooleanType): | |
| if isinstance(value, str): | |
| return value.lower() in ("true", "t", "1", "yes", "y") | |
| return bool(value) |
There was a problem hiding this comment.
Good catch, fixed in a20e89c. Confirmed the bug: bool("false") is True, and so is bool("maybe").
I took a stricter line than the suggestion, in a _to_boolean helper. Mapping every unrecognised string to False trades one silent wrong answer for another: "maybe" would become a confident False. Instead a value that does not clearly denote a boolean becomes null, which is already how the rest of _coerce reports an answer that does not fit the requested schema. true/t/yes/y/1 and false/f/no/n/0 are accepted, bool is checked before int since it is a subclass, and everything else is null.
Covered by test_a_boolean_field_reads_the_answer.
| if prompt is None: | ||
| return True | ||
| if isinstance(prompt, float) and math.isnan(prompt): | ||
| return True | ||
| return not isinstance(prompt, str) |
There was a problem hiding this comment.
In Spark SQL, non-string columns (like integers) passed to the UDF are currently treated as missing because _is_missing returns True for any non-string. This causes them to be silently nulled without any error or status message. Update _is_missing to only check for actual nulls (None and NaN), and handle string coercion in _generate_one.
| if prompt is None: | |
| return True | |
| if isinstance(prompt, float) and math.isnan(prompt): | |
| return True | |
| return not isinstance(prompt, str) | |
| if prompt is None: | |
| return True | |
| if isinstance(prompt, float) and math.isnan(prompt): | |
| return True | |
| return False |
There was a problem hiding this comment.
Agreed and fixed in a20e89c. _is_missing now only detects real nulls: None, NaN, and pandas.NA (which a nullable column can produce and which neither of the other two checks catches).
Worth recording why this was reachable at all: ai_generate does cast its prompt column to STRING on the driver, so the DataFrame path was safe. The gap was the Spark SQL path, where ai_generate_udf() is registered and applied to the column as declared with no cast inserted. That is now covered end to end by test_a_non_string_column_works_in_spark_sql.
| if _is_missing(prompt): | ||
| # A null prompt produces a null result without calling the model. | ||
| return _row(schema, result=None, full_response=None, status="") | ||
|
|
||
| for attempt in range(_MAX_ATTEMPTS): | ||
| try: | ||
| response = await generate_content( | ||
| model=endpoint, contents=prompt, config=config | ||
| ) |
There was a problem hiding this comment.
Coerce the prompt to a string if it is not already one, to support non-string columns (e.g., integers) passed via Spark SQL without silently nulling them.
| if _is_missing(prompt): | |
| # A null prompt produces a null result without calling the model. | |
| return _row(schema, result=None, full_response=None, status="") | |
| for attempt in range(_MAX_ATTEMPTS): | |
| try: | |
| response = await generate_content( | |
| model=endpoint, contents=prompt, config=config | |
| ) | |
| if _is_missing(prompt): | |
| # A null prompt produces a null result without calling the model. | |
| return _row(schema, result=None, full_response=None, status="") | |
| prompt_str = str(prompt) if not isinstance(prompt, str) else prompt | |
| for attempt in range(_MAX_ATTEMPTS): | |
| try: | |
| response = await generate_content( | |
| model=endpoint, contents=prompt_str, config=config | |
| ) |
There was a problem hiding this comment.
Fixed in a20e89c, though not quite as suggested, because writing the test for this turned up a second bug in the fix itself.
A plain str(prompt) renders a BIGINT 2 as "2.0". Arrow delivers a nullable integer column to pandas as float64, since int64 cannot hold NaN. So the prompt text depended on whether some unrelated row in the same Arrow batch happened to be null: identical inputs, different prompt. The suggested patch would have shipped that.
Rendering now goes through _to_text, which drops the fraction from a whole number, making the output stable regardless of neighbouring nulls:
n BIGINT, no nulls in batch -> 2 -> "2"
n BIGINT, null in batch -> 2.0 -> "2" (was "2.0")
n DOUBLE -> 3.5 -> "3.5"
Covered by test_a_non_string_prompt_is_rendered_as_text (unit) and test_a_non_string_column_works_in_spark_sql (through Spark, with a null row).
| for bigquery_name, spark_name in _TYPE_ALIASES.items(): | ||
| normalized = re.sub( | ||
| rf"\b{bigquery_name}\b", | ||
| spark_name, | ||
| normalized, | ||
| flags=re.IGNORECASE, | ||
| ) |
There was a problem hiding this comment.
Using a simple word boundary \b to replace BigQuery type names will corrupt column names that happen to match those type names (e.g., a column named int64 or bool). For example, "int64 INT64" will be incorrectly replaced with "BIGINT BIGINT", which fails to parse. Use lookarounds to ensure we only replace type names when they are acting as types.
| for bigquery_name, spark_name in _TYPE_ALIASES.items(): | |
| normalized = re.sub( | |
| rf"\b{bigquery_name}\b", | |
| spark_name, | |
| normalized, | |
| flags=re.IGNORECASE, | |
| ) | |
| for bigquery_name, spark_name in _TYPE_ALIASES.items(): | |
| normalized = re.sub( | |
| rf"(?<=[\s:<]){bigquery_name}(?=\s*(?:,|>|\)|$))", | |
| spark_name, | |
| normalized, | |
| flags=re.IGNORECASE, | |
| ) |
There was a problem hiding this comment.
Real bug, fixed in a20e89c — but the suggested regex introduces a regression, so I used a variant.
Confirming the bug first, and it is worse than a parse failure. "int64 INT64" currently normalises to "BIGINT BIGINT", which parses cleanly as struct<BIGINT:bigint>. The user silently gets a field named BIGINT.
The proposed lookahead (?=\s*(?:,|>|\)|$)) fixes that but drops the modifiers Spark accepts after a type. Checked against real StructType.fromDDL:
| input | current | suggested | shipped |
|---|---|---|---|
int64 INT64 |
struct<BIGINT:bigint> 🐛 |
struct<int64:bigint> |
struct<int64:bigint> |
s STRUCT<bool: BOOL> |
struct<s:struct<BOOLEAN:boolean>> 🐛 |
ok | ok |
a INT64 NOT NULL |
ok | fails to parse 🐛 | ok |
a INT64 COMMENT 'how many' |
ok | fails to parse 🐛 | ok |
scores ARRAY<INT64> |
ok | ok | ok |
The shipped pattern keeps the lookbehind and extends the lookahead to allow those modifiers:
(?<=[\s:<])INT64(?=\s*(?:,|>|\)|$|NOT\s+NULL|COMMENT\b))
Both halves are now regression tested by test_a_field_named_after_a_type_keeps_its_name and test_a_type_keeps_its_modifiers.
Three fixes from the review bot on GoogleCloudDataproc#8, each with a regression test: - Only rewrite a BigQuery type name where a type may appear. Matching the bare name also rewrote a field named after a type, so "int64 INT64" parsed as a field actually called BIGINT. The narrower match still accepts the NOT NULL and COMMENT modifiers that Spark allows after a type. - Treat only nulls as a missing prompt. Any other value is now rendered as text, so a non-string column no longer arrives as a whole column of nulls when the function is registered and called from Spark SQL, where no cast to STRING is inserted. A whole number keeps its integer spelling: Arrow delivers a nullable BIGINT column as floats, so otherwise a row's prompt would depend on whether an unrelated row in the same batch happened to be null. - Read a boolean field explicitly rather than through Python truthiness, which recorded the answer "false" as True.
1 participant
What this adds
A new
google.cloud.dataproc_ml.sqlmodule whose entry point,ai_generate,is the Spark counterpart of BigQuery's
AI.GENERATE.It calls a Gemini model once per row and returns a struct, so a query written
for BigQuery can move to Spark with little change.
Because it is a plain pandas UDF, the same implementation serves Spark SQL:
This covers the
AI.GENERATErow of Milestone 1.The argument surface is deliberately BigQuery's
BigQuery's
AI.GENERATEtakesprompt,connection_id,endpoint,model_paramsandoutput_schema. This function takes the same set, withprojectandlocationin place ofconnection_id, and nothing else.That is a deliberate constraint rather than an oversight. Concurrency limits,
request pacing and retry timing are all implemented, but they are internal
constants, exactly as they are in BigQuery. Exposing them is easy to do later
and impossible to undo, so it is left to a follow-up once there is evidence of
what users actually need to tune. See "Follow-ups" below.
Two deliberate divergences from BigQuery
1. No
connection_id; authentication is Application Default Credentials.This is not a shortcut, it is a constraint. A BigQuery
CLOUD_RESOURCEconnection owns a hidden service account that only the BigQuery service can
authenticate as. Its credentials cannot be exported, and an external workload
such as a Spark executor cannot impersonate it merely by knowing the
connection ID. There is therefore no way for this function to honour a
connection_id.Instead the call uses ADC, which on Dataproc, Dataproc Serverless and GKE is
the cluster's own identity.
projectandlocationplay the role that theproject and location components of a connection ID play in BigQuery: they
decide which project is billed and which Vertex AI endpoint is called.
Consequence for reviewers: a query moved verbatim from BigQuery will need
its
connection_idargument removed, and the cluster's service account needsroles/aiplatform.user.2. Different default model and location.
BigQuery currently defaults to
gemini-2.5-flash. This defaults togemini-3.6-flash, and consequently to thegloballocation rather than aregion, because the Gemini 3.x models are not served from regional endpoints
(verified:
gemini-3.6-flashreturnsNOT_FOUNDinus-central1).The location also deliberately ignores
GOOGLE_CLOUD_LOCATIONandGOOGLE_CLOUD_REGION. A Dataproc cluster sets those to the region it runs in,which is not a statement about where a model should be called. Honouring them
would have routed the default model to a region that does not serve it, and
because a failed row reports its error in
statusrather than raising, thesymptom would have been a silent column of NULLs. Only an explicit
location=argument selects a region.Behaviour, and why it looks like BigQuery
statusandresultisNULL.statusis the empty string on success, not"OK".This is the opposite of the existing
GenAiModelHandler, which raises. Onlyerrors that are genuinely the remote service's —
google.genai.errors.APIErrorand transport errors — are reported this way. Anything else is a bug in this
library and is allowed to fail the task rather than hide in a column.
output_schemafields come back sorted by name, not in declarationorder, and
resultdisappears.full_responseandstatusstay, in thatorder, at the end. Both BigQuery type names (
INT64,FLOAT64,BOOL) andtheir Spark spellings are accepted.
full_responseisVARIANTon Spark 4, the closest equivalent ofBigQuery's
JSON, falling back to a JSON string on older Spark. Its keys aresnake_case, matching BigQuery's serialization.
candidate. A response blocked by a safety filter is not an error:
statusis empty and
resultisNULL.Implementation notes for the reviewer
Everything is validated on the driver, inside
ai_generate_udf, before astage is scheduled, so a typo fails in milliseconds rather than after a
cluster has spun up.
_output_schema.pydoes not hand-roll a DDL parser. It rewrites theBigQuery type names and hands the string to
StructType.fromDDL._model_params.pyflattens thegenerateContentrequest body onto theflat
types.GenerateContentConfig, so either the nested{"generation_config": {...}}form or the flat form works. camelCase isconverted at the top level only, deliberately, so user property names inside
response_schemaandtoolsare left alone._client.pycaches the client and one asyncio loop per worker process.One long-lived loop on a daemon thread avoids building and tearing down a
loop per Arrow batch.
Retries use full jitter,
uniform(0, min(cap, base * 2**attempt)).Test seam: both public functions accept a private
_generate_contentargument that replaces the network call, so the unit tests exercise the real
Spark path — serialization, Arrow batching, struct construction — with no
credentials and no network.
Why every module here is underscored, unlike
inference/. The publicsurface of this package is exactly the two names in
__all__; everythingelse should stay free to change without that being a breaking change. Two
things make
sql/differ from its sibling:inference/exports classes whose names differ from their modules(
GenAiModelHandlerfromgen_ai_model_handler.py), and it has a realextension point in
BaseModelHandler, which is not in__all__and sois reached through its module path. That path is part of its API and
must not be underscored.
sql/exports functions whose names would equal their module names.sql/ai_generate.pyplus a function calledai_generatemakesfrom ...sql import ai_generateambiguous, and which one wins depends onimport order.
A non-underscored
generate.pywould also have dodged that collision, sothis is a choice rather than a necessity. Happy to rename if you would
rather the two packages match.
Review findings addressed
This went through an independent review pass. Two substantive fixes came out
of it:
were cast to STRING, so
ai_generate(col("user_id"))reached the worker asan int, was treated as a null prompt, and produced a whole column of NULLs
with an empty
status— no error anywhere. Every prompt path now casts, anda regression test covers all three forms.
request_typewas dead API surface — it accepted onlyUNSPECIFIEDandraised for everything else. Removed until provisioned throughput is actually
implemented.
One reported "critical" finding was investigated and rejected: the claim that
constructing the
genai.Clientoff the event loop breaks itshttpxtransport. Reproduced directly against Vertex AI — off-loop construction
followed by sequential and concurrent calls on the shared loop all succeed,
because
httpx.AsyncClientbinds lazily on first use, not at construction.From the review bot, in
a20e89cAll three findings were real. Each is fixed with a regression test, and each
fix differs from the suggested patch — details in the review threads.
"int64 INT64"normalised to"BIGINT BIGINT"and parsed cleanly asstruct<BIGINT:bigint>, so thefield silently came back called
BIGINT. Type-name rewriting is nowanchored to type position. The suggested lookahead also fixed this but
stopped
NOT NULLandCOMMENTfrom parsing at all, so the shippedpattern allows those too.
ai_generatecovered the DataFrame path, but a UDF registered withspark.udf.registeris applied to the column as declared, with no castinserted.
_is_missingnow detects only nulls and other values arerendered as text. Writing the test for this exposed a second bug in the
fix: a plain
str()renders BIGINT2as"2.0", because Arrow deliversa nullable integer column as
float64. A row's prompt would havedepended on whether an unrelated row in the same batch was null. Whole
numbers now keep their integer spelling.
bool("false")isTrue. A boolean field answered as text wasrecorded as True. Booleans are now parsed explicitly, and an answer that
denotes no boolean at all becomes null rather than being forced to
False.Tests
tests/unit/sql/runs without credentials.tests/integration/sql/callsVertex AI for real, and was run against a live project in
us-central1withgemini-3.6-flashon theglobalendpoint.Lint, per
contributing.md:pylinton the package is silent. The single warning above is pre-existing ina file this PR does not touch.
The docs build (
sphinx-build -b html docs/ docs/_build/html) succeeds, withno new warnings from this module.
Full suite, with
TEST_GCS_BUCKETset so that theinference/integrationtests can run:
All three failures are in
tests/integration/inference/, which this PR doesnot touch. Nothing in this PR imports or modifies
google/cloud/dataproc_ml/inference/, and none of those tests importdataproc_ml.sql.test_vertex_endpoint_handler.pyfails reproducibly and cannot passoutside the environment it was written in: it hardcodes a Vertex AI
endpoint ID and carries a matching
TODO.test_gen_ai.py::test_prompt_templateandtest_pytorch.py::test_invalid_model_path_missing_componentsare flaky.Both pass when rerun on their own against the same project:
They also passed inside a full-suite run earlier the same day, on the same
code for
inference/.A note for whoever runs the linter: with the
pyinkversion resolved bypip install ".[dev]"today,pyink .also reformatsinference/base_model_handler.pyandinference/vertex_endpoint_handler.py,which this PR does not touch. That churn was reverted here to keep the diff
focused, so a local
pyink .will show those two files as modified. It lookslike the pinned formatter has drifted from whatever produced the current
formatting.
Dependency
Adds
google-genai>=1.38.0, <3.0.0. This is the current, non-deprecated GeminiSDK; the existing
GenAiModelHandlerstill uses the deprecatedvertexai.generative_modelsand is untouched here. The1.38.0floor wasverified by installing it and checking every API this module relies on.
Known limitations
One Vertex AI request per row, and cluster-wide concurrency is uncapped.
This function makes one
generateContentcall per non-null row, the samerow-wise model as BigQuery's
AI.GENERATE. A null prompt short-circuitswithout a request; a row that hits a retryable error costs up to
_MAX_ATTEMPTSrequests. So a one million row table is at least one millionrequests, and up to five million in the worst case.
Concurrency is bounded at five requests in flight per Python worker, and
Spark runs one worker per core. There is no cluster-wide limit, so a 100-core
cluster can have roughly 500 requests in flight against Vertex AI. On a
project whose quota is lower than that, the sequence is:
429responses,exponential backoff, and for any row that exhausts its attempts, an error
recorded in
status.The failure mode is worth stating plainly, because it follows from the
BigQuery-compatible soft-fail semantics: a quota-limited run does not fail.
It returns, with some rows carrying a
RESOURCE_EXHAUSTEDmessage instatusand a NULL result. A user who does not inspectstatuswill readthat as missing data rather than as an error.
Two mitigations exist today, both manual: reduce parallelism with
spark.dynamicAllocation.maxExecutorsorcoalesce, and checkWHERE status <> ''after the query. A per-worker rate limit is the properfix and is the first follow-up below.
Follow-ups
Intentionally left out of this PR:
timing are internal constants today. A cluster large enough to exhaust a
Vertex AI quota will want them, and a token bucket implementation already
exists in the history of this branch, but exposing arguments BigQuery does
not have deserves its own review rather than riding along here.
request_type/DEDICATEDprovisioned throughput routing.AI.GENERATE_BOOL/_INT/_DOUBLE,AI.EMBED,AI.COUNT_TOKENS,AI.GENERATE_TABLE,AI.IF/SCORE/CLASSIFY/AGG.max_error_ratio, prompt-hashcaching, metrics, and
OPTIONS(...)insideoutput_schema.