query

pykmhelpers.pipeline.query

KmindexQuery

Source code in pykmhelpers/pipeline/query.py

class KmindexQuery:
    def __init__(self, path: str = "", sequence: Optional[Sequence] = None) -> None:
        if not path and sequence is None:
            raise ValueError("Either path or sequence must be provided")
        self._sequence = sequence
        self._path = path
        if sequence:
            if path:
                path = Toolbox.get_canonical_path(path)
                if os.path.isfile(path):
                    raise FileExistsError(f"Sequence file already exists: {path}")
                os.makedirs(os.path.dirname(path), exist_ok=True)
                with open(path, "w") as f:
                    f.write(sequence.to_fasta())
        else:
            if not os.path.isfile(path):
                raise FileNotFoundError(f"Query file not found: {path}")

    def execute(
        self,
        registry_path: str,
        output_dir="query",
        index_ids: Optional[list[str]] = None,
        z: int = 6,
        threshold=0.01,
        single_query: Optional[str] = None,
        aggregate: bool = False,
        threads: int = 1,
        fast: bool = True,
        is_compressed: bool = False,
        method: str = "seq",
        vec: bool = False,
    ):
        """Run a query against the kmindex registry.

        Args:
            registry_path (str): Path to the kmindex registry.
            output_dir (str): Output directory for query results.
            index_ids (list[str]): Index IDs to query against; empty list queries all.
            z (int): Z-value (error rate parameter) for kmindex.
            threshold (float): Minimum score threshold for reported hits.
            single_query (str, optional): Query identifier; treats all sequences as one query.
            aggregate (bool): Whether to aggregate batch results into a single file.
            threads (int): Number of threads to use.
            fast (bool): Enable fast mode (disabled automatically when `is_compressed` is True).
            is_compressed (bool): Whether the index is stored in compressed form.
            method (str): Query method passed to kmindex (e.g. ``"seq"``).
            vec (bool): Use ``jsonl_vec`` output format instead of ``jsonl``.
        """
        index_ids = index_ids if index_ids is not None else []
        result_dir = os.path.join(output_dir, KMINDEX_QUERY_OUTPUT)
        os.makedirs(output_dir, exist_ok=True)

        query_path = os.path.join(output_dir, os.path.basename(self._path))
        shutil.copy(self._path, query_path)

        output = KmindexWrapper().query(
            input_registry=registry_path,
            query_file=query_path,
            output_dir=result_dir,
            names=index_ids,
            single_query=single_query,
            aggregate=aggregate,
            threads=threads,
            zvalue=z,
            is_compressed=is_compressed,
            fast=fast and not is_compressed,
            threshold=threshold,
            method=method,
            format="jsonl_vec" if vec else "jsonl",
        )

        # Save result to info.yaml
        info_file = os.path.join(output_dir, "info.yaml")
        with open(info_file, "w") as f:
            yaml.safe_dump(output, f)

        result = []

        for f in os.listdir(result_dir):
            fpath = os.path.join(result_dir, f)
            if os.path.isfile(fpath) and f.endswith(".jsonl"):
                try:
                    result.append(KmindexQueryResult(fpath))
                except Exception as e:
                    logger.warning(f"Could not read result from {fpath}: {e}")

        return result

execute(registry_path, output_dir='query', index_ids=None, z=6, threshold=0.01, single_query=None, aggregate=False, threads=1, fast=True, is_compressed=False, method='seq', vec=False)

Run a query against the kmindex registry.

Parameters:

Name	Type	Description	Default
registry_path	str	Path to the kmindex registry.	required
output_dir	str	Output directory for query results.	'query'
index_ids	list[str]	Index IDs to query against; empty list queries all.	None
z	int	Z-value (error rate parameter) for kmindex.	6
threshold	float	Minimum score threshold for reported hits.	0.01
single_query	str	Query identifier; treats all sequences as one query.	None
aggregate	bool	Whether to aggregate batch results into a single file.	False
threads	int	Number of threads to use.	1
fast	bool	Enable fast mode (disabled automatically when is_compressed is True).	True
is_compressed	bool	Whether the index is stored in compressed form.	False
method	str	Query method passed to kmindex (e.g. "seq").	'seq'
vec	bool	Use jsonl_vec output format instead of jsonl.	False

