tensorflow-io 0.23.1读取BigQuery表报错及cardinality返回-2问题

参考官方示例读取BigQuery表数据，根据官方文档说明，parallel_read_rows() 方法的返回值应为tf.dataset类型，但实际调用dataset.cardinality().numpy()时返回值为-2。
实现代码如下：

from tensorflow_io.bigquery import BigQueryClient

def read_bigquery(table_name):
    bigquery_client = BigQueryClient()
    project, dataset, table = table_name.split(".")
    read_session = bigquery_client.read_session(
        "projects/" + project,
        project,
        table,
        dataset,
        {
            "label": {"output_type": dtypes.int64, "mode": BigQueryClient.FieldMode.NULLABLE},
            "features": {"mode": BigQueryClient.FieldMode.REPEATED, "output_type": dtypes.float32},
        },
        requested_streams=2,
    )
    return read_session.parallel_read_rows()

dataset = read_bigquery(my_table)

读取数据集过程中出现异常日志，返回的数据集类型为<class 'tensorflow.python.data.ops.dataset_ops.ParallelInterleaveDataset'>，具体报错日志如下：

returned ds type: <class 'tensorflow.python.data.ops.dataset_ops.ParallelInterleaveDataset'>
2022-06-30 15:18:47.292204: E tensorflow/core/framework/dataset.cc:577] UNIMPLEMENTED: Cannot compute input sources for dataset of type IO>BigQueryDataset, because the dataset does not implement `InputDatasets`.
2022-06-30 15:18:47.292243: E tensorflow/core/framework/dataset.cc:581] UNIMPLEMENTED: Cannot merge options for dataset of type IO>BigQueryDataset, because the dataset does not implement `InputDatasets`.
2022-06-30 15:18:47.292455: E tensorflow/core/framework/dataset.cc:577] UNIMPLEMENTED: Cannot compute input sources for dataset of type IO>BigQueryDataset, because the dataset does not implement `InputDatasets`.
2022-06-30 15:18:47.292482: E tensorflow/core/framework/dataset.cc:581] UNIMPLEMENTED: Cannot merge options for dataset of type IO>BigQueryDataset, because the dataset does not implement `InputDatasets`.

当前运行环境使用的tensorflow-io版本为tensorflow-io==0.23.1，需要排查解决上述问题。

• 首先明确两个现象不属于代码逻辑错误：
  • dataset.cardinality().numpy()返回-2是TensorFlow Dataset的标准返回值，对应常量UNKNOWN_CARDINALITY，代表框架无法在不遍历全量数据的前提下提前确定数据集总条数，不是返回值异常。
  • parallel_read_rows()返回ParallelInterleaveDataset是符合预期的，该类型本身就是tf.data.Dataset的子类，和文档描述的返回类型没有冲突。
• 日志中打印的UNIMPLEMENTED提示是tensorflow-io 0.23.1版本的已知兼容问题：BigQuery自定义数据集没有实现InputDatasets接口，这部分是底层框架的提示日志，不会阻断正常的数据读取流程，只要后续迭代数据集能正常拿到数据就不需要特殊处理。
• 可根据需求选择对应修复方案：
  • 如果需要消除该报错日志，可以升级tensorflow和tensorflow-io到严格匹配的稳定版本，两个包版本必须一一对应（例如tensorflow 2.10.x匹配tensorflow-io 0.25.x），高版本已经修复了该接口未实现的问题。
  • 如果需要获取数据集总条数，不要依赖cardinality()方法，可以直接在BigQuery侧执行count查询拿到总行数，或者遍历全量数据集手动计数。
  • 验证数据集读取逻辑是否正常，可以直接取小批量数据测试，例如执行for batch in dataset.take(1): print(batch)，如果能正常输出特征和标签字段，说明读取逻辑完全可用，不需要额外处理底层提示日志。

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