Catalogs and Tables#

def create_namespace(self, identifier: Identifier | str) -> None:
    """Creates a namespace in this catalog.

    Args:
        identifier (Identifier | str): namespace identifier
    """
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    self._create_namespace(identifier)

def create_namespace_if_not_exists(self, identifier: Identifier | str) -> None:
    """Creates a namespace in this catalog if it does not already exist.

    Args:
        identifier (Identifier | str): namespace identifier
    """
    if not self.has_namespace(identifier):
        self.create_namespace(identifier)

def create_table(
    self,
    identifier: Identifier | str,
    source: Schema | DataFrame,
    properties: Properties | None = None,
) -> Table:
    """Creates a table in this catalog.

    Args:
        identifier (Identifier | str): table identifier
        source (Schema | DataFrame): table source object such as a Schema or DataFrame.

    Returns:
        Table: new table instance.
    """
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)
    schema = source.schema() if isinstance(source, DataFrame) else source

    table = self._create_table(identifier, schema, properties)
    if isinstance(source, DataFrame):
        table.append(source)

    return table

def create_table_if_not_exists(
    self,
    identifier: Identifier | str,
    source: Schema | DataFrame,
    properties: Properties | None = None,
) -> Table:
    """Creates a table in this catalog if it does not already exist.

    Args:
        identifier (Identifier | str): table identifier
        source (Schema | DataFrame): table source object such as a Schema or DataFrame.

    Returns:
        Table: the existing table (if exists) or the new table instance.
    """
    if self.has_table(identifier):
        return self.get_table(identifier)
    else:
        return self.create_table(identifier, source, properties)

def drop_namespace(self, identifier: Identifier | str) -> None:
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    self._drop_namespace(identifier)

def drop_table(self, identifier: Identifier | str) -> None:
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    self._drop_table(identifier)

@staticmethod
def from_glue(
    name: str,
    client: object | None = None,
    session: object | None = None,
) -> Catalog:
    """Creates a Daft Catalog backed by the AWS Glue service, with optional client or session.

    Terms:
        - AWS Glue          -> Daft Catalog
        - AWS Glue Database -> Daft Namespace
        - AWS Glue Table    -> Daft Table

    Args:
        name (str): glue database name
        client: optional boto3 client
        session: optional boto3 session

    Returns:
        Catalog: new daft catalog instance backed by AWS Glue.
    """
    try:
        from daft.catalog.__glue import GlueCatalog

        if client is not None and session is not None:
            raise ValueError("Can provide either a client or session but not both.")
        elif client is not None:
            return GlueCatalog.from_client(name, client)
        elif session is not None:
            return GlueCatalog.from_session(name, session)
        else:
            raise ValueError("Must provide either a client or session.")
    except ImportError:
        raise ImportError("AWS Glue support not installed: pip install -U 'daft[aws]'")

@staticmethod
def from_iceberg(catalog: object) -> Catalog:
    """Create a Daft Catalog from a PyIceberg catalog object.

    Args:
        catalog (object): a PyIceberg catalog instance

    Returns:
        Catalog: a new Catalog instance backed by the PyIceberg catalog.

    Examples:
        >>> from pyiceberg.catalog import load_catalog
        >>> iceberg_catalog = load_catalog("my_iceberg_catalog")
        >>> catalog = Catalog.from_iceberg(iceberg_catalog)
    """
    try:
        from daft.catalog.__iceberg import IcebergCatalog

        return IcebergCatalog._from_obj(catalog)
    except ImportError:
        raise ImportError("Iceberg support not installed: pip install -U 'daft[iceberg]'")

@staticmethod
def from_pydict(tables: dict[Identifier | str, object], name: str = "default") -> Catalog:
    """Returns an in-memory catalog from a dictionary of table-like objects.

    The table-like objects can be pydicts, dataframes, or a Table implementation.
    For qualified tables, namespaces are created if necessary.

    Args:
        tables (dict[str,object]): a dictionary of table-like objects (pydicts, dataframes, and tables)

    Returns:
        Catalog: new catalog instance with name 'default'

    Examples:
        >>> import daft
        >>> from daft.catalog import Catalog, Table
        >>>
        >>> dictionary = {"x": [1, 2, 3]}
        >>> dataframe = daft.from_pydict(dictionary)
        >>> table = Table.from_df("temp", dataframe)
        >>>
        >>> catalog = Catalog.from_pydict(
        ...     {
        ...         "R": dictionary,
        ...         "S": dataframe,
        ...         "T": table,
        ...     }
        ... )
        >>> catalog.list_tables()
        ['R', 'S', 'T']

    """
    from daft.catalog.__internal import MemoryCatalog

    catalog = MemoryCatalog._new(name)

    for ident, source in tables.items():
        ident = ident if isinstance(ident, Identifier) else Identifier.from_str(ident)

        # has namespace, create one if it doesn't exist
        if len(ident) > 1:
            namespace = Identifier(*ident[:-1])
            catalog.create_namespace_if_not_exists(namespace)

        df: DataFrame
        if isinstance(source, Table):
            df = source.read()
        elif isinstance(source, DataFrame):
            df = source
        elif isinstance(source, dict):
            df = DataFrame._from_pydict(source)
        else:
            raise ValueError(f"Unsupported table source {type(source)}")

        catalog.create_table(ident, df)

    return catalog