Source code in pykmhelpers/pipeline/query.py

def execute(
    self,
    registry_path: str,
    output_dir="query",
    index_ids: Optional[list[str]] = None,
    z: int = 6,
    threshold=0.01,
    single_query: Optional[str] = None,
    aggregate: bool = False,
    threads: int = 1,
    fast: bool = True,
    is_compressed: bool = False,
    method: str = "seq",
    vec: bool = False,
):
    """Run a query against the kmindex registry.

    Args:
        registry_path (str): Path to the kmindex registry.
        output_dir (str): Output directory for query results.
        index_ids (list[str]): Index IDs to query against; empty list queries all.
        z (int): Z-value (error rate parameter) for kmindex.
        threshold (float): Minimum score threshold for reported hits.
        single_query (str, optional): Query identifier; treats all sequences as one query.
        aggregate (bool): Whether to aggregate batch results into a single file.
        threads (int): Number of threads to use.
        fast (bool): Enable fast mode (disabled automatically when `is_compressed` is True).
        is_compressed (bool): Whether the index is stored in compressed form.
        method (str): Query method passed to kmindex (e.g. ``"seq"``).
        vec (bool): Use ``jsonl_vec`` output format instead of ``jsonl``.
    """
    index_ids = index_ids if index_ids is not None else []
    result_dir = os.path.join(output_dir, KMINDEX_QUERY_OUTPUT)
    os.makedirs(output_dir, exist_ok=True)

    query_path = os.path.join(output_dir, os.path.basename(self._path))
    shutil.copy(self._path, query_path)

    output = KmindexWrapper().query(
        input_registry=registry_path,
        query_file=query_path,
        output_dir=result_dir,
        names=index_ids,
        single_query=single_query,
        aggregate=aggregate,
        threads=threads,
        zvalue=z,
        is_compressed=is_compressed,
        fast=fast and not is_compressed,
        threshold=threshold,
        method=method,
        format="jsonl_vec" if vec else "jsonl",
    )

    # Save result to info.yaml
    info_file = os.path.join(output_dir, "info.yaml")
    with open(info_file, "w") as f:
        yaml.safe_dump(output, f)

    result = []

    for f in os.listdir(result_dir):
        fpath = os.path.join(result_dir, f)
        if os.path.isfile(fpath) and f.endswith(".jsonl"):
            try:
                result.append(KmindexQueryResult(fpath))
            except Exception as e:
                logger.warning(f"Could not read result from {fpath}: {e}")

    return result

QueryRunnerConfig dataclass

Configuration for a QueryRunner instance.

Attributes:

Name	Type	Description
registry_path	str	Path to the kmindex registry directory.
output_dir	str	Root output directory; per-query subdirectories are created here.
index_ids	list[str]	Index IDs to query against. Empty means all indices.
zvalue	int	Z-value for the findere false-positive filter.
threshold	float	Score threshold applied when filtering results.
threads	int	Number of threads passed to kmindex.
single_query	Optional[str]	When set, all sequences are merged under this identifier.
batch	bool	Concatenate all input files into one query before running.
aggregate	bool	Aggregate batch results into a single output file.
compressed	bool	Whether the index is stored in compressed form.
output_format	str	Output format for result conversion (json, yaml, md, html, tsv).
timestamp	bool	Append a YYYYmmdd_HHMMSS suffix to each per-query output directory.
on_existing	str	Behaviour when the output directory already exists (skip, fail, delete, new-name).
parallel	str	Parallelisation strategy passed to kmindex (seq or sub). Forced to sub when compressed is True.
force	bool	Skip confirmation prompts (e.g. when on_existing="delete").
print_output	bool	Write converted results to stdout instead of saving to file. Only meaningful when format is not json.
on_result	Optional[Callable[[list[KmindexQueryResult]], None]]	Optional callback invoked with each per-query result list as it completes. Useful for streaming results to the caller without waiting for the full run to finish.

Source code in pykmhelpers/pipeline/query.py

