From f0a2ca1323942b0cea13f85dcee8e20fd27195ae Mon Sep 17 00:00:00 2001 From: 2646799270 <151953254+2646799270@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:46:50 +0800 Subject: [PATCH 01/12] Update chinese_convert_mapper.py --- data_engine/ops/mapper/chinese_convert_mapper.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/data_engine/ops/mapper/chinese_convert_mapper.py b/data_engine/ops/mapper/chinese_convert_mapper.py index 76ed134..6e075bd 100644 --- a/data_engine/ops/mapper/chinese_convert_mapper.py +++ b/data_engine/ops/mapper/chinese_convert_mapper.py @@ -26,8 +26,10 @@ def prepare_converter(mode): class ChineseConvertMapper(Mapper): """Mapper to convert Chinese between Traditional Chinese, Simplified Chinese and Japanese Kanji.""" + _supports_streaming = True def __init__(self, mode: str = 's2t', *args, **kwargs): + """ Initialization method. From 39787920309aa6ac62472ff85f2efe58294cea16 Mon Sep 17 00:00:00 2001 From: 2646799270 <151953254+2646799270@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:22:21 +0800 Subject: [PATCH 02/12] Update run_dataflow_task.py --- run_dataflow_task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/run_dataflow_task.py b/run_dataflow_task.py index b1d0cbe..1d22480 100644 --- a/run_dataflow_task.py +++ b/run_dataflow_task.py @@ -8,7 +8,7 @@ from datetime import datetime, timezone from loguru import logger - +s # Suppress Pydantic V2 namespace warnings from third-party packages # (e.g. AWS SDK models using "model_arn" which conflicts with Pydantic V2 "model_" namespace) From ecba7f79315992c00f14e0a20b59c97d9ab0480e Mon Sep 17 00:00:00 2001 From: 2646799270 <151953254+2646799270@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:23:05 +0800 Subject: [PATCH 03/12] Update run_dataflow_task.py --- run_dataflow_task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/run_dataflow_task.py b/run_dataflow_task.py index 1d22480..b1d0cbe 100644 --- a/run_dataflow_task.py +++ b/run_dataflow_task.py @@ -8,7 +8,7 @@ from datetime import datetime, timezone from loguru import logger -s + # Suppress Pydantic V2 namespace warnings from third-party packages # (e.g. AWS SDK models using "model_arn" which conflicts with Pydantic V2 "model_" namespace) From 7a9ca0c0853bc07dd992eb321938805386cde5d8 Mon Sep 17 00:00:00 2001 From: 2646799270 <151953254+2646799270@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:57:57 +0800 Subject: [PATCH 04/12] Update config.py --- data_engine/config/config.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/data_engine/config/config.py b/data_engine/config/config.py index 9d8cc4c..53c6871 100644 --- a/data_engine/config/config.py +++ b/data_engine/config/config.py @@ -143,6 +143,17 @@ def init_configs(args=None,redirect=True): type=PositiveInt, default=4, help='Number of processes to process dataset.') + parser.add_argument( + '--use_streaming', + type=bool, + default=True, + help='Whether to use streaming mode for dataset loading and processing. ' + 'Streaming mode significantly reduces memory usage by processing ' + 'data iteratively without loading the entire dataset into memory. ' + 'However, it disables caching, checkpointing, and multi-process ' + 'parallelism. Only simple mappers and filters support streaming mode. ' + 'Recommended for datasets larger than available memory (>10GB). ' + 'When enabled, automatically sets batch_size=100 and enables sample counting.') parser.add_argument( '--text_keys', type=Union[str, List[str]], @@ -869,4 +880,4 @@ def get_init_configs(cfg): with open(temp_file, 'w') as f: json.dump(cfg, f) inited_dj_cfg = init_configs(['--config', temp_file]) - return inited_dj_cfg \ No newline at end of file + return inited_dj_cfg From 58754b985dfaa82e077fee4a3fb966601473622a Mon Sep 17 00:00:00 2001 From: 2646799270 <151953254+2646799270@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:02:26 +0800 Subject: [PATCH 05/12] Add files via upload --- data_engine/core/streaming_data.py | 449 +++++++++++++++++++++++++++++ 1 file changed, 449 insertions(+) create mode 100644 data_engine/core/streaming_data.py diff --git a/data_engine/core/streaming_data.py b/data_engine/core/streaming_data.py new file mode 100644 index 0000000..39663d6 --- /dev/null +++ b/data_engine/core/streaming_data.py @@ -0,0 +1,449 @@ +""" +Streaming Dataset Wrapper for Low Memory Processing + +This module provides a wrapper around HuggingFace's IterableDataset +to enable streaming mode processing with minimal memory footprint. +""" + +import traceback +from loguru import logger +from datasets import IterableDataset +from tqdm import tqdm + + +class StreamingDataset: + """ + Wrapper class for IterableDataset to provide a similar interface + to NestedDataset while operating in streaming mode. + + Key differences from NestedDataset: + - No caching or fingerprint management + - No multi-process parallelism (num_proc is ignored) + - No random access (len(), select() not supported) + - Significantly lower memory usage + - Optional batch processing for efficiency + - Optional tqdm progress bar support with total sample count + + + """ + + def __init__(self, iterable_dataset, batch_size=1, total_samples=None): + """ + Initialize StreamingDataset + + Args: + iterable_dataset: HuggingFace IterableDataset instance + batch_size: Number of samples to process in each batch (default: 1) + Larger batch_size can improve efficiency but uses more memory + total_samples: Total number of samples in dataset (optional) + If provided, enables progress bar with percentage and ETA + """ + if not isinstance(iterable_dataset, IterableDataset): + raise TypeError( + f"Expected IterableDataset, got {type(iterable_dataset)}. " + f"Make sure to load dataset with streaming=True" + ) + + self.dataset = iterable_dataset + self._supports_streaming = True + self.batch_size = batch_size + self.total_samples = total_samples + + if batch_size > 1: + logger.info(f'StreamingDataset initialized with batch_size={batch_size} ' + f'for improved processing efficiency') + + if total_samples is not None: + logger.info(f'StreamingDataset initialized with total_samples={total_samples:,} ' + f'(progress bar will be displayed)') + + def map(self, function=None, **kwargs): + """ + Apply a function to each sample (or batch) in the dataset. + + Args: + function: Function to apply to each sample or batch + **kwargs: Additional arguments + - batched: If True, function receives a batch dict with lists of values + - batch_size: Override instance batch_size for this operation + - num_proc: Ignored (streaming is single-process) + - with_rank: Ignored (no GPU parallelism in streaming) + - desc: Description for progress display + - disable_progress_bar: If True, disable progress bar even if total_samples is set + + Returns: + StreamingDataset: New streaming dataset with function applied + """ + # Extract parameters + desc = kwargs.get('desc', 'Processing') + batched = kwargs.get('batched', False) + operation_batch_size = kwargs.get('batch_size', self.batch_size) + disable_progress_bar = kwargs.get('disable_progress_bar', False) + + # Warn about ignored parameters + ignored_params = ['num_proc', 'with_rank', 'new_fingerprint'] + for param in ignored_params: + if param in kwargs and kwargs[param] is not None: + logger.debug(f"Parameter '{param}' is ignored in streaming mode") + + if function is None: + function = lambda x: x + + # Wrap function with progress bar if total_samples is available + if self.total_samples is not None and not disable_progress_bar: + wrapped_function = self._wrap_with_progress_bar( + function, + desc=desc, + total=self.total_samples, + batched=batched, + batch_size=operation_batch_size + ) + else: + wrapped_function = function + + # Apply function using IterableDataset's map + if batched: + # Explicitly requested batched processing + mapped_ds = self.dataset.map( + wrapped_function, + batched=True, + batch_size=operation_batch_size + ) + logger.debug(f'{desc}: Applied batched function with batch_size={operation_batch_size} in streaming mode') + else: + # Single sample processing + mapped_ds = self.dataset.map(wrapped_function) + logger.debug(f'{desc}: Applied function in streaming mode') + + # Preserve total_samples for chained operations + return StreamingDataset(mapped_ds, batch_size=self.batch_size, total_samples=self.total_samples) + + def filter(self, function=None, **kwargs): + """ + Filter samples based on a predicate function. + + Args: + function: Predicate function that returns True to keep sample + - If batched=False: function(sample: Dict) -> bool + - If batched=True: function(batch: Dict[str, List]) -> List[bool] + **kwargs: Additional arguments + - batched: If True, function receives batches and returns List[bool] (default: False) + - batch_size: Batch size for filtering (uses instance default if not specified) + - desc: Description for progress display + - disable_progress_bar: If True, disable progress bar even if total_samples is set + - num_proc: Ignored (streaming is single-process) + + Returns: + StreamingDataset: Filtered streaming dataset + + + """ + desc = kwargs.get('desc', 'Filtering') + batched = kwargs.get('batched', False) + operation_batch_size = kwargs.get('batch_size', self.batch_size) + disable_progress_bar = kwargs.get('disable_progress_bar', False) + + # Warn about ignored parameters + if 'num_proc' in kwargs and kwargs['num_proc'] is not None: + logger.debug("Parameter 'num_proc' is ignored in streaming mode") + + # Default function + if function is None: + if batched: + # For batched mode, return list of True with same length as batch + function = lambda batch: [True] * len(batch[next(iter(batch))]) + else: + # For single sample mode, return True + function = lambda x: True + + # Wrap function with progress bar if total_samples is available + if self.total_samples is not None and not disable_progress_bar: + wrapped_function = self._wrap_with_progress_bar( + function, + desc=desc, + total=self.total_samples, + batched=batched, + batch_size=operation_batch_size, + is_filter=True + ) + else: + wrapped_function = function + + # Apply filter using IterableDataset's filter + # HuggingFace IterableDataset.filter DOES support batched parameter + if batched: + filtered_ds = self.dataset.filter( + wrapped_function, + batched=True, + batch_size=operation_batch_size + ) + logger.debug(f'{desc}: Applied batched filter with batch_size={operation_batch_size} in streaming mode') + else: + filtered_ds = self.dataset.filter(wrapped_function) + logger.debug(f'{desc}: Applied filter in streaming mode') + + # Note: After filtering, total_samples is no longer accurate, set to None + return StreamingDataset(filtered_ds, batch_size=self.batch_size, total_samples=None) + + def _wrap_with_progress_bar(self, function, desc, total, batched=False, batch_size=1, is_filter=False): + """ + Wrap a function with tqdm progress bar. + + Args: + function: Function to wrap + desc: Description for progress bar + total: Total number of samples + batched: Whether function is batched + batch_size: Batch size for batched processing + is_filter: Whether this is a filter operation + + Returns: + Wrapped function with progress bar + """ + # Create a closure to maintain progress bar state + pbar = {'bar': None, 'count': 0} + + def wrapped_function(sample): + # Initialize progress bar on first call + if pbar['bar'] is None: + pbar['bar'] = tqdm( + total=total, + desc=desc, + unit='samples', + dynamic_ncols=True, + colour='green' + ) + + # Call original function + result = function(sample) + + # Update progress bar + if batched: + # For batched processing, get the batch size from the sample + if isinstance(sample, dict): + # Get the first key to determine batch size + first_key = next(iter(sample.keys())) + current_batch_size = len(sample[first_key]) + else: + current_batch_size = batch_size + pbar['bar'].update(current_batch_size) + pbar['count'] += current_batch_size + else: + # For single sample processing + pbar['bar'].update(1) + pbar['count'] += 1 + + # Close progress bar when done + if pbar['count'] >= total: + pbar['bar'].close() + + return result + + return wrapped_function + + + def process(self, operators, *, exporter=None, checkpointer=None, tracer=None): + """ + Process dataset through a list of operators. + + Args: + operators: List of operator instances to apply + exporter: Exporter for saving results (optional) + checkpointer: Ignored in streaming mode (no checkpoints) + tracer: Ignored in streaming mode (no tracing) + + Returns: + StreamingDataset: Processed dataset + + Raises: + ValueError: If any operator doesn't support streaming mode + """ + if operators is None: + return self + + if not isinstance(operators, list): + operators = [operators] + + # Warn about unsupported features + if checkpointer is not None: + logger.warning( + "Checkpointing is not supported in streaming mode and will be ignored" + ) + if tracer is not None: + logger.warning( + "Tracing is not supported in streaming mode and will be ignored" + ) + + dataset = self + processed_count = 0 + + for op in operators: + # Check if operator supports streaming + if not getattr(op, '_supports_streaming', False): + raise ValueError( + f"Operator [{op._name}] does not support streaming mode. " + f"Please either:\n" + f" 1. Remove this operator from the pipeline, or\n" + f" 2. Use normal mode (set use_streaming=False)" + ) + + logger.info(f'Processing with operator [{op._name}] in streaming mode...') + + try: + # Run operator in streaming mode + dataset = op.run(dataset, exporter=exporter, tracer=None) + processed_count += 1 + logger.info( + f'OP [{op._name}] completed in streaming mode ' + f'({processed_count}/{len(operators)})' + ) + except Exception as e: + logger.error( + f'An error occurred during Op [{op._name}] in streaming mode: {e}' + ) + raise + + return dataset + + + def __iter__(self): + """ + Iterate over samples in the dataset. + + Returns: + Iterator over dataset samples + """ + return iter(self.dataset) + + def __repr__(self): + return f"StreamingDataset({self.dataset})" + + def take(self, n): + """ + Take first n samples from the dataset. + + Args: + n: Number of samples to take + + Returns: + StreamingDataset: Dataset with first n samples + """ + # Update total_samples to n if it was set + new_total = min(n, self.total_samples) if self.total_samples is not None else n + return StreamingDataset(self.dataset.take(n), batch_size=self.batch_size, total_samples=new_total) + + def skip(self, n): + """ + Skip first n samples from the dataset. + + Args: + n: Number of samples to skip + + Returns: + StreamingDataset: Dataset with first n samples skipped + """ + # Update total_samples if it was set + new_total = max(0, self.total_samples - n) if self.total_samples is not None else None + return StreamingDataset(self.dataset.skip(n), batch_size=self.batch_size, total_samples=new_total) + + def shuffle(self, seed=None, buffer_size=1000): + """ + Shuffle the dataset using a buffer. + + Note: This is not a true shuffle but uses a buffer-based approach. + For true random shuffle, use normal mode instead. + + Args: + seed: Random seed + buffer_size: Size of shuffle buffer + + Returns: + StreamingDataset: Shuffled dataset + """ + return StreamingDataset( + self.dataset.shuffle(seed=seed, buffer_size=buffer_size), + batch_size=self.batch_size, + total_samples=self.total_samples + ) + + +def is_streaming_dataset(dataset): + """ + Check if a dataset is in streaming mode. + + Args: + dataset: Dataset to check + + Returns: + bool: True if dataset is StreamingDataset, False otherwise + """ + return isinstance(dataset, StreamingDataset) + + +def catch_streaming_exception(method): + """ + Exception handler for streaming mode operations. + Converts error samples to dict with empty lists, which will be filtered out. + + This matches catch_map_single_exception in base_op.py for consistency: + - Returns {key: []} for all keys to maintain schema + - Logs error details for debugging + - Empty list values will be filtered out by filter_empty_samples + """ + from functools import wraps + from data_engine.utils.constant import Fields + + @wraps(method) + def wrapper(sample, *args, **kwargs): + try: + return method(sample, *args, **kwargs) + except Exception as e: + logger.error( + f'An error occurred in streaming operation when processing ' + f'sample, {type(e)}: {e}' + ) + traceback.print_exc() + # Return dict with empty lists (matches normal mode behavior) + # This maintains the schema structure while marking sample as invalid + ret = {key: [] for key in sample.keys()} if sample else {} + ret[Fields.stats] = [] + ret[Fields.source_file] = [] + return ret + return wrapper + + +def filter_empty_samples(dataset): + """ + Filter out samples with empty list values that resulted from exceptions. + Should be applied after map operations that might return {key: []} on error. + + Args: + dataset: StreamingDataset that might contain error samples + + Returns: + StreamingDataset with error samples removed + """ + from data_engine.utils.constant import Fields + + def is_not_empty(sample): + # Filter out samples where all values are empty lists (error samples) + # Check if sample is empty or if it's an error sample from catch_streaming_exception + if not sample: + return False + + # Check if this is an error sample: all values are empty lists + # Skip checking Fields.stats and Fields.source_file as they're added by exception handler + data_keys = [k for k in sample.keys() if k not in [Fields.stats, Fields.source_file]] + if not data_keys: + # Only has stats/source_file, likely an error sample + return False + + # Check if all data values are empty lists + all_empty = all( + isinstance(sample[key], list) and len(sample[key]) == 0 + for key in data_keys + ) + return not all_empty + + # Apply filter using dataset's filter method (preserves StreamingDataset wrapper) + return dataset.filter(is_not_empty, desc='filter_empty_samples') From 75060dc38a6aa5ef7e8c5bc45c0dcaf7a4312630 Mon Sep 17 00:00:00 2001 From: shenren123 <2646799270@qq.com> Date: Tue, 18 Aug 2026 15:21:01 +0800 Subject: [PATCH 06/12] =?UTF-8?q?=E6=B5=81=E5=BC=8F=E5=8A=A0=E8=BD=BD?= =?UTF-8?q?=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data_engine/exporter/base_exporter.py | 171 +++++++++- data_engine/exporter/csghub_exporter.py | 14 +- data_engine/format/formatter.py | 225 +++++++++++++- data_engine/ops/base_op.py | 294 ++++++++++++++++++ data_engine/ops/filter/alphanumeric_filter.py | 2 + .../ops/filter/average_line_length_filter.py | 2 + .../ops/filter/character_repetition_filter.py | 3 + .../ops/filter/flagged_words_filter.py | 2 + .../ops/filter/language_id_score_filter.py | 2 + .../ops/filter/maximum_line_length_filter.py | 2 + .../ops/filter/multi_keyword_filter.py | 2 + .../ops/filter/special_characters_filter.py | 3 + .../ops/filter/specified_field_filter.py | 3 + .../filter/specified_numeric_field_filter.py | 3 + data_engine/ops/filter/stopwords_filter.py | 2 + data_engine/ops/filter/suffix_filter.py | 3 + .../ops/filter/text_high_score_filter.py | 2 + data_engine/ops/filter/text_length_filter.py | 2 + .../ops/filter/word_repetition_filter.py | 2 + data_engine/ops/filter/words_num_filter.py | 2 + .../ops/mapper/chinese_convert_mapper.py | 5 +- .../ops/mapper/clean_copyright_mapper.py | 2 + data_engine/ops/mapper/clean_email_mapper.py | 2 + data_engine/ops/mapper/clean_html_mapper.py | 2 + data_engine/ops/mapper/clean_ip_mapper.py | 3 + data_engine/ops/mapper/clean_links_mapper.py | 3 + data_engine/ops/mapper/expand_macro_mapper.py | 3 + data_engine/ops/mapper/fix_unicode_mapper.py | 2 + data_engine/ops/mapper/nlpaug_en_mapper.py | 1 + data_engine/ops/mapper/nlpcda_zh_mapper.py | 1 + .../punctuation_normalization_mapper.py | 3 + .../ops/mapper/remove_bibliography_mapper.py | 3 + .../ops/mapper/remove_comments_mapper.py | 3 + .../ops/mapper/remove_header_mapper.py | 3 + .../ops/mapper/remove_long_words_mapper.py | 3 + .../remove_non_chinese_character_mapper.py | 3 + .../mapper/remove_repeat_sentences_mapper.py | 3 + .../mapper/remove_specific_chars_mapper.py | 3 + .../ops/mapper/remove_table_text_mapper.py | 3 + ..._words_with_incorrect_substrings_mapper.py | 2 + .../ops/mapper/replace_content_mapper.py | 2 + .../ops/mapper/sentence_split_mapper.py | 2 + .../mapper/whitespace_normalization_mapper.py | 2 + data_server/pod/common_tasks.py | 103 +++++- 44 files changed, 882 insertions(+), 21 deletions(-) diff --git a/data_engine/exporter/base_exporter.py b/data_engine/exporter/base_exporter.py index 84a3fe1..dc7d3e9 100644 --- a/data_engine/exporter/base_exporter.py +++ b/data_engine/exporter/base_exporter.py @@ -52,7 +52,7 @@ def __init__( """ if Path(export_path).is_dir() and not path_is_dir: export_path = os.path.join(export_path, "x.jsonl") - + self.export_path = export_path self.export_shard_size = export_shard_size self.export_in_parallel = export_in_parallel @@ -125,7 +125,13 @@ def _export_impl(self, dataset, export_path, suffix, export_stats=True): :param export_stats: whether to export stats of dataset. :return: """ - if Fields.stats in dataset.features and export_stats: + # 流式模式:不支持 features 属性和 select_columns 方法 + from data_engine.core.streaming_data import StreamingDataset + from datasets import IterableDataset + + is_streaming = isinstance(dataset, (StreamingDataset, IterableDataset)) + + if not is_streaming and Fields.stats in dataset.features and export_stats: # export stats of datasets into a single file. logger.info('Exporting computed stats into a single file...') ds_stats = dataset.select_columns(Fields.stats) @@ -164,10 +170,21 @@ def _export_impl(self, dataset, export_path, suffix, export_stats=True): HashKeys.is_duplicate, HashKeys.similarity_hash, }) - feature_fields = set(dataset.features.keys()) - removed_fields = fields_to_remove.intersection(feature_fields) - if removed_fields: - dataset = dataset.remove_columns(removed_fields) + # 流式模式:使用 map 移除内部字段 + if not is_streaming: + feature_fields = set(dataset.features.keys()) + removed_fields = fields_to_remove.intersection(feature_fields) + if removed_fields: + dataset = dataset.remove_columns(removed_fields) + else: + # 流式模式:通过 map 过滤字段 + logger.info('Streaming mode: filtering internal fields before export...') + + def remove_fields(sample): + return {k: v for k, v in sample.items() if k not in fields_to_remove} + + dataset = dataset.map(remove_fields) + export_method = Exporter._router()[suffix] dirname = os.path.join(os.path.dirname(os.path.abspath(self.export_path)), "_data") @@ -234,13 +251,149 @@ def export_large_folder(self): def export(self, dataset): """ Export method for a dataset. + Supports both NestedDataset and StreamingDataset. - :param dataset: the dataset to export. - :return: + :param dataset: the dataset to export (NestedDataset or StreamingDataset) + :return: empty string or branch name if pushing to repo """ - self._export_impl(dataset, self.export_path, self.suffix, self.export_stats) + from data_engine.core.streaming_data import is_streaming_dataset + + if is_streaming_dataset(dataset): + return self._export_streaming(dataset) + else: + self._export_impl(dataset, self.export_path, self.suffix, self.export_stats) + return "" + + def _export_streaming(self, dataset): + """ + Export StreamingDataset to file(s) with low memory usage. + + :param dataset: StreamingDataset to export + :return: empty string + """ + import json + import tempfile + import shutil + from tqdm import tqdm + + logger.info('Exporting dataset in STREAMING mode (low memory)...') + + # Determine export directory and filename + export_dir = os.path.dirname(os.path.abspath(self.export_path)) + basename = os.path.basename(self.export_path) + + # Create directory for data files + data_dir = os.path.join(export_dir, "_data") + os.makedirs(data_dir, exist_ok=True) + + output_file = os.path.join(data_dir, basename) + + # Use temporary file to avoid overwriting input file during streaming + # (fixes issue where input and output paths are the same) + temp_fd, temp_file = tempfile.mkstemp( + suffix='.jsonl', + prefix='_df_tmp_', + dir=data_dir, + text=True + ) + + try: + # Write samples to temporary file iteratively + sample_count = 0 + with os.fdopen(temp_fd, 'w', encoding='utf-8') as f: + # Use tqdm for progress display (without total count) + pbar = tqdm(desc='Exporting samples', unit=' samples') + + for sample in dataset: + # Skip empty samples (error samples from exception handling) + # Error samples have all values as empty lists: {key: []} + if not sample: + continue + + # Check if this is an error sample (all values are empty lists) + is_error_sample = False + if isinstance(sample, dict): + # Get non-internal keys + from data_engine.utils.constant import Fields + data_keys = [k for k in sample.keys() if k not in [Fields.stats, Fields.source_file]] + if data_keys: + # Check if all data values are empty lists + is_error_sample = all( + isinstance(sample[key], list) and len(sample[key]) == 0 + for key in data_keys + ) + + if is_error_sample: + continue + + # Remove internal fields before export + sample = self._clean_sample_for_export(sample) + + # Write as JSON line + f.write(json.dumps(sample, ensure_ascii=False) + '\n') + sample_count += 1 + pbar.update(1) + + pbar.close() + + # Move temporary file to final location (atomic operation on same filesystem) + shutil.move(temp_file, output_file) + logger.info(f'Exported {sample_count} samples to {output_file} in streaming mode') + + except Exception as e: + # Clean up temporary file on error + if os.path.exists(temp_file): + try: + os.remove(temp_file) + except OSError: + pass + raise e + return "" + def _clean_sample_for_export(self, sample): + """ + Remove internal fields from sample before export. + + :param sample: sample dict + :return: cleaned sample dict + """ + # Fields to remove + fields_to_remove = set() + + if not self.keep_stats_in_res_ds: + fields_to_remove.add(Fields.stats) + + if not self.keep_hashes_in_res_ds: + fields_to_remove.update({ + HashKeys.hash, + HashKeys.minhash, + HashKeys.simhash, + HashKeys.imagehash, + HashKeys.videohash, + }) + + # Other internal fields + fields_to_remove.update({ + Fields.suffix, + Fields.context, + Fields.meta, + Fields.source_file, + Fields.video_frame_tags, + Fields.video_audio_tags, + Fields.multimodal_data_output_dir, + HashKeys.is_duplicate, + HashKeys.similarity_hash, + }) + + # Remove fields that exist in sample + cleaned_sample = { + k: v for k, v in sample.items() + if k not in fields_to_remove + } + + return cleaned_sample + def export_compute_stats(self, dataset, export_path): """ Export method for saving compute status in filters diff --git a/data_engine/exporter/csghub_exporter.py b/data_engine/exporter/csghub_exporter.py index e60125e..83ed0e6 100644 --- a/data_engine/exporter/csghub_exporter.py +++ b/data_engine/exporter/csghub_exporter.py @@ -136,10 +136,22 @@ def export_large_folder(self): def export(self, dataset): """ Export method for a dataset. + Supports both NestedDataset and StreamingDataset. :param dataset: the dataset to export. - :return: + :return: branch name if pushing to repo, empty string otherwise """ + from data_engine.core.streaming_data import is_streaming_dataset + + # Export dataset (streaming or normal mode) + if is_streaming_dataset(dataset): + # Call parent's streaming export method + super()._export_streaming(dataset) + else: + # Call parent's normal export method + self._export_impl(dataset, self.export_path, self.suffix, self.export_stats) + + # After export, push to repo if repo_id is configured self._export_impl(dataset, self.export_path, self.suffix, self.export_stats) self.upload_path = os.path.join(self.work_dir, "_data") self.repo_work_dir = os.path.join(self.work_dir, "_git") diff --git a/data_engine/format/formatter.py b/data_engine/format/formatter.py index bbf7be8..23ced30 100644 --- a/data_engine/format/formatter.py +++ b/data_engine/format/formatter.py @@ -68,7 +68,23 @@ def load_dataset(self, num_proc: int = 1, global_cfg=None) -> Dataset: :param num_proc: number of processes when loading the dataset :param global_cfg: global cfg used in consequent processes, - :return: formatted dataset + :return: formatted dataset (NestedDataset or StreamingDataset) + """ + # Check if streaming mode is enabled + use_streaming = getattr(global_cfg, 'use_streaming', False) + + if use_streaming: + return self._load_dataset_streaming(global_cfg) + else: + return self._load_dataset_normal(num_proc, global_cfg) + + def _load_dataset_normal(self, num_proc: int = 1, global_cfg=None) -> Dataset: + """ + Load dataset in normal mode (original implementation). + + :param num_proc: number of processes when loading the dataset + :param global_cfg: global cfg used in consequent processes + :return: NestedDataset """ from datasets.exceptions import DatasetGenerationError @@ -999,6 +1015,213 @@ def fill_missing_fields(sample): from datasets import Dataset return Dataset.from_list(data_list, features=features) + def _load_dataset_streaming(self, global_cfg=None): + """ + Load dataset in streaming mode for low memory usage. + + :param global_cfg: global cfg used in consequent processes + :return: StreamingDataset + """ + from data_engine.core.streaming_data import StreamingDataset + from datasets import interleave_datasets + + logger.info('Loading dataset in STREAMING mode (low memory)...') + + # First, find data files + self.data_files = find_files_with_suffix(self.dataset_path, self.suffixes) + + # Streaming mode: Fixed configuration + # - Always pre-scan sample count (for progress bar) + # - Fixed batch_size=100 (memory-efficient batch processing) + estimated_total_samples = None + batch_size = 1000 # Fixed batch size for streaming mode (balances memory and performance) + + if global_cfg and global_cfg.use_streaming: + # Pre-scan sample count (mandatory in streaming mode) + estimated_total_samples = self._prescan_sample_count() + + # Escape glob special characters in file paths + escaped_data_files = { + key: [escape_glob_chars(f) for f in files] + for key, files in self.data_files.items() + } + + # Load with streaming=True + datasets = load_dataset( + self.type, + data_files={ + key.strip('.'): escaped_data_files[key] + for key in escaped_data_files + }, + streaming=True, # Key parameter for streaming mode + # Note: num_proc is not supported in streaming mode + **self.kwargs + ) + + # Merge multiple splits using interleave_datasets + dataset_list = [ds for _, ds in datasets.items()] + if len(dataset_list) > 1: + logger.info(f'Merging {len(dataset_list)} dataset splits in streaming mode...') + ds = interleave_datasets(dataset_list) + else: + ds = dataset_list[0] + + # Apply format unification in streaming mode + ds = self._unify_format_streaming(ds, global_cfg) + + # Handle add_suffix in streaming mode + if self.add_suffix: + logger.info('Adding suffix info to streaming dataset...') + + def add_suffix_field(sample): + if Fields.suffix not in sample: + # Add suffix based on file type + sample[Fields.suffix] = f'.{self.type}' + return sample + + ds = ds.map(add_suffix_field) + + logger.info(f'Dataset loaded successfully in streaming mode (batch_size={batch_size})') + + # Create StreamingDataset with total_samples for progress bar support + streaming_dataset = StreamingDataset( + ds, + batch_size=batch_size, + total_samples=estimated_total_samples # Pass to constructor for progress bar + ) + + return streaming_dataset + + def _unify_format_streaming(self, dataset, global_cfg): + """ + Unify dataset format in streaming mode. + Simplified version without multi-process support. + + :param dataset: IterableDataset + :param global_cfg: global config + :return: formatted IterableDataset + """ + # Get text_keys from config + text_keys = self.text_keys if self.text_keys else ['text'] + if isinstance(text_keys, str): + text_keys = [text_keys] + + # Define format unification function + def unify_sample(sample): + # Ensure text field exists + if text_keys[0] not in sample: + # Try to find alternative text field + for key in text_keys: + if key in sample: + sample[text_keys[0]] = sample[key] + break + else: + # No text field found, set to empty string + sample[text_keys[0]] = "" + + return sample + + # Apply unification + dataset = dataset.map(unify_sample) + + # Filter out samples with empty or None text (same as normal mode) + logger.info('Filtering empty text samples in streaming mode...') + + def non_empty_text(sample): + """Filter function to remove samples with None text fields""" + for target_key in text_keys: + if target_key in sample and sample[target_key] is None: + # Filter out samples with None in any text column + return False + return True + + # Apply filter + dataset = dataset.filter(non_empty_text) + logger.info('Empty text filtering applied (streaming mode)') + + return dataset + + def _prescan_sample_count(self): + """ + Pre-scan data files to count total samples (low memory, O(1) complexity). + + Uses fast line counting (binary read, no parsing) with minimal memory (~9KB). + Supports line-based formats: JSONL, CSV, TXT. + + :return: Total sample count or None if failed + """ + logger.info('Pre-scanning sample count (fast line counting)...') + + total_samples = 0 + file_count = 0 + + try: + # Collect all data files + all_files = [] + for suffix, files in self.data_files.items(): + all_files.extend(files) + + if not all_files: + logger.warning('No data files found for pre-scanning') + return None + + # Count samples in each file + for file_path in all_files: + try: + # Check if file format supports line counting + if not self._is_line_based_format(file_path): + logger.warning( + f'Skipping {file_path}: Not a line-based format (only JSONL, CSV, TXT supported)') + continue + + # Fast line counting (binary read, O(1) memory) + file_samples = self._count_lines_fast(file_path) + total_samples += file_samples + file_count += 1 + + # Log per-file count + logger.info(f' {os.path.basename(file_path)}: {file_samples:,} samples') + + except Exception as e: + logger.warning(f'Failed to count samples in {file_path}: {e}') + continue + + if file_count > 0: + logger.info(f'Pre-scan complete: {total_samples:,} total samples across {file_count} file(s)') + return total_samples + else: + logger.warning('No files were successfully scanned') + return None + + except Exception as e: + logger.warning(f'Pre-scan failed: {e}') + return None + + @staticmethod + def _count_lines_fast(file_path): + """ + Fast line counting using binary read (minimal memory: ~9KB, O(1) complexity). + + :param file_path: Path to file + :return: Number of lines + """ + count = 0 + with open(file_path, 'rb') as f: + for _ in f: # Binary iteration, no decoding + count += 1 + return count + + @staticmethod + def _is_line_based_format(file_path): + """ + Check if file format is line-based (supports fast line counting). + + :param file_path: Path to file + :return: True if line-based format + """ + ext = os.path.splitext(file_path)[1].lower() + line_based_formats = {'.jsonl', '.csv', '.txt', '.tsv'} + return ext in line_based_formats class RemoteFormatter(BaseFormatter): """The class is used to load a dataset from repository of huggingface diff --git a/data_engine/ops/base_op.py b/data_engine/ops/base_op.py index bc254d3..3150f51 100644 --- a/data_engine/ops/base_op.py +++ b/data_engine/ops/base_op.py @@ -227,6 +227,7 @@ def init_params(cls): pass class Mapper(OP): + _supports_streaming = False # Default: does not support streaming def __init__(self, *args, **kwargs): """ @@ -261,6 +262,30 @@ def process(self, sample): raise NotImplementedError def run(self, dataset, *, exporter=None, tracer=None): + """ + Run mapper on dataset, supporting both normal and streaming modes. + + :param dataset: NestedDataset or StreamingDataset + :param exporter: exporter instance + :param tracer: tracer instance + :return: processed dataset + """ + from data_engine.core.streaming_data import is_streaming_dataset + + if is_streaming_dataset(dataset): + # Streaming mode + if not self._supports_streaming: + raise ValueError( + f"Mapper [{self._name}] does not support streaming mode. " + f"Please either remove this operator or use normal mode." + ) + return self.run_streaming(dataset, exporter=exporter) + else: + # Normal mode (original implementation) + return self._run_normal(dataset, exporter=exporter, tracer=tracer) + + def _run_normal(self, dataset, *, exporter=None, tracer=None): + """Original run implementation for normal mode.""" insert_pipline_job_run_task_log_info(self.job_uid, f"starting mapper job", operator_name=self._name, operator_index=self.pipline_index) set_pipline_job_operator_status(self.job_uid,OperatorStatusEnum.Processing,self._name,self.pipline_index) @@ -286,8 +311,127 @@ def run(self, dataset, *, exporter=None, tracer=None): finally: insert_pipline_job_run_task_log_info(self.job_uid,"ending mapper job",operator_name=self._name,operator_index=self.pipline_index) + def run_streaming(self, dataset, *, exporter=None): + """ + Run mapper in streaming mode (low memory). + + :param dataset: StreamingDataset + :param exporter: exporter instance + :return: processed StreamingDataset + """ + insert_pipline_job_run_task_log_info( + self.job_uid, + f"Starting mapper job in STREAMING mode", + operator_name=self._name, + operator_index=self.pipline_index + ) + set_pipline_job_operator_status( + self.job_uid, + OperatorStatusEnum.Processing, + self._name, + self.pipline_index + ) + + try: + # Get batch_size from dataset + batch_size = getattr(dataset, 'batch_size', 1) + use_batched = batch_size > 1 + + if use_batched: + logger.info(f'Processing {self._name} in streaming mode (batched with batch_size={batch_size})...') + else: + logger.info(f'Processing {self._name} in streaming mode (single-process)...') + + # Import streaming exception handler + from data_engine.core.streaming_data import catch_streaming_exception, filter_empty_samples + + if use_batched: + # Define batched processing function + def safe_process_batched(batch): + """Process batch with exception handling""" + # Get batch size + first_key = next(iter(batch.keys())) + current_batch_size = len(batch[first_key]) + + # Initialize result batch + result_batch = {key: [] for key in batch.keys()} + + # Process each sample in batch + for i in range(current_batch_size): + try: + # Extract single sample from batch + sample = {key: batch[key][i] for key in batch.keys()} + + # Process sample + processed_sample = self.process(sample) + + # Add to result batch + for key in processed_sample.keys(): + if key not in result_batch: + result_batch[key] = [] + result_batch[key].append(processed_sample[key]) + except Exception as e: + logger.error( + f'An error occurred in mapper {self._name} when processing sample {i}: {e}' + ) + import traceback + traceback.print_exc() + # Return empty dict for error samples + for key in batch.keys(): + if key not in result_batch: + result_batch[key] = [] + result_batch[key].append({}) + + return result_batch + + # Apply batched processing + new_dataset = dataset.map( + safe_process_batched, + batched=True, + batch_size=batch_size, + desc=self._name + '_process' + ) + else: + # Wrap process method with exception handler (similar to normal mode) + safe_process = catch_streaming_exception(self.process) + + # Apply safe process in single-sample mode + new_dataset = dataset.map(safe_process, desc=self._name + '_process') + + # Filter out empty dict samples that resulted from exceptions + new_dataset = filter_empty_samples(new_dataset) + + set_pipline_job_operator_status( + self.job_uid, + OperatorStatusEnum.SUCCESS, + self._name, + self.pipline_index + ) + return new_dataset + except Exception as e: + set_pipline_job_operator_status( + self.job_uid, + OperatorStatusEnum.ERROR, + self._name, + self.pipline_index + ) + insert_pipline_job_run_task_log_error( + self.job_uid, + f"An error occurred during streaming data mapping: {e}", + operator_name=self._name, + operator_index=self.pipline_index + ) + raise + finally: + insert_pipline_job_run_task_log_info( + self.job_uid, + "Ending mapper job (streaming mode)", + operator_name=self._name, + operator_index=self.pipline_index + ) class Filter(OP): + _supports_streaming = False # Default: does not support streaming def __init__(self, *args, **kwargs): """ @@ -334,6 +478,30 @@ def process(self, sample): raise NotImplementedError def run(self, dataset, *, exporter=None, tracer=None): + """ + Run filter on dataset, supporting both normal and streaming modes. + + :param dataset: NestedDataset or StreamingDataset + :param exporter: exporter instance + :param tracer: tracer instance + :return: filtered dataset + """ + from data_engine.core.streaming_data import is_streaming_dataset + + if is_streaming_dataset(dataset): + # Streaming mode + if not self._supports_streaming: + raise ValueError( + f"Filter [{self._name}] does not support streaming mode. " + f"Please either remove this operator or use normal mode." + ) + return self.run_streaming(dataset, exporter=exporter) + else: + # Normal mode (original implementation) + return self._run_normal(dataset, exporter=exporter, tracer=tracer) + + def _run_normal(self, dataset, *, exporter=None, tracer=None): + """Original run implementation for normal mode (3-stage processing).""" insert_pipline_job_run_task_log_info(self.job_uid, f"starting filter job", operator_name=self._name, operator_index=self.pipline_index) set_pipline_job_operator_status(self.job_uid, OperatorStatusEnum.Processing, self._name, self.pipline_index) @@ -376,6 +544,132 @@ def run(self, dataset, *, exporter=None, tracer=None): insert_pipline_job_run_task_log_info(self.job_uid, "ending filter job", operator_name=self._name, operator_index=self.pipline_index) + def run_streaming(self, dataset, *, exporter=None): + """ + Run filter in streaming mode (low memory). + Combines compute_stats and process into single-pass filtering. + + :param dataset: StreamingDataset + :param exporter: exporter instance + :return: filtered StreamingDataset + """ + insert_pipline_job_run_task_log_info( + self.job_uid, + f"Starting filter job in STREAMING mode", + operator_name=self._name, + operator_index=self.pipline_index + ) + set_pipline_job_operator_status( + self.job_uid, + OperatorStatusEnum.Processing, + self._name, + self.pipline_index + ) + + try: + # Get batch_size from dataset + batch_size = getattr(dataset, 'batch_size', 1) + use_batched = batch_size > 1 + + if use_batched: + logger.info( + f'Processing {self._name} in streaming mode (batched filtering with batch_size={batch_size})...') + else: + logger.info(f'Processing {self._name} in streaming mode (single-pass filtering)...') + + # Define combined filter function with exception handling + if use_batched: + def safe_combined_filter_batched(batch): + """Compute stats and filter in batch mode with exception handling""" + results = [] + + # Get batch size + first_key = next(iter(batch.keys())) + current_batch_size = len(batch[first_key]) + + # Process each sample in batch + for i in range(current_batch_size): + try: + # Extract single sample from batch + sample = {key: batch[key][i] for key in batch.keys()} + + # Compute stats (compute_stats is wrapped with exception handler) + sample = self.compute_stats(sample) + + # Apply filter + keep = self.process(sample) + results.append(keep) + except Exception as e: + logger.error( + f'An error occurred in filter {self._name} when processing sample {i}: {e}' + ) + import traceback + traceback.print_exc() + # Return False to filter out error samples + results.append(False) + + return results + + # Apply combined filter in batched mode + new_dataset = dataset.filter( + safe_combined_filter_batched, + batched=True, + batch_size=batch_size, + desc=self._name + '_process' + ) + else: + def safe_combined_filter(sample): + """Compute stats and filter in single pass with exception handling""" + try: + # Compute stats (compute_stats is wrapped with exception handler) + sample = self.compute_stats(sample) + + # Apply filter + return self.process(sample) + except Exception as e: + logger.error( + f'An error occurred in filter {self._name} when processing sample: {e}' + ) + import traceback + traceback.print_exc() + # Return False to filter out error samples + return False + + # Apply combined filter in single-sample mode + new_dataset = dataset.filter( + safe_combined_filter, + desc=self._name + '_process' + ) + + set_pipline_job_operator_status( + self.job_uid, + OperatorStatusEnum.SUCCESS, + self._name, + self.pipline_index + ) + return new_dataset + except Exception as e: + set_pipline_job_operator_status( + self.job_uid, + OperatorStatusEnum.ERROR, + self._name, + self.pipline_index + ) + insert_pipline_job_run_task_log_error( + self.job_uid, + f"An error occurred during streaming data filter: {e}", + operator_name=self._name, + operator_index=self.pipline_index + ) + raise + finally: + insert_pipline_job_run_task_log_info( + self.job_uid, + "Ending filter job (streaming mode)", + operator_name=self._name, + operator_index=self.pipline_index + ) + def _log_filter_details(self, original_dataset, filtered_dataset): """ Generate detailed logging for filter operations. diff --git a/data_engine/ops/filter/alphanumeric_filter.py b/data_engine/ops/filter/alphanumeric_filter.py index b2be218..e10ca94 100644 --- a/data_engine/ops/filter/alphanumeric_filter.py +++ b/data_engine/ops/filter/alphanumeric_filter.py @@ -19,6 +19,8 @@ class AlphanumericFilter(Filter): """Filter to keep samples with alphabet/numeric ratio within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, tokenization: bool = False, diff --git a/data_engine/ops/filter/average_line_length_filter.py b/data_engine/ops/filter/average_line_length_filter.py index f9306dd..c29fc72 100644 --- a/data_engine/ops/filter/average_line_length_filter.py +++ b/data_engine/ops/filter/average_line_length_filter.py @@ -13,6 +13,8 @@ class AverageLineLengthFilter(Filter): """Filter to keep samples with average line length within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, min_len: PositiveInt = 10, diff --git a/data_engine/ops/filter/character_repetition_filter.py b/data_engine/ops/filter/character_repetition_filter.py index b8c2431..5c918df 100644 --- a/data_engine/ops/filter/character_repetition_filter.py +++ b/data_engine/ops/filter/character_repetition_filter.py @@ -14,6 +14,9 @@ class CharacterRepetitionFilter(Filter): """Filter to keep samples with char-level n-gram repetition ratio within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, rep_len: PositiveInt = 10, diff --git a/data_engine/ops/filter/flagged_words_filter.py b/data_engine/ops/filter/flagged_words_filter.py index 3ccf26a..09ba0de 100644 --- a/data_engine/ops/filter/flagged_words_filter.py +++ b/data_engine/ops/filter/flagged_words_filter.py @@ -32,6 +32,8 @@ class FlaggedWordFilter(Filter): """Filter to keep samples with flagged-word ratio less than a specific max value.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, lang: str = 'en', diff --git a/data_engine/ops/filter/language_id_score_filter.py b/data_engine/ops/filter/language_id_score_filter.py index 3043d25..c83b5bc 100644 --- a/data_engine/ops/filter/language_id_score_filter.py +++ b/data_engine/ops/filter/language_id_score_filter.py @@ -19,6 +19,8 @@ class LanguageIDScoreFilter(Filter): """Filter to keep samples in a specific language with confidence score larger than a specific min value.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, lang: Union[str, List[str], Tuple[str]] = '', diff --git a/data_engine/ops/filter/maximum_line_length_filter.py b/data_engine/ops/filter/maximum_line_length_filter.py index 04e22f1..c6cac07 100644 --- a/data_engine/ops/filter/maximum_line_length_filter.py +++ b/data_engine/ops/filter/maximum_line_length_filter.py @@ -13,6 +13,8 @@ class MaximumLineLengthFilter(Filter): """Filter to keep samples with maximum line length within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, min_len: PositiveInt = 10, diff --git a/data_engine/ops/filter/multi_keyword_filter.py b/data_engine/ops/filter/multi_keyword_filter.py index c6d8b9e..dea59a1 100644 --- a/data_engine/ops/filter/multi_keyword_filter.py +++ b/data_engine/ops/filter/multi_keyword_filter.py @@ -10,6 +10,8 @@ @OPERATORS.register_module(OP_NAME) class MultiKeywordFilter(Filter): """Filter to remove samples that contain any of the specified keywords.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, keywords: Union[str, List[str], Tuple[str]] = [], diff --git a/data_engine/ops/filter/special_characters_filter.py b/data_engine/ops/filter/special_characters_filter.py index 960675c..39ab1f1 100644 --- a/data_engine/ops/filter/special_characters_filter.py +++ b/data_engine/ops/filter/special_characters_filter.py @@ -14,6 +14,9 @@ class SpecialCharactersFilter(Filter): """Filter to keep samples with special-char ratio within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, min_ratio: ClosedUnitInterval = 0.0, diff --git a/data_engine/ops/filter/specified_field_filter.py b/data_engine/ops/filter/specified_field_filter.py index 53c69a9..5faa40c 100644 --- a/data_engine/ops/filter/specified_field_filter.py +++ b/data_engine/ops/filter/specified_field_filter.py @@ -11,6 +11,9 @@ class SpecifiedFieldFilter(Filter): If the specified field information in the sample is not within the specified target value, the sample will be filtered. """ + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, field_key: str = '', diff --git a/data_engine/ops/filter/specified_numeric_field_filter.py b/data_engine/ops/filter/specified_numeric_field_filter.py index b067a74..acde9b7 100644 --- a/data_engine/ops/filter/specified_numeric_field_filter.py +++ b/data_engine/ops/filter/specified_numeric_field_filter.py @@ -21,6 +21,9 @@ class SpecifiedNumericFieldFilter(Filter): If the specified numeric information in the sample is not within the specified range, the sample will be filtered. """ + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, field_key: str = '', diff --git a/data_engine/ops/filter/stopwords_filter.py b/data_engine/ops/filter/stopwords_filter.py index 581bbe3..5ed32f2 100644 --- a/data_engine/ops/filter/stopwords_filter.py +++ b/data_engine/ops/filter/stopwords_filter.py @@ -25,6 +25,8 @@ class StopWordsFilter(Filter): """Filter to keep samples with stopword ratio larger than a specific min value.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, lang: str = 'en', diff --git a/data_engine/ops/filter/suffix_filter.py b/data_engine/ops/filter/suffix_filter.py index 1192348..ca460e1 100644 --- a/data_engine/ops/filter/suffix_filter.py +++ b/data_engine/ops/filter/suffix_filter.py @@ -15,6 +15,9 @@ @OPERATORS.register_module('suffix_filter') class SuffixFilter(Filter): """Filter to keep samples with specified suffix.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, suffixes: Union[str, List[str], Tuple[str]] = [], diff --git a/data_engine/ops/filter/text_high_score_filter.py b/data_engine/ops/filter/text_high_score_filter.py index e631d6e..7290302 100644 --- a/data_engine/ops/filter/text_high_score_filter.py +++ b/data_engine/ops/filter/text_high_score_filter.py @@ -7,6 +7,8 @@ @OPERATORS.register_module('text_high_score_filter') class TextHighScoreFilter(Filter): + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, score_field: str = 'text_score', min_score: float = 0.0, diff --git a/data_engine/ops/filter/text_length_filter.py b/data_engine/ops/filter/text_length_filter.py index d1418f1..5666182 100644 --- a/data_engine/ops/filter/text_length_filter.py +++ b/data_engine/ops/filter/text_length_filter.py @@ -11,6 +11,8 @@ class TextLengthFilter(Filter): """Filter to keep samples with total text length within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, min_len: PositiveInt = 10, diff --git a/data_engine/ops/filter/word_repetition_filter.py b/data_engine/ops/filter/word_repetition_filter.py index 9608a10..6e687f3 100644 --- a/data_engine/ops/filter/word_repetition_filter.py +++ b/data_engine/ops/filter/word_repetition_filter.py @@ -24,6 +24,8 @@ class WordRepetitionFilter(Filter): """Filter to keep samples with word-level n-gram repetition ratio within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, lang: str = 'en', diff --git a/data_engine/ops/filter/words_num_filter.py b/data_engine/ops/filter/words_num_filter.py index a7a354d..2805bd8 100644 --- a/data_engine/ops/filter/words_num_filter.py +++ b/data_engine/ops/filter/words_num_filter.py @@ -22,6 +22,8 @@ class WordsNumFilter(Filter): """Filter to keep samples with total words number within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, lang: str = 'en', diff --git a/data_engine/ops/mapper/chinese_convert_mapper.py b/data_engine/ops/mapper/chinese_convert_mapper.py index 6e075bd..a7e025e 100644 --- a/data_engine/ops/mapper/chinese_convert_mapper.py +++ b/data_engine/ops/mapper/chinese_convert_mapper.py @@ -26,10 +26,11 @@ def prepare_converter(mode): class ChineseConvertMapper(Mapper): """Mapper to convert Chinese between Traditional Chinese, Simplified Chinese and Japanese Kanji.""" - _supports_streaming = True + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, mode: str = 's2t', *args, **kwargs): - + """ Initialization method. diff --git a/data_engine/ops/mapper/clean_copyright_mapper.py b/data_engine/ops/mapper/clean_copyright_mapper.py index 0de8cc3..4e8c0e4 100644 --- a/data_engine/ops/mapper/clean_copyright_mapper.py +++ b/data_engine/ops/mapper/clean_copyright_mapper.py @@ -14,6 +14,8 @@ @OPERATORS.register_module('clean_copyright_mapper') class CleanCopyrightMapper(Mapper): """Remove copyright notices from documents (PDF->MD->JSONL).""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, matching_rules: Union[str, List[str], Tuple[str]] = [], diff --git a/data_engine/ops/mapper/clean_email_mapper.py b/data_engine/ops/mapper/clean_email_mapper.py index 9c54056..6f25aab 100644 --- a/data_engine/ops/mapper/clean_email_mapper.py +++ b/data_engine/ops/mapper/clean_email_mapper.py @@ -6,6 +6,8 @@ @OPERATORS.register_module('clean_email_mapper') class CleanEmailMapper(Mapper): """Mapper to clean email in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, pattern: str = None, repl: str = '', *args, **kwargs): """ diff --git a/data_engine/ops/mapper/clean_html_mapper.py b/data_engine/ops/mapper/clean_html_mapper.py index 2e73088..d11f364 100644 --- a/data_engine/ops/mapper/clean_html_mapper.py +++ b/data_engine/ops/mapper/clean_html_mapper.py @@ -15,6 +15,8 @@ @OPERATORS.register_module(OP_NAME) class CleanHtmlMapper(Mapper): """Mapper to clean html code in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, *args, **kwargs): """ diff --git a/data_engine/ops/mapper/clean_ip_mapper.py b/data_engine/ops/mapper/clean_ip_mapper.py index 3a15bd9..8aaf86e 100644 --- a/data_engine/ops/mapper/clean_ip_mapper.py +++ b/data_engine/ops/mapper/clean_ip_mapper.py @@ -6,6 +6,9 @@ @OPERATORS.register_module('clean_ip_mapper') class CleanIpMapper(Mapper): """Mapper to clean ipv4 and ipv6 address in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, pattern: str = None, repl: str = '', *args, **kwargs): """ diff --git a/data_engine/ops/mapper/clean_links_mapper.py b/data_engine/ops/mapper/clean_links_mapper.py index 846f5de..e619542 100644 --- a/data_engine/ops/mapper/clean_links_mapper.py +++ b/data_engine/ops/mapper/clean_links_mapper.py @@ -9,6 +9,9 @@ @OPERATORS.register_module('clean_links_mapper') class CleanLinksMapper(Mapper): """Mapper to clean links like http/https/ftp in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, pattern: str = None, repl: str = '', *args, **kwargs): """ diff --git a/data_engine/ops/mapper/expand_macro_mapper.py b/data_engine/ops/mapper/expand_macro_mapper.py index 7248e69..67d4fbf 100644 --- a/data_engine/ops/mapper/expand_macro_mapper.py +++ b/data_engine/ops/mapper/expand_macro_mapper.py @@ -11,6 +11,9 @@ class ExpandMacroMapper(Mapper): """Mapper to expand macro definitions in the document body of Latex samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, *args, **kwargs): """ diff --git a/data_engine/ops/mapper/fix_unicode_mapper.py b/data_engine/ops/mapper/fix_unicode_mapper.py index b634f9f..64422a8 100644 --- a/data_engine/ops/mapper/fix_unicode_mapper.py +++ b/data_engine/ops/mapper/fix_unicode_mapper.py @@ -11,6 +11,8 @@ @OPERATORS.register_module(OP_NAME) class FixUnicodeMapper(Mapper): """Mapper to fix unicode errors in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, normalization: str = None, *args, **kwargs): """ diff --git a/data_engine/ops/mapper/nlpaug_en_mapper.py b/data_engine/ops/mapper/nlpaug_en_mapper.py index e8f62d6..34ee5c1 100644 --- a/data_engine/ops/mapper/nlpaug_en_mapper.py +++ b/data_engine/ops/mapper/nlpaug_en_mapper.py @@ -20,6 +20,7 @@ class NlpaugEnMapper(Mapper): """Mapper to simply augment samples in English based on nlpaug library.""" _batched_op = True + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, sequential: bool = False, diff --git a/data_engine/ops/mapper/nlpcda_zh_mapper.py b/data_engine/ops/mapper/nlpcda_zh_mapper.py index ab6e40f..ef2f796 100644 --- a/data_engine/ops/mapper/nlpcda_zh_mapper.py +++ b/data_engine/ops/mapper/nlpcda_zh_mapper.py @@ -18,6 +18,7 @@ class NlpcdaZhMapper(Mapper): """Mapper to simply augment samples in Chinese based on nlpcda library.""" _batched_op = True + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, sequential: bool = False, diff --git a/data_engine/ops/mapper/punctuation_normalization_mapper.py b/data_engine/ops/mapper/punctuation_normalization_mapper.py index 845f9e3..74406f2 100644 --- a/data_engine/ops/mapper/punctuation_normalization_mapper.py +++ b/data_engine/ops/mapper/punctuation_normalization_mapper.py @@ -9,6 +9,9 @@ class PunctuationNormalizationMapper(Mapper): """Mapper to normalize unicode punctuations to English punctuations in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, *args, **kwargs): """ diff --git a/data_engine/ops/mapper/remove_bibliography_mapper.py b/data_engine/ops/mapper/remove_bibliography_mapper.py index 96167fd..660b780 100644 --- a/data_engine/ops/mapper/remove_bibliography_mapper.py +++ b/data_engine/ops/mapper/remove_bibliography_mapper.py @@ -11,6 +11,9 @@ class RemoveBibliographyMapper(Mapper): """Mapper to remove bibliography at the end of documents in Latex samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, *args, **kwargs): """ diff --git a/data_engine/ops/mapper/remove_comments_mapper.py b/data_engine/ops/mapper/remove_comments_mapper.py index e24a477..593d7fd 100644 --- a/data_engine/ops/mapper/remove_comments_mapper.py +++ b/data_engine/ops/mapper/remove_comments_mapper.py @@ -16,6 +16,9 @@ class RemoveCommentsMapper(Mapper): Only support 'tex' for now. """ + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, doc_type: Union[str, List[str]] = 'tex', diff --git a/data_engine/ops/mapper/remove_header_mapper.py b/data_engine/ops/mapper/remove_header_mapper.py index 16e8f39..e2549bf 100644 --- a/data_engine/ops/mapper/remove_header_mapper.py +++ b/data_engine/ops/mapper/remove_header_mapper.py @@ -11,6 +11,9 @@ class RemoveHeaderMapper(Mapper): """Mapper to remove headers at the beginning of documents in Latex samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, drop_no_head: bool = True, *args, **kwargs): """ diff --git a/data_engine/ops/mapper/remove_long_words_mapper.py b/data_engine/ops/mapper/remove_long_words_mapper.py index 8ee2032..9d22a14 100644 --- a/data_engine/ops/mapper/remove_long_words_mapper.py +++ b/data_engine/ops/mapper/remove_long_words_mapper.py @@ -14,6 +14,9 @@ @OPERATORS.register_module('remove_long_words_mapper') class RemoveLongWordsMapper(Mapper): """Mapper to remove long words within a specific range.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, min_len: PositiveInt = 1, diff --git a/data_engine/ops/mapper/remove_non_chinese_character_mapper.py b/data_engine/ops/mapper/remove_non_chinese_character_mapper.py index 2e9d31a..50bd738 100644 --- a/data_engine/ops/mapper/remove_non_chinese_character_mapper.py +++ b/data_engine/ops/mapper/remove_non_chinese_character_mapper.py @@ -6,6 +6,9 @@ @OPERATORS.register_module('remove_non_chinese_character_mapper') class RemoveNonChineseCharacterlMapper(Mapper): """Mapper to remove non chinese Character in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, keep_alphabet: bool = True, diff --git a/data_engine/ops/mapper/remove_repeat_sentences_mapper.py b/data_engine/ops/mapper/remove_repeat_sentences_mapper.py index 4e428f4..610cb7c 100644 --- a/data_engine/ops/mapper/remove_repeat_sentences_mapper.py +++ b/data_engine/ops/mapper/remove_repeat_sentences_mapper.py @@ -14,6 +14,9 @@ def split_sentence(text): @OPERATORS.register_module('remove_repeat_sentences_mapper') class RemoveRepeatSentencesMapper(Mapper): """Mapper to remove repeat sentences in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, lowercase: bool = False, diff --git a/data_engine/ops/mapper/remove_specific_chars_mapper.py b/data_engine/ops/mapper/remove_specific_chars_mapper.py index 71dcf56..443ab6a 100644 --- a/data_engine/ops/mapper/remove_specific_chars_mapper.py +++ b/data_engine/ops/mapper/remove_specific_chars_mapper.py @@ -8,6 +8,9 @@ @OPERATORS.register_module('remove_specific_chars_mapper') class RemoveSpecificCharsMapper(Mapper): """Mapper to clean specific chars in text samples.""" + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, chars_to_remove: Union[str, List[str]] = '◆●■►▼▲▴∆▻▷❖♡□', diff --git a/data_engine/ops/mapper/remove_table_text_mapper.py b/data_engine/ops/mapper/remove_table_text_mapper.py index 781e089..55d603e 100644 --- a/data_engine/ops/mapper/remove_table_text_mapper.py +++ b/data_engine/ops/mapper/remove_table_text_mapper.py @@ -15,6 +15,9 @@ class RemoveTableTextMapper(Mapper): Regular expression is used to remove tables in the range of column number of tables. """ + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, min_col: from_2_to_20 = 2, diff --git a/data_engine/ops/mapper/remove_words_with_incorrect_substrings_mapper.py b/data_engine/ops/mapper/remove_words_with_incorrect_substrings_mapper.py index 5075e71..5d99436 100644 --- a/data_engine/ops/mapper/remove_words_with_incorrect_substrings_mapper.py +++ b/data_engine/ops/mapper/remove_words_with_incorrect_substrings_mapper.py @@ -17,6 +17,8 @@ @OPERATORS.register_module(OP_NAME) class RemoveWordsWithIncorrectSubstringsMapper(Mapper): """Mapper to remove words with incorrect substrings.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, lang: str = 'en', diff --git a/data_engine/ops/mapper/replace_content_mapper.py b/data_engine/ops/mapper/replace_content_mapper.py index 8a14401..7872e32 100644 --- a/data_engine/ops/mapper/replace_content_mapper.py +++ b/data_engine/ops/mapper/replace_content_mapper.py @@ -10,6 +10,8 @@ class ReplaceContentMapper(Mapper): """Mapper to replace all content in the text that matches a specific regular expression pattern with a designated replacement string.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, pattern: Union[str, List[str]] = None, diff --git a/data_engine/ops/mapper/sentence_split_mapper.py b/data_engine/ops/mapper/sentence_split_mapper.py index 2a24314..cad2c13 100644 --- a/data_engine/ops/mapper/sentence_split_mapper.py +++ b/data_engine/ops/mapper/sentence_split_mapper.py @@ -13,6 +13,8 @@ @OPERATORS.register_module(OP_NAME) class SentenceSplitMapper(Mapper): """Mapper to split text samples to sentences.""" + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, lang: str = 'en', *args, **kwargs): """ diff --git a/data_engine/ops/mapper/whitespace_normalization_mapper.py b/data_engine/ops/mapper/whitespace_normalization_mapper.py index e986d4b..61b1efb 100644 --- a/data_engine/ops/mapper/whitespace_normalization_mapper.py +++ b/data_engine/ops/mapper/whitespace_normalization_mapper.py @@ -15,6 +15,8 @@ class WhitespaceNormalizationMapper(Mapper): Different kinds of whitespaces can be found here: https://en.wikipedia.org/wiki/Whitespace_character """ + + _supports_streaming = True # Supports streaming mode for low memory usage def __init__(self, *args, **kwargs): """ diff --git a/data_server/pod/common_tasks.py b/data_server/pod/common_tasks.py index e0ec993..a8d0fb0 100644 --- a/data_server/pod/common_tasks.py +++ b/data_server/pod/common_tasks.py @@ -485,15 +485,102 @@ def run_operator_execute(task_params: dict): except OSError: logger.warning("Failed to remove temp operator config {}", temp_path) - formatter = load_formatter( - recipe.dataset_path, - cfg.generated_dataset_config, - cfg.text_keys, - cfg.suffixes, - cfg.add_suffix, - ) - dataset = formatter.load_dataset(cfg.np, cfg) + # Load operators first to check streaming compatibility ops = load_ops(cfg.process, cfg.op_fusion, job_uid=str(task_params.get("job_id") or "")) + + # Smart formatter selection based on streaming mode and dataset requirements + use_streaming = getattr(cfg, 'use_streaming', False) + + if use_streaming: + # Check if dataset path contains weights or requires mixing + dataset_path_tokens = recipe.dataset_path.split() + has_weights = False + + for token in dataset_path_tokens: + try: + float(token) + has_weights = True + break + except ValueError: + continue + + needs_mixture = has_weights or (hasattr(cfg, 'max_samples') and cfg.max_samples is not None) + + if needs_mixture: + # Dataset requires MixtureFormatter (weights/sampling), must use normal mode + logger.warning("Dataset path contains weights or max_samples - MixtureFormatter required") + logger.warning("Automatically switching to NORMAL mode") + cfg.use_streaming = False + + formatter = load_formatter( + recipe.dataset_path, + cfg.generated_dataset_config, + cfg.text_keys, + cfg.suffixes, + cfg.add_suffix, + ) + + job_uid = str(task_params.get("job_id") or "") + if job_uid: + from data_server.pod.pod_logger import log_task_info + log_task_info(job_uid, "⚠️ MixtureFormatter required, streaming disabled") + else: + # Simple dataset path, check if operators support streaming + unsupported_ops = [] + for op in ops: + if not getattr(op, '_supports_streaming', False): + unsupported_ops.append(op._name) + + if unsupported_ops: + # Operators don't support streaming, fallback to normal mode with MixtureFormatter + logger.warning(f"Operators [{', '.join(unsupported_ops)}] not compatible with streaming") + logger.warning("Automatically switching to NORMAL mode") + cfg.use_streaming = False + + formatter = load_formatter( + recipe.dataset_path, + cfg.generated_dataset_config, + cfg.text_keys, + cfg.suffixes, + cfg.add_suffix, + ) + + job_uid = str(task_params.get("job_id") or "") + if job_uid: + from data_server.pod.pod_logger import log_task_info + log_task_info(job_uid, + f"⚠️ Streaming disabled: operators [{', '.join(unsupported_ops)}] not compatible") + else: + # All operators support streaming, use smart formatter + from data_engine.format.formatter import load_formatter as smart_load_formatter + formatter = smart_load_formatter( + dataset_path=recipe.dataset_path, + text_keys=cfg.text_keys, + suffixes=cfg.suffixes, + add_suffix=cfg.add_suffix, + ) + + logger.info(f"✓ Streaming mode enabled: using {formatter.__class__.__name__}") + logger.info(f"✓ All {len(ops)} operator(s) support streaming mode") + + job_uid = str(task_params.get("job_id") or "") + if job_uid: + from data_server.pod.pod_logger import log_task_info + log_task_info(job_uid, f"✓ Streaming mode: {formatter.__class__.__name__} + {len(ops)} operator(s)") + else: + # Normal mode: use MixtureFormatter (existing behavior) + formatter = load_formatter( + recipe.dataset_path, + cfg.generated_dataset_config, + cfg.text_keys, + cfg.suffixes, + cfg.add_suffix, + ) + logger.info(f"Normal mode: using {formatter.__class__.__name__}") + + # Load dataset with appropriate mode + dataset = formatter.load_dataset(cfg.np, cfg) + exporter = load_exporter( recipe.export_path, cfg.export_shard_size, From 80525ebfedcec17dffa534380624b5d7874dbd36 Mon Sep 17 00:00:00 2001 From: shenren123 <2646799270@qq.com> Date: Mon, 24 Aug 2026 12:19:33 +0800 Subject: [PATCH 07/12] =?UTF-8?q?=E6=96=B0=E5=A2=9E=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E7=AE=97=E5=AD=90=E6=94=AF=E6=8C=81=E6=B5=81=E5=BC=8F=EF=BC=8C?= =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=B5=81=E5=BC=8F=E5=8A=9F=E8=83=BD=E4=BB=A3?= =?UTF-8?q?=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data_engine/config/config.py | 10 + data_engine/exporter/base_exporter.py | 143 +++++-------- data_engine/format/formatter.py | 5 +- data_engine/ops/base_op.py | 196 +++++++----------- .../filter/annotate_edu_train_bert_scorer.py | 3 + data_engine/ops/filter/perplexity_filter.py | 1 + data_engine/ops/filter/text_action_filter.py | 1 + .../filter/text_entity_dependency_filter.py | 1 + data_engine/ops/filter/token_num_filter.py | 2 + data_engine/ops/mapper/extract_qa_mapper.py | 1 + .../mapper/generate_code_qa_pair_mapper.py | 2 + .../ops/mapper/optimize_instruction_mapper.py | 2 + .../ops/mapper/text_make_cosmopedia.py | 1 + 13 files changed, 155 insertions(+), 213 deletions(-) diff --git a/data_engine/config/config.py b/data_engine/config/config.py index 53c6871..eabe3e5 100644 --- a/data_engine/config/config.py +++ b/data_engine/config/config.py @@ -154,6 +154,16 @@ def init_configs(args=None,redirect=True): 'parallelism. Only simple mappers and filters support streaming mode. ' 'Recommended for datasets larger than available memory (>10GB). ' 'When enabled, automatically sets batch_size=100 and enables sample counting.') + parser.add_argument( + '--streaming_batch_size', + type=PositiveInt, + default=100, + help='Batch size for streaming mode processing. Larger values improve ' + 'processing efficiency but consume more memory. This parameter ' + 'controls how many samples are processed together in each batch. ' + 'Recommended range: 100-5000. For text data: 1000-2000;for simple operators: ' + '2000-5000; for complex operators (LLM inference): 50-200. ' + 'Default: 100. Only effective when use_streaming=True.') parser.add_argument( '--text_keys', type=Union[str, List[str]], diff --git a/data_engine/exporter/base_exporter.py b/data_engine/exporter/base_exporter.py index dc7d3e9..212d6b0 100644 --- a/data_engine/exporter/base_exporter.py +++ b/data_engine/exporter/base_exporter.py @@ -147,43 +147,11 @@ def _export_impl(self, dataset, export_path, suffix, export_stats=True): if self.export_ds: # Collect all internal fields to remove before export. # intersection() ensures only existing columns are removed, no error if absent. - fields_to_remove = set() - if not self.keep_stats_in_res_ds: - fields_to_remove.add(Fields.stats) - if not self.keep_hashes_in_res_ds: - fields_to_remove.update({ - HashKeys.hash, - HashKeys.minhash, - HashKeys.simhash, - HashKeys.imagehash, - HashKeys.videohash, - }) - # Other internal __dj__ fields that should not appear in export - fields_to_remove.update({ - Fields.suffix, - Fields.context, - Fields.meta, - Fields.source_file, - Fields.video_frame_tags, - Fields.video_audio_tags, - Fields.multimodal_data_output_dir, - HashKeys.is_duplicate, - HashKeys.similarity_hash, - }) - # 流式模式:使用 map 移除内部字段 - if not is_streaming: - feature_fields = set(dataset.features.keys()) - removed_fields = fields_to_remove.intersection(feature_fields) - if removed_fields: - dataset = dataset.remove_columns(removed_fields) - else: - # 流式模式:通过 map 过滤字段 - logger.info('Streaming mode: filtering internal fields before export...') - - def remove_fields(sample): - return {k: v for k, v in sample.items() if k not in fields_to_remove} - - dataset = dataset.map(remove_fields) + fields_to_remove = self._get_fields_to_remove() + feature_fields = set(dataset.features.keys()) + removed_fields = fields_to_remove.intersection(feature_fields) + if removed_fields: + dataset = dataset.remove_columns(removed_fields) export_method = Exporter._router()[suffix] @@ -248,6 +216,42 @@ def export_from_files(self, upload_path: Path): def export_large_folder(self): pass + def _get_fields_to_remove(self): + """ + Get set of internal fields to remove before export. + Shared by both normal and streaming modes. + + :return: set of field names to remove + """ + fields_to_remove = set() + + if not self.keep_stats_in_res_ds: + fields_to_remove.add(Fields.stats) + + if not self.keep_hashes_in_res_ds: + fields_to_remove.update({ + HashKeys.hash, + HashKeys.minhash, + HashKeys.simhash, + HashKeys.imagehash, + HashKeys.videohash, + }) + + # Other internal fields + fields_to_remove.update({ + Fields.suffix, + Fields.context, + Fields.meta, + Fields.source_file, + Fields.video_frame_tags, + Fields.video_audio_tags, + Fields.multimodal_data_output_dir, + HashKeys.is_duplicate, + HashKeys.similarity_hash, + }) + + return fields_to_remove + def export(self, dataset): """ Export method for a dataset. @@ -275,9 +279,13 @@ def _export_streaming(self, dataset): import tempfile import shutil from tqdm import tqdm + from data_engine.core.streaming_data import filter_empty_samples logger.info('Exporting dataset in STREAMING mode (low memory)...') + # Filter out error samples before export (reuse streaming_data.py logic) + dataset = filter_empty_samples(dataset) + # Determine export directory and filename export_dir = os.path.dirname(os.path.abspath(self.export_path)) basename = os.path.basename(self.export_path) @@ -305,28 +313,7 @@ def _export_streaming(self, dataset): pbar = tqdm(desc='Exporting samples', unit=' samples') for sample in dataset: - # Skip empty samples (error samples from exception handling) - # Error samples have all values as empty lists: {key: []} - if not sample: - continue - - # Check if this is an error sample (all values are empty lists) - is_error_sample = False - if isinstance(sample, dict): - # Get non-internal keys - from data_engine.utils.constant import Fields - data_keys = [k for k in sample.keys() if k not in [Fields.stats, Fields.source_file]] - if data_keys: - # Check if all data values are empty lists - is_error_sample = all( - isinstance(sample[key], list) and len(sample[key]) == 0 - for key in data_keys - ) - - if is_error_sample: - continue - - # Remove internal fields before export + # Remove internal fields before export (reuse shared method) sample = self._clean_sample_for_export(sample) # Write as JSON line @@ -354,45 +341,13 @@ def _export_streaming(self, dataset): def _clean_sample_for_export(self, sample): """ Remove internal fields from sample before export. + Reuses the shared field list from _get_fields_to_remove(). :param sample: sample dict :return: cleaned sample dict """ - # Fields to remove - fields_to_remove = set() - - if not self.keep_stats_in_res_ds: - fields_to_remove.add(Fields.stats) - - if not self.keep_hashes_in_res_ds: - fields_to_remove.update({ - HashKeys.hash, - HashKeys.minhash, - HashKeys.simhash, - HashKeys.imagehash, - HashKeys.videohash, - }) - - # Other internal fields - fields_to_remove.update({ - Fields.suffix, - Fields.context, - Fields.meta, - Fields.source_file, - Fields.video_frame_tags, - Fields.video_audio_tags, - Fields.multimodal_data_output_dir, - HashKeys.is_duplicate, - HashKeys.similarity_hash, - }) - - # Remove fields that exist in sample - cleaned_sample = { - k: v for k, v in sample.items() - if k not in fields_to_remove - } - - return cleaned_sample + fields_to_remove = self._get_fields_to_remove() + return {k: v for k, v in sample.items() if k not in fields_to_remove} def export_compute_stats(self, dataset, export_path): """ diff --git a/data_engine/format/formatter.py b/data_engine/format/formatter.py index 23ced30..9b223e0 100644 --- a/data_engine/format/formatter.py +++ b/data_engine/format/formatter.py @@ -1030,11 +1030,10 @@ def _load_dataset_streaming(self, global_cfg=None): # First, find data files self.data_files = find_files_with_suffix(self.dataset_path, self.suffixes) - # Streaming mode: Fixed configuration # - Always pre-scan sample count (for progress bar) - # - Fixed batch_size=100 (memory-efficient batch processing) estimated_total_samples = None - batch_size = 1000 # Fixed batch size for streaming mode (balances memory and performance) + # Get batch_size from config (user-controllable) + batch_size = getattr(global_cfg, 'streaming_batch_size', 100) if global_cfg else 100 if global_cfg and global_cfg.use_streaming: # Pre-scan sample count (mandatory in streaming mode) diff --git a/data_engine/ops/base_op.py b/data_engine/ops/base_op.py index 3150f51..bd63227 100644 --- a/data_engine/ops/base_op.py +++ b/data_engine/ops/base_op.py @@ -335,68 +335,57 @@ def run_streaming(self, dataset, *, exporter=None): try: # Get batch_size from dataset batch_size = getattr(dataset, 'batch_size', 1) - use_batched = batch_size > 1 - if use_batched: - logger.info(f'Processing {self._name} in streaming mode (batched with batch_size={batch_size})...') - else: - logger.info(f'Processing {self._name} in streaming mode (single-process)...') + logger.info(f'Processing {self._name} in streaming mode (batch_size={batch_size})...') # Import streaming exception handler - from data_engine.core.streaming_data import catch_streaming_exception, filter_empty_samples - - if use_batched: - # Define batched processing function - def safe_process_batched(batch): - """Process batch with exception handling""" - # Get batch size - first_key = next(iter(batch.keys())) - current_batch_size = len(batch[first_key]) - - # Initialize result batch - result_batch = {key: [] for key in batch.keys()} - - # Process each sample in batch - for i in range(current_batch_size): - try: - # Extract single sample from batch - sample = {key: batch[key][i] for key in batch.keys()} - - # Process sample - processed_sample = self.process(sample) - - # Add to result batch - for key in processed_sample.keys(): - if key not in result_batch: - result_batch[key] = [] - result_batch[key].append(processed_sample[key]) - except Exception as e: - logger.error( - f'An error occurred in mapper {self._name} when processing sample {i}: {e}' - ) - import traceback - traceback.print_exc() - # Return empty dict for error samples - for key in batch.keys(): - if key not in result_batch: - result_batch[key] = [] - result_batch[key].append({}) - - return result_batch - - # Apply batched processing - new_dataset = dataset.map( - safe_process_batched, - batched=True, - batch_size=batch_size, - desc=self._name + '_process' - ) - else: - # Wrap process method with exception handler (similar to normal mode) - safe_process = catch_streaming_exception(self.process) + from data_engine.core.streaming_data import filter_empty_samples + + # Batched processing function (supports any batch_size including 1) + def safe_process_batched(batch): + """Process batch with exception handling""" + # Get batch size + first_key = next(iter(batch.keys())) + current_batch_size = len(batch[first_key]) + + # Initialize result batch + result_batch = {key: [] for key in batch.keys()} + + # Process each sample in batch + for i in range(current_batch_size): + try: + # Extract single sample from batch + sample = {key: batch[key][i] for key in batch.keys()} + + # Process sample + processed_sample = self.process(sample) + + # Add to result batch + for key in processed_sample.keys(): + if key not in result_batch: + result_batch[key] = [] + result_batch[key].append(processed_sample[key]) + except Exception as e: + logger.error( + f'An error occurred in mapper {self._name} when processing sample {i}: {e}' + ) + import traceback + traceback.print_exc() + # Mark error sample: add empty list (will be filtered by filter_empty_samples) + for key in batch.keys(): + if key not in result_batch: + result_batch[key] = [] + result_batch[key].append([]) - # Apply safe process in single-sample mode - new_dataset = dataset.map(safe_process, desc=self._name + '_process') + return result_batch + + # Apply batched processing + new_dataset = dataset.map( + safe_process_batched, + batched=True, + batch_size=batch_size, + desc=self._name + '_process' + ) # Filter out empty dict samples that resulted from exceptions new_dataset = filter_empty_samples(new_dataset) @@ -569,77 +558,52 @@ def run_streaming(self, dataset, *, exporter=None): try: # Get batch_size from dataset batch_size = getattr(dataset, 'batch_size', 1) - use_batched = batch_size > 1 - if use_batched: - logger.info( - f'Processing {self._name} in streaming mode (batched filtering with batch_size={batch_size})...') - else: - logger.info(f'Processing {self._name} in streaming mode (single-pass filtering)...') - - # Define combined filter function with exception handling - if use_batched: - def safe_combined_filter_batched(batch): - """Compute stats and filter in batch mode with exception handling""" - results = [] - - # Get batch size - first_key = next(iter(batch.keys())) - current_batch_size = len(batch[first_key]) - - # Process each sample in batch - for i in range(current_batch_size): - try: - # Extract single sample from batch - sample = {key: batch[key][i] for key in batch.keys()} - - # Compute stats (compute_stats is wrapped with exception handler) - sample = self.compute_stats(sample) - - # Apply filter - keep = self.process(sample) - results.append(keep) - except Exception as e: - logger.error( - f'An error occurred in filter {self._name} when processing sample {i}: {e}' - ) - import traceback - traceback.print_exc() - # Return False to filter out error samples - results.append(False) - - return results - - # Apply combined filter in batched mode - new_dataset = dataset.filter( - safe_combined_filter_batched, - batched=True, - batch_size=batch_size, - desc=self._name + '_process' - ) - else: - def safe_combined_filter(sample): - """Compute stats and filter in single pass with exception handling""" + logger.info(f'Processing {self._name} in streaming mode (batch_size={batch_size})...') + + # Batched filter function (supports any batch_size including 1) + def safe_combined_filter_batched(batch): + """Compute stats and filter in batch mode with exception handling""" + results = [] + + # Get batch size + first_key = next(iter(batch.keys())) + current_batch_size = len(batch[first_key]) + + # Process each sample in batch + for i in range(current_batch_size): try: + # Extract single sample from batch + sample = {key: batch[key][i] for key in batch.keys()} + + # 🔧 FIX: Ensure Fields.stats exists and is a dict (not None) + if Fields.stats not in sample or sample[Fields.stats] is None: + sample[Fields.stats] = {} + # Compute stats (compute_stats is wrapped with exception handler) sample = self.compute_stats(sample) # Apply filter - return self.process(sample) + keep = self.process(sample) + results.append(keep) except Exception as e: logger.error( - f'An error occurred in filter {self._name} when processing sample: {e}' + f'An error occurred in filter {self._name} when processing sample {i}: {e}' ) import traceback traceback.print_exc() # Return False to filter out error samples - return False + results.append(False) - # Apply combined filter in single-sample mode - new_dataset = dataset.filter( - safe_combined_filter, - desc=self._name + '_process' - ) + return results + + # Apply batched filter + new_dataset = dataset.filter( + safe_combined_filter_batched, + batched=True, + batch_size=batch_size, + desc=self._name + '_process' + ) set_pipline_job_operator_status( self.job_uid, diff --git a/data_engine/ops/filter/annotate_edu_train_bert_scorer.py b/data_engine/ops/filter/annotate_edu_train_bert_scorer.py index f1709cd..8309fb7 100644 --- a/data_engine/ops/filter/annotate_edu_train_bert_scorer.py +++ b/data_engine/ops/filter/annotate_edu_train_bert_scorer.py @@ -21,6 +21,9 @@ @OPERATORS.register_module(OP_NAME) @LOADED_AUDIOS.register_module(OP_NAME) class AnnotateEduTrainBertScorer(Mapper): + + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, auth_token: str = "", model_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1", diff --git a/data_engine/ops/filter/perplexity_filter.py b/data_engine/ops/filter/perplexity_filter.py index ca48816..7cff4e9 100644 --- a/data_engine/ops/filter/perplexity_filter.py +++ b/data_engine/ops/filter/perplexity_filter.py @@ -41,6 +41,7 @@ class PerplexityFilter(Filter): """Filter to keep samples with perplexity score less than a specific max value. Uses LLM API to evaluate text quality.""" + _supports_streaming = True # Supports streaming mode for low memory usage _accelerator = 'cpu' def __init__(self, diff --git a/data_engine/ops/filter/text_action_filter.py b/data_engine/ops/filter/text_action_filter.py index 44091d2..3bc65d2 100644 --- a/data_engine/ops/filter/text_action_filter.py +++ b/data_engine/ops/filter/text_action_filter.py @@ -31,6 +31,7 @@ class TextActionFilter(Filter): Uses remote LLM API to detect actions. """ + _supports_streaming = True # Supports streaming mode for low memory usage _accelerator = 'cpu' def __init__(self, diff --git a/data_engine/ops/filter/text_entity_dependency_filter.py b/data_engine/ops/filter/text_entity_dependency_filter.py index 5b4004d..34fa835 100644 --- a/data_engine/ops/filter/text_entity_dependency_filter.py +++ b/data_engine/ops/filter/text_entity_dependency_filter.py @@ -44,6 +44,7 @@ class TextEntityDependencyFilter(Filter): Uses remote LLM API to detect entity dependencies. """ + _supports_streaming = True # Supports streaming mode for low memory usage _accelerator = 'cpu' def __init__(self, diff --git a/data_engine/ops/filter/token_num_filter.py b/data_engine/ops/filter/token_num_filter.py index b229203..b7bcfde 100644 --- a/data_engine/ops/filter/token_num_filter.py +++ b/data_engine/ops/filter/token_num_filter.py @@ -20,6 +20,8 @@ class TokenNumFilter(Filter): """Filter to keep samples with total token number within a specific range.""" + _supports_streaming = True # Supports streaming mode for low memory usage + def __init__(self, hf_tokenizer: str = 'EleutherAI/pythia-6.9b-deduped', min_num: PositiveInt = 10, diff --git a/data_engine/ops/mapper/extract_qa_mapper.py b/data_engine/ops/mapper/extract_qa_mapper.py index babe6cf..6a39248 100644 --- a/data_engine/ops/mapper/extract_qa_mapper.py +++ b/data_engine/ops/mapper/extract_qa_mapper.py @@ -40,6 +40,7 @@ class ExtractQAMapper(Mapper): Supports OpenAI-compatible API formats including Qwen, DeepSeek, GPT, etc. """ + _supports_streaming = True # Supports streaming mode for low memory usage _accelerator = 'cpu' def __init__(self, diff --git a/data_engine/ops/mapper/generate_code_qa_pair_mapper.py b/data_engine/ops/mapper/generate_code_qa_pair_mapper.py index a3297b5..e66259f 100644 --- a/data_engine/ops/mapper/generate_code_qa_pair_mapper.py +++ b/data_engine/ops/mapper/generate_code_qa_pair_mapper.py @@ -23,6 +23,8 @@ class GenerateCodeQAPairMapper(Mapper): Mapper to generate code QA pairs using remote LLM API. Supports OpenAI-compatible API formats including Qwen, DeepSeek, GPT, etc. """ + + _supports_streaming = True # Supports streaming mode for low memory usage _accelerator = 'cpu' def __init__(self, diff --git a/data_engine/ops/mapper/optimize_instruction_mapper.py b/data_engine/ops/mapper/optimize_instruction_mapper.py index 2837e88..dcc2767 100644 --- a/data_engine/ops/mapper/optimize_instruction_mapper.py +++ b/data_engine/ops/mapper/optimize_instruction_mapper.py @@ -28,6 +28,8 @@ @OPERATORS.register_module(OP_NAME) class OptimizeInstructionMapper(Mapper): + + _supports_streaming = True # Supports streaming mode for low memory usage _accelerator = 'cpu' def __init__(self, diff --git a/data_engine/ops/mapper/text_make_cosmopedia.py b/data_engine/ops/mapper/text_make_cosmopedia.py index 0b66f7d..76b2688 100644 --- a/data_engine/ops/mapper/text_make_cosmopedia.py +++ b/data_engine/ops/mapper/text_make_cosmopedia.py @@ -11,6 +11,7 @@ class MakeCosmopediaMapper(Mapper): """Mapper to generate synthetic tutorial data from seed text samples.""" + _supports_streaming = True # Supports streaming mode for low memory usage # _batched_op = False def __init__(self, *args, **kwargs): From ba13d3b6a0852745238dcb62a61df743c4d97e1c Mon Sep 17 00:00:00 2001 From: shenren123 <2646799270@qq.com> Date: Fri, 28 Aug 2026 11:39:27 +0800 Subject: [PATCH 08/12] =?UTF-8?q?=E4=BF=AE=E5=A4=8Dbug,=E4=BC=98=E5=8C=96?= =?UTF-8?q?=E4=BB=A3=E7=A0=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data_engine/config/config.py | 2 +- data_engine/exporter/csghub_exporter.py | 1 - data_engine/format/formatter.py | 49 +++++++++++++++---- data_engine/ops/filter/token_num_filter.py | 2 - .../ops/mapper/text_make_cosmopedia.py | 1 - data_server/logic/models.py | 5 ++ 6 files changed, 46 insertions(+), 14 deletions(-) diff --git a/data_engine/config/config.py b/data_engine/config/config.py index eabe3e5..2e5c00d 100644 --- a/data_engine/config/config.py +++ b/data_engine/config/config.py @@ -146,7 +146,7 @@ def init_configs(args=None,redirect=True): parser.add_argument( '--use_streaming', type=bool, - default=True, + default=False, help='Whether to use streaming mode for dataset loading and processing. ' 'Streaming mode significantly reduces memory usage by processing ' 'data iteratively without loading the entire dataset into memory. ' diff --git a/data_engine/exporter/csghub_exporter.py b/data_engine/exporter/csghub_exporter.py index 83ed0e6..e9f61f9 100644 --- a/data_engine/exporter/csghub_exporter.py +++ b/data_engine/exporter/csghub_exporter.py @@ -152,7 +152,6 @@ def export(self, dataset): self._export_impl(dataset, self.export_path, self.suffix, self.export_stats) # After export, push to repo if repo_id is configured - self._export_impl(dataset, self.export_path, self.suffix, self.export_stats) self.upload_path = os.path.join(self.work_dir, "_data") self.repo_work_dir = os.path.join(self.work_dir, "_git") self._export_common() diff --git a/data_engine/format/formatter.py b/data_engine/format/formatter.py index 9b223e0..33c8e9b 100644 --- a/data_engine/format/formatter.py +++ b/data_engine/format/formatter.py @@ -1145,14 +1145,14 @@ def _prescan_sample_count(self): Pre-scan data files to count total samples (low memory, O(1) complexity). Uses fast line counting (binary read, no parsing) with minimal memory (~9KB). - Supports line-based formats: JSONL, CSV, TXT. - + Supports line-based formats: JSONL, CSV, TSV, TXT. :return: Total sample count or None if failed """ logger.info('Pre-scanning sample count (fast line counting)...') total_samples = 0 file_count = 0 + skipped_files = [] try: # Collect all data files @@ -1168,13 +1168,25 @@ def _prescan_sample_count(self): for file_path in all_files: try: # Check if file format supports line counting - if not self._is_line_based_format(file_path): - logger.warning( - f'Skipping {file_path}: Not a line-based format (only JSONL, CSV, TXT supported)') + format_info = self._is_line_based_format(file_path) + if format_info is None: + skipped_files.append((file_path, 'not line-based format')) + continue + + is_supported, has_header = format_info + if not is_supported: + skipped_files.append((file_path, 'not line-based format')) continue # Fast line counting (binary read, O(1) memory) - file_samples = self._count_lines_fast(file_path) + line_count = self._count_lines_fast(file_path) + + # Adjust for header row if present (CSV/TSV) + if has_header and line_count > 0: + file_samples = line_count - 1 + else: + file_samples = line_count + total_samples += file_samples file_count += 1 @@ -1183,8 +1195,15 @@ def _prescan_sample_count(self): except Exception as e: logger.warning(f'Failed to count samples in {file_path}: {e}') + skipped_files.append((file_path, str(e))) continue + # Log skipped files with reasons + if skipped_files: + logger.warning(f'Skipped {len(skipped_files)} file(s) during pre-scan:') + for file_path, reason in skipped_files: + logger.warning(f' - {os.path.basename(file_path)}: {reason}') + if file_count > 0: logger.info(f'Pre-scan complete: {total_samples:,} total samples across {file_count} file(s)') return total_samples @@ -1216,11 +1235,23 @@ def _is_line_based_format(file_path): Check if file format is line-based (supports fast line counting). :param file_path: Path to file - :return: True if line-based format + :return: Tuple (is_supported, has_header) or None if not supported + - is_supported: True if line-based format + - has_header: True if format has header row (CSV/TSV) """ ext = os.path.splitext(file_path)[1].lower() - line_based_formats = {'.jsonl', '.csv', '.txt', '.tsv'} - return ext in line_based_formats + + # Line-based formats without header + no_header_formats = {'.jsonl', '.txt'} + # Line-based formats with header row + header_formats = {'.csv', '.tsv'} + + if ext in no_header_formats: + return (True, False) # Supported, no header + elif ext in header_formats: + return (True, True) # Supported, has header + else: + return None # Not supported class RemoteFormatter(BaseFormatter): """The class is used to load a dataset from repository of huggingface diff --git a/data_engine/ops/filter/token_num_filter.py b/data_engine/ops/filter/token_num_filter.py index b7bcfde..b229203 100644 --- a/data_engine/ops/filter/token_num_filter.py +++ b/data_engine/ops/filter/token_num_filter.py @@ -20,8 +20,6 @@ class TokenNumFilter(Filter): """Filter to keep samples with total token number within a specific range.""" - _supports_streaming = True # Supports streaming mode for low memory usage - def __init__(self, hf_tokenizer: str = 'EleutherAI/pythia-6.9b-deduped', min_num: PositiveInt = 10, diff --git a/data_engine/ops/mapper/text_make_cosmopedia.py b/data_engine/ops/mapper/text_make_cosmopedia.py index 76b2688..0b66f7d 100644 --- a/data_engine/ops/mapper/text_make_cosmopedia.py +++ b/data_engine/ops/mapper/text_make_cosmopedia.py @@ -11,7 +11,6 @@ class MakeCosmopediaMapper(Mapper): """Mapper to generate synthetic tutorial data from seed text samples.""" - _supports_streaming = True # Supports streaming mode for low memory usage # _batched_op = False def __init__(self, *args, **kwargs): diff --git a/data_server/logic/models.py b/data_server/logic/models.py index 4706504..4be2384 100644 --- a/data_server/logic/models.py +++ b/data_server/logic/models.py @@ -161,6 +161,11 @@ class Recipe(BaseModelExtended): keep_stats_in_res_ds: bool = False keep_hashes_in_res_ds: bool = False + # streaming mode control + use_streaming: bool = False + # Batch size for streaming mode processing + streaming_batch_size: Optional[int] = 100 + # for distributed processing executor_type: Union[Literal["default"], Literal["ray"]] = "default" ray_address: str = "auto" From 46a590f5d9d790f0378abf50061497d8a83f1aec Mon Sep 17 00:00:00 2001 From: shenren123 <2646799270@qq.com> Date: Fri, 28 Aug 2026 19:23:54 +0800 Subject: [PATCH 09/12] =?UTF-8?q?=E5=89=8D=E7=AB=AF=E6=8E=A5=E5=8F=97?= =?UTF-8?q?=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data_server/api/endpoints/job.py | 22 ++++++++++++++++++++++ data_server/job/JobsManager.py | 9 +++++++++ 2 files changed, 31 insertions(+) diff --git a/data_server/api/endpoints/job.py b/data_server/api/endpoints/job.py index 9af114f..492de3e 100644 --- a/data_server/api/endpoints/job.py +++ b/data_server/api/endpoints/job.py @@ -600,6 +600,28 @@ async def run_pipline_job( space_resource_id=data.get("space_resource_id"), storage_size=data.get("storage_size"), ) + + # Update streaming parameters in yaml_config + use_streaming = data.get("use_streaming") + streaming_batch_size = data.get("streaming_batch_size") + + if use_streaming is not None or streaming_batch_size is not None: + try: + yaml_config_dict = yaml.safe_load(job.yaml_config) if job.yaml_config else {} + + if use_streaming is not None: + yaml_config_dict["use_streaming"] = use_streaming + + if use_streaming and streaming_batch_size is not None: + yaml_config_dict["streaming_batch_size"] = streaming_batch_size + elif use_streaming is False: + yaml_config_dict.pop("streaming_batch_size", None) + + job.yaml_config = yaml.dump(yaml_config_dict, sort_keys=False, default_flow_style=False, indent=2, + width=float("inf")) + except Exception as e: + logger.warning(f"Failed to update streaming parameters: {e}") + session.commit() ok, msg = execute_job( diff --git a/data_server/job/JobsManager.py b/data_server/job/JobsManager.py index d3828ba..dfc1af9 100644 --- a/data_server/job/JobsManager.py +++ b/data_server/job/JobsManager.py @@ -934,6 +934,15 @@ def parse_yaml_config(yaml_string: str,config): "trace_num": '1', } + # Add streaming mode configuration if provided + if hasattr(config, 'use_streaming'): + fields_to_insert["use_streaming"] = config.use_streaming + + # Add streaming_batch_size only if streaming mode is enabled + if hasattr(config, 'use_streaming') and config.use_streaming: + if hasattr(config, 'streaming_batch_size') and config.streaming_batch_size is not None: + fields_to_insert["streaming_batch_size"] = config.streaming_batch_size + dsl_data = yaml.safe_load(yaml_string) dsl_data.update(fields_to_insert) From 8fa87aa451aa3f33621c8905d2bcfa6b84beb068 Mon Sep 17 00:00:00 2001 From: shenren123 <2646799270@qq.com> Date: Mon, 31 Aug 2026 20:23:08 +0800 Subject: [PATCH 10/12] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E5=8F=82=E6=95=B0?= =?UTF-8?q?=E7=99=BD=E5=90=8D=E5=8D=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data_server/utils/csghub_pipeline_config.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/data_server/utils/csghub_pipeline_config.py b/data_server/utils/csghub_pipeline_config.py index 5720174..57b2806 100644 --- a/data_server/utils/csghub_pipeline_config.py +++ b/data_server/utils/csghub_pipeline_config.py @@ -28,6 +28,8 @@ "percentiles", "export_original_dataset", "save_stats_in_one_file", + "use_streaming", + "streaming_batch_size", ) From 898d09939facc197db5afc4100481471934c997e Mon Sep 17 00:00:00 2001 From: shenren123 <2646799270@qq.com> Date: Wed, 9 Sep 2026 18:36:43 +0800 Subject: [PATCH 11/12] =?UTF-8?q?=E6=A0=BC=E5=BC=8F=E8=BD=AC=E6=8D=A2?= =?UTF-8?q?=E6=B5=81=E5=BC=8F=E5=8A=A0=E8=BD=BD=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data_server/pod/common_tasks.py | 45 ++ data_server/pod/formatify_helpers.py | 793 ++++++++++++++++++++++++++- 2 files changed, 831 insertions(+), 7 deletions(-) diff --git a/data_server/pod/common_tasks.py b/data_server/pod/common_tasks.py index a8d0fb0..08597ec 100644 --- a/data_server/pod/common_tasks.py +++ b/data_server/pod/common_tasks.py @@ -1002,6 +1002,14 @@ def _select_convert_func(from_type, to_type): def _run_convert_func(convert_func, file_path: str, task_uid: str, task_params: dict): + """ + Run conversion function with appropriate parameters. + + Supports streaming mode for Excel/CSV conversions: + - use_streaming: bool (default: False) + - chunk_size: int (default: 50000) + """ + # PDF to Markdown has special parameters if convert_func is convert_pdf_to_markdown: return convert_func( file_path, @@ -1009,6 +1017,43 @@ 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"), ) + + # Check if conversion function supports streaming mode + # (Excel/CSV conversions support use_streaming and chunk_size parameters) + streaming_supported_funcs = ( + convert_excel_to_csv, + convert_excel_to_json, + convert_excel_to_parquet, + convert_csv_to_excel, + ) + + if convert_func in streaming_supported_funcs: + # Extract streaming parameters from task_params + # Default: use_streaming=True (streaming mode enabled by default for better memory efficiency) + use_streaming = task_params.get("use_streaming", False) + chunk_size = task_params.get("chunk_size", 50000) + + # Convert to appropriate types + if isinstance(use_streaming, str): + use_streaming = use_streaming.lower() in ("true", "1", "yes") + use_streaming = bool(use_streaming) + + if isinstance(chunk_size, str): + try: + chunk_size = int(chunk_size) + except (ValueError, TypeError): + chunk_size = 50000 + chunk_size = int(chunk_size) if chunk_size else 50000 + + # Call with streaming parameters + return convert_func( + file_path, + task_uid, + use_streaming=use_streaming, + chunk_size=chunk_size + ) + + # Other conversion functions (Word/PPT/TXT/HTML to Markdown) 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..6b944bc 100644 --- a/data_server/pod/formatify_helpers.py +++ b/data_server/pod/formatify_helpers.py @@ -1,4 +1,5 @@ import json +import mmap import os import re import shutil @@ -28,6 +29,46 @@ 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 _count_csv_rows_fast(file_path: str) -> int: + """Fast CSV row counting without loading into memory.""" + try: + # Method 1: Use memory mapping (fastest for large files) + with open(file_path, 'r+b') as f: + mmapped = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) + count = 0 + while mmapped.readline(): + count += 1 + mmapped.close() + return count - 1 # Exclude header + except Exception: + # Method 2: Fallback to simple line counting + count = 0 + with open(file_path, 'rb') as f: + for _ in f: + count += 1 + return count - 1 # Exclude header + + 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): @@ -41,7 +82,27 @@ def _non_conflicting_output_path(file_path: str) -> str: return candidate -def convert_excel_to_csv(file_path: str, task_uid) -> Optional[Dict[str, str]]: +def convert_excel_to_csv(file_path: str, task_uid, use_streaming: bool = False, chunk_size: int = 50000) -> Optional[Dict[str, str]]: + """ + Convert Excel/CSV to CSV format. + + Args: + file_path: Input file path + task_uid: Task identifier for logging + use_streaming: Use streaming mode for large files (default: False) + chunk_size: Number of rows per chunk in streaming mode (default: 50000) + + Returns: + Conversion result dictionary + """ + if use_streaming: + return _convert_excel_to_csv_streaming(file_path, task_uid, chunk_size) + else: + return _convert_excel_to_csv_legacy(file_path, task_uid) + + +def _convert_excel_to_csv_legacy(file_path: str, task_uid) -> Optional[Dict[str, str]]: + """Original non-streaming implementation.""" if file_path.lower().endswith(".csv"): log_task_info(task_uid, f"CSV source file will be copied without conversion: {file_path}") return { @@ -53,9 +114,14 @@ def convert_excel_to_csv(file_path: str, task_uid) -> Optional[Dict[str, str]]: if file_path.lower().endswith((".xlsx", ".xls")): log_task_info(task_uid, f"Source file address:{file_path}") try: - xls = pd.ExcelFile(file_path) - sheet_names = xls.sheet_names + # Use openpyxl to avoid file locking issues on Windows + from openpyxl import load_workbook + + # Get sheet names + wb_info = load_workbook(file_path, read_only=True, data_only=True) + sheet_names = wb_info.sheetnames sheet_count = len(sheet_names) + wb_info.close() log_task_info(task_uid, f"Found {sheet_count} sheet(s) in Excel file") @@ -82,7 +148,17 @@ def convert_excel_to_csv(file_path: str, task_uid) -> Optional[Dict[str, str]]: log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") continue - os.remove(file_path) + # Remove source file + try: + os.remove(file_path) + except PermissionError: + # Retry after short delay for Windows file locking + import time + time.sleep(0.5) + try: + os.remove(file_path) + except Exception as e: + log_task_error(task_uid, f"Warning: Could not delete source file: {e}") if len(result_files) == 0: return {"from": file_path, "to": None, "status": "failure", "error": "No sheets converted"} @@ -100,7 +176,220 @@ def convert_excel_to_csv(file_path: str, task_uid) -> Optional[Dict[str, str]]: return None -def convert_excel_to_json(file_path: str, task_uid) -> Optional[Dict[str, str]]: +def _read_excel_sheet_in_chunks(file_path: str, sheet_name: str, chunk_size: int): + """ + Read Excel sheet in chunks using openpyxl. + + Note: pandas.read_excel() does NOT support chunksize parameter. + This is a workaround using openpyxl's read_only mode. + + Yields: + DataFrame: Chunk of data as pandas DataFrame + + Important: + This function does NOT load all rows into memory at once. + It yields chunks as it reads, providing true streaming behavior. + """ + wb = None + try: + from openpyxl import load_workbook + except ImportError: + # Fallback: if openpyxl not available, read entire sheet + df = pd.read_excel(file_path, sheet_name=sheet_name) + yield df + return + + try: + # Use read_only mode to save memory + wb = load_workbook(file_path, read_only=True, data_only=True) + ws = wb[sheet_name] + + # Get header (first row) + rows_iter = ws.iter_rows(values_only=True) + header = next(rows_iter, None) + + if header is None: + return + + # Convert header to list and clean up + header = [str(cell) if cell is not None else f"Column_{i}" for i, cell in enumerate(header)] + + # Stream processing: yield chunks without loading all rows + chunk_data = [] + for row in rows_iter: + chunk_data.append(row) + + if len(chunk_data) >= chunk_size: + # Convert to DataFrame and yield + df = pd.DataFrame(chunk_data, columns=header) + yield df + chunk_data = [] + del df + + # Yield remaining data + if chunk_data: + df = pd.DataFrame(chunk_data, columns=header) + yield df + del df + + except Exception as e: + # Fallback to reading entire sheet + df = pd.read_excel(file_path, sheet_name=sheet_name) + yield df + finally: + # Ensure workbook is closed + if wb is not None: + try: + wb.close() + except Exception: + pass + + +def _convert_excel_to_csv_streaming(file_path: str, task_uid, chunk_size: int) -> Optional[Dict[str, str]]: + """ + Streaming implementation with chunked reading and writing. + + Note: pandas.read_excel() does NOT support chunksize parameter. + We use openpyxl's read_only mode as a workaround. + """ + if file_path.lower().endswith(".csv"): + log_task_info(task_uid, f"CSV source file will be copied without conversion: {file_path}") + return { + "from": file_path, + "to": file_path, + "to_files": [file_path], + "status": "success", + } + + if file_path.lower().endswith((".xlsx", ".xls")): + 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") + + try: + # Use openpyxl directly to avoid file handle issues + from openpyxl import load_workbook + + # First pass: get sheet names + wb_info = load_workbook(file_path, read_only=True, data_only=True) + sheet_names = wb_info.sheetnames + sheet_count = len(sheet_names) + wb_info.close() + + log_task_info(task_uid, f"Found {sheet_count} sheet(s) in Excel file") + + result_files = [] + base_name = os.path.splitext(file_path)[0] + + for idx, sheet_name in enumerate(sheet_names, 1): + new_file = None + try: + log_task_info(task_uid, f"Processing sheet {idx}/{sheet_count}: '{sheet_name}'") + + # Generate output filename + safe_sheet_name = re.sub(r'[<>:"/\\|?*]', '_', sheet_name) + if sheet_count == 1: + new_file = f"{base_name}.csv" + else: + new_file = f"{base_name}_{safe_sheet_name}.csv" + new_file = _non_conflicting_output_path(new_file) + + # Count total rows for progress reporting using openpyxl + log_task_info(task_uid, f"Counting total rows in sheet '{sheet_name}'...") + wb_count = load_workbook(file_path, read_only=True, data_only=True) + ws_count = wb_count[sheet_name] + total_rows = ws_count.max_row - 1 # Exclude header + wb_count.close() + log_task_info(task_uid, f"Total rows: {total_rows:,}") + + # Stream processing using openpyxl + first_chunk = True + rows_processed = 0 + + for chunk in _read_excel_sheet_in_chunks(file_path, sheet_name, chunk_size): + chunk.to_csv( + new_file, + mode='w' if first_chunk else 'a', + header=first_chunk, + index=False, + encoding='utf-8-sig' + ) + first_chunk = False + rows_processed += len(chunk) + + # Progress reporting + progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 + log_task_info( + task_uid, + f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" + ) + + # Release memory + del chunk + + result_files.append(new_file) + log_task_info(task_uid, f"Sheet '{sheet_name}' converted successfully to {new_file}") + + except Exception as sheet_error: + log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") + # Rollback: delete partial file + if new_file and os.path.exists(new_file): + try: + os.remove(new_file) + log_task_info(task_uid, f"Rolled back partial file: {new_file}") + except Exception: + pass + continue + + # Remove source file only if at least one sheet succeeded + if len(result_files) > 0: + try: + os.remove(file_path) + except PermissionError: + # File might be locked, try again after a short delay + import time + time.sleep(0.5) + try: + os.remove(file_path) + except Exception as e: + log_task_error(task_uid, f"Warning: Could not delete source file: {e}") + + if len(result_files) == 0: + return {"from": file_path, "to": None, "status": "failure", "error": "No sheets converted"} + + return { + "from": file_path, + "to": result_files[0] if len(result_files) == 1 else result_files, + "to_files": result_files, + "status": "success", + "sheets_count": len(result_files) + } + except Exception as e: + log_task_error(task_uid, f"convert file {file_path} error: {e}") + return {"from": file_path, "to": None, "status": "failure", "error": str(e)} + return None + + +def convert_excel_to_json(file_path: str, task_uid, use_streaming: bool = False, chunk_size: int = 50000) -> Optional[Dict[str, str]]: + """ + Convert Excel/CSV to JSON format. + + Args: + file_path: Input file path + task_uid: Task identifier for logging + use_streaming: Use streaming mode for large files (default: False) + chunk_size: Number of rows per chunk in streaming mode (default: 50000) + + Returns: + Conversion result dictionary + """ + if use_streaming: + return _convert_excel_to_json_streaming(file_path, task_uid, chunk_size) + else: + return _convert_excel_to_json_legacy(file_path, task_uid) + + +def _convert_excel_to_json_legacy(file_path: str, task_uid) -> Optional[Dict[str, str]]: + """Original non-streaming implementation.""" if file_path.lower().endswith(".csv"): log_task_info(task_uid, f"Source file address: {file_path}") try: @@ -166,7 +455,200 @@ def convert_excel_to_json(file_path: str, task_uid) -> Optional[Dict[str, str]]: return None -def convert_excel_to_parquet(file_path: str, task_uid) -> Optional[Dict[str, str]]: +def _convert_excel_to_json_streaming(file_path: str, task_uid, chunk_size: int) -> Optional[Dict[str, str]]: + """ + Streaming implementation with chunked reading and writing. + + Note: pandas.read_excel() does NOT support chunksize parameter. + We use openpyxl's read_only mode as a workaround. + """ + if file_path.lower().endswith(".csv"): + 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 = _non_conflicting_output_path( + f"{os.path.splitext(file_path)[0]}.json" + ) + + # Count total rows + log_task_info(task_uid, "Counting total rows...") + total_rows = _count_csv_rows_fast(file_path) + log_task_info(task_uid, f"Total rows: {total_rows:,}") + + # Stream processing + first_chunk = True + rows_processed = 0 + + with open(new_file, 'w', encoding='utf-8') as json_file: + json_file.write('[') + + for chunk in _read_csv_chunked(file_path, chunk_size): + json_str = chunk.to_json(orient='records', force_ascii=False, indent=None) + json_str = json_str[1:-1] # Remove [ ] + + if json_str: + if not first_chunk: + json_file.write(',') + json_file.write(json_str) + first_chunk = False + + rows_processed += len(chunk) + progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 + log_task_info( + task_uid, + f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" + ) + + del chunk + + 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}") + # Rollback + 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)} + + if file_path.lower().endswith((".xlsx", ".xls")): + 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") + + try: + from openpyxl import load_workbook + + # First pass: get sheet names + wb_info = load_workbook(file_path, read_only=True, data_only=True) + sheet_names = wb_info.sheetnames + sheet_count = len(sheet_names) + wb_info.close() + + log_task_info(task_uid, f"Found {sheet_count} sheet(s) in Excel file") + + result_files = [] + base_name = os.path.splitext(file_path)[0] + + for idx, sheet_name in enumerate(sheet_names, 1): + new_file = None + try: + log_task_info(task_uid, f"Processing sheet {idx}/{sheet_count}: '{sheet_name}'") + + safe_sheet_name = re.sub(r'[<>:"/\\|?*]', '_', sheet_name) + if sheet_count == 1: + new_file = f"{base_name}.json" + else: + new_file = f"{base_name}_{safe_sheet_name}.json" + new_file = _non_conflicting_output_path(new_file) + + # Count total rows using openpyxl + log_task_info(task_uid, f"Counting total rows in sheet '{sheet_name}'...") + wb_count = load_workbook(file_path, read_only=True, data_only=True) + ws_count = wb_count[sheet_name] + total_rows = ws_count.max_row - 1 # Exclude header + wb_count.close() + log_task_info(task_uid, f"Total rows: {total_rows:,}") + + # Stream processing + first_chunk = True + rows_processed = 0 + + with open(new_file, 'w', encoding='utf-8') as json_file: + json_file.write('[') + + for chunk in _read_excel_sheet_in_chunks(file_path, sheet_name, chunk_size): + json_str = chunk.to_json(orient='records', force_ascii=False, indent=None) + json_str = json_str[1:-1] + + if json_str: + if not first_chunk: + json_file.write(',') + json_file.write(json_str) + first_chunk = False + + rows_processed += len(chunk) + progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 + log_task_info( + task_uid, + f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" + ) + + del chunk + + json_file.write(']') + + result_files.append(new_file) + log_task_info(task_uid, f"Sheet '{sheet_name}' converted successfully to {new_file}") + + except Exception as sheet_error: + log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") + # Rollback + if new_file and os.path.exists(new_file): + try: + os.remove(new_file) + log_task_info(task_uid, f"Rolled back partial file: {new_file}") + except Exception: + pass + continue + + if len(result_files) > 0: + try: + os.remove(file_path) + except PermissionError: + import time + time.sleep(0.5) + try: + os.remove(file_path) + except Exception as e: + log_task_error(task_uid, f"Warning: Could not delete source file: {e}") + + if len(result_files) == 0: + return {"from": file_path, "to": None, "status": "failure", "error": "No sheets converted"} + + return { + "from": file_path, + "to": result_files[0] if len(result_files) == 1 else result_files, + "to_files": result_files, + "status": "success", + "sheets_count": len(result_files) + } + except Exception as e: + log_task_error(task_uid, f"convert file {file_path} error: {e}") + return {"from": file_path, "to": None, "status": "failure", "error": str(e)} + return None + + +def convert_excel_to_parquet(file_path: str, task_uid, use_streaming: bool = False, chunk_size: int = 50000) -> Optional[Dict[str, str]]: + """ + Convert Excel/CSV to Parquet format. + + Args: + file_path: Input file path + task_uid: Task identifier for logging + use_streaming: Use streaming mode for large files (default: False) + chunk_size: Number of rows per chunk in streaming mode (default: 50000) + + Returns: + Conversion result dictionary + """ + if use_streaming: + return _convert_excel_to_parquet_streaming(file_path, task_uid, chunk_size) + else: + return _convert_excel_to_parquet_legacy(file_path, task_uid) + + +def _convert_excel_to_parquet_legacy(file_path: str, task_uid) -> Optional[Dict[str, str]]: + """Original non-streaming implementation.""" if file_path.lower().endswith(".csv"): log_task_info(task_uid, f"Source file address: {file_path}") try: @@ -286,7 +768,215 @@ 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]]: +def _convert_excel_to_parquet_streaming(file_path: str, task_uid, chunk_size: int) -> Optional[Dict[str, str]]: + """ + Streaming implementation with chunked reading and writing. + + Note: pandas.read_excel() does NOT support chunksize parameter. + We use openpyxl's read_only mode as a workaround. + """ + 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"} + + if file_path.lower().endswith(".csv"): + 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 = _non_conflicting_output_path( + f"{os.path.splitext(file_path)[0]}.parquet" + ) + + # Count total rows + log_task_info(task_uid, "Counting total rows...") + total_rows = _count_csv_rows_fast(file_path) + log_task_info(task_uid, f"Total rows: {total_rows:,}") + + # Stream processing + writer = None + rows_processed = 0 + + for chunk in _read_csv_chunked(file_path, chunk_size): + # Data type processing + for col in chunk.columns: + if chunk[col].dtype == "object": + chunk[col] = chunk[col].astype(str).replace("nan", None) + elif pd.api.types.is_integer_dtype(chunk[col]) and chunk[col].isna().any(): + chunk[col] = chunk[col].astype(str) + + # Convert to Arrow Table + table = pa.Table.from_pandas(chunk) + + # Initialize or append + if writer is None: + writer = pq.ParquetWriter(new_file, table.schema) + writer.write_table(table) + + rows_processed += len(chunk) + progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 + log_task_info( + task_uid, + f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" + ) + + del chunk, table + + 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}") + # Rollback + 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)} + + if file_path.lower().endswith((".xlsx", ".xls")): + 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") + + try: + from openpyxl import load_workbook + + # First pass: get sheet names + wb_info = load_workbook(file_path, read_only=True, data_only=True) + sheet_names = wb_info.sheetnames + sheet_count = len(sheet_names) + wb_info.close() + + log_task_info(task_uid, f"Found {sheet_count} sheet(s) in Excel file") + + result_files = [] + base_name = os.path.splitext(file_path)[0] + + for idx, sheet_name in enumerate(sheet_names, 1): + new_file = None + writer = None + try: + log_task_info(task_uid, f"Processing sheet {idx}/{sheet_count}: '{sheet_name}'") + + safe_sheet_name = re.sub(r'[<>:"/\\|?*]', '_', 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) + + # Count total rows using openpyxl + log_task_info(task_uid, f"Counting total rows in sheet '{sheet_name}'...") + wb_count = load_workbook(file_path, read_only=True, data_only=True) + ws_count = wb_count[sheet_name] + total_rows = ws_count.max_row - 1 # Exclude header + wb_count.close() + log_task_info(task_uid, f"Total rows: {total_rows:,}") + + # Stream processing + rows_processed = 0 + + for chunk in _read_excel_sheet_in_chunks(file_path, sheet_name, chunk_size): + # Data type processing + for col in chunk.columns: + if chunk[col].dtype == "object": + chunk[col] = chunk[col].astype(str).replace("nan", None) + elif pd.api.types.is_integer_dtype(chunk[col]) and chunk[col].isna().any(): + chunk[col] = chunk[col].astype(str) + + # Convert to Arrow Table + table = pa.Table.from_pandas(chunk) + + # Initialize or append + if writer is None: + writer = pq.ParquetWriter(new_file, table.schema) + writer.write_table(table) + + rows_processed += len(chunk) + progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 + log_task_info( + task_uid, + f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" + ) + + del chunk, table + + if writer: + writer.close() + + result_files.append(new_file) + log_task_info(task_uid, f"Sheet '{sheet_name}' converted successfully to {new_file}") + + except Exception as sheet_error: + log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") + # Rollback + if new_file and os.path.exists(new_file): + try: + os.remove(new_file) + log_task_info(task_uid, f"Rolled back partial file: {new_file}") + except Exception: + pass + continue + + if len(result_files) > 0: + try: + os.remove(file_path) + except PermissionError: + import time + time.sleep(0.5) + try: + os.remove(file_path) + except Exception as e: + log_task_error(task_uid, f"Warning: Could not delete source file: {e}") + + if len(result_files) == 0: + return {"from": file_path, "to": None, "status": "failure", "error": "No sheets converted"} + + return { + "from": file_path, + "to": result_files[0] if len(result_files) == 1 else result_files, + "to_files": result_files, + "status": "success", + "sheets_count": len(result_files) + } + except Exception as e: + log_task_error(task_uid, f"convert file {file_path} error: {e}") + return {"from": file_path, "to": None, "status": "failure", "error": str(e)} + return None + + +def convert_csv_to_excel(file_path: str, task_uid, use_streaming: bool = False, chunk_size: int = 50000) -> Optional[Dict[str, str]]: + """ + Convert CSV to Excel format. + + Args: + file_path: Input file path + task_uid: Task identifier for logging + use_streaming: Use streaming mode for large files (default: False) + chunk_size: Number of rows per chunk in streaming mode (default: 50000) + + Returns: + Conversion result dictionary + """ + if use_streaming: + return _convert_csv_to_excel_streaming(file_path, task_uid, chunk_size) + else: + return _convert_csv_to_excel_legacy(file_path, task_uid) + + +def _convert_csv_to_excel_legacy(file_path: str, task_uid) -> Optional[Dict[str, str]]: + """Original non-streaming implementation.""" if not file_path.lower().endswith(".csv"): return None @@ -307,6 +997,95 @@ def convert_csv_to_excel(file_path: str, task_uid) -> Optional[Dict[str, str]]: 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) -> Optional[Dict[str, str]]: + """Streaming implementation using xlsxwriter's constant_memory mode.""" + 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") + + 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"} + + new_file = None + try: + new_file = _non_conflicting_output_path( + f"{os.path.splitext(file_path)[0]}.xlsx" + ) + + # Count total rows + log_task_info(task_uid, "Counting total rows...") + total_rows = _count_csv_rows_fast(file_path) + log_task_info(task_uid, f"Total rows: {total_rows:,}") + + # Create workbook with constant_memory mode + workbook = xlsxwriter.Workbook(new_file, { + 'constant_memory': True, + 'use_zip64': True, + 'strings_to_numbers': False, + 'strings_to_urls': False + }) + worksheet = workbook.add_worksheet() + + # Stream processing + current_row = 0 + header_written = False + rows_processed = 0 + + for chunk in _read_csv_chunked(file_path, chunk_size): + # Write header (first chunk only) + 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 + + # Write data + for _, data_row in chunk.iterrows(): + for col_idx, value in enumerate(data_row): + 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) + progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 + log_task_info( + task_uid, + f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" + ) + + del chunk + + # Finalize workbook + log_task_info(task_uid, "Finalizing Excel file (building ZIP structure)...") + workbook.close() + + log_task_info(task_uid, f"Conversion completed: {new_file}") + + 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}") + # Rollback + if new_file and os.path.exists(new_file): + try: + os.remove(new_file) + log_task_info(task_uid, f"Rolled back partial file: {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'([^<]+)' From 0aa25bf93713e0679e7560bc150703b47b30fce0 Mon Sep 17 00:00:00 2001 From: shenren123 <2646799270@qq.com> Date: Fri, 11 Sep 2026 14:40:52 +0800 Subject: [PATCH 12/12] =?UTF-8?q?=E4=BC=98=E5=8C=96=E4=BB=A3=E7=A0=81?= =?UTF-8?q?=E9=80=BB=E8=BE=91=EF=BC=8C=E5=8E=BB=E9=99=A4excl=E6=B5=81?= =?UTF-8?q?=E5=BC=8F=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data_server/formatify/FormatifyManager.py | 22 + data_server/formatify/schemas.py | 4 + data_server/pod/common_tasks.py | 91 ++- data_server/pod/formatify_helpers.py | 859 ++++------------------ 4 files changed, 237 insertions(+), 739 deletions(-) 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 08597ec..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,13 +1038,7 @@ def _select_convert_func(from_type, to_type): def _run_convert_func(convert_func, file_path: str, task_uid: str, task_params: dict): - """ - Run conversion function with appropriate parameters. - - Supports streaming mode for Excel/CSV conversions: - - use_streaming: bool (default: False) - - chunk_size: int (default: 50000) - """ + # PDF to Markdown has special parameters if convert_func is convert_pdf_to_markdown: return convert_func( @@ -1018,42 +1048,23 @@ def _run_convert_func(convert_func, file_path: str, task_uid: str, task_params: task_params.get("mineru_backend"), ) - # Check if conversion function supports streaming mode - # (Excel/CSV conversions support use_streaming and chunk_size parameters) - streaming_supported_funcs = ( - convert_excel_to_csv, - convert_excel_to_json, - convert_excel_to_parquet, - convert_csv_to_excel, + # 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_supported_funcs: - # Extract streaming parameters from task_params - # Default: use_streaming=True (streaming mode enabled by default for better memory efficiency) - use_streaming = task_params.get("use_streaming", False) + 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 - # Convert to appropriate types - if isinstance(use_streaming, str): - use_streaming = use_streaming.lower() in ("true", "1", "yes") - use_streaming = bool(use_streaming) - - if isinstance(chunk_size, str): - try: - chunk_size = int(chunk_size) - except (ValueError, TypeError): - chunk_size = 50000 - chunk_size = int(chunk_size) if chunk_size else 50000 - - # Call with streaming parameters - return convert_func( - file_path, - task_uid, - use_streaming=use_streaming, - chunk_size=chunk_size - ) + return convert_func(file_path, task_uid, chunk_size=chunk_size) - # Other conversion functions (Word/PPT/TXT/HTML to Markdown) + # 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 6b944bc..51afb50 100644 --- a/data_server/pod/formatify_helpers.py +++ b/data_server/pod/formatify_helpers.py @@ -1,5 +1,5 @@ +import gc import json -import mmap import os import re import shutil @@ -49,26 +49,6 @@ def _read_csv_chunked(file_path: str, chunk_size: int): return pd.read_csv(file_path, sep=None, engine="python", chunksize=chunk_size, iterator=True) -def _count_csv_rows_fast(file_path: str) -> int: - """Fast CSV row counting without loading into memory.""" - try: - # Method 1: Use memory mapping (fastest for large files) - with open(file_path, 'r+b') as f: - mmapped = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) - count = 0 - while mmapped.readline(): - count += 1 - mmapped.close() - return count - 1 # Exclude header - except Exception: - # Method 2: Fallback to simple line counting - count = 0 - with open(file_path, 'rb') as f: - for _ in f: - count += 1 - return count - 1 # Exclude header - - 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): @@ -82,27 +62,7 @@ def _non_conflicting_output_path(file_path: str) -> str: return candidate -def convert_excel_to_csv(file_path: str, task_uid, use_streaming: bool = False, chunk_size: int = 50000) -> Optional[Dict[str, str]]: - """ - Convert Excel/CSV to CSV format. - - Args: - file_path: Input file path - task_uid: Task identifier for logging - use_streaming: Use streaming mode for large files (default: False) - chunk_size: Number of rows per chunk in streaming mode (default: 50000) - - Returns: - Conversion result dictionary - """ - if use_streaming: - return _convert_excel_to_csv_streaming(file_path, task_uid, chunk_size) - else: - return _convert_excel_to_csv_legacy(file_path, task_uid) - - -def _convert_excel_to_csv_legacy(file_path: str, task_uid) -> Optional[Dict[str, str]]: - """Original non-streaming implementation.""" +def convert_excel_to_csv(file_path: str, task_uid) -> Optional[Dict[str, str]]: if file_path.lower().endswith(".csv"): log_task_info(task_uid, f"CSV source file will be copied without conversion: {file_path}") return { @@ -114,14 +74,9 @@ def _convert_excel_to_csv_legacy(file_path: str, task_uid) -> Optional[Dict[str, if file_path.lower().endswith((".xlsx", ".xls")): log_task_info(task_uid, f"Source file address:{file_path}") try: - # Use openpyxl to avoid file locking issues on Windows - from openpyxl import load_workbook - - # Get sheet names - wb_info = load_workbook(file_path, read_only=True, data_only=True) - sheet_names = wb_info.sheetnames + xls = pd.ExcelFile(file_path) + sheet_names = xls.sheet_names sheet_count = len(sheet_names) - wb_info.close() log_task_info(task_uid, f"Found {sheet_count} sheet(s) in Excel file") @@ -139,7 +94,7 @@ def _convert_excel_to_csv_legacy(file_path: str, task_uid) -> Optional[Dict[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) @@ -148,210 +103,7 @@ def _convert_excel_to_csv_legacy(file_path: str, task_uid) -> Optional[Dict[str, log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") continue - # Remove source file - try: - os.remove(file_path) - except PermissionError: - # Retry after short delay for Windows file locking - import time - time.sleep(0.5) - try: - os.remove(file_path) - except Exception as e: - log_task_error(task_uid, f"Warning: Could not delete source file: {e}") - - if len(result_files) == 0: - return {"from": file_path, "to": None, "status": "failure", "error": "No sheets converted"} - - return { - "from": file_path, - "to": result_files[0] if len(result_files) == 1 else result_files, - "to_files": result_files, - "status": "success", - "sheets_count": len(result_files) - } - except Exception as e: - log_task_error(task_uid, f"convert file {file_path} error: {e}") - return {"from": file_path, "to": None, "status": "failure", "error": str(e)} - return None - - -def _read_excel_sheet_in_chunks(file_path: str, sheet_name: str, chunk_size: int): - """ - Read Excel sheet in chunks using openpyxl. - - Note: pandas.read_excel() does NOT support chunksize parameter. - This is a workaround using openpyxl's read_only mode. - - Yields: - DataFrame: Chunk of data as pandas DataFrame - - Important: - This function does NOT load all rows into memory at once. - It yields chunks as it reads, providing true streaming behavior. - """ - wb = None - try: - from openpyxl import load_workbook - except ImportError: - # Fallback: if openpyxl not available, read entire sheet - df = pd.read_excel(file_path, sheet_name=sheet_name) - yield df - return - - try: - # Use read_only mode to save memory - wb = load_workbook(file_path, read_only=True, data_only=True) - ws = wb[sheet_name] - - # Get header (first row) - rows_iter = ws.iter_rows(values_only=True) - header = next(rows_iter, None) - - if header is None: - return - - # Convert header to list and clean up - header = [str(cell) if cell is not None else f"Column_{i}" for i, cell in enumerate(header)] - - # Stream processing: yield chunks without loading all rows - chunk_data = [] - for row in rows_iter: - chunk_data.append(row) - - if len(chunk_data) >= chunk_size: - # Convert to DataFrame and yield - df = pd.DataFrame(chunk_data, columns=header) - yield df - chunk_data = [] - del df - - # Yield remaining data - if chunk_data: - df = pd.DataFrame(chunk_data, columns=header) - yield df - del df - - except Exception as e: - # Fallback to reading entire sheet - df = pd.read_excel(file_path, sheet_name=sheet_name) - yield df - finally: - # Ensure workbook is closed - if wb is not None: - try: - wb.close() - except Exception: - pass - - -def _convert_excel_to_csv_streaming(file_path: str, task_uid, chunk_size: int) -> Optional[Dict[str, str]]: - """ - Streaming implementation with chunked reading and writing. - - Note: pandas.read_excel() does NOT support chunksize parameter. - We use openpyxl's read_only mode as a workaround. - """ - if file_path.lower().endswith(".csv"): - log_task_info(task_uid, f"CSV source file will be copied without conversion: {file_path}") - return { - "from": file_path, - "to": file_path, - "to_files": [file_path], - "status": "success", - } - - if file_path.lower().endswith((".xlsx", ".xls")): - 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") - - try: - # Use openpyxl directly to avoid file handle issues - from openpyxl import load_workbook - - # First pass: get sheet names - wb_info = load_workbook(file_path, read_only=True, data_only=True) - sheet_names = wb_info.sheetnames - sheet_count = len(sheet_names) - wb_info.close() - - log_task_info(task_uid, f"Found {sheet_count} sheet(s) in Excel file") - - result_files = [] - base_name = os.path.splitext(file_path)[0] - - for idx, sheet_name in enumerate(sheet_names, 1): - new_file = None - try: - log_task_info(task_uid, f"Processing sheet {idx}/{sheet_count}: '{sheet_name}'") - - # Generate output filename - safe_sheet_name = re.sub(r'[<>:"/\\|?*]', '_', sheet_name) - if sheet_count == 1: - new_file = f"{base_name}.csv" - else: - new_file = f"{base_name}_{safe_sheet_name}.csv" - new_file = _non_conflicting_output_path(new_file) - - # Count total rows for progress reporting using openpyxl - log_task_info(task_uid, f"Counting total rows in sheet '{sheet_name}'...") - wb_count = load_workbook(file_path, read_only=True, data_only=True) - ws_count = wb_count[sheet_name] - total_rows = ws_count.max_row - 1 # Exclude header - wb_count.close() - log_task_info(task_uid, f"Total rows: {total_rows:,}") - - # Stream processing using openpyxl - first_chunk = True - rows_processed = 0 - - for chunk in _read_excel_sheet_in_chunks(file_path, sheet_name, chunk_size): - chunk.to_csv( - new_file, - mode='w' if first_chunk else 'a', - header=first_chunk, - index=False, - encoding='utf-8-sig' - ) - first_chunk = False - rows_processed += len(chunk) - - # Progress reporting - progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 - log_task_info( - task_uid, - f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" - ) - - # Release memory - del chunk - - result_files.append(new_file) - log_task_info(task_uid, f"Sheet '{sheet_name}' converted successfully to {new_file}") - - except Exception as sheet_error: - log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") - # Rollback: delete partial file - if new_file and os.path.exists(new_file): - try: - os.remove(new_file) - log_task_info(task_uid, f"Rolled back partial file: {new_file}") - except Exception: - pass - continue - - # Remove source file only if at least one sheet succeeded - if len(result_files) > 0: - try: - os.remove(file_path) - except PermissionError: - # File might be locked, try again after a short delay - import time - time.sleep(0.5) - try: - os.remove(file_path) - except Exception as e: - log_task_error(task_uid, f"Warning: Could not delete source file: {e}") + os.remove(file_path) if len(result_files) == 0: return {"from": file_path, "to": None, "status": "failure", "error": "No sheets converted"} @@ -369,27 +121,7 @@ def _convert_excel_to_csv_streaming(file_path: str, task_uid, chunk_size: int) - return None -def convert_excel_to_json(file_path: str, task_uid, use_streaming: bool = False, chunk_size: int = 50000) -> Optional[Dict[str, str]]: - """ - Convert Excel/CSV to JSON format. - - Args: - file_path: Input file path - task_uid: Task identifier for logging - use_streaming: Use streaming mode for large files (default: False) - chunk_size: Number of rows per chunk in streaming mode (default: 50000) - - Returns: - Conversion result dictionary - """ - if use_streaming: - return _convert_excel_to_json_streaming(file_path, task_uid, chunk_size) - else: - return _convert_excel_to_json_legacy(file_path, task_uid) - - -def _convert_excel_to_json_legacy(file_path: str, task_uid) -> Optional[Dict[str, str]]: - """Original non-streaming implementation.""" +def convert_excel_to_json(file_path: str, task_uid) -> Optional[Dict[str, str]]: if file_path.lower().endswith(".csv"): log_task_info(task_uid, f"Source file address: {file_path}") try: @@ -455,200 +187,7 @@ def _convert_excel_to_json_legacy(file_path: str, task_uid) -> Optional[Dict[str return None -def _convert_excel_to_json_streaming(file_path: str, task_uid, chunk_size: int) -> Optional[Dict[str, str]]: - """ - Streaming implementation with chunked reading and writing. - - Note: pandas.read_excel() does NOT support chunksize parameter. - We use openpyxl's read_only mode as a workaround. - """ - if file_path.lower().endswith(".csv"): - 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 = _non_conflicting_output_path( - f"{os.path.splitext(file_path)[0]}.json" - ) - - # Count total rows - log_task_info(task_uid, "Counting total rows...") - total_rows = _count_csv_rows_fast(file_path) - log_task_info(task_uid, f"Total rows: {total_rows:,}") - - # Stream processing - first_chunk = True - rows_processed = 0 - - with open(new_file, 'w', encoding='utf-8') as json_file: - json_file.write('[') - - for chunk in _read_csv_chunked(file_path, chunk_size): - json_str = chunk.to_json(orient='records', force_ascii=False, indent=None) - json_str = json_str[1:-1] # Remove [ ] - - if json_str: - if not first_chunk: - json_file.write(',') - json_file.write(json_str) - first_chunk = False - - rows_processed += len(chunk) - progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 - log_task_info( - task_uid, - f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" - ) - - del chunk - - 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}") - # Rollback - 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)} - - if file_path.lower().endswith((".xlsx", ".xls")): - 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") - - try: - from openpyxl import load_workbook - - # First pass: get sheet names - wb_info = load_workbook(file_path, read_only=True, data_only=True) - sheet_names = wb_info.sheetnames - sheet_count = len(sheet_names) - wb_info.close() - - log_task_info(task_uid, f"Found {sheet_count} sheet(s) in Excel file") - - result_files = [] - base_name = os.path.splitext(file_path)[0] - - for idx, sheet_name in enumerate(sheet_names, 1): - new_file = None - try: - log_task_info(task_uid, f"Processing sheet {idx}/{sheet_count}: '{sheet_name}'") - - safe_sheet_name = re.sub(r'[<>:"/\\|?*]', '_', sheet_name) - if sheet_count == 1: - new_file = f"{base_name}.json" - else: - new_file = f"{base_name}_{safe_sheet_name}.json" - new_file = _non_conflicting_output_path(new_file) - - # Count total rows using openpyxl - log_task_info(task_uid, f"Counting total rows in sheet '{sheet_name}'...") - wb_count = load_workbook(file_path, read_only=True, data_only=True) - ws_count = wb_count[sheet_name] - total_rows = ws_count.max_row - 1 # Exclude header - wb_count.close() - log_task_info(task_uid, f"Total rows: {total_rows:,}") - - # Stream processing - first_chunk = True - rows_processed = 0 - - with open(new_file, 'w', encoding='utf-8') as json_file: - json_file.write('[') - - for chunk in _read_excel_sheet_in_chunks(file_path, sheet_name, chunk_size): - json_str = chunk.to_json(orient='records', force_ascii=False, indent=None) - json_str = json_str[1:-1] - - if json_str: - if not first_chunk: - json_file.write(',') - json_file.write(json_str) - first_chunk = False - - rows_processed += len(chunk) - progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 - log_task_info( - task_uid, - f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" - ) - - del chunk - - json_file.write(']') - - result_files.append(new_file) - log_task_info(task_uid, f"Sheet '{sheet_name}' converted successfully to {new_file}") - - except Exception as sheet_error: - log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") - # Rollback - if new_file and os.path.exists(new_file): - try: - os.remove(new_file) - log_task_info(task_uid, f"Rolled back partial file: {new_file}") - except Exception: - pass - continue - - if len(result_files) > 0: - try: - os.remove(file_path) - except PermissionError: - import time - time.sleep(0.5) - try: - os.remove(file_path) - except Exception as e: - log_task_error(task_uid, f"Warning: Could not delete source file: {e}") - - if len(result_files) == 0: - return {"from": file_path, "to": None, "status": "failure", "error": "No sheets converted"} - - return { - "from": file_path, - "to": result_files[0] if len(result_files) == 1 else result_files, - "to_files": result_files, - "status": "success", - "sheets_count": len(result_files) - } - except Exception as e: - log_task_error(task_uid, f"convert file {file_path} error: {e}") - return {"from": file_path, "to": None, "status": "failure", "error": str(e)} - return None - - -def convert_excel_to_parquet(file_path: str, task_uid, use_streaming: bool = False, chunk_size: int = 50000) -> Optional[Dict[str, str]]: - """ - Convert Excel/CSV to Parquet format. - - Args: - file_path: Input file path - task_uid: Task identifier for logging - use_streaming: Use streaming mode for large files (default: False) - chunk_size: Number of rows per chunk in streaming mode (default: 50000) - - Returns: - Conversion result dictionary - """ - if use_streaming: - return _convert_excel_to_parquet_streaming(file_path, task_uid, chunk_size) - else: - return _convert_excel_to_parquet_legacy(file_path, task_uid) - - -def _convert_excel_to_parquet_legacy(file_path: str, task_uid) -> Optional[Dict[str, str]]: - """Original non-streaming implementation.""" +def convert_excel_to_parquet(file_path: str, task_uid) -> Optional[Dict[str, str]]: if file_path.lower().endswith(".csv"): log_task_info(task_uid, f"Source file address: {file_path}") try: @@ -688,7 +227,7 @@ def _convert_excel_to_parquet_legacy(file_path: str, task_uid) -> Optional[Dict[ 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) @@ -703,18 +242,18 @@ def _convert_excel_to_parquet_legacy(file_path: str, task_uid) -> Optional[Dict[ 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) @@ -724,7 +263,7 @@ def _convert_excel_to_parquet_legacy(file_path: str, task_uid) -> Optional[Dict[ 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 @@ -756,7 +295,7 @@ def _convert_excel_to_parquet_legacy(file_path: str, task_uid) -> Optional[Dict[ "status": "success", "sheets_count": len(result_files) } - + except Exception as e: log_task_error(task_uid, f"convert file {file_path} error: {e}") return { @@ -767,225 +306,147 @@ def _convert_excel_to_parquet_legacy(file_path: str, task_uid) -> Optional[Dict[ } 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 + + log_task_info(task_uid, f"Source file address: {file_path}") + try: + new_file = _non_conflicting_output_path( + f"{os.path.splitext(file_path)[0]}.xlsx" + ) + _read_csv(file_path).to_excel(new_file, index=False, engine="openpyxl") + 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}") + return {"from": file_path, "to": None, "status": "failure", "error": str(e)} -def _convert_excel_to_parquet_streaming(file_path: str, task_uid, chunk_size: int) -> Optional[Dict[str, str]]: +def convert_csv_to_json_streaming(file_path: str, task_uid, chunk_size: int = 50000) -> Optional[Dict[str, str]]: """ - Streaming implementation with chunked reading and writing. + Convert CSV to JSON using streaming mode (low memory footprint). - Note: pandas.read_excel() does NOT support chunksize parameter. - We use openpyxl's read_only mode as a workaround. + 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 """ - 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"} + if not file_path.lower().endswith(".csv"): + return None - if file_path.lower().endswith(".csv"): - 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") + 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" - new_file = None - try: - new_file = _non_conflicting_output_path( - f"{os.path.splitext(file_path)[0]}.parquet" - ) - - # Count total rows - log_task_info(task_uid, "Counting total rows...") - total_rows = _count_csv_rows_fast(file_path) - log_task_info(task_uid, f"Total rows: {total_rows:,}") - - # Stream processing - writer = None - rows_processed = 0 + 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): - # Data type processing - for col in chunk.columns: - if chunk[col].dtype == "object": - chunk[col] = chunk[col].astype(str).replace("nan", None) - elif pd.api.types.is_integer_dtype(chunk[col]) and chunk[col].isna().any(): - chunk[col] = chunk[col].astype(str) - - # Convert to Arrow Table - table = pa.Table.from_pandas(chunk) - - # Initialize or append - if writer is None: - writer = pq.ParquetWriter(new_file, table.schema) - writer.write_table(table) + columns = chunk.columns.tolist() - rows_processed += len(chunk) - progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 - log_task_info( - task_uid, - f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" - ) + # 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, table + del chunk + 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}") - # Rollback - 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)} - - if file_path.lower().endswith((".xlsx", ".xls")): - 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") + json_file.write(']') - try: - from openpyxl import load_workbook - - # First pass: get sheet names - wb_info = load_workbook(file_path, read_only=True, data_only=True) - sheet_names = wb_info.sheetnames - sheet_count = len(sheet_names) - wb_info.close() - - log_task_info(task_uid, f"Found {sheet_count} sheet(s) in Excel file") - - result_files = [] - base_name = os.path.splitext(file_path)[0] - - for idx, sheet_name in enumerate(sheet_names, 1): - new_file = None - writer = None - try: - log_task_info(task_uid, f"Processing sheet {idx}/{sheet_count}: '{sheet_name}'") - - safe_sheet_name = re.sub(r'[<>:"/\\|?*]', '_', 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) - - # Count total rows using openpyxl - log_task_info(task_uid, f"Counting total rows in sheet '{sheet_name}'...") - wb_count = load_workbook(file_path, read_only=True, data_only=True) - ws_count = wb_count[sheet_name] - total_rows = ws_count.max_row - 1 # Exclude header - wb_count.close() - log_task_info(task_uid, f"Total rows: {total_rows:,}") - - # Stream processing - rows_processed = 0 - - for chunk in _read_excel_sheet_in_chunks(file_path, sheet_name, chunk_size): - # Data type processing - for col in chunk.columns: - if chunk[col].dtype == "object": - chunk[col] = chunk[col].astype(str).replace("nan", None) - elif pd.api.types.is_integer_dtype(chunk[col]) and chunk[col].isna().any(): - chunk[col] = chunk[col].astype(str) - - # Convert to Arrow Table - table = pa.Table.from_pandas(chunk) - - # Initialize or append - if writer is None: - writer = pq.ParquetWriter(new_file, table.schema) - writer.write_table(table) - - rows_processed += len(chunk) - progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 - log_task_info( - task_uid, - f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" - ) - - del chunk, table - - if writer: - writer.close() - - result_files.append(new_file) - log_task_info(task_uid, f"Sheet '{sheet_name}' converted successfully to {new_file}") - - except Exception as sheet_error: - log_task_error(task_uid, f"Failed to convert sheet '{sheet_name}': {sheet_error}") - # Rollback - if new_file and os.path.exists(new_file): - try: - os.remove(new_file) - log_task_info(task_uid, f"Rolled back partial file: {new_file}") - except Exception: - pass - continue - - if len(result_files) > 0: - try: - os.remove(file_path) - except PermissionError: - import time - time.sleep(0.5) - try: - os.remove(file_path) - except Exception as e: - log_task_error(task_uid, f"Warning: Could not delete source file: {e}") - - if len(result_files) == 0: - return {"from": file_path, "to": None, "status": "failure", "error": "No sheets converted"} - - return { - "from": file_path, - "to": result_files[0] if len(result_files) == 1 else result_files, - "to_files": result_files, - "status": "success", - "sheets_count": len(result_files) - } - except Exception as e: - log_task_error(task_uid, f"convert file {file_path} error: {e}") - return {"from": file_path, "to": None, "status": "failure", "error": str(e)} - return None + 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(file_path: str, task_uid, use_streaming: bool = False, chunk_size: int = 50000) -> Optional[Dict[str, str]]: +def convert_csv_to_parquet_streaming(file_path: str, task_uid, chunk_size: int = 50000) -> Optional[Dict[str, str]]: """ - Convert CSV to Excel format. + Convert CSV to Parquet using streaming mode (low memory footprint). Args: - file_path: Input file path + file_path: CSV file path task_uid: Task identifier for logging - use_streaming: Use streaming mode for large files (default: False) - chunk_size: Number of rows per chunk in streaming mode (default: 50000) + 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 use_streaming: - return _convert_csv_to_excel_streaming(file_path, task_uid, chunk_size) - else: - return _convert_csv_to_excel_legacy(file_path, task_uid) - - -def _convert_csv_to_excel_legacy(file_path: str, task_uid) -> Optional[Dict[str, str]]: - """Original non-streaming implementation.""" if not file_path.lower().endswith(".csv"): return None - - log_task_info(task_uid, f"Source file address: {file_path}") + try: - new_file = _non_conflicting_output_path( - f"{os.path.splitext(file_path)[0]}.xlsx" - ) - _read_csv(file_path).to_excel(new_file, index=False, engine="openpyxl") + 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, @@ -994,35 +455,42 @@ def _convert_csv_to_excel_legacy(file_path: str, task_uid) -> Optional[Dict[str, } 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) -> Optional[Dict[str, str]]: - """Streaming implementation using xlsxwriter's constant_memory mode.""" +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 - 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") - 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 = _non_conflicting_output_path( - f"{os.path.splitext(file_path)[0]}.xlsx" - ) - - # Count total rows - log_task_info(task_uid, "Counting total rows...") - total_rows = _count_csv_rows_fast(file_path) - log_task_info(task_uid, f"Total rows: {total_rows:,}") + new_file = f"{os.path.splitext(file_path)[0]}.xlsx" - # Create workbook with constant_memory mode workbook = xlsxwriter.Workbook(new_file, { 'constant_memory': True, 'use_zip64': True, @@ -1031,22 +499,21 @@ def _convert_csv_to_excel_streaming(file_path: str, task_uid, chunk_size: int) - }) worksheet = workbook.add_worksheet() - # Stream processing current_row = 0 header_written = False rows_processed = 0 for chunk in _read_csv_chunked(file_path, chunk_size): - # Write header (first chunk only) 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 - # Write data - for _, data_row in chunk.iterrows(): - for col_idx, value in enumerate(data_row): + 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: @@ -1054,20 +521,16 @@ def _convert_csv_to_excel_streaming(file_path: str, task_uid, chunk_size: int) - current_row += 1 rows_processed += len(chunk) - progress = (rows_processed / total_rows * 100) if total_rows > 0 else 0 - log_task_info( - task_uid, - f"Progress: {rows_processed:,}/{total_rows:,} rows ({progress:.1f}%)" - ) del chunk + del chunk_values + + if rows_processed % (chunk_size * 5) == 0: + gc.collect() - # Finalize workbook - log_task_info(task_uid, "Finalizing Excel file (building ZIP structure)...") + log_task_info(task_uid, "Finalizing Excel file...") workbook.close() - log_task_info(task_uid, f"Conversion completed: {new_file}") - return { "from": file_path, "to": new_file, @@ -1076,11 +539,9 @@ def _convert_csv_to_excel_streaming(file_path: str, task_uid, chunk_size: int) - } except Exception as e: log_task_error(task_uid, f"convert file {file_path} error: {e}") - # Rollback if new_file and os.path.exists(new_file): try: os.remove(new_file) - log_task_info(task_uid, f"Rolled back partial file: {new_file}") except Exception: pass return {"from": file_path, "to": None, "status": "failure", "error": str(e)}