Merge pull request #320 from shijinpjlab/dev_parquet

feat: support parquet file

Author: shijinpjlab

dingo/config/__init__.py

@@ -1,2 +1,2 @@
-from dingo.config.input_args import (DatasetArgs, DatasetCsvArgs, DatasetExcelArgs, DatasetFieldArgs, DatasetHFConfigArgs, DatasetS3ConfigArgs, DatasetSqlArgs, EvalPipline,  # noqa E402.
-                                     EvalPiplineConfig, EvaluatorLLMArgs, EvaluatorRuleArgs, ExecutorArgs, ExecutorResultSaveArgs, InputArgs)
+from dingo.config.input_args import (DatasetArgs, DatasetCsvArgs, DatasetExcelArgs, DatasetFieldArgs, DatasetHFConfigArgs, DatasetParquetArgs, DatasetS3ConfigArgs, DatasetSqlArgs,  # noqa E402.
+                                     EvalPipline, EvalPiplineConfig, EvaluatorLLMArgs, EvaluatorRuleArgs, ExecutorArgs, ExecutorResultSaveArgs, InputArgs)

dingo/config/input_args.py

@@ -40,6 +40,11 @@ class DatasetCsvArgs(BaseModel):
     quotechar: str = '"'  # 引号字符，默认双引号
 
 
+class DatasetParquetArgs(BaseModel):
+    batch_size: int = 10000  # 每次读取的行数，用于流式读取大文件
+    columns: Optional[List[str]] = None  # 指定读取的列，None 表示读取所有列
+
+
 class DatasetFieldArgs(BaseModel):
     id: str = ''
     prompt: str = ''
@@ -58,6 +63,7 @@ class DatasetArgs(BaseModel):
     sql_config: DatasetSqlArgs = DatasetSqlArgs()
     excel_config: DatasetExcelArgs = DatasetExcelArgs()
     csv_config: DatasetCsvArgs = DatasetCsvArgs()
+    parquet_config: DatasetParquetArgs = DatasetParquetArgs()
 
 
 class ExecutorResultSaveArgs(BaseModel):

dingo/data/converter/base.py

@@ -299,6 +299,25 @@ def _convert(raw: Union[str, Dict]):
         return _convert
 
 
+@BaseConverter.register("parquet")
+class ParquetConverter(BaseConverter):
+    """Parquet file converter."""
+
+    def __init__(self):
+        super().__init__()
+
+    @classmethod
+    def convertor(cls, input_args: InputArgs) -> Callable:
+        def _convert(raw: Union[str, Dict]):
+            j = raw
+            if isinstance(raw, str):
+                j = json.loads(raw)
+            data_dict = j
+            return Data(**data_dict)
+
+        return _convert
+
+
 @BaseConverter.register("listjson")
 class ListJsonConverter(BaseConverter):
     """List json file converter."""

dingo/data/datasource/local.py

@@ -142,6 +142,75 @@ def _load_excel_file_xlsx(self, path: str) -> Generator[str, None, None]:
             if wb:
                 wb.close()
 
+    def _load_parquet_file(self, path: str) -> Generator[str, None, None]:
+        """
+        Load a Parquet file and return its contents row by row as JSON strings.
+        Supports streaming for large files to avoid memory overflow.
+
+        Args:
+            path (str): The path to the Parquet file.
+
+        Returns:
+            Generator[str]: Each row as a JSON string with column keys.
+        """
+        try:
+            import pyarrow.parquet as pq
+        except ImportError:
+            raise RuntimeError(
+                "pyarrow is required to read Parquet files. "
+                "Please install it using: pip install pyarrow"
+            )
+
+        # 获取 Parquet 配置
+        batch_size = self.input_args.dataset.parquet_config.batch_size
+        columns = self.input_args.dataset.parquet_config.columns
+
+        try:
+            # 打开 Parquet 文件
+            parquet_file = pq.ParquetFile(path)
+
+            # 使用流式读取，分批次处理
+            for batch in parquet_file.iter_batches(batch_size=batch_size, columns=columns):
+                # 将 batch 转换为字典格式
+                batch_dict = batch.to_pydict()
+
+                # 获取批次中的行数
+                num_rows = len(next(iter(batch_dict.values()))) if batch_dict else 0
+
+                # 逐行处理
+                for i in range(num_rows):
+                    # 构建每一行的字典
+                    row_dict = {col: batch_dict[col][i] for col in batch_dict}
+
+                    # 处理特殊类型的值
+                    for key, value in row_dict.items():
+                        # 处理 None 值
+                        if value is None:
+                            row_dict[key] = ""
+                        # 处理 bytes 类型
+                        elif isinstance(value, bytes):
+                            try:
+                                row_dict[key] = value.decode('utf-8')
+                            except UnicodeDecodeError:
+                                row_dict[key] = str(value)
+                        # 处理其他不可 JSON 序列化的类型
+                        elif not isinstance(value, (str, int, float, bool, list, dict)):
+                            row_dict[key] = str(value)
+
+                    # 转换为 JSON 字符串并 yield
+                    yield json.dumps(row_dict, ensure_ascii=False) + '\n'
+
+        except ImportError as ie:
+            raise RuntimeError(
+                f'Failed to load required library for Parquet: {str(ie)}. '
+                f'Please install pyarrow using: pip install pyarrow'
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f'Failed to read Parquet file "{path}": {str(e)}. '
+                f'Please ensure the file is a valid Parquet file.'
+            )
+
     def _load_csv_file(self, path: str) -> Generator[str, None, None]:
         """
         Load a CSV file and return its contents row by row as JSON strings.
@@ -334,6 +403,11 @@ def _load_local_file(self) -> Generator[str, None, None]:
                 if self.input_args.dataset.format != 'csv':
                     raise RuntimeError(f'CSV file "{f}" is not supported. Please set dataset.format to "csv" to read CSV files.')
                 yield from self._load_csv_file(f)
+            # Check if file is Parquet
+            elif f.endswith('.parquet'):
+                if self.input_args.dataset.format != 'parquet':
+                    raise RuntimeError(f'Parquet file "{f}" is not supported. Please set dataset.format to "parquet" to read Parquet files.')
+                yield from self._load_parquet_file(f)
             # Check if file is Excel
             elif f.endswith('.xlsx'):
                 if self.input_args.dataset.format != 'excel':