From 5c95279dac918bf3475126f0c2c4f99b626da5c5 Mon Sep 17 00:00:00 2001 From: Joshua Klein Date: Thu, 5 Dec 2024 22:50:45 -0500 Subject: [PATCH] Address last closing commentary --- quantmsio/core/fragpipe.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/quantmsio/core/fragpipe.py b/quantmsio/core/fragpipe.py index 7fae0d5..7fbe0b3 100644 --- a/quantmsio/core/fragpipe.py +++ b/quantmsio/core/fragpipe.py @@ -143,6 +143,7 @@ def to_arrow_single(self): fields.append({"position": pos, "localization_probability": 1.0}) return {"name": str(self.identifier), "fields": fields} + @classmethod def to_arrow(cls, batch: Iterator[List["MzTabModification"]]): modifications = [] for block in batch: @@ -375,19 +376,20 @@ def write_psms_to_parquet( writer = None file_metadata = [] - - for i, batch in enumerate(self.convert_psms(file_path, batch_size=batch_size)): - logger.debug("Converting batch %d with %d entries", i, batch.num_rows) - if writer is None: - logger.debug("Initializing ParquetWriter with schema %r", batch.schema) - writer = pq.ParquetWriter(output_path, schema=batch.schema, metadata_collector=file_metadata) - writer.add_key_value_metadata(metadata) - - writer.write_batch(batch) - if writer is not None: - writer.close() - else: - logger.warning("No PSMs found. Not writing PSM parquet file") + try: + for i, batch in enumerate(self.convert_psms(file_path, batch_size=batch_size)): + logger.debug("Converting batch %d with %d entries", i, batch.num_rows) + if writer is None: + logger.debug("Initializing ParquetWriter with schema %r", batch.schema) + writer = pq.ParquetWriter(output_path, schema=batch.schema, metadata_collector=file_metadata) + writer.add_key_value_metadata(metadata) + + writer.write_batch(batch) + finally: + if writer is not None: + writer.close() + else: + logger.warning("No PSMs found. Not writing PSM parquet file") return file_metadata def convert_psms(