Source code for zephon.build_index

# Copyright 2025 DatologyAI
# SPDX-License-Identifier: Apache-2.0

"""Build an ``index.json`` for a dataset directory (fast shard discovery).

CLI: ``python -m zephon.build_index <format> <dataset_dir>``.
"""

from __future__ import annotations

import importlib
import sys
from pathlib import Path

from zephon._internal.io.index.index_builder import create_index as _create_index

# A format's index module registers its builder on import but pulls an optional
# dep (pyarrow / vortex), so import it on demand to keep this module dep-free.
_FORMAT_INDEX_MODULES = {
    "jsonl": "zephon._internal.io.index.jsonl_index",
    "parquet": "zephon._internal.io.index.parquet_index",
    "vortex": "zephon._internal.io.index.vortex_index",
}




def build_index(
    format_name: str,
    dataset_dir: str | Path,
    *,
    output_path: str | Path | None = None,
    progress: bool = True,
    max_workers: int = 8,
) -> str:
    """Build an ``index.json`` for a JSONL, Parquet, or Vortex dataset.

    Writes to ``dataset_dir/index.json`` unless *output_path* is given, printing
    progress when *progress* is true. Local paths and remote roots supported by
    Zephon's storage backend are accepted. Returns the output path or URL as a
    string. At most ``max_workers`` shards are inspected concurrently; set it to
    1 for sequential indexing.
    """
    module = _FORMAT_INDEX_MODULES.get(format_name)
    if module is not None:
        importlib.import_module(module)
    return _create_index(
        format_name,
        dataset_dir,
        output_path=output_path,
        progress=progress,
        max_workers=max_workers,
    )



def _cli() -> None:
    if len(sys.argv) != 3:
        formats = ", ".join(sorted(_FORMAT_INDEX_MODULES))
        print("Usage: python -m zephon.build_index <format> <dataset_dir>")
        print()
        print("Creates index.json for fast dataset discovery.")
        print(f"Indexable formats: {formats}")
        sys.exit(1)
    try:
        build_index(sys.argv[1], sys.argv[2])
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    _cli()


__all__ = ["build_index"]