处理Hugging Face Dataset时遇pyarrow长度不匹配错误求助

错误日志 #

Traceback (most recent call last):
  File "C:\Users\conno\LegalAIDataset\LegalAIDataset\main.py", line 39, in <module>
    new_dataset = dataset.map(process_row, batched=True, batch_size=1, remove_columns=None)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\conno\LegalAIDataset\LegalAIDataset\.venv\Lib\site-packages\datasets\arrow_dataset.py", line 602, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\conno\LegalAIDataset\LegalAIDataset\.venv\Lib\site-packages\datasets\arrow_dataset.py", line 567, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\conno\LegalAIDataset\LegalAIDataset\.venv\Lib\site-packages\datasets\arrow_dataset.py", line 3156, in map
    for rank, done, content in Dataset._map_single(**dataset_kwargs):
  File "C:\Users\conno\LegalAIDataset\LegalAIDataset\.venv\Lib\site-packages\datasets\arrow_dataset.py", line 3570, in _map_single
    writer.write_batch(batch)
  File "C:\Users\conno\LegalAIDataset\LegalAIDataset\.venv\Lib\site-packages\datasets\arrow_writer.py", line 571, in write_batch
    pa_table = pa.Table.from_arrays(arrays, schema=schema)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "pyarrow\table.pxi", line 4642, in pyarrow.lib.Table.from_arrays
  File "pyarrow\table.pxi", line 3922, in pyarrow.lib.Table.validate
  File "pyarrow\error.pxi", line 91, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Column 1 named type expected length 44 but got length 21

复现代码 #

import datasets
from datasets import load_dataset, Dataset
from semantic_text_splitter import TextSplitter


# Step 1: Load the existing dataset
dataset = load_dataset('HF_Dataset')
# Slice the 'train' split of the dataset
sliced_data = dataset['train'][:100]

# Convert the sliced data back into a Dataset object
dataset = Dataset.from_dict(sliced_data)


def chunk_text(text_list, metadata):
    splitter = TextSplitter(1000)
    chunks = [chunk for text in text_list for chunk in splitter.chunks(text)]
    return {"text_chunks": chunks, **metadata}
# Define a global executor
#executor = ThreadPoolExecutor(max_workers=1)

def process_row(batch):
    # Initialize a dictionary to store the results
    results = {k: [] for k in batch.keys()}
    results['text_chunks'] = []  # Add 'text_chunks' key to the results dictionary

    # Process each row in the batch
    for i in range(len(batch['text'])):
        # Apply the chunk_text function to the text
        chunks = chunk_text(batch['text'][i], {k: v[i] for k, v in batch.items() if k != 'text'})
        # Add the results to the dictionary
        for k, v in chunks.items():
            results[k].extend(v)

    # Return the results
    return results

# Apply the function to the dataset
new_dataset = dataset.map(process_row, batched=True, batch_size=1, remove_columns=None)

# Save and upload the new dataset
new_dataset.to_json('dataset.jsonl')
dataset_dict = datasets.DatasetDict({"split": new_dataset})


# dataset_dict.save_to_disk("", format="json")
# dataset_dict.upload_to_hub("", "This is a test dataset")

问题原因 #

错误核心是列长度不匹配：单条文本拆分成多个chunk后，元数据（如type列）只保留原单条数据的1份，但text_chunks列有多个元素，导致所有列长度不一致，Arrow无法生成合法表格。

解决方案 #

修改代码逻辑，确保每个文本chunk对应一份完整的原数据元数据，保证所有列长度一致。修正后的代码如下：

import datasets
from datasets import load_dataset, Dataset
from semantic_text_splitter import TextSplitter

# 加载数据集并切片
dataset = load_dataset('HF_Dataset')
sliced_data = dataset['train'][:100]
dataset = Dataset.from_dict(sliced_data)

def chunk_text(text, metadata):
    splitter = TextSplitter(1000)
    chunks = splitter.chunks(text)
    # 为每个chunk复制一份原元数据，确保一一对应
    chunked_rows = []
    for chunk in chunks:
        chunked_rows.append({
            "text_chunks": chunk,
            **metadata
        })
    return chunked_rows

def process_row(batch):
    all_chunked_rows = []
    # 遍历批次中的每条数据
    for i in range(len(batch['text'])):
        # 提取当前数据的元数据（排除text字段）
        current_metadata = {k: v[i] for k, v in batch.items() if k != 'text'}
        # 生成带元数据的chunk列表
        chunked_rows = chunk_text(batch['text'][i], current_metadata)
        all_chunked_rows.extend(chunked_rows)
    
    # 将列表转换为Dataset要求的字典格式（每个键对应所有行的该字段值列表）
    if not all_chunked_rows:
        return {}
    return {key: [row[key] for row in all_chunked_rows] for key in all_chunked_rows[0].keys()}

# 应用处理函数，可选择移除原text字段（避免冗余）
new_dataset = dataset.map(process_row, batched=True, batch_size=1, remove_columns=['text'])

# 保存为jsonl文件
new_dataset.to_json('dataset.jsonl')
dataset_dict = datasets.DatasetDict({"split": new_dataset})

关键修改点 #

1. 元数据复制：chunk_text函数为每个生成的chunk复制一份原数据的元数据，确保每个chunk都带有完整上下文信息。
2. 结果格式转换：process_row先收集所有带元数据的chunk行，再转换为Dataset要求的键值对格式，保证所有列长度完全一致。
3. 冗余字段处理：移除原text字段（可选），避免数据冗余，也可根据需求保留（需确保长度匹配）。

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