@dataclass
class QueryRunnerConfig:
    """Configuration for a ``QueryRunner`` instance.

    Attributes:
        registry_path: Path to the kmindex registry directory.
        output_dir: Root output directory; per-query subdirectories are created here.
        index_ids: Index IDs to query against.  Empty means all indices.
        zvalue: Z-value for the findere false-positive filter.
        threshold: Score threshold applied when filtering results.
        threads: Number of threads passed to kmindex.
        single_query: When set, all sequences are merged under this identifier.
        batch: Concatenate all input files into one query before running.
        aggregate: Aggregate batch results into a single output file.
        compressed: Whether the index is stored in compressed form.
        output_format: Output format for result conversion (``json``, ``yaml``, ``md``, ``html``, ``tsv``).
        timestamp: Append a ``YYYYmmdd_HHMMSS`` suffix to each per-query output directory.
        on_existing: Behaviour when the output directory already exists
            (``skip``, ``fail``, ``delete``, ``new-name``).
        parallel: Parallelisation strategy passed to kmindex (``seq`` or ``sub``).
            Forced to ``sub`` when ``compressed`` is ``True``.
        force: Skip confirmation prompts (e.g. when ``on_existing="delete"``).
        print_output: Write converted results to stdout instead of saving to
            file.  Only meaningful when ``format`` is not ``json``.
        on_result: Optional callback invoked with each per-query result list as
            it completes.  Useful for streaming results to the caller without
            waiting for the full run to finish.
    """

    registry_path: str
    output_dir: str
    index_ids: list[str] = field(default_factory=list)
    zvalue: int = 6
    threshold: float = 0.05
    threads: int = 1
    single_query: Optional[str] = None
    batch: bool = False
    aggregate: bool = False
    compressed: bool = False
    output_format: str = "json"
    timestamp: bool = False
    on_existing: str = "skip"
    parallel: str = "seq"
    force: bool = False
    vec: bool = False
    print_output: bool = False
    on_result: Optional[Callable[[list["KmindexQueryResult"]], None]] = None

QueryRunner

Orchestrates one or more kmindex query operations.

Accepts a list of query file paths (or "-" for stdin), resolves them to concrete files, optionally batches them, and runs each query via KmindexQuery. Output-directory conflict resolution, format conversion, and temp-file cleanup are all handled internally.

Parameters:

Name	Type	Description	Default
config	QueryRunnerConfig	Runtime configuration. See QueryRunnerConfig.	required

Source code in pykmhelpers/pipeline/query.py

