schema_parser.py

"""Schema parser for converting schema.txt to DeepLake SchemaConfig."""

import re
from dataclasses import dataclass
from typing import List, Tuple


@dataclass
class SchemaConfig:
    """Dynamic schema configuration parsed from schema file."""

    STRING_COLUMNS: List[str]
    INT_COLUMNS: List[str]
    LONG_COLUMNS: List[str]
    FLOAT_COLUMNS: List[str]
    DOUBLE_COLUMNS: List[str]
    BOOL_COLUMNS: List[str]
    TIMESTAMP_COLUMNS: List[str]
    ARRAY_COLUMNS: List[str]  # Simple arrays (e.g., embeddings with primitive types)
    DICT_COLUMNS: List[str]  # Complex arrays with struct/nested types
    EMBEDDING_COLUMN: str
    EMBEDDING_SIZE: int
    INVERTED_COLUMN: str
    UNIQUE_ID_COLUMN: str
    BATCH_SIZE: int = 200_000  # Default batch size for processing

    @property
    def all_columns(self) -> List[str]:
        """Return all column names in processing order."""
        return (
            self.STRING_COLUMNS
            + self.LONG_COLUMNS
            + self.INT_COLUMNS
            + self.FLOAT_COLUMNS
            + self.DOUBLE_COLUMNS
            + self.BOOL_COLUMNS
            + self.TIMESTAMP_COLUMNS
            + self.ARRAY_COLUMNS
            + self.DICT_COLUMNS
        )


class SchemaParser:
    """Parses schema.txt files in Spark schema format."""

    TYPE_MAPPING = {
        "string": "STRING",
        "long": "LONG",
        "integer": "INT",
        "float": "FLOAT",
        "double": "DOUBLE",
        "boolean": "BOOL",
        "timestamp": "TIMESTAMP",
    }

    def __init__(
        self,
        schema_path: str,
        embedding_column: str = "title_emb",
        embedding_size: int = 768,
        inverted_column: str = "title1",
        unique_id_column: str = "video_id",
    ):
        """
        Initialize schema parser.

        Args:
            schema_path: Path to schema.txt file
            embedding_column: Name of the embedding column
            embedding_size: Size of embedding vectors
            inverted_column: Name of the column to index with inverted index
            unique_id_column: Name of the column to use as unique ID
        """
        self.schema_path = schema_path
        self.embedding_column = embedding_column
        self.embedding_size = embedding_size
        self.inverted_column = inverted_column
        self.unique_id_column = unique_id_column

    def parse(self) -> SchemaConfig:
        """Parse schema file and return SchemaConfig."""
        columns_by_type = {
            "STRING": [],
            "INT": [],
            "LONG": [],
            "FLOAT": [],
            "DOUBLE": [],
            "BOOL": [],
            "TIMESTAMP": [],
            "ARRAY": [],
            "DICT": [],
        }

        with open(self.schema_path, "r") as f:
            lines = f.readlines()

        for i, line in enumerate(lines):
            parsed = self._parse_line(line)
            if parsed:
                column_name, column_type = parsed

                if column_type == "ARRAY":
                    if i + 1 < len(lines):
                        next_line = lines[i + 1]
                        if "struct" in next_line:
                            # Complex array -> treat as dict
                            columns_by_type["DICT"].append(column_name)
                        else:
                            # Simple array (e.g., float elements)
                            columns_by_type["ARRAY"].append(column_name)
                elif column_type in columns_by_type:
                    columns_by_type[column_type].append(column_name)

        return SchemaConfig(
            STRING_COLUMNS=columns_by_type["STRING"],
            INT_COLUMNS=columns_by_type["INT"],
            LONG_COLUMNS=columns_by_type["LONG"],
            FLOAT_COLUMNS=columns_by_type["FLOAT"],
            DOUBLE_COLUMNS=columns_by_type["DOUBLE"],
            BOOL_COLUMNS=columns_by_type["BOOL"],
            TIMESTAMP_COLUMNS=columns_by_type["TIMESTAMP"],
            ARRAY_COLUMNS=columns_by_type["ARRAY"],
            DICT_COLUMNS=columns_by_type["DICT"],
            EMBEDDING_COLUMN=self.embedding_column,
            EMBEDDING_SIZE=self.embedding_size,
            INVERTED_COLUMN=self.inverted_column,
            UNIQUE_ID_COLUMN=self.unique_id_column,
            BATCH_SIZE=200_000,
        )

    def _parse_line(self, line: str) -> Tuple[str, str] | None:
        """
        Parse a single schema line.

        Format: " |-- column_name: type (nullable = true)"

        Returns:
            Tuple of (column_name, type_category) or None if not a column definition
        """
        match = re.match(r"\s*\|--\s+(\w+):\s+(\w+)", line)
        if not match:
            return None

        column_name = match.group(1)
        column_type = match.group(2)

        if column_type == "array":
            return (column_name, "ARRAY")

        type_category = self.TYPE_MAPPING.get(column_type)
        if type_category:
            return (column_name, type_category)

        return None


def parse_schema_file(
    schema_path: str,
    embedding_column: str = "title_emb",
    embedding_size: int = 768,
    inverted_column: str = "title1",
    unique_id_column: str = "video_id",
) -> SchemaConfig:
    """
    Convenience function to parse a schema file.

    Args:
        schema_path: Path to schema.txt file
        embedding_column: Name of the embedding column
        embedding_size: Size of embedding vectors
        inverted_column: Name of the column to index with inverted index
        unique_id_column: Name of the column to use as unique ID

    Returns:
        SchemaConfig object
    """
    parser = SchemaParser(
        schema_path, embedding_column, embedding_size, inverted_column, unique_id_column
    )
    return parser.parse()