BigQuery Storage API读取会话：获取流元数据及指定流行数方法

一、获取每个流的数据量元数据 #

BigQuery Storage API 创建会话后，不会直接返回每个流的行数或页数，可通过以下两种方式获取相关估算信息：

• 预读取批次估算：从每个流读取第一批数据，结合批次行数和剩余数据预估量推算总行数：

from google.cloud import bigquery_storage_v1

# 假设已完成read_session和bqstorageclient的初始化
for stream in read_session.streams:
    reader = bqstorageclient.read_rows(stream.name)
    first_batch = next(reader.rows())
    batch_rows = first_batch.num_rows
    # remaining_rows为API返回的预估剩余行数，仅作参考
    estimated_total = batch_rows + reader.remaining_rows
    print(f"流 {stream.name} 预估总行数: {estimated_total}")

• 表元数据均分估算：先通过BigQuery常规API获取表总行数，再按流数量均分（注意实际流拆分可能不均，仅为近似值）：

from google.cloud import bigquery

bq_client = bigquery.Client()
table = bq_client.get_table("your-project.your-dataset.your-table")
total_rows = table.num_rows
stream_count = len(read_session.streams)
estimated_per_stream = total_rows // stream_count
print(f"每个流预估行数: {estimated_per_stream}")

二、设置流大小为9998行 #

BigQuery Storage API 不支持直接指定流的行数，但可通过读取时控行或自定义分片实现近似效果：

• 读取时限制行数：读取流数据时累计行数，达到9998行即停止，合并为DataFrame：

import pyarrow as pa
import pandas as pd

def read_stream_with_row_limit(stream, row_limit=9998):
    reader = bqstorageclient.read_rows(stream.name)
    rows_iter = reader.rows()
    collected_batches = []
    current_row_count = 0
    
    for batch in rows_iter:
        if current_row_count + batch.num_rows > row_limit:
            # 截取所需剩余行数
            remaining_rows = row_limit - current_row_count
            collected_batches.append(batch.slice(0, remaining_rows))
            current_row_count += remaining_rows
            break
        collected_batches.append(batch)
        current_row_count += batch.num_rows
        if current_row_count >= row_limit:
            break
    
    # 合并批次为DataFrame
    arrow_table = pa.concat_tables(collected_batches)
    return arrow_table.to_pandas()

# 使用示例
target_stream = read_session.streams[3]
limited_df = read_stream_with_row_limit(target_stream, 9998)

• 自定义分片控制：通过ReadSession的read_options指定行过滤条件，将数据拆分到单流中，间接控制流的行数：

read_options = types.ReadSession.TableReadOptions(
    # 根据表中字段分布设置过滤条件，确保结果行数接近9998
    row_restriction="your_sortable_column BETWEEN 'start_value' AND 'end_value'"
)
requested_session = types.ReadSession(
    table=table,
    data_format=types.DataFormat.ARROW,
    read_options=read_options,
)
# 创建单流会话，该流的数据量由过滤条件控制
read_session = bqstorageclient.create_read_session(
    parent=parent,
    read_session=requested_session,
    max_stream_count=1
)

内容的提问来源于stack exchange，提问作者hello