class QueryRunner:
    """Orchestrates one or more kmindex query operations.

    Accepts a list of query file paths (or ``"-"`` for stdin), resolves them to
    concrete files, optionally batches them, and runs each query via
    ``KmindexQuery``.  Output-directory conflict resolution, format conversion,
    and temp-file cleanup are all handled internally.

    Args:
        config: Runtime configuration.  See ``QueryRunnerConfig``.
    """

    def __init__(self, config: QueryRunnerConfig) -> None:
        self._config = config
        if self._config.compressed and self._config.parallel != "sub":
            logger.warning(
                "--compressed requires sub parallelization strategy; forcing parallel=sub"
            )
            self._config.parallel = "sub"

    @property
    def config(self) -> QueryRunnerConfig:
        return self._config

    def run(self, query_files: Iterable[str]) -> list[list["KmindexQueryResult"]]:
        """Run queries for all provided input paths.

        Args:
            query_files: Paths to FASTA/FASTQ files or directories.  Pass
                ``"-"`` to read from stdin.

        Returns:
            A list of per-query result lists, in the same order as the resolved
            input files (one entry per executed query).
        """
        os.makedirs(self._config.output_dir, exist_ok=True)

        all_results: list[list[KmindexQueryResult]] = []

        resolved, temp_files = self._resolve_files(query_files)
        errors: list[str] = []
        try:

            if self._config.batch:
                batch_path = os.path.join(tempfile.gettempdir(), "kmhelpers_batch.fa")
                temp_files.append(batch_path)
                with open(batch_path, "wb") as fout:
                    for qfile in resolved:
                        with open(qfile, "rb") as fin:
                            data = fin.read()
                        fout.write(data)
                        if not data.endswith(b"\n"):
                            fout.write(b"\n")
                logger.info(f"Batching {len(resolved)} file(s) into a single query...")
                result = self._run_single_safe(batch_path, total=1, idx=1)
                if result:
                    all_results.append(result)
                else:
                    errors.append(f"{os.path.basename(batch_path)}")
            else:
                total = len(resolved)
                for idx, qfile in enumerate(resolved, 1):
                    result = self._run_single_safe(qfile, total=total, idx=idx)
                    if result:
                        all_results.append(result)
                    else:
                        errors.append(f"{os.path.basename(qfile)}")
        finally:
            for tmp in temp_files:
                try:
                    os.unlink(tmp)
                except OSError:
                    pass

        if errors:
            raise RuntimeError(
                f"{len(errors)} query file(s) failed: " + "; ".join(errors)
            )

        return all_results

    # ---
    # PRIVATE METHODS

    def _resolve_files(self, query_files: Iterable[str]) -> tuple[list[str], list[str]]:
        resolved: list[str] = []
        temp_files: list[str] = []
        for qfile in query_files:
            if qfile == "-":
                tmp = tempfile.NamedTemporaryFile(mode="wb", suffix=".fa", delete=False)
                tmp.write(sys.stdin.buffer.read())
                tmp.close()
                resolved.append(tmp.name)
                temp_files.append(tmp.name)
            elif os.path.isdir(qfile):
                for root, _, files in os.walk(qfile):
                    for fname in sorted(files):
                        if any(fname.endswith(ext) for ext in DATA_EXT):
                            resolved.append(os.path.join(root, fname))
            else:
                if not os.path.isfile(qfile):
                    raise FileNotFoundError(f"Query file not found: {qfile}")
                resolved.append(qfile)
        return resolved, temp_files

    def _run_single_safe(self, qfile: str, total: int, idx: int):
        try:
            result = self._run_single(qfile, total=total, idx=idx)
            return result
        except Exception as e:
            logger.error(f"[{os.path.basename(qfile)}] {e}")
            return None

    def _run_single(
        self, qfile: str, total: int, idx: int
    ) -> list["KmindexQueryResult"]:
        start = time.time()
        cfg = self._config

        stem = os.path.splitext(os.path.basename(qfile))[0]
        query_output = os.path.join(cfg.output_dir, stem)

        if cfg.timestamp:
            query_output = f"{query_output}_{time.strftime('%Y%m%d_%H%M%S')}"

        query_output = self._resolve_existing(query_output, stem)
        if query_output is None:
            return []

        logger.info(f"[{idx}/{total}] Querying: {stem}...")

        kq = KmindexQuery(path=qfile)
        results = kq.execute(
            registry_path=cfg.registry_path,
            output_dir=query_output,
            index_ids=cfg.index_ids,
            z=cfg.zvalue,
            single_query=cfg.single_query,
            aggregate=cfg.aggregate,
            threads=cfg.threads,
            is_compressed=cfg.compressed,
            fast=not cfg.compressed,
            threshold=cfg.threshold,
            method=cfg.parallel,
            vec=cfg.vec,
        )

        elapsed = time.time() - start
        result_dir = os.path.join(query_output, KMINDEX_QUERY_OUTPUT)
        logger.debug(f"kmindex output dir: {result_dir}")
        logger.info(f"Time: {elapsed:.2f}s")

        if cfg.output_format:
            self._convert_results(result_dir)

        if cfg.on_result is not None:
            cfg.on_result(results)

        return results

    def _resolve_existing(self, output_path: str, label: str) -> Optional[str]:
        """Return the (possibly adjusted) output path, or ``None`` to skip."""
        if not os.path.exists(output_path):
            return output_path

        strategy = self._config.on_existing

        if strategy == "skip":
            logger.warning(f"Skipping {label}: output directory already exists")
            return None
        elif strategy == "fail":
            raise FileExistsError(f"Output directory already exists: {output_path}")
        elif strategy == "delete":
            if not self._config.force:
                raise PermissionError(
                    f"Output directory exists and force=False: {output_path}. "
                    "Set force=True to delete automatically."
                )
            logger.debug(f"Deleting existing output directory: {output_path}")
            shutil.rmtree(output_path)
            return output_path
        elif strategy == "new-name":
            new_path = f"{output_path}_{time.strftime('%Y%m%d_%H%M%S')}"
            logger.debug(f"Output directory renamed to: {new_path}")
            return new_path

        return output_path

    def _convert_results(self, result_dir: str) -> None:
        fmt = self._config.output_format
        out_file = os.path.join(os.path.dirname(result_dir), f"results.{fmt}")
        logger.debug(f"Merge results to {out_file}...")
        threshold = self._config.threshold
        merged = KmindexQueryResult()
        for fname in sorted(os.listdir(result_dir)):
            if not fname.endswith(".jsonl"):
                continue
            json_path = os.path.join(result_dir, fname)
            try:
                merged.load_jsonl(json_path)
            except Exception as e:
                logger.warning(f"Failed to read {fname}: {e}")
        if not merged.items:
            logger.info(f"No match")
            return
        converted = merged.convert(format=fmt, threshold=threshold)
        if self._config.print_output:
            sys.stdout.write(f"{converted}\n")
        else:
            with open(out_file, "w") as f:
                f.write(converted)
            logger.info(f"Results: {out_file}")

        # Other formats embed coverage, a matrix TSV cannot hold a second table
        if fmt == "tsv" and merged.has_vectors:
            coverage = merged.generate_coverage_tsv(threshold)
            if self._config.print_output:
                sys.stdout.write(f"{coverage}\n")
            else:
                coverage_file = os.path.join(
                    os.path.dirname(result_dir), "coverage.tsv"
                )
                with open(coverage_file, "w") as f:
                    f.write(coverage)
                logger.info(f"Coverage: {coverage_file}")

