Phase 4: Parquet I/O and crema convert CLI subcommand

crema/__init__.py

@@ -28,19 +28,23 @@
 from .parsers.pepxml import read_pepxml
 from .confidence import TdcConfidence, DuckdbTdcConfidence, assign_confidence
 from .writers.txt import to_txt
+from .writers.parquet import to_parquet
+from .parsers.parquet import read_parquet
 
 __all__ = [
+    "DuckdbTdcConfidence",
     "PsmDataset",
     "TdcConfidence",
-    "DuckdbTdcConfidence",
     "assign_confidence",
     "read_comet",
     "read_msamanda",
     "read_msfragger",
     "read_msgf",
     "read_mztab",
+    "read_parquet",
     "read_pepxml",
     "read_tide",
     "read_txt",
+    "to_parquet",
     "to_txt",
 ]

crema/confidence.py

@@ -13,6 +13,7 @@
 from . import utils
 
 from .writers.txt import to_txt
+from .writers.parquet import to_parquet as _to_parquet
 
 np.random.seed(0)
 
@@ -385,6 +386,34 @@ def to_txt(self, output_dir=None, file_root=None, sep="\t", decoys=False):
             decoys=decoys,
         )
 
+    def to_parquet(self, output_dir=None, file_root=None, decoys=False):
+        """Save confidence estimates to Parquet files.
+
+        Requires the ``pyarrow`` package (``pip install crema[fast]``).
+
+        Parameters
+        ----------
+        output_dir : str or None, optional
+            The directory in which to save the files. ``None`` uses the
+            current working directory.
+        file_root : str or None, optional
+            An optional prefix for the output file names. Files are always
+            named ``[file_root.]crema.{level}.parquet``.
+        decoys : bool, optional
+            Save decoy confidence estimates as well?
+
+        Returns
+        -------
+        list of str
+            The paths to the saved files.
+        """
+        return _to_parquet(
+            self,
+            output_dir=output_dir,
+            file_root=file_root,
+            decoys=decoys,
+        )
+
 
 class TdcConfidence(Confidence):
     """Assign confidence estimates using target decoy competition.

crema/crema.py

@@ -6,6 +6,7 @@
 import sys
 import time
 import logging
+from pathlib import Path
 
 from .parsers.tide import read_tide
 from .parsers.msamanda import read_msamanda
@@ -14,32 +15,40 @@
 from .parsers.comet import read_comet
 from .parsers.mztab import read_mztab
 from .parsers.pepxml import read_pepxml
+from .parsers.txt import read_txt
 from .params import Params
 
 
 def main():
     """The CLI entry point"""
     start_time = time.time()
 
-    # Creates the parser for parse args and reads in command line arguments
     args = Params()
 
-    # Set up logging files
+    if args.command is None:
+        args.parser.print_help()
+        return
+
+    if args.command == "assign-confidence":
+        _run_assign_confidence(args, start_time)
+    elif args.command == "convert":
+        _run_convert(args, start_time)
+
+
+def _setup_logging(output_dir, file_root):
+    """Configure logging to file and stderr."""
     log_file = "crema.log.txt"
-    if args.file_root is not None:
-        log_file = args.file_root + "." + log_file
-    if args.output_dir is None:
-        args.output_dir = os.getcwd()
+    if file_root is not None:
+        log_file = file_root + "." + log_file
+    if output_dir is None:
+        output_dir = os.getcwd()
 
-    # Configure logging
     logging.basicConfig(
-        filename=os.path.join(args.output_dir, log_file),
+        filename=os.path.join(output_dir, log_file),
         filemode="w+",
         level=logging.INFO,
         format="[%(levelname)s] %(message)s",
     )
-
-    # Write logs to stderr as well
     logging.getLogger().addHandler(logging.StreamHandler())
 
     logging.info("crema")
@@ -55,7 +64,19 @@ def main():
     logging.info("Starting Analysis")
     logging.info("=================")
 
-    # Create dataset object
+    return output_dir
+
+
+def _auto_read(psm_files):
+    """Try each parser in turn; return the first that succeeds."""
+    # Parquet is tried first via extension check to avoid slow fallback
+    files = psm_files if isinstance(psm_files, list) else [psm_files]
+    if all(str(f).endswith(".parquet") for f in files):
+        raise ValueError(
+            "Parquet input requires explicit column arguments. "
+            "Use crema.read_parquet() in Python to load Parquet files."
+        )
+
     readers = [
         read_tide,
         read_msgf,
@@ -64,12 +85,13 @@ def main():
         read_msfragger,
         read_pepxml,
         read_mztab,
+        read_txt,
     ]
 
     psms = None
     for read_fn in readers:
         try:
-            psms = read_fn(args.psm_files)
+            psms = read_fn(psm_files)
             break
         except Exception as exc:
             logging.debug("%s failed: %s", read_fn.__name__, exc)
@@ -78,21 +100,90 @@ def main():
     if psms is None:
         raise ValueError("Unrecognized file type.")
 
+    return psms
+
+
+def _run_assign_confidence(args, start_time):
+    """Run the assign-confidence subcommand."""
+    output_dir = _setup_logging(
+        getattr(args, "output_dir", None), getattr(args, "file_root", None)
+    )
+
+    # Check if inputs are Parquet
+    files = args.psm_files
+    if all(str(f).endswith(".parquet") for f in files):
+        raise ValueError(
+            "Parquet input to assign-confidence is not yet supported via the "
+            "CLI. Load the Parquet file with crema.read_parquet() in Python."
+        )
+
+    psms = _auto_read(args.psm_files)
+
+    score = args.score
+    desc = {"True": True, "False": False, "None": None}[args.desc]
     conf = psms.assign_confidence(
-        score_column=args.score,
+        score_column=score,
+        threshold=args.threshold,
+        pep_fdr_type=args.pep_fdr_type,
+        prot_fdr_type=args.prot_fdr_type,
+        desc=desc,
         eval_fdr=args.eval_fdr,
         method=args.method,
     )
 
-    # Write result to file
     logging.info("Writing results...")
-    conf.to_txt(output_dir=args.output_dir, file_root=args.file_root)
+    if getattr(args, "parquet", False):
+        conf.to_parquet(output_dir=output_dir, file_root=args.file_root)
+    else:
+        conf.to_txt(output_dir=output_dir, file_root=args.file_root)
+
+    end_time = time.time()
+    logging.info("==== DONE! =====")
+    logging.info("Wall Time: %.2fs", end_time - start_time)
+
+
+def _run_convert(args, start_time):
+    """Run the convert subcommand."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format="[%(levelname)s] %(message)s",
+    )
+
+    logging.info("crema convert")
+    logging.info("Command issued:")
+    logging.info("%s", " ".join(sys.argv))
+    logging.info("")
+
+    logging.info("Reading PSMs...")
+    psms = read_txt(
+        args.psm_files,
+        target_column=args.target_column,
+        spectrum_columns=args.spectrum_columns,
+        score_columns=args.score_columns,
+        peptide_column=args.peptide_column,
+        protein_column=args.protein_column,
+        protein_delim=args.protein_delim,
+    )
+
+    if args.output is not None:
+        out_path = args.output
+    else:
+        out_path = str(Path(args.psm_files[0]).with_suffix(".parquet"))
+
+    try:
+        import pyarrow  # noqa: F401
+    except ImportError as exc:
+        raise ImportError(
+            "The 'pyarrow' package is required to write Parquet files. "
+            "Install it with: pip install crema[fast]"
+        ) from exc
+
+    logging.info("Writing Parquet to %s...", out_path)
+    psms.data.to_parquet(out_path, index=False)
 
-    # Calculate how long the confidence estimation took
     end_time = time.time()
-    total_time = end_time - start_time
     logging.info("==== DONE! =====")
-    logging.info("Wall Time: %.2fs", total_time)
+    logging.info("Wall Time: %.2fs", end_time - start_time)
 
 
 if __name__ == "__main__":