为了使用Petastorm从存储在S3中的Parquet创建TensorFlow数据集，实现更高效和可维护的实现，我尝试使用tf.data.Dataset.from_generator方法。这允许定义产生数据示例和标签的生成器函数。
下面是一个如何使用tf.data.dataset.from_generator修改代码的示例：

import tensorflow as tf
from petastorm import make_batch_reader, make_petastorm_dataset

cols = ['col1_nm', 'col2_nm', ...]

def generator_fn():
    with make_batch_reader(s3_paths, schema_fields=cols + ['target']) as reader:
        for e in reader:
            x_features = [getattr(e, c) for c in cols]
            y = getattr(e, 'target')
            yield x_features, y

def parse(x_features, y):
    X = tf.stack(x_features, axis=1)
    return X, y

dataset = tf.data.Dataset.from_generator(generator_fn, output_signature=(
    tf.TensorSpec(shape=(len(cols),), dtype=tf.float32),
    tf.TensorSpec(shape=(), dtype=tf.float32)
))

parsed_dataset = dataset.map(parse)

for e in parsed_dataset.take(3):
    print(e)

1.generator_fn函数被定义为生成数据示例（x_features）和标签（y）的生成器函数。在generator函数中调用make_batch_reader来创建读取器并迭代数据。

1. parse函数将x_features和y作为参数，并处理它们以创建所需的X和y格式。
   1.tf.data.Dataset.from_generator用于从generator_fn创建数据集。output_signature参数用于指定生成元素的形状和数据类型。
   1.然后使用parse函数Map结果数据集，对每个元素应用必要的处理。
   与直接在Petastorm数据集上使用map函数相比，这种方法应该可以对数据加载过程提供更好的控制和灵活性，并且可以更有效和更可维护。
   P.S：我没有测试代码。但它可能会给予你一些想法。