This repository was archived by the owner on Apr 1, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 67
Expand file tree
/
Copy pathexecutor.py
More file actions
352 lines (293 loc) · 11.6 KB
/
executor.py
File metadata and controls
352 lines (293 loc) · 11.6 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
# Copyright 2024 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
#
# http://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.
from __future__ import annotations
import abc
import dataclasses
import functools
import itertools
from typing import Iterator, Literal, Optional, Sequence, Union
from google.cloud import bigquery, bigquery_storage_v1
import google.cloud.bigquery.table as bq_table
import pandas as pd
import pyarrow
import pyarrow as pa
import bigframes
import bigframes.core
from bigframes.core import bq_data, local_data, pyarrow_utils
import bigframes.core.schema
import bigframes.dtypes
import bigframes.session._io.pandas as io_pandas
import bigframes.session.execution_spec as ex_spec
_ROW_LIMIT_EXCEEDED_TEMPLATE = (
"Execution has downloaded {result_rows} rows so far, which exceeds the "
"limit of {maximum_result_rows}. You can adjust this limit by setting "
"`bpd.options.compute.maximum_result_rows`."
)
class ResultsIterator(Iterator[pa.RecordBatch]):
"""
Iterator for query results, with some extra metadata attached.
"""
def __init__(
self,
batches: Iterator[pa.RecordBatch],
schema: bigframes.core.schema.ArraySchema,
total_rows: Optional[int] = 0,
total_bytes: Optional[int] = 0,
):
self._batches = batches
self._schema = schema
self._total_rows = total_rows
self._total_bytes = total_bytes
@property
def approx_total_rows(self) -> Optional[int]:
return self._total_rows
@property
def approx_total_bytes(self) -> Optional[int]:
return self._total_bytes
def __next__(self) -> pa.RecordBatch:
return next(self._batches)
@property
def arrow_batches(self) -> Iterator[pyarrow.RecordBatch]:
result_rows = 0
for batch in self._batches:
result_rows += batch.num_rows
maximum_result_rows = bigframes.options.compute.maximum_result_rows
if maximum_result_rows is not None and result_rows > maximum_result_rows:
message = bigframes.exceptions.format_message(
_ROW_LIMIT_EXCEEDED_TEMPLATE.format(
result_rows=result_rows,
maximum_result_rows=maximum_result_rows,
)
)
raise bigframes.exceptions.MaximumResultRowsExceeded(message)
yield batch
def to_arrow_table(self, limit: Optional[int] = None) -> pyarrow.Table:
# Need to provide schema if no result rows, as arrow can't infer
# If ther are rows, it is safest to infer schema from batches.
# Any discrepencies between predicted schema and actual schema will produce errors.
batches = iter(self.arrow_batches)
peek_it = itertools.islice(batches, 0, 1)
peek_value = list(peek_it)
# TODO: Enforce our internal schema on the table for consistency
if len(peek_value) > 0:
batches = itertools.chain(peek_value, batches) # reconstruct
if limit:
batches = pyarrow_utils.truncate_pyarrow_iterable(
batches, max_results=limit
)
return pyarrow.Table.from_batches(batches)
else:
try:
return self._schema.to_pyarrow().empty_table()
except pa.ArrowNotImplementedError:
# Bug with some pyarrow versions, empty_table only supports base storage types, not extension types.
return self._schema.to_pyarrow(use_storage_types=True).empty_table()
def to_pandas(self, limit: Optional[int] = None) -> pd.DataFrame:
return io_pandas.arrow_to_pandas(self.to_arrow_table(limit=limit), self._schema)
def to_pandas_batches(
self, page_size: Optional[int] = None, max_results: Optional[int] = None
) -> Iterator[pd.DataFrame]:
assert (page_size is None) or (page_size > 0)
assert (max_results is None) or (max_results > 0)
batch_iter: Iterator[
Union[pyarrow.Table, pyarrow.RecordBatch]
] = self.arrow_batches
if max_results is not None:
batch_iter = pyarrow_utils.truncate_pyarrow_iterable(
batch_iter, max_results
)
if page_size is not None:
batches_iter = pyarrow_utils.chunk_by_row_count(batch_iter, page_size)
batch_iter = map(
lambda batches: pyarrow.Table.from_batches(batches), batches_iter
)
yield from map(
functools.partial(io_pandas.arrow_to_pandas, schema=self._schema),
batch_iter,
)
def to_py_scalar(self):
columns = list(self.to_arrow_table().to_pydict().values())
if len(columns) != 1:
raise ValueError(
f"Expected single column result, got {len(columns)} columns."
)
column = columns[0]
if len(column) != 1:
raise ValueError(f"Expected single row result, got {len(column)} rows.")
return column[0]
class ExecuteResult(abc.ABC):
@property
@abc.abstractmethod
def execution_metadata(self) -> ExecutionMetadata:
...
@property
@abc.abstractmethod
def schema(self) -> bigframes.core.schema.ArraySchema:
...
@abc.abstractmethod
def batches(self, sample_rate: Optional[float] = None) -> ResultsIterator:
...
@property
def query_job(self) -> Optional[bigquery.QueryJob]:
return self.execution_metadata.query_job
@property
def total_bytes_processed(self) -> Optional[int]:
return self.execution_metadata.bytes_processed
@dataclasses.dataclass(frozen=True)
class ExecutionMetadata:
query_job: Optional[bigquery.QueryJob] = None
bytes_processed: Optional[int] = None
@classmethod
def from_iterator_and_job(
cls, iterator: bq_table.RowIterator, job: Optional[bigquery.QueryJob]
) -> ExecutionMetadata:
return cls(query_job=job, bytes_processed=iterator.total_bytes_processed)
class LocalExecuteResult(ExecuteResult):
def __init__(
self,
data: pa.Table,
bf_schema: bigframes.core.schema.ArraySchema,
execution_metadata: ExecutionMetadata = ExecutionMetadata(),
):
self._data = local_data.ManagedArrowTable.from_pyarrow(data, bf_schema)
self._execution_metadata = execution_metadata
@property
def execution_metadata(self) -> ExecutionMetadata:
return self._execution_metadata
@property
def schema(self) -> bigframes.core.schema.ArraySchema:
return self._data.schema
def batches(self, sample_rate: Optional[float] = None) -> ResultsIterator:
return ResultsIterator(
iter(self._data.to_arrow(sample_rate=sample_rate)[1]),
self.schema,
self._data.metadata.row_count,
self._data.metadata.total_bytes,
)
class EmptyExecuteResult(ExecuteResult):
def __init__(
self,
bf_schema: bigframes.core.schema.ArraySchema,
execution_metadata: ExecutionMetadata = ExecutionMetadata(),
):
self._schema = bf_schema
self._execution_metadata = execution_metadata
@property
def execution_metadata(self) -> ExecutionMetadata:
return self._execution_metadata
@property
def schema(self) -> bigframes.core.schema.ArraySchema:
return self._schema
def batches(self, sample_rate: Optional[float] = None) -> ResultsIterator:
return ResultsIterator(iter([]), self.schema, 0, 0)
class BQTableExecuteResult(ExecuteResult):
def __init__(
self,
data: bq_data.BigqueryDataSource,
storage_client: bigquery_storage_v1.BigQueryReadClient,
project_id: str,
*,
execution_metadata: ExecutionMetadata = ExecutionMetadata(),
limit: Optional[int] = None,
selected_fields: Optional[Sequence[tuple[str, str]]] = None,
):
self._data = data
self._project_id = project_id
self._execution_metadata = execution_metadata
self._storage_client = storage_client
self._limit = limit
self._selected_fields = selected_fields or [
(name, name) for name in data.schema.names
]
@property
def execution_metadata(self) -> ExecutionMetadata:
return self._execution_metadata
@property
@functools.cache
def schema(self) -> bigframes.core.schema.ArraySchema:
source_ids = [selection[0] for selection in self._selected_fields]
return self._data.schema.select(source_ids).rename(dict(self._selected_fields))
def batches(self, sample_rate: Optional[float] = None) -> ResultsIterator:
read_batches = bq_data.get_arrow_batches(
self._data,
[x[0] for x in self._selected_fields],
self._storage_client,
self._project_id,
sample_rate=sample_rate,
)
arrow_batches: Iterator[pa.RecordBatch] = map(
functools.partial(
pyarrow_utils.rename_batch, names=list(self.schema.names)
),
read_batches.iter,
)
approx_bytes: Optional[int] = read_batches.approx_bytes
approx_rows: Optional[int] = self._data.n_rows or read_batches.approx_rows
if self._limit is not None:
if approx_rows is not None:
approx_rows = min(approx_rows, self._limit)
arrow_batches = pyarrow_utils.truncate_pyarrow_iterable(
arrow_batches, self._limit
)
if self._data.sql_predicate:
approx_bytes = None
approx_rows = None
return ResultsIterator(arrow_batches, self.schema, approx_rows, approx_bytes)
@dataclasses.dataclass(frozen=True)
class HierarchicalKey:
columns: tuple[str, ...]
@dataclasses.dataclass(frozen=True)
class CacheConfig(abc.ABC):
optimize_for: Union[Literal["auto", "head"], HierarchicalKey] = "auto"
if_cached: Literal["reuse-strict", "reuse-any", "replace"] = "reuse-any"
class Executor(abc.ABC):
"""
Interface for an executor, which compiles and executes ArrayValue objects.
"""
def to_sql(
self,
array_value: bigframes.core.ArrayValue,
offset_column: Optional[str] = None,
ordered: bool = False,
enable_cache: bool = True,
) -> str:
"""
Convert an ArrayValue to a sql query that will yield its value.
"""
raise NotImplementedError("to_sql not implemented for this executor")
@abc.abstractmethod
def execute(
self,
array_value: bigframes.core.ArrayValue,
execution_spec: ex_spec.ExecutionSpec,
) -> ExecuteResult:
"""
Execute the ArrayValue.
"""
...
def dry_run(
self, array_value: bigframes.core.ArrayValue, ordered: bool = True
) -> bigquery.QueryJob:
"""
Dry run executing the ArrayValue.
Does not actually execute the data but will get stats and indicate any invalid query errors.
"""
raise NotImplementedError("dry_run not implemented for this executor")
def cached(
self,
array_value: bigframes.core.ArrayValue,
*,
config: CacheConfig,
) -> None:
raise NotImplementedError("cached not implemented for this executor")