Lance IO performance is 20x worse than TFRecord

[neuyilan]

I plan to train a model using the TensorFlow training framework, using either Lance or TFRecord for data. I ran a test and found that Lance's pure I/O performance is significantly worse than TFRecord's.
TFRecord's performance is approximately 20 times faster that of Lance.
Below is the results and attached is the test code.

I don't know Lance very well, is this as expected?

dataloader_benchmark.py.txt

[Xuanwo]

Hi, it looks like you are in a very slow loop:

dataset = tf.data.Dataset.from_generator(
    lambda: _lance_generator_full_parse(data_path, feature_configs),
    output_signature=output_signature
)

Please try using lance's own tf integration:

ds = lance.tf.data.from_lance(
    "s3://my-bucket/my-dataset",
    columns=["image", "label"],
    filter="split = 'train' AND collected_time > timestamp '2020-01-01'",
    batch_size=256)

See also https://lancedb.github.io/lance/integrations/tensorflow/

[neuyilan]

@Xuanwo Thanks for advice
I have tried the method mentioned above, but it reports the following error. The code is the example code.

ValueError: Can't convert non-rectangular Python sequence to Tensor.
 [[{{node PyFunc}}]] [Op:IteratorGetNext] name: 

The test codes like below:

import tensorflow as tf
from lance.tf.data import from_lance, lance_fragments

dataset_uri = "dataset.lance"

ds = from_lance(
    dataset_uri,
    batch_size=32
    )
for batch in ds:
    print(batch)

[neuyilan]

I have revised the code for reading tfrecord and lance, mainly by simulating the reading of each row of data and calculating the md5 of each row to avoid the performance impact of lazy loading of tfrecord. The code is shown below, and the result show that TFRecord is 5x faster than Lance

Lance

import argparse
import csv
import hashlib
import time
from typing import Dict, List

import tensorflow as tf
import lance
import numpy as np
import pandas as pd
import pyarrow.dataset as ds


# --- 特征配置类 (无变化) ---

class FeatureType:
    SPARSE = "SPARSE"
    SEQUENCE = "SEQUENCE"
    EMBEDDING = "EMBEDDING"
    AD_SEQUENCE = "AD_SEQUENCE"
    LABEL = "LABEL"

    @classmethod
    def from_str(cls, type_str: str) -> str:
        type_upper = type_str.upper()
        if type_upper in [cls.SPARSE, cls.SEQUENCE, cls.EMBEDDING, cls.AD_SEQUENCE, cls.LABEL]:
            return type_upper
        raise ValueError(f"无法解析特征类型: {type_str}")

class FeatureConfig:
    def __init__(self, name: str, value_dim: int, dtype: str, feature_type: str):
        self.name = name
        self.shape = [value_dim] if value_dim > 0 else []
        self.dtype = self._parse_dtype(dtype)
        self.feature_type = FeatureType.from_str(feature_type)
        self.ftype = self._get_ftype(feature_type)  # 'sparse', 'dense', 'scalar'

    def _parse_dtype(self, dtype_str: str) -> tf.DType:
        mapping = {'int64': tf.int64, 'float32': tf.float32, 'string': tf.string}
        lower_str = dtype_str.lower()
        if lower_str in mapping:
            return mapping[lower_str]
        raise ValueError(f"不支持的数据类型: {dtype_str}")

    def _get_ftype(self, feature_type_str: str) -> str:
        if FeatureType.from_str(feature_type_str) in [FeatureType.SEQUENCE, FeatureType.AD_SEQUENCE]:
            return 'sparse'
        elif FeatureType.from_str(feature_type_str) == FeatureType.EMBEDDING:
            return 'dense'
        else:
            return 'scalar'


class CSVConfigParser:
    def parse(self, file_path: str) -> Dict[str, FeatureConfig]:
        features: Dict[str, FeatureConfig] = {}
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                for row in reader:
                    if not row or row[0].strip().startswith('#'):
                        continue
                    self._process_row(row, features)
        except Exception as e:
            raise RuntimeError(f"解析特征配置文件失败: {file_path}, 错误: {e}")
        return features

    def _process_row(self, row: List[str], features: Dict[str, FeatureConfig]):
        try:
            name, _, dim, dtype, _, feature_type, *_ = row
            config = FeatureConfig(
                name=name.strip(),
                value_dim=int(dim),
                dtype=dtype.strip(),
                feature_type=feature_type.strip(),
            )
            features[config.name] = config
        except Exception as e:
            raise ValueError(f"无效的配置行: {row}, 错误: {e}")


