util.py

import os
import glob
from typing import List
import pyarrow.fs as pafs
import logging

logger = logging.getLogger(__name__)


def _is_s3_path(path: str) -> bool:
    """Check if a path is a S3 path."""
    return path.startswith("s3://")


def check_path_exists(path: str) -> bool:
    """Check if a path exists (supports both S3 and local paths)."""
    if _is_s3_path(path):
        return check_s3_path_exists(path)
    else:
        return check_local_path_exists(path)


def check_local_path_exists(path: str) -> bool:
    """Check if a local path exists."""
    return os.path.exists(path)


def check_s3_path_exists(path: str) -> bool:
    """Check if a S3 path exists."""
    fs = pafs.S3FileSystem(region="us-east-2")

    # Remove s3:// prefix if present
    if path.startswith("s3://"):
       path = path[5:]

    # Remove trailing slash
    path = path.rstrip("/")

    logger.info(f"Listing parquet files in s3://{path}")

    # Use FileSelector with recursive=True
    selector = pafs.FileSelector(path, recursive=True)
    file_infos = fs.get_file_info(selector)
    return len(file_infos) > 0


def list_parquet_files(path: str) -> List[str]:
    """
    List all parquet files in a directory recursively (supports both GCS and local paths).

    Args:
        path: Path to directory (GCS path like gs://bucket/path/ or local path)

    Returns:
        List of full paths to parquet files
    """
    if _is_s3_path(path):
        return list_s3_parquet_files(path)
    else:
        return list_local_parquet_files(path)


def list_local_parquet_files(path: str) -> List[str]:
    """
    List all parquet files in a local directory recursively.

    Args:
        path: Local path to directory

    Returns:
        List of full local paths to parquet files
    """
    path = path.rstrip("/")
    logger.info(f"Listing parquet files in {path}")

    parquet_files = []

    if os.path.isfile(path):
        # Single file
        if path.endswith('.parquet'):
            parquet_files.append(path)
    elif os.path.isdir(path):
        # Directory - search recursively
        pattern = os.path.join(path, "**", "*.parquet")
        parquet_files = glob.glob(pattern, recursive=True)
    else:
        # Could be a glob pattern
        parquet_files = [f for f in glob.glob(path, recursive=True) if f.endswith('.parquet')]

    logger.info(f"Found {len(parquet_files)} parquet files")
    return parquet_files


def list_s3_parquet_files(path: str) -> List[str]:
    """
    List all parquet files in a S3 directory recursively using PyArrow.

    Args:
        path: S3 path (e.g., s3://bucket/path/)

    Returns:
        List of full S3 paths to parquet files
    """
    # Create filesystem
    fs = pafs.S3FileSystem(region="us-east-2")

    # Remove gs:// prefix if present
    if path.startswith("s3://"):
        path = path[5:]

    # Remove trailing slash
    path = path.rstrip("/")

    logger.info(f"Listing parquet files in s3://{path}")

    # Use FileSelector with recursive=True
    selector = pafs.FileSelector(path, recursive=True)
    file_infos = fs.get_file_info(selector)

    # Filter for parquet files
    parquet_files = []
    for file_info in file_infos:
        if file_info.type == pafs.FileType.File and file_info.path.endswith('.parquet'):
            parquet_files.append(f"s3://{file_info.path}")

    logger.info(f"Found {len(parquet_files)} parquet files")

    return parquet_files