diff --git a/data_server/formatify/FormatifyManager.py b/data_server/formatify/FormatifyManager.py index f6be30a..0217bcb 100644 --- a/data_server/formatify/FormatifyManager.py +++ b/data_server/formatify/FormatifyManager.py @@ -63,6 +63,8 @@ def _submit_formatify_task_to_csghub( namespace: str, user_name: str | None = None, task_run_time: str | None = None, + use_streaming: bool | None = None, + chunk_size: int | None = None, ): flow_id = build_job_flow_id("formatify", formatify_task.id) formatify_task.flow_id = flow_id @@ -73,6 +75,12 @@ def _submit_formatify_task_to_csghub( task_params["flow_id"] = flow_id if task_run_time: task_params["execute_time"] = task_run_time + + # Add streaming mode parameters (not stored in database, passed via task_params) + if use_streaming is not None: + task_params["use_streaming"] = use_streaming + if chunk_size is not None: + task_params["chunk_size"] = chunk_size dag_tasks = build_formatify_dag(flow_id, task_params) payload = build_csghub_payload( job_id=flow_id, @@ -124,6 +132,10 @@ def create_formatify_task( # Prepare skip_meta value (use provided value or default to False) skip_meta_value = dataFormatTask.skip_meta if dataFormatTask.skip_meta is not None else False + # Extract streaming mode parameters (not stored in database) + use_streaming = dataFormatTask.use_streaming + chunk_size = dataFormatTask.chunk_size + data_format_task_db = DataFormatTask(name=dataFormatTask.name, des=dataFormatTask.des, from_csg_hub_dataset_name=dataFormatTask.from_csg_hub_dataset_name, @@ -163,6 +175,8 @@ def create_formatify_task( user_token, nu, user_name=user_name, + use_streaming=use_streaming, + chunk_size=chunk_size, ) data_format_task_db.task_status = DataFormatTaskStatusEnum.WAITING.value except Exception as e: @@ -409,6 +423,8 @@ def execute_formatify_task( user_name: str, user_token: str, task_run_time: str | None = None, + use_streaming: bool | None = None, + chunk_size: int | None = None, ): """Waiting and not yet submitted to CSGHub: first submit this record.""" try: @@ -425,6 +441,8 @@ def execute_formatify_task( nu, user_name=user_name, task_run_time=task_run_time, + use_streaming=use_streaming, + chunk_size=chunk_size, ) formatify_task.task_status = DataFormatTaskStatusEnum.WAITING.value db_session.commit() @@ -442,6 +460,8 @@ def execute_new_formatify_task( user_name: str, user_token: str, task_run_time: str | None = None, + use_streaming: bool | None = None, + chunk_size: int | None = None, ): """List "Execute": copy new task and submit to CSGHub; do not re-run old record.""" try: @@ -462,6 +482,8 @@ def execute_new_formatify_task( nu, user_name=user_name, task_run_time=task_run_time, + use_streaming=use_streaming, + chunk_size=chunk_size, ) new_task.task_status = DataFormatTaskStatusEnum.WAITING.value db_session.commit() diff --git a/data_server/formatify/schemas.py b/data_server/formatify/schemas.py index f00b8ac..8769bc0 100644 --- a/data_server/formatify/schemas.py +++ b/data_server/formatify/schemas.py @@ -26,6 +26,10 @@ class DataFormatTaskRequest(BaseModel): storage_size: Optional[str] = None namespace_uuid: Optional[str] = None namespace_type: str = "personal" + + # Streaming mode parameters (not stored in database, only passed to conversion tasks) + use_streaming: Optional[bool] = None + chunk_size: Optional[int] = None @field_validator("storage_size", mode="before") @classmethod diff --git a/data_server/pod/common_tasks.py b/data_server/pod/common_tasks.py index a8d0fb0..35ee3de 100644 --- a/data_server/pod/common_tasks.py +++ b/data_server/pod/common_tasks.py @@ -38,6 +38,9 @@ convert_excel_to_csv, convert_excel_to_json, convert_excel_to_parquet, + convert_csv_to_json_streaming, + convert_csv_to_parquet_streaming, + convert_csv_to_excel_streaming, convert_word_to_markdown, convert_txt_to_markdown, convert_html_to_markdown, @@ -828,7 +831,17 @@ def run_format_conversion(task_params: dict): if not found: raise ValueError("No source files found for format conversion") - convert_func = _select_convert_func(task_params.get("from_data_type"), task_params.get("to_data_type")) + # Extract streaming mode parameters + use_streaming = task_params.get("use_streaming", False) + if isinstance(use_streaming, str): + use_streaming = use_streaming.lower() in ("true", "1", "yes") + use_streaming = bool(use_streaming) + + convert_func = _select_convert_func( + task_params.get("from_data_type"), + task_params.get("to_data_type"), + use_streaming=use_streaming + ) if convert_func is None: raise ValueError("Unsupported format conversion") @@ -983,7 +996,30 @@ def run_format_conversion(task_params: dict): } -def _select_convert_func(from_type, to_type): +def _select_convert_func(from_type, to_type, use_streaming=False): + """ + Select conversion function based on format types and mode. + + Args: + from_type: Source format type + to_type: Target format type + use_streaming: If True, use streaming mode for CSV conversions (when available) + + Returns: + Conversion function or None + """ + # For CSV source with streaming mode enabled, use streaming functions + if use_streaming and from_type == DataFormatTypeEnum.Csv.value: + streaming_mapping = { + DataFormatTypeEnum.Excel.value: convert_csv_to_excel_streaming, + DataFormatTypeEnum.Json.value: convert_csv_to_json_streaming, + DataFormatTypeEnum.Parquet.value: convert_csv_to_parquet_streaming, + } + streaming_func = streaming_mapping.get(to_type) + if streaming_func: + return streaming_func + + # Standard (non-streaming) mode mapping mapping = { (DataFormatTypeEnum.Excel.value, DataFormatTypeEnum.Csv.value): convert_excel_to_csv, (DataFormatTypeEnum.Excel.value, DataFormatTypeEnum.Json.value): convert_excel_to_json, @@ -1002,6 +1038,8 @@ def _select_convert_func(from_type, to_type): def _run_convert_func(convert_func, file_path: str, task_uid: str, task_params: dict): + + # PDF to Markdown has special parameters if convert_func is convert_pdf_to_markdown: return convert_func( file_path, @@ -1009,6 +1047,24 @@ def _run_convert_func(convert_func, file_path: str, task_uid: str, task_params: task_params.get("mineru_api_url"), task_params.get("mineru_backend"), ) + + # Streaming conversion functions require chunk_size + streaming_funcs = ( + convert_csv_to_json_streaming, + convert_csv_to_parquet_streaming, + convert_csv_to_excel_streaming, + ) + + if convert_func in streaming_funcs: + chunk_size = task_params.get("chunk_size", 50000) + try: + chunk_size = int(chunk_size) + except (ValueError, TypeError): + chunk_size = 50000 + + return convert_func(file_path, task_uid, chunk_size=chunk_size) + + # Standard conversion functions return convert_func(file_path, task_uid) diff --git a/data_server/pod/formatify_helpers.py b/data_server/pod/formatify_helpers.py index 2b612cb..51afb50 100644 --- a/data_server/pod/formatify_helpers.py +++ b/data_server/pod/formatify_helpers.py @@ -1,3 +1,4 @@ +import gc import json import os import re @@ -28,6 +29,26 @@ def _read_csv(file_path: str) -> pd.DataFrame: return pd.read_csv(file_path, sep=None, engine="python") +def _read_csv_chunked(file_path: str, chunk_size: int): + """Read CSV in chunks with proper encoding detection.""" + last_error = None + for encoding in ("utf-8-sig", "utf-8", "gb18030"): + try: + return pd.read_csv( + file_path, + encoding=encoding, + sep=None, + engine="python", + chunksize=chunk_size, + iterator=True + ) + except UnicodeDecodeError as error: + last_error = error + if last_error is not None: + raise last_error + return pd.read_csv(file_path, sep=None, engine="python", chunksize=chunk_size, iterator=True) + + def _non_conflicting_output_path(file_path: str) -> str: """Avoid overwriting an existing source/target file in the raw stage directory.""" if not os.path.exists(file_path): @@ -73,7 +94,7 @@ def convert_excel_to_csv(file_path: str, task_uid) -> Optional[Dict[str, str]]: else: new_file = f"{base_name}_{safe_sheet_name}.csv" new_file = _non_conflicting_output_path(new_file) - + # Use utf-8-sig encoding to ensure Excel can open the CSV correctly df.to_csv(new_file, index=False, encoding='utf-8-sig') result_files.append(new_file) @@ -206,7 +227,7 @@ def convert_excel_to_parquet(file_path: str, task_uid) -> Optional[Dict[str, str for idx, sheet_name in enumerate(sheet_names, 1): try: log_task_info(task_uid, f"Processing sheet {idx}/{sheet_count}: '{sheet_name}'") - + # Read the sheet df = pd.read_excel(file_path, sheet_name=sheet_name) @@ -221,18 +242,18 @@ def convert_excel_to_parquet(file_path: str, task_uid) -> Optional[Dict[str, str elif pd.api.types.is_float_dtype(df[col]): if df[col].isna().any(): pass - + # Generate output file name # Clean sheet name to remove invalid file system characters safe_sheet_name = re.sub(r'[<>:"/\\|?*]', '_', sheet_name) - + # If only one sheet, use simple naming; otherwise include sheet name if sheet_count == 1: new_file = f"{base_name}.parquet" else: new_file = f"{base_name}_{safe_sheet_name}.parquet" new_file = _non_conflicting_output_path(new_file) - + # Save to parquet df.to_parquet(new_file, index=False, engine="pyarrow") result_files.append(new_file) @@ -242,7 +263,7 @@ def convert_excel_to_parquet(file_path: str, task_uid) -> Optional[Dict[str, str f"Sheet '{sheet_name}' converted successfully: {new_file} " f"({len(df)} rows, {len(df.columns)} columns)" ) - + except Exception as sheet_error: log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") # Continue processing other sheets even if one fails @@ -274,7 +295,7 @@ def convert_excel_to_parquet(file_path: str, task_uid) -> Optional[Dict[str, str "status": "success", "sheets_count": len(result_files) } - + except Exception as e: log_task_error(task_uid, f"convert file {file_path} error: {e}") return { @@ -285,7 +306,6 @@ def convert_excel_to_parquet(file_path: str, task_uid) -> Optional[Dict[str, str } return None - def convert_csv_to_excel(file_path: str, task_uid) -> Optional[Dict[str, str]]: if not file_path.lower().endswith(".csv"): return None @@ -306,6 +326,226 @@ def convert_csv_to_excel(file_path: str, task_uid) -> Optional[Dict[str, str]]: log_task_error(task_uid, f"convert file {file_path} error: {e}") return {"from": file_path, "to": None, "status": "failure", "error": str(e)} +def convert_csv_to_json_streaming(file_path: str, task_uid, chunk_size: int = 50000) -> Optional[Dict[str, str]]: + """ + Convert CSV to JSON using streaming mode (low memory footprint). + + Args: + file_path: CSV file path + task_uid: Task identifier for logging + chunk_size: Number of rows per chunk (default: 50000) + + Returns: + Conversion result dictionary + """ + if not file_path.lower().endswith(".csv"): + return None + + log_task_info(task_uid, f"[Streaming Mode] Source file address: {file_path}") + log_task_info(task_uid, f"[Streaming Mode] Chunk size: {chunk_size:,} rows") + + new_file = None + try: + new_file = f"{os.path.splitext(file_path)[0]}.json" + + first_row = True + + with open(new_file, 'w', encoding='utf-8') as json_file: + json_file.write('[') + + for chunk in _read_csv_chunked(file_path, chunk_size): + columns = chunk.columns.tolist() + + # Use itertuples for better performance (faster than iloc) + for row in chunk.itertuples(index=False, name=None): + if not first_row: + json_file.write(',') + first_row = False + + # Build dict for single row + row_dict = {} + for col, value in zip(columns, row): + row_dict[col] = None if pd.isna(value) else value + + json_file.write(json.dumps(row_dict, ensure_ascii=False)) + + del chunk + gc.collect() + + json_file.write(']') + + return { + "from": file_path, + "to": new_file, + "to_files": [new_file], + "status": "success", + } + except Exception as e: + log_task_error(task_uid, f"convert file {file_path} error: {e}") + if new_file and os.path.exists(new_file): + try: + os.remove(new_file) + except Exception: + pass + return {"from": file_path, "to": None, "status": "failure", "error": str(e)} + + +def convert_csv_to_parquet_streaming(file_path: str, task_uid, chunk_size: int = 50000) -> Optional[Dict[str, str]]: + """ + Convert CSV to Parquet using streaming mode (low memory footprint). + + Args: + file_path: CSV file path + task_uid: Task identifier for logging + chunk_size: Number of rows per chunk (default: 50000) + + Returns: + Conversion result dictionary + + Note: + Data type processing is consistent with convert_excel_to_parquet() + """ + if not file_path.lower().endswith(".csv"): + return None + + try: + import pyarrow as pa + import pyarrow.parquet as pq + except ImportError: + log_task_error(task_uid, "PyArrow is required for streaming Parquet conversion") + return {"from": file_path, "to": None, "status": "failure", "error": "PyArrow not installed"} + + log_task_info(task_uid, f"[Streaming Mode] Source file address: {file_path}") + log_task_info(task_uid, f"[Streaming Mode] Chunk size: {chunk_size:,} rows") + + new_file = None + try: + new_file = f"{os.path.splitext(file_path)[0]}.parquet" + + writer = None + + for chunk in _read_csv_chunked(file_path, chunk_size): + # Data type processing (consistent with original logic) + for col in chunk.columns: + col_dtype = chunk[col].dtype + if col_dtype == "object": + # Convert to string and replace "nan" with None (consistent with original) + chunk[col] = chunk[col].astype(str).replace("nan", None) + elif pd.api.types.is_integer_dtype(col_dtype) and chunk[col].isna().any(): + chunk[col] = chunk[col].astype(str) + + table = pa.Table.from_pandas(chunk, preserve_index=False) + + if writer is None: + writer = pq.ParquetWriter(new_file, table.schema) + writer.write_table(table) + + del chunk + del table + gc.collect() + + if writer: + writer.close() + + return { + "from": file_path, + "to": new_file, + "to_files": [new_file], + "status": "success", + } + except Exception as e: + log_task_error(task_uid, f"convert file {file_path} error: {e}") + if new_file and os.path.exists(new_file): + try: + os.remove(new_file) + except Exception: + pass + return {"from": file_path, "to": None, "status": "failure", "error": str(e)} + + +def convert_csv_to_excel_streaming(file_path: str, task_uid, chunk_size: int = 50000) -> Optional[Dict[str, str]]: + """ + Convert CSV to Excel using streaming mode (constant memory). + + Args: + file_path: CSV file path + task_uid: Task identifier for logging + chunk_size: Number of rows per chunk (default: 50000) + + Returns: + Conversion result dictionary + """ + if not file_path.lower().endswith(".csv"): + return None + + try: + import xlsxwriter + except ImportError: + log_task_error(task_uid, "xlsxwriter is required for streaming Excel conversion") + return {"from": file_path, "to": None, "status": "failure", "error": "xlsxwriter not installed"} + + log_task_info(task_uid, f"[Streaming Mode] Source file address: {file_path}") + log_task_info(task_uid, f"[Streaming Mode] Chunk size: {chunk_size:,} rows") + + new_file = None + try: + new_file = f"{os.path.splitext(file_path)[0]}.xlsx" + + workbook = xlsxwriter.Workbook(new_file, { + 'constant_memory': True, + 'use_zip64': True, + 'strings_to_numbers': False, + 'strings_to_urls': False + }) + worksheet = workbook.add_worksheet() + + current_row = 0 + header_written = False + rows_processed = 0 + + for chunk in _read_csv_chunked(file_path, chunk_size): + if not header_written: + for col_idx, col_name in enumerate(chunk.columns): + worksheet.write(0, col_idx, col_name) + current_row = 1 + header_written = True + + chunk_values = chunk.values + for row_idx in range(len(chunk_values)): + for col_idx in range(len(chunk_values[row_idx])): + value = chunk_values[row_idx][col_idx] + if pd.isna(value): + worksheet.write_blank(current_row, col_idx, None) + else: + worksheet.write(current_row, col_idx, value) + current_row += 1 + + rows_processed += len(chunk) + + del chunk + del chunk_values + + if rows_processed % (chunk_size * 5) == 0: + gc.collect() + + log_task_info(task_uid, "Finalizing Excel file...") + workbook.close() + + return { + "from": file_path, + "to": new_file, + "to_files": [new_file], + "status": "success", + } + except Exception as e: + log_task_error(task_uid, f"convert file {file_path} error: {e}") + if new_file and os.path.exists(new_file): + try: + os.remove(new_file) + except Exception: + pass + return {"from": file_path, "to": None, "status": "failure", "error": str(e)} + def fix_email_links_in_html(html_content: str) -> str: pattern1 = r'([^<]+)'