@staticmethod
def from_s3tables(
    table_bucket_arn: str,
    client: object | None = None,
    session: object | None = None,
) -> Catalog:
    """Creates a Daft Catalog from S3 Tables bucket ARN, with optional client or session.

    If neither a boto3 client nor session is provided, the Iceberg REST
    client will be used under the hood.

    Args:
        table_bucket_arn (str): ARN of the S3 Tables bucket
        client (object, optional): a boto3 client
        session (object, optional): a boto3 session

    Returns:
        Catalog: a new Catalog instance backed by S3 Tables.

    Examples:
        >>> arn = "arn:aws:s3:::my-s3tables-bucket"
        >>> catalog = Catalog.from_s3tables(arn)
        >>> catalog.list_tables()
    """
    try:
        from daft.catalog.__s3tables import S3Catalog

        if client is not None and session is not None:
            raise ValueError("Can provide either a client or session but not both.")
        elif client is not None:
            return S3Catalog.from_client(table_bucket_arn, client)
        elif session is not None:
            return S3Catalog.from_session(table_bucket_arn, session)
        else:
            return S3Catalog.from_arn(table_bucket_arn)
    except ImportError:
        raise ImportError("S3 Tables support not installed: pip install -U 'daft[aws]'")

@staticmethod
def from_unity(catalog: object) -> Catalog:
    """Create a Daft Catalog from a Unity Catalog client.

    Args:
        catalog (object): a Unity Catalog client instance

    Returns:
        Catalog: a new Catalog instance backed by the Unity catalog.

    Examples:
        >>> from unity_sdk import UnityCatalogClient
        >>> unity_client = UnityCatalogClient(...)
        >>> catalog = Catalog.from_unity(unity_client)

    """
    try:
        from daft.catalog.__unity import UnityCatalog

        return UnityCatalog._from_obj(catalog)
    except ImportError:
        raise ImportError("Unity support not installed: pip install -U 'daft[unity]'")

def get_table(self, identifier: Identifier | str) -> Table:
    """Get a table by its identifier or raises if the table does not exist.

    Args:
        identifier (Identifier|str): table identifier

    Returns:
        Table: matched table or raises if the table does not exist.
    """
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    return self._get_table(identifier)

def has_namespace(self, identifier: Identifier | str) -> bool:
    """Returns True if the namespace exists, otherwise False."""
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    return self._has_namespace(identifier)

def has_table(self, identifier: Identifier | str) -> bool:
    """Returns True if the table exists, otherwise False."""
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    return self._has_table(identifier)

def list_namespaces(self, pattern: str | None = None) -> list[Identifier]:
    """List namespaces in the catalog which match the given pattern.

    Args:
        pattern (str): pattern to match such as a namespace prefix

    Returns:
        list[Identifier]: list of namespace identifiers matching the pattern.
    """
    return self._list_namespaces(pattern)

def list_tables(self, pattern: str | None = None) -> list[Identifier]:
    """List tables in the catalog which match the given pattern.

    Args:
        pattern (str): pattern to match such as a namespace prefix

    Returns:
        list[str]: list of table identifiers matching the pattern.
    """
    return self._list_tables(pattern)

def read_table(self, identifier: Identifier | str, **options: dict[str, Any]) -> DataFrame:
    """Returns the table as a DataFrame or raises an exception if it does not exist."""
    return self.get_table(identifier).read(**options)

def write_table(
    self,
    identifier: Identifier | str,
    df: DataFrame,
    mode: Literal["append", "overwrite"] = "append",
    **options: dict[str, Any],
) -> None:
    return self.get_table(identifier).write(df, mode=mode, **options)

def drop(self, n: int = 1) -> Identifier:
    """Returns a new Identifier with the first n parts removed.

    Args:
        n (int): Number of parts to drop from the beginning. Defaults to 1.

    Returns:
        Identifier: A new Identifier with the first n parts removed.

    Raises:
        ValueError: If dropping n parts would result in an empty Identifier.
    """
    if n <= 0:
        return Identifier(*self)
    if n >= len(self):
        raise ValueError(f"Cannot drop {n} parts from Identifier with {len(self)} parts")
    parts = tuple(self)
    return Identifier(*parts[n:])