def _lance_generator_no_parse(data_path: str, batch_size: int, feature_configs: Dict[str, FeatureConfig]):
    """生成器，用于模拟 Lance I/O，不进行解析"""
    columns_to_read = list(feature_configs.keys())
    dataset = lance.dataset(data_path)
    for batch_data in dataset.to_batches(batch_size=batch_size, columns=columns_to_read, batch_readahead=batch_size):
        batch_df = batch_data.to_pandas()
        yield from _parse_batch_no_parse(batch_df, feature_configs)

def _safe_array(value, dtype, default_shape=(0,)):
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return np.empty(default_shape, dtype=dtype)
    return np.array(value, dtype=dtype)


def _parse_batch(batch_df: pd.DataFrame, feature_configs: Dict[str, FeatureConfig]):
    batch_size = len(batch_df)
    if batch_size == 0: return

    final_schemas = [{} for _ in range(batch_size)]
    for name, config in feature_configs.items():
        if name not in batch_df.columns: continue

        column_data = batch_df[name]
        processed_column = [None] * batch_size

        if config.ftype == 'sparse':
            np_dtype = np.int64 if config.dtype == tf.int64 else np.float32
            processed_column = [tf.sparse.from_dense(_safe_array(val, dtype=np_dtype)) for val in column_data]
        elif config.ftype == 'dense':
            np_dtype = np.int64 if config.dtype == tf.int64 else np.float32
            processed_column = [_safe_array(val, dtype=np_dtype) for val in column_data]
        else:  # scalar
            np_dtype = np.int64 if config.dtype == tf.int64 else np.float32 if config.dtype == tf.float32 else str
            processed_column = column_data.astype(np_dtype).to_numpy()

        for i in range(batch_size):
            final_schemas[i][name] = processed_column[i]

    # print("final_schemas = ", final_schemas)
    yield from final_schemas


def _parse_batch_no_parse(batch_df: pd.DataFrame, feature_configs: Dict[str, FeatureConfig]):
    batch_size = len(batch_df)
    if batch_size == 0:
        return

    # 创建空字典结构的样本，保持output_signature的形状一致性
    final_schemas = [{} for _ in range(batch_size)]

    # 为所有特征填充0占位符，并计算每行的MD5
    for i in range(batch_size):
        # 提取当前行所有特征列的值并生成MD5
        row_values = []
        for name in feature_configs.keys():
            value = batch_df.at[i, name]
            md5_hash = hashlib.md5(value).hexdigest()
            # print(md5_hash)
            # 填充占位符
            final_schemas[i][name] = 0
    yield from final_schemas

def _lance_generator(lance_file: str, batch_size: int, feature_configs: Dict[str, FeatureConfig]):
    columns_to_read = list(feature_configs.keys())
    dataset = lance.dataset(lance_file)
    for batch_data in dataset.to_batches(batch_size=batch_size, columns=columns_to_read, batch_readahead=batch_size):
        batch_df = batch_data.to_pandas()
        yield from _parse_batch(batch_df, feature_configs)


def create_lance_dataset(data_path: str, feature_configs: Dict[str, FeatureConfig], batch_size: int,no_parse: bool
                         ) -> tf.data.Dataset:
    if no_parse:
        print("构建 Lance 数据集管道 (无解析)...")
        output_signature = {
            name: tf.TensorSpec(shape=(), dtype=tf.int32)  # 使用int32作为占位符类型
            for name in feature_configs.keys()
        }
        dataset = tf.data.Dataset.from_generator(
            lambda: _lance_generator_no_parse(data_path, batch_size ,feature_configs),
            output_signature=output_signature
        )
        # For generator-based datasets, batching should happen after the generator yields individual items.
        dataset = dataset.batch(batch_size, drop_remainder=True)
        return dataset
    else:
        output_signature = {}
        for name, config in feature_configs.items():
            if config.ftype == 'sparse':
                output_signature[name] = tf.SparseTensorSpec(shape=(None,), dtype=config.dtype)
            elif config.ftype == 'dense':
                output_signature[name] = tf.TensorSpec(shape=config.shape, dtype=config.dtype)
            else:  # scalar
                output_signature[name] = tf.TensorSpec(shape=(), dtype=config.dtype)

        dataset = tf.data.Dataset.from_generator(
            lambda: _lance_generator(data_path, batch_size, feature_configs),
            output_signature=output_signature
        )
        # For generator-based datasets, batching should happen after the generator yields individual items.
        dataset = dataset.batch(batch_size, drop_remainder=True)
        return dataset

