enhance: pickle the iterator if something failed

Signed-off-by: yangxuan <xuan.yang@zilliz.com>

Author: XuanYang-cn

.gitignore

@@ -11,4 +11,8 @@ venv/
 .venv/
 .idea/
 results/
-logs/
+logs/
+
+# AI rules
+CLAUDE.md
+AGENTS.md

vectordb_bench/backend/dataset.py

@@ -6,8 +6,8 @@
 
 import logging
 import pathlib
-import typing
 from enum import Enum
+from typing import Any, NamedTuple
 
 import pandas as pd
 import polars as pl
@@ -25,7 +25,7 @@
 log = logging.getLogger(__name__)
 
 
-class SizeLabel(typing.NamedTuple):
+class SizeLabel(NamedTuple):
     size: int
     label: str
     file_count: int
@@ -405,6 +405,17 @@ def __init__(self, dataset: DatasetManager):
         self._cur = None
         self._sub_idx = [0 for i in range(len(self._ds.train_files))]  # iter num for each file
 
+    def __getstate__(self):
+        """Custom pickle support to handle unpicklable generator."""
+        state = self.__dict__.copy()
+        # Remove the unpicklable generator from ParquetFile.iter_batches()
+        state["_cur"] = None
+        return state
+
+    def __setstate__(self, state: Any):
+        """Restore state after unpickling."""
+        self.__dict__.update(state)
+
     def __iter__(self):
         return self