@staticmethod
def from_sql(input: str, normalize: bool = False) -> Identifier:
    """Parses an Identifier from an SQL string, normalizing to lowercase if specified.

    Args:
        input (str): input sql string
        normalize (bool): flag to case-normalize the identifier text

    Returns:
        Identifier: new identifier instance

    Examples:
        >>> from daft.catalog import Identifier
        >>> Identifier.from_sql("namespace.table") == Identifier("namespace", "table")
        >>> Identifier.from_sql('"a.b"') == Identifier('"a.b."')
        >>> Identifier.from_sql('ABC."xYz"', normalize=True) == Identifier("abc", "xYz")

    """
    i = Identifier.__new__(Identifier)
    i._ident = PyIdentifier.from_sql(input, normalize)
    return i

@staticmethod
def from_str(input: str) -> Identifier:
    """Parses an Identifier from a dot-delimited Python string without normalization.

    Args:
        input (str): input identifier string

    Returns:
        Identifier: new identifier instance

    Examples:
        >>> from daft.catalog import Identifier
        >>> Identifier.from_str("namespace.table") == Identifier("namespace", "table")

    """
    return Identifier(*input.split("."))

@abstractmethod
def append(self, df: DataFrame, **options: Any) -> None:
    """Appends the DataFrame to this table.

    Args:
        df (DataFrame): dataframe to append
        **options (Any): additional format-dependent write options
    """

@staticmethod
def from_df(name: str, dataframe: DataFrame) -> Table:
    """Returns a read-only table backed by the DataFrame.

    Args:
        name (str): table name
        dataframe (DataFrame): table source dataframe

    Returns:
        Table: new table instance

    Examples:
        >>> import daft
        >>> from daft.catalog import Table
        >>> Table.from_df("my_table", daft.from_pydict({"x": [1, 2, 3]}))

    """
    from daft.catalog.__internal import MemoryTable

    table = MemoryTable._new(name, dataframe.schema())
    table.append(dataframe)

    return table

@staticmethod
def from_iceberg(table: object) -> Table:
    """Creates a Daft Table instance from an Iceberg table.

    Args:
        table (object): a pyiceberg table

    Returns:
        Table: new daft table instance
    """
    try:
        from daft.catalog.__iceberg import IcebergTable

        return IcebergTable._from_obj(table)
    except ImportError:
        raise ImportError("Iceberg support not installed: pip install -U 'daft[iceberg]'")

@staticmethod
def from_pydict(name: str, data: dict[str, InputListType]) -> Table:
    """Returns a read-only table backed by the given data.

    Args:
        name (str): table table
        data dict[str,object]: keys are column names and the values are python lists, numpy arrays, or arrow arrays.

    Returns:
        DataFrame: new read-only table instance

    Examples:
        >>> from daft.catalog import Table
        >>> table = Table.from_pydict({"foo": [1, 2]})
        >>> table.show()
        ╭───────╮
        │ foo   │
        │ ---   │
        │ Int64 │
        ╞═══════╡
        │ 1     │
        ├╌╌╌╌╌╌╌┤
        │ 2     │
        ╰───────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)

    """
    df = DataFrame._from_pydict(data)
    return Table.from_df(name, df)

@staticmethod
def from_unity(table: object) -> Table:
    """Returns a Daft Table instance from a Unity table.

    Args:
        table (object): unity table instance.
    """
    try:
        from daft.catalog.__unity import UnityTable

        return UnityTable._from_obj(table)
    except ImportError:
        raise ImportError("Unity support not installed: pip install -U 'daft[unity]'")

@abstractmethod
def overwrite(self, df: DataFrame, **options: Any) -> None:
    """Overwrites this table with the given DataFrame.

    Args:
        df (DataFrame): dataframe to overwrite this table with
        **options (Any): additional format-dependent write options
    """

@abstractmethod
def read(self, **options: Any) -> DataFrame:
    """Creates a new DataFrame from this table.

    Args:
        **options (Any): additional format-dependent read options

    Returns:
        DataFrame: new DataFrame instance
    """

@abstractmethod
def schema(self) -> Schema:
    """Returns the table's schema."""

def select(self, *columns: ColumnInputType) -> DataFrame:
    """Creates a new DataFrame from the table applying the provided expressions.

    Args:
        *columns (Expression|str): columns to select from the current DataFrame

    Returns:
        DataFrame: new DataFrame instance with the select columns
    """
    return self.read().select(*columns)

def show(self, n: int = 8) -> None:
    """Shows the first n rows from this table.

    Args:
        n (int): number of rows to show

    Returns:
        None
    """
    return self.read().show(n)

def write(self, df: DataFrame, mode: Literal["append", "overwrite"] = "append", **options: Any) -> None:
    """Writes the DataFrame to this table.

    Args:
        df (DataFrame): datafram to write
        mode (str): write mode such as 'append' or 'overwrite'
        **options (Any): additional format-dependent write options
    """
    if mode == "append":
        return self.append(df, **options)