def run_benchmark(args):
    mode_str = "无解析 (I/O-only)" if args.no_parse else "有解析 (Full)"
    print(f"--- 开始 {args.format.upper()} 数据加载性能测试 ({mode_str}) ---")
    print(f"数据文件路径: {args.data_path}")
    print(f"特征配置文件: {args.feature_config}")
    print(f"批次大小 (Batch Size): {args.batch_size}")
    print(f"迭代轮数 (Epochs): {args.epochs}")
    print("-" * 40)

    config_parser = CSVConfigParser()
    feature_configs = config_parser.parse(args.feature_config)
    print(f"成功解析 {len(feature_configs)} 个特征配置。")

    if args.format == 'lance':
        dataset = create_lance_dataset(args.data_path, feature_configs, args.batch_size, args.no_parse)
    else:
        raise ValueError(f"不支持的数据格式: {args.format}")

    dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
    print("数据集管道构建完成。开始迭代...")

    total_batches = 0
    total_examples = 0
    start_time = time.time()

    for epoch in range(args.epochs):
        epoch_batches = 0
        # 必须要遍历，才能够正确的遍历tfrecord的数据集。
        count = 0
        for row in dataset:
            print("the dict size =", len(row))
            if not args.no_parse:
                print("the dict keys =", row.keys())
            total_batches += 1
            epoch_batches += 1

        total_examples += epoch_batches * args.batch_size
        print(f"Epoch {epoch + 1}/{args.epochs} 完成, 处理了 {epoch_batches} 个批次。"
              f"处理了 {epoch_batches * args.batch_size} 个样本。")

    duration = time.time() - start_time

    if duration > 0:
        batches_per_second = total_batches / duration
        examples_per_second = total_examples / duration
    else:
        batches_per_second = float('inf')
        examples_per_second = float('inf')

    print(f"\n--- {args.format.upper()} 性能测试报告 ({mode_str}) ---")
    print(f"总耗时: {duration:.2f} 秒")
    print(f"总处理批次数: {total_batches}")
    print(f"总处理样本数: {total_examples} (约)")
    print(f"平均每秒处理批次数: {batches_per_second:.2f} batches/sec")
    print(f"平均每秒处理样本数: {examples_per_second:.2f} examples/sec")
    print("-" * 40)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Dataloader Performance Benchmark for TFRecord and Lance")
    parser.add_argument(
        "--data_path", type=str, required=True,
        help="数据文件的路径。TFRecord支持通配符; Lance指向 .lance 目录。"
    )
    parser.add_argument(
        "--format", type=str, required=True, choices=['tfrecord', 'lance', 'parquet'],
        help="要测试的数据格式。"
    )
    parser.add_argument(
        "--feature_config", type=str, default="model_tf2/config/feature_config_example.csv",
        help="特征配置CSV文件的路径。"
    )
    parser.add_argument(
        "--batch_size", type=int, default=512,
        help="每个批次的样本数量。"
    )
    parser.add_argument(
        "--epochs", type=int, default=3,
        help="完整遍历数据集的次数。"
    )
    parser.add_argument(
        "--no_parse", action='store_true',
        help="如果设置此标志，将运行I/O-only的无解析性能测试。"
    )
    
    args = parser.parse_args()
    

    run_benchmark(args)

TFRecord

import argparse
import csv
import time
from typing import Dict, List

import tensorflow as tf
import lance
import numpy as np
import pandas as pd
import pyarrow.dataset as ds


# --- 特征配置类 (无变化) ---

