Ray含list<double>数据集转TensorFlow格式异常，如何转为(25,2)形状？

我的Ray数据集结构如下：

Column  Type
------  ----
x_1     list<element: double not null>
x_2     list<element: double not null>
y       int32

尝试用data.to_tf(feature_columns=["x_1", "x_2"], label_columns="y")转换为TensorFlow对象时，得到如下结果：

<_OptionsDataset element_spec=({'x_1': TensorSpec(shape=(None,), dtype=tf.string, name='x_1'), 'x_2': TensorSpec(shape=(None,), dtype=tf.string, name='x_2')}, TensorSpec(shape=(None,), dtype=tf.int32, name='y'))>

可见特征x_1和x_2被默认转为string类型，我需要将特征转换为形状为(25, 2)的格式（25是列表长度，2是特征数量），请问该如何操作？

方法1：批量预处理合并特征 #

先对Ray数据集做批量预处理，将两个列表特征合并为二维张量，再转为TensorFlow数据集：

import tensorflow as tf
from ray.data import Dataset

def preprocess_batch(batch: dict) -> dict:
    # 将x_1和x_2堆叠成形状为(批次大小, 25, 2)的张量
    features = tf.stack([
        tf.convert_to_tensor(batch["x_1"], dtype=tf.float32),
        tf.convert_to_tensor(batch["x_2"], dtype=tf.float32)
    ], axis=-1)
    return {"features": features, "label": batch["y"]}

# 预处理后转为TF数据集
processed_data = data.map_batches(preprocess_batch, batch_format="pandas")
tf_dataset = processed_data.to_tf(label_columns="label", feature_columns="features")

处理后单条样本的特征形状为(25, 2)，批次维度为None。

方法2：指定张量规格后合并特征 #

直接在to_tf中指定特征的张量规格，再通过map合并特征：

import tensorflow as tf
from ray.data import Dataset

# 为每个列表特征指定固定长度的张量规格
tensor_specs = {
    "x_1": tf.TensorSpec(shape=(25,), dtype=tf.float32),
    "x_2": tf.TensorSpec(shape=(25,), dtype=tf.float32)
}

# 转换后合并特征
tf_dataset = data.to_tf(
    feature_columns=["x_1", "x_2"],
    label_columns="y",
    tensor_spec=tensor_specs
).map(lambda features, label: (tf.stack([features["x_1"], features["x_2"]], axis=-1), label))

前置注意事项 #

如果样本的x_1/x_2列表长度不固定，需要先做截断或填充处理：

from ray.data.preprocessors import PadOrTrim

# 将x_1和x_2统一为25长度
pad_trim_x1 = PadOrTrim(column="x_1", target_length=25)
data = pad_trim_x1.fit_transform(data)

pad_trim_x2 = PadOrTrim(column="x_2", target_length=25)
data = pad_trim_x2.fit_transform(data)

内容的提问来源于stack exchange，提问作者Prabhjot Singh Rai