|
28 | 28 | import com.google.api.gax.paging.Page; |
29 | 29 | import com.google.api.gax.rpc.HeaderProvider; |
30 | 30 | import com.google.api.gax.rpc.NoHeaderProvider; |
| 31 | +import com.google.api.gax.rpc.ServerStream; |
31 | 32 | import com.google.api.services.bigquery.model.ErrorProto; |
32 | 33 | import com.google.api.services.bigquery.model.GetQueryResultsResponse; |
33 | 34 | import com.google.api.services.bigquery.model.ProjectList; |
|
54 | 55 | import com.google.cloud.bigquery.storage.v1.BigQueryReadSettings; |
55 | 56 | import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest; |
56 | 57 | import com.google.cloud.bigquery.storage.v1.DataFormat; |
| 58 | +import com.google.cloud.bigquery.storage.v1.ReadRowsRequest; |
| 59 | +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; |
57 | 60 | import com.google.cloud.bigquery.storage.v1.ReadSession; |
58 | 61 | import com.google.common.annotations.VisibleForTesting; |
59 | 62 | import com.google.common.base.Function; |
|
68 | 71 | import io.opentelemetry.api.trace.Span; |
69 | 72 | import io.opentelemetry.context.Scope; |
70 | 73 | import java.io.IOException; |
| 74 | +import java.util.ArrayDeque; |
71 | 75 | import java.util.ArrayList; |
72 | 76 | import java.util.Collections; |
| 77 | +import java.util.Iterator; |
73 | 78 | import java.util.List; |
74 | 79 | import java.util.Map; |
| 80 | +import java.util.Queue; |
75 | 81 | import java.util.concurrent.Callable; |
76 | 82 | import java.util.concurrent.ConcurrentHashMap; |
77 | 83 | import java.util.regex.Matcher; |
@@ -278,6 +284,192 @@ public Page<FieldValueList> getNextPage() { |
278 | 284 | } |
279 | 285 | } |
280 | 286 |
|
| 287 | + /** |
| 288 | + * NextPageFetcher implementation for queries returning results in Arrow format. Reads subsequent |
| 289 | + * pages from the job's default gRPC storage read stream. Rows are buffered and converted on |
| 290 | + * demand from Arrow record batches into FieldValueList instances in batches up to {@code |
| 291 | + * pageSize}. |
| 292 | + * |
| 293 | + * <p>Note: Neither {@link Page} nor {@link TableResult} implements {@link AutoCloseable}. The |
| 294 | + * underlying gRPC stream is automatically canceled and resources released when iteration reaches |
| 295 | + * the end (or maximum results requested) or when an error occurs. Callers that do not iterate to |
| 296 | + * completion rely on server-side stream timeouts and garbage collection to release stream |
| 297 | + * resources. |
| 298 | + */ |
| 299 | + static class ArrowQueryPageFetcher implements NextPageFetcher<FieldValueList> { |
| 300 | + private static final long serialVersionUID = 1L; |
| 301 | + private static final long DEFAULT_PAGE_SIZE = 10000L; |
| 302 | + |
| 303 | + private final JobId jobId; |
| 304 | + private final Schema schema; |
| 305 | + private final byte[] arrowSchemaBytes; |
| 306 | + private final BigQueryOptions serviceOptions; |
| 307 | + private final long maxResults; |
| 308 | + private final Map<BigQueryRpc.Option, ?> optionsMap; |
| 309 | + |
| 310 | + private transient org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo; |
| 311 | + private transient BigQueryReadClient bqReadClient; |
| 312 | + private transient ServerStream<ReadRowsResponse> stream; |
| 313 | + private transient Iterator<ReadRowsResponse> streamIterator; |
| 314 | + private transient Queue<FieldValueList> buffer = new ArrayDeque<>(); |
| 315 | + private long totalRowsReturned = 0L; |
| 316 | + private boolean streamClosed = false; |
| 317 | + |
| 318 | + ArrowQueryPageFetcher( |
| 319 | + JobId jobId, |
| 320 | + Schema schema, |
| 321 | + byte[] arrowSchemaBytes, |
| 322 | + org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo, |
| 323 | + BigQueryOptions serviceOptions, |
| 324 | + long initialRowOffset, |
| 325 | + Long maxResults, |
| 326 | + Map<BigQueryRpc.Option, ?> optionsMap) { |
| 327 | + this.jobId = jobId; |
| 328 | + this.schema = schema; |
| 329 | + this.arrowSchemaBytes = arrowSchemaBytes; |
| 330 | + this.arrowSchemaPojo = arrowSchemaPojo; |
| 331 | + this.serviceOptions = serviceOptions; |
| 332 | + this.totalRowsReturned = initialRowOffset; |
| 333 | + this.maxResults = maxResults != null ? maxResults : Long.MAX_VALUE; |
| 334 | + this.optionsMap = optionsMap; |
| 335 | + } |
| 336 | + |
| 337 | + @Override |
| 338 | + public Page<FieldValueList> getNextPage() { |
| 339 | + if (buffer == null) { |
| 340 | + buffer = new ArrayDeque<>(); |
| 341 | + } |
| 342 | + if (streamClosed || totalRowsReturned >= maxResults) { |
| 343 | + closeClient(); |
| 344 | + return null; |
| 345 | + } |
| 346 | + |
| 347 | + Number optionPageSize = |
| 348 | + optionsMap != null ? (Number) optionsMap.get(BigQueryRpc.Option.MAX_RESULTS) : null; |
| 349 | + long pageSize = |
| 350 | + optionPageSize != null && optionPageSize.longValue() > 0 |
| 351 | + ? optionPageSize.longValue() |
| 352 | + : DEFAULT_PAGE_SIZE; |
| 353 | + List<FieldValueList> rowBatch = new ArrayList<>((int) Math.min(pageSize, 10000L)); |
| 354 | + |
| 355 | + try { |
| 356 | + // Resolve job location in order: JobId location -> BigQueryOptions location -> "global" |
| 357 | + // default. |
| 358 | + // The Storage Read API stream resource name requires a location component (e.g. |
| 359 | + // projects/{project}/locations/{location}/jobs/{job}/streams/_default). If no specific |
| 360 | + // location |
| 361 | + // was provided on the job or service options, defaulting to "global" allows queries created |
| 362 | + // without an explicit location to still stream results without failing. |
| 363 | + String location = jobId.getLocation(); |
| 364 | + if (location == null) { |
| 365 | + location = serviceOptions.getLocation(); |
| 366 | + } |
| 367 | + if (location == null) { |
| 368 | + location = "global"; |
| 369 | + } |
| 370 | + |
| 371 | + if (streamIterator == null) { |
| 372 | + if (bqReadClient == null) { |
| 373 | + BigQuery service = serviceOptions.getService(); |
| 374 | + if (!(service instanceof BigQueryImpl)) { |
| 375 | + // Arrow pagination relies on BigQueryImpl to manage and cache the underlying |
| 376 | + // BigQueryReadClient across page fetches. Custom BigQuery implementations that do |
| 377 | + // not extend BigQueryImpl cannot provide the managed gRPC storage read client. |
| 378 | + throw new IllegalStateException( |
| 379 | + "Arrow query result pagination requires an instance of BigQueryImpl to manage BigQueryReadClient lifecycle"); |
| 380 | + } |
| 381 | + bqReadClient = ((BigQueryImpl) service).getBigQueryReadClient(location); |
| 382 | + } |
| 383 | + |
| 384 | + // Construct the default stream path for reading job query results via Storage Read API. |
| 385 | + String streamName = |
| 386 | + String.format( |
| 387 | + "projects/%s/locations/%s/jobs/%s/streams/_default", |
| 388 | + jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(), |
| 389 | + location, |
| 390 | + jobId.getJob()); |
| 391 | + |
| 392 | + ReadRowsRequest readRowsRequest = |
| 393 | + ReadRowsRequest.newBuilder() |
| 394 | + .setReadStream(streamName) |
| 395 | + .setOffset(totalRowsReturned) |
| 396 | + .build(); |
| 397 | + |
| 398 | + // Open the server-streaming RPC and obtain the response iterator. |
| 399 | + stream = bqReadClient.readRowsCallable().call(readRowsRequest); |
| 400 | + streamIterator = stream.iterator(); |
| 401 | + } |
| 402 | + |
| 403 | + // Lazily deserialize the Arrow schema from serialized bytes on the first page fetch (e.g., |
| 404 | + // after fetcher deserialization) and cache the POJO schema for subsequent batches. |
| 405 | + if (arrowSchemaPojo == null && arrowSchemaBytes != null) { |
| 406 | + arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes); |
| 407 | + } |
| 408 | + |
| 409 | + boolean hasMore = |
| 410 | + ArrowDeserializer.loadArrowRows( |
| 411 | + streamIterator, |
| 412 | + arrowSchemaPojo, |
| 413 | + schema, |
| 414 | + rowBatch, |
| 415 | + buffer, |
| 416 | + pageSize, |
| 417 | + totalRowsReturned, |
| 418 | + maxResults); |
| 419 | + |
| 420 | + // If no rows were read (e.g., EOF reached on an empty stream or after all results |
| 421 | + // were exhausted), terminate pagination. Note: when hasMore is false, rowBatch may |
| 422 | + // still contain the final page of data to return, so we only return null if rowBatch |
| 423 | + // is empty. |
| 424 | + if (rowBatch.isEmpty()) { |
| 425 | + streamClosed = true; |
| 426 | + closeClient(); |
| 427 | + return null; |
| 428 | + } |
| 429 | + |
| 430 | + totalRowsReturned += rowBatch.size(); |
| 431 | + |
| 432 | + // Generate nextPageToken only if there are more rows in the stream and maxResults |
| 433 | + // has not been reached. Otherwise, close the stream early to free resources. |
| 434 | + String nextPageToken = null; |
| 435 | + if (hasMore && totalRowsReturned < maxResults) { |
| 436 | + nextPageToken = String.valueOf(totalRowsReturned); |
| 437 | + } else { |
| 438 | + streamClosed = true; |
| 439 | + closeClient(); |
| 440 | + } |
| 441 | + |
| 442 | + return new PageImpl<>(this, nextPageToken, rowBatch); |
| 443 | + |
| 444 | + } catch (BigQueryException e) { |
| 445 | + streamClosed = true; |
| 446 | + closeClient(); |
| 447 | + throw e; |
| 448 | + } catch (Exception e) { |
| 449 | + streamClosed = true; |
| 450 | + closeClient(); |
| 451 | + throw new BigQueryException(0, "Failed to read Arrow rows from storage stream", e); |
| 452 | + } |
| 453 | + } |
| 454 | + |
| 455 | + /** |
| 456 | + * Cancels the active Storage Read API server stream (if any) and clears transient references so |
| 457 | + * gRPC stream resources can be reclaimed immediately. |
| 458 | + */ |
| 459 | + private void closeClient() { |
| 460 | + if (stream != null) { |
| 461 | + try { |
| 462 | + stream.cancel(); |
| 463 | + } catch (Exception e) { |
| 464 | + // Ignore cancellation exceptions during teardown |
| 465 | + } |
| 466 | + } |
| 467 | + bqReadClient = null; |
| 468 | + streamIterator = null; |
| 469 | + stream = null; |
| 470 | + } |
| 471 | + } |
| 472 | + |
281 | 473 | private static final int MAX_CACHED_READ_CLIENTS = 100; |
282 | 474 |
|
283 | 475 | private transient ConcurrentHashMap<String, BigQueryReadClient> bqReadClients; |
|
0 commit comments