class FeatureType:
    SPARSE = "SPARSE"
    SEQUENCE = "SEQUENCE"
    EMBEDDING = "EMBEDDING"
    AD_SEQUENCE = "AD_SEQUENCE"
    LABEL = "LABEL"

    @classmethod
    def from_str(cls, type_str: str) -> str:
        type_upper = type_str.upper()
        if type_upper in [cls.SPARSE, cls.SEQUENCE, cls.EMBEDDING, cls.AD_SEQUENCE, cls.LABEL]:
            return type_upper
        raise ValueError(f"无法解析特征类型: {type_str}")

class FeatureConfig:
    def __init__(self, name: str, value_dim: int, dtype: str, feature_type: str):
        self.name = name
        self.shape = [value_dim] if value_dim > 0 else []
        self.dtype = self._parse_dtype(dtype)
        self.feature_type = FeatureType.from_str(feature_type)
        self.ftype = self._get_ftype(feature_type)  # 'sparse', 'dense', 'scalar'

    def _parse_dtype(self, dtype_str: str) -> tf.DType:
        mapping = {'int64': tf.int64, 'float32': tf.float32, 'string': tf.string}
        lower_str = dtype_str.lower()
        if lower_str in mapping:
            return mapping[lower_str]
        raise ValueError(f"不支持的数据类型: {dtype_str}")

    def _get_ftype(self, feature_type_str: str) -> str:
        if FeatureType.from_str(feature_type_str) in [FeatureType.SEQUENCE, FeatureType.AD_SEQUENCE]:
            return 'sparse'
        elif FeatureType.from_str(feature_type_str) == FeatureType.EMBEDDING:
            return 'dense'
        else:
            return 'scalar'


class CSVConfigParser:
    def parse(self, file_path: str) -> Dict[str, FeatureConfig]:
        features: Dict[str, FeatureConfig] = {}
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                for row in reader:
                    if not row or row[0].strip().startswith('#'):
                        continue
                    self._process_row(row, features)
        except Exception as e:
            raise RuntimeError(f"解析特征配置文件失败: {file_path}, 错误: {e}")
        return features

    def _process_row(self, row: List[str], features: Dict[str, FeatureConfig]):
        try:
            name, _, dim, dtype, _, feature_type, *_ = row
            config = FeatureConfig(
                name=name.strip(),
                value_dim=int(dim),
                dtype=dtype.strip(),
                feature_type=feature_type.strip(),
            )
            features[config.name] = config
        except Exception as e:
            raise ValueError(f"无效的配置行: {row}, 错误: {e}")

# --- 数据加载逻辑 ---

def create_tfrecord_dataset(data_path: str, feature_configs: Dict[str, FeatureConfig], batch_size: int, no_parse: bool) -> tf.data.Dataset:
    """为TFRecord创建数据集，可通过 no_parse 控制是否解析"""
    if no_parse:
        print("构建 TFRecord 数据集管道 (无解析)...")
        files = tf.data.Dataset.list_files(data_path, shuffle=False)
        dataset = tf.data.TFRecordDataset(files, num_parallel_reads=tf.data.AUTOTUNE)
        # 添加 map 操作以确保每个样本被处理（即使不做任何解析）
        dataset = dataset.batch(batch_size, drop_remainder=True)
        dataset = dataset.map(lambda x: x, num_parallel_calls=tf.data.AUTOTUNE)
    else:
        print("构建 TFRecord 数据集管道 (有解析)...")
        tf_schema = {
            name: tf.io.VarLenFeature(dtype=config.dtype) 
            if config.feature_type in [FeatureType.SEQUENCE, FeatureType.AD_SEQUENCE, FeatureType.LABEL] 
            else tf.io.FixedLenFeature(shape=config.shape, dtype=config.dtype)
            for name, config in feature_configs.items()
        }
        def _parse_fn(example_proto):
            return tf.io.parse_example(example_proto, tf_schema)

        files = tf.data.Dataset.list_files(data_path, shuffle=False)
        dataset = tf.data.TFRecordDataset(files, num_parallel_reads=tf.data.AUTOTUNE)
        dataset = dataset.batch(batch_size, drop_remainder=True)
        dataset = dataset.map(_parse_fn, num_parallel_calls=tf.data.AUTOTUNE)
    return dataset

def _safe_array(value, dtype, default_shape=(0,)):
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return np.empty(default_shape, dtype=dtype)
    return np.array(value, dtype=dtype)

