Skip to content

Commit 34ca90e

Browse files
sycaigemini-code-assist[bot]
sycai
and
gemini-code-assist[bot]
authored
chore(bigframes): move some ai accessor tests outside the SQLGlot compiler dir (googleapis#17318)
These tests do not generate golden SQLs, so it makes less sense to place them under the SQLGlot test directory. This is a follow-up to googleapis#17302 --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 parent 1281790 commit 34ca90e

3 files changed

Lines changed: 290 additions & 292 deletions

File tree

‎packages/bigframes/tests/unit/core/compile/sqlglot/test_dataframe_accessor.py‎

Lines changed: 0 additions & 292 deletions
Original file line numberDiff line numberDiff line change
@@ -58,295 +58,3 @@ def test_bigframes_sql_scalar(scalar_types_df: bpd.DataFrame, snapshot):
5858
# Bigframes implementation returns a bigframes.series.Series
5959
sql, _, _ = result.to_frame()._to_sql_query(include_index=True)
6060
snapshot.assert_match(sql, "out.sql")
61-
62-
63-
def test_ai_forecast(snapshot, monkeypatch):
64-
import bigframes.bigquery.ai
65-
import bigframes.session
66-
67-
session = mock.create_autospec(bigframes.session.Session)
68-
bf_df = mock.create_autospec(bpd.DataFrame)
69-
session.read_pandas.return_value = bf_df
70-
71-
def mock_ai_forecast(df, **kwargs):
72-
assert df is bf_df
73-
result_df = mock.create_autospec(bpd.DataFrame)
74-
result_df.to_pandas.return_value = kwargs
75-
return result_df
76-
77-
import bigframes.bigquery.ai
78-
79-
monkeypatch.setattr(bigframes.bigquery.ai, "forecast", mock_ai_forecast)
80-
81-
df = pd.DataFrame({"date": ["2020-01-01"], "value": [1.0]})
82-
result = df.bigquery.ai.forecast(
83-
timestamp_col="date",
84-
data_col="value",
85-
horizon=5,
86-
session=session,
87-
)
88-
89-
session.read_pandas.assert_called_once()
90-
assert result == {
91-
"timestamp_col": "date",
92-
"data_col": "value",
93-
"model": "TimesFM 2.0",
94-
"id_cols": None,
95-
"horizon": 5,
96-
"confidence_level": 0.95,
97-
"context_window": None,
98-
"output_historical_time_series": False,
99-
}
100-
101-
102-
def test_bigframes_ai_forecast(snapshot, monkeypatch):
103-
import bigframes.bigquery.ai
104-
import bigframes.session
105-
106-
session = mock.create_autospec(bigframes.session.Session)
107-
bf_df = mock.create_autospec(bpd.DataFrame)
108-
109-
def mock_ai_forecast(df, **kwargs):
110-
assert df is bf_df
111-
result_df = mock.create_autospec(bpd.DataFrame)
112-
return result_df
113-
114-
monkeypatch.setattr(bigframes.bigquery.ai, "forecast", mock_ai_forecast)
115-
116-
result = bf_df.bigquery.ai.forecast(
117-
timestamp_col="date",
118-
data_col="value",
119-
horizon=5,
120-
session=session,
121-
)
122-
123-
session.read_pandas.assert_not_called()
124-
# BigFrames accessor returns the bf_df directly without calling to_pandas
125-
assert result is not None
126-
127-
128-
def test_ai_generate(monkeypatch):
129-
import bigframes.bigquery.ai
130-
131-
def mock_generate(prompt, **kwargs):
132-
result_series = mock.create_autospec(bpd.Series)
133-
result_series.to_pandas.return_value = (prompt, kwargs)
134-
return result_series
135-
136-
monkeypatch.setattr(bigframes.bigquery.ai, "generate", mock_generate)
137-
138-
df = pd.DataFrame({"text_input": ["Is this a positive review?"]})
139-
result = df.bigquery.ai.generate(
140-
df["text_input"],
141-
connection_id="conn",
142-
endpoint="endpoint",
143-
request_type="dedicated",
144-
model_params={"temp": 0.5},
145-
output_schema={"res": "STRING"},
146-
)
147-
148-
assert result == (
149-
df["text_input"],
150-
{
151-
"connection_id": "conn",
152-
"endpoint": "endpoint",
153-
"request_type": "dedicated",
154-
"model_params": {"temp": 0.5},
155-
"output_schema": {"res": "STRING"},
156-
},
157-
)
158-
159-
160-
def test_bigframes_ai_generate(scalar_types_df: bpd.DataFrame, monkeypatch):
161-
import bigframes.bigquery.ai
162-
import bigframes.session
163-
164-
session = mock.create_autospec(bigframes.session.Session)
165-
bf_series = mock.create_autospec(bpd.Series)
166-
167-
def mock_generate(prompt, **kwargs):
168-
assert prompt is bf_series
169-
result_series = mock.create_autospec(bpd.Series)
170-
return result_series
171-
172-
monkeypatch.setattr(bigframes.bigquery.ai, "generate", mock_generate)
173-
174-
result = scalar_types_df.bigquery.ai.generate(
175-
bf_series,
176-
connection_id="conn",
177-
endpoint="endpoint",
178-
request_type="dedicated",
179-
model_params={"temp": 0.5},
180-
output_schema={"res": "STRING"},
181-
)
182-
183-
session.read_pandas.assert_not_called()
184-
assert result is not None
185-
186-
187-
def test_ai_generate_bool(monkeypatch):
188-
import bigframes.bigquery.ai
189-
190-
def mock_generate_bool(prompt, **kwargs):
191-
result_series = mock.create_autospec(bpd.Series)
192-
result_series.to_pandas.return_value = (prompt, kwargs)
193-
return result_series
194-
195-
monkeypatch.setattr(bigframes.bigquery.ai, "generate_bool", mock_generate_bool)
196-
197-
df = pd.DataFrame({"text_input": ["Is this a positive review?"]})
198-
result = df.bigquery.ai.generate_bool(
199-
df["text_input"],
200-
connection_id="conn",
201-
endpoint="endpoint",
202-
request_type="dedicated",
203-
model_params={"temp": 0.5},
204-
)
205-
206-
assert result == (
207-
df["text_input"],
208-
{
209-
"connection_id": "conn",
210-
"endpoint": "endpoint",
211-
"request_type": "dedicated",
212-
"model_params": {"temp": 0.5},
213-
},
214-
)
215-
216-
217-
def test_bigframes_ai_generate_bool(scalar_types_df: bpd.DataFrame, monkeypatch):
218-
import bigframes.bigquery.ai
219-
import bigframes.session
220-
221-
session = mock.create_autospec(bigframes.session.Session)
222-
bf_series = mock.create_autospec(bpd.Series)
223-
224-
def mock_generate_bool(prompt, **kwargs):
225-
assert prompt is bf_series
226-
result_series = mock.create_autospec(bpd.Series)
227-
return result_series
228-
229-
monkeypatch.setattr(bigframes.bigquery.ai, "generate_bool", mock_generate_bool)
230-
231-
result = scalar_types_df.bigquery.ai.generate_bool(
232-
bf_series,
233-
connection_id="conn",
234-
endpoint="endpoint",
235-
request_type="dedicated",
236-
model_params={"temp": 0.5},
237-
)
238-
239-
session.read_pandas.assert_not_called()
240-
assert result is not None
241-
242-
243-
def test_ai_generate_int(monkeypatch):
244-
import bigframes.bigquery.ai
245-
246-
def mock_generate_int(prompt, **kwargs):
247-
result_series = mock.create_autospec(bpd.Series)
248-
result_series.to_pandas.return_value = (prompt, kwargs)
249-
return result_series
250-
251-
monkeypatch.setattr(bigframes.bigquery.ai, "generate_int", mock_generate_int)
252-
253-
df = pd.DataFrame({"text_input": ["How many legs?"]})
254-
result = df.bigquery.ai.generate_int(
255-
df["text_input"],
256-
connection_id="conn",
257-
endpoint="endpoint",
258-
request_type="dedicated",
259-
model_params={"temp": 0.5},
260-
)
261-
262-
assert result == (
263-
df["text_input"],
264-
{
265-
"connection_id": "conn",
266-
"endpoint": "endpoint",
267-
"request_type": "dedicated",
268-
"model_params": {"temp": 0.5},
269-
},
270-
)
271-
272-
273-
def test_bigframes_ai_generate_int(scalar_types_df: bpd.DataFrame, monkeypatch):
274-
import bigframes.bigquery.ai
275-
import bigframes.session
276-
277-
session = mock.create_autospec(bigframes.session.Session)
278-
bf_series = mock.create_autospec(bpd.Series)
279-
280-
def mock_generate_int(prompt, **kwargs):
281-
assert prompt is bf_series
282-
result_series = mock.create_autospec(bpd.Series)
283-
return result_series
284-
285-
monkeypatch.setattr(bigframes.bigquery.ai, "generate_int", mock_generate_int)
286-
287-
result = scalar_types_df.bigquery.ai.generate_int(
288-
bf_series,
289-
connection_id="conn",
290-
endpoint="endpoint",
291-
request_type="dedicated",
292-
model_params={"temp": 0.5},
293-
)
294-
295-
session.read_pandas.assert_not_called()
296-
assert result is not None
297-
298-
299-
def test_ai_generate_double(monkeypatch):
300-
import bigframes.bigquery.ai
301-
302-
def mock_generate_double(prompt, **kwargs):
303-
result_series = mock.create_autospec(bpd.Series)
304-
result_series.to_pandas.return_value = (prompt, kwargs)
305-
return result_series
306-
307-
monkeypatch.setattr(bigframes.bigquery.ai, "generate_double", mock_generate_double)
308-
309-
df = pd.DataFrame({"text_input": ["How tall?"]})
310-
result = df.bigquery.ai.generate_double(
311-
df["text_input"],
312-
connection_id="conn",
313-
endpoint="endpoint",
314-
request_type="dedicated",
315-
model_params={"temp": 0.5},
316-
)
317-
318-
assert result == (
319-
df["text_input"],
320-
{
321-
"connection_id": "conn",
322-
"endpoint": "endpoint",
323-
"request_type": "dedicated",
324-
"model_params": {"temp": 0.5},
325-
},
326-
)
327-
328-
329-
def test_bigframes_ai_generate_double(scalar_types_df: bpd.DataFrame, monkeypatch):
330-
import bigframes.bigquery.ai
331-
import bigframes.session
332-
333-
session = mock.create_autospec(bigframes.session.Session)
334-
bf_series = mock.create_autospec(bpd.Series)
335-
336-
def mock_generate_double(prompt, **kwargs):
337-
assert prompt is bf_series
338-
result_series = mock.create_autospec(bpd.Series)
339-
return result_series
340-
341-
monkeypatch.setattr(bigframes.bigquery.ai, "generate_double", mock_generate_double)
342-
343-
result = scalar_types_df.bigquery.ai.generate_double(
344-
bf_series,
345-
connection_id="conn",
346-
endpoint="endpoint",
347-
request_type="dedicated",
348-
model_params={"temp": 0.5},
349-
)
350-
351-
session.read_pandas.assert_not_called()
352-
assert result is not None

‎packages/bigframes/tests/unit/extensions/core/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

0 commit comments

Comments
 (0)