Skip to content

Commit

Permalink
Address last closing commentary
Browse files Browse the repository at this point in the history
  • Loading branch information
mobiusklein committed Dec 6, 2024
1 parent 5c0eb6c commit 5c95279
Showing 1 changed file with 15 additions and 13 deletions.
28 changes: 15 additions & 13 deletions quantmsio/core/fragpipe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand Down

0 comments on commit 5c95279

Please sign in to comment.