# --- 基准测试主逻辑 ---

def run_benchmark(args):
    mode_str = "无解析 (I/O-only)" if args.no_parse else "有解析 (Full)"
    print(f"--- 开始 {args.format.upper()} 数据加载性能测试 ({mode_str}) ---")
    print(f"数据文件路径: {args.data_path}")
    print(f"特征配置文件: {args.feature_config}")
    print(f"批次大小 (Batch Size): {args.batch_size}")
    print(f"迭代轮数 (Epochs): {args.epochs}")
    print("-" * 40)

    config_parser = CSVConfigParser()
    feature_configs = config_parser.parse(args.feature_config)
    print(f"成功解析 {len(feature_configs)} 个特征配置。")

    if args.format == 'tfrecord':
        dataset = create_tfrecord_dataset(args.data_path, feature_configs, args.batch_size, args.no_parse)
    else:
        raise ValueError(f"不支持的数据格式: {args.format}")

    dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
    print("数据集管道构建完成。开始迭代...")

    total_batches = 0
    total_examples = 0
    start_time = time.time()

    for epoch in range(args.epochs):
        epoch_batches = 0
        # 必须要遍历，才能够正确的遍历tfrecord的数据集。
        count = 0
        for rows in dataset:
            import hashlib
            # print("the dict row =", row)
            print("the dict size =", len(rows))
            for row in rows:
                md5_hash = hashlib.md5(str(row).encode('utf-8')).hexdigest()
                # print("the md5_hash =", md5_hash)
                # print("the row =", row)
            if not args.no_parse:
                print("the dict keys =", rows.keys())
            total_batches += 1
            epoch_batches += 1

        total_examples += epoch_batches * args.batch_size
        print(f"Epoch {epoch + 1}/{args.epochs} 完成, 处理了 {epoch_batches} 个批次。"
              f"处理了 {epoch_batches * args.batch_size} 个样本。")

    duration = time.time() - start_time

    if duration > 0:
        batches_per_second = total_batches / duration
        examples_per_second = total_examples / duration
    else:
        batches_per_second = float('inf')
        examples_per_second = float('inf')

    print(f"\n--- {args.format.upper()} 性能测试报告 ({mode_str}) ---")
    print(f"总耗时: {duration:.2f} 秒")
    print(f"总处理批次数: {total_batches}")
    print(f"总处理样本数: {total_examples} (约)")
    print(f"平均每秒处理批次数: {batches_per_second:.2f} batches/sec")
    print(f"平均每秒处理样本数: {examples_per_second:.2f} examples/sec")
    print("-" * 40)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Dataloader Performance Benchmark for TFRecord and Lance")
    parser.add_argument(
        "--data_path", type=str, required=True,
        help="数据文件的路径。TFRecord支持通配符; Lance指向 .lance 目录。"
    )
    parser.add_argument(
        "--format", type=str, required=True, choices=['tfrecord', 'lance', 'parquet'],
        help="要测试的数据格式。"
    )
    parser.add_argument(
        "--feature_config", type=str, default="model_tf2/config/feature_config_example.csv",
        help="特征配置CSV文件的路径。"
    )
    parser.add_argument(
        "--batch_size", type=int, default=512,
        help="每个批次的样本数量。"
    )
    parser.add_argument(
        "--epochs", type=int, default=3,
        help="完整遍历数据集的次数。"
    )
    parser.add_argument(
        "--no_parse", action='store_true',
        help="如果设置此标志，将运行I/O-only的无解析性能测试。"
    )
    
    args = parser.parse_args()

    run_benchmark(args)

[neuyilan]

@Xuanwo Thanks for advice I have tried the method mentioned above, but it reports the following error. The code is the example code.

ValueError: Can't convert non-rectangular Python sequence to Tensor.
 [[{{node PyFunc}}]] [Op:IteratorGetNext] name: 

The test codes like below:

import tensorflow as tf
from lance.tf.data import from_lance, lance_fragments

dataset_uri = "dataset.lance"

ds = from_lance(
    dataset_uri,
    batch_size=32
    )
for batch in ds:
    print(batch)

The reason for this issue is that there are arrays of varying lengths in the data of a certain column. Lance failed to handle this situation correctly.