run(query_files)

Run queries for all provided input paths.

Parameters:

Name	Type	Description	Default
query_files	Iterable[str]	Paths to FASTA/FASTQ files or directories. Pass "-" to read from stdin.	required

Returns:

Type	Description
list[list[KmindexQueryResult]]	A list of per-query result lists, in the same order as the resolved
list[list[KmindexQueryResult]]	input files (one entry per executed query).

Source code in pykmhelpers/pipeline/query.py

def run(self, query_files: Iterable[str]) -> list[list["KmindexQueryResult"]]:
    """Run queries for all provided input paths.

    Args:
        query_files: Paths to FASTA/FASTQ files or directories.  Pass
            ``"-"`` to read from stdin.

    Returns:
        A list of per-query result lists, in the same order as the resolved
        input files (one entry per executed query).
    """
    os.makedirs(self._config.output_dir, exist_ok=True)

    all_results: list[list[KmindexQueryResult]] = []

    resolved, temp_files = self._resolve_files(query_files)
    errors: list[str] = []
    try:

        if self._config.batch:
            batch_path = os.path.join(tempfile.gettempdir(), "kmhelpers_batch.fa")
            temp_files.append(batch_path)
            with open(batch_path, "wb") as fout:
                for qfile in resolved:
                    with open(qfile, "rb") as fin:
                        data = fin.read()
                    fout.write(data)
                    if not data.endswith(b"\n"):
                        fout.write(b"\n")
            logger.info(f"Batching {len(resolved)} file(s) into a single query...")
            result = self._run_single_safe(batch_path, total=1, idx=1)
            if result:
                all_results.append(result)
            else:
                errors.append(f"{os.path.basename(batch_path)}")
        else:
            total = len(resolved)
            for idx, qfile in enumerate(resolved, 1):
                result = self._run_single_safe(qfile, total=total, idx=idx)
                if result:
                    all_results.append(result)
                else:
                    errors.append(f"{os.path.basename(qfile)}")
    finally:
        for tmp in temp_files:
            try:
                os.unlink(tmp)
            except OSError:
                pass

    if errors:
        raise RuntimeError(
            f"{len(errors)} query file(s) failed: " + "; ".join(errors)
        )

    return all_results