Sessions#

def attach(self, object: Catalog | Provider | Table | UDF | DataFrame, alias: str | None = None) -> None:
    """Attaches a known attachable object like a Catalog, Table, UDF, or DataFrame.

    Args:
        object (Catalog|Table|UDF|DataFrame): object which is attachable to a session

    Returns:
        None
    """
    if isinstance(object, Catalog):
        self.attach_catalog(object, alias)
    elif isinstance(object, Provider):
        self.attach_provider(object, alias)
    elif isinstance(object, Table):
        self.attach_table(object, alias)
    elif isinstance(object, UDF):
        self.attach_function(object, alias)
    elif isinstance(object, DataFrame):
        if alias is None:
            raise ValueError("Cannot attach a DataFrame without an alias. Please provide `alias`.")
        self.attach_view(object, alias)
    else:
        raise ValueError(f"Cannot attach object with type {type(object)}")

def attach_catalog(self, catalog: Catalog | object, alias: str | None = None) -> Catalog:
    """Attaches an external catalog to this session.

    Args:
        catalog (object): catalog instance or supported catalog object
        alias (str|None): optional alias for name resolution

    Returns:
        Catalog: new daft catalog instance
    """
    c = catalog if isinstance(catalog, Catalog) else Catalog._from_obj(catalog)
    a = alias if alias else c.name
    self._session.attach_catalog(c, a)
    return c

def attach_function(self, function: UDF, alias: str | None = None) -> None:
    """Attaches a Python function as a UDF in the current session."""
    self._session.attach_function(function, alias)

def attach_provider(self, provider: Provider, alias: str | None = None) -> Provider:
    """Attaches a provider instance to this session.

    Args:
        provider (Provider): provider instance
        alias (str | None): optional alias for name resolution

    Returns:
        Provider: the provider instance
    """
    p = provider  # TODO: support attaching provider-like objects e.g. OpenAI client.
    a = alias if alias else p.name
    self._session.attach_provider(p, a)
    return p

def attach_table(self, table: Table | object, alias: str | None = None) -> Table:
    """Attaches an external table instance to this session.

    Args:
        table (Table | object): table instance or supported table object
        alias (str | None): optional alias for name resolution

    Returns:
        Table: new daft table instance
    """
    t = table if isinstance(table, Table) else Table._from_obj(table)
    a = alias if alias else t.name
    self._session.attach_table(t, a)
    return t

def attach_view(self, view: DataFrame, alias: str) -> Table:
    """Attaches a DataFrame as a non-materialized temporary view for SQL resolution.

    Unlike ``attach_table(Table.from_df(...))``, this does not materialize data into a MemoryTable.
    """
    py_source = self._to_py_table_source(view)
    return self._create_temp_table_with_source(alias, py_source, replace=False)

def create_namespace(self, identifier: Identifier | str) -> None:
    """Creates a namespace in the current catalog."""
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    if resolved := self._resolve_catalog(identifier):
        cat, identifier = resolved
        return cat.create_namespace(identifier)

    if not (catalog := self.current_catalog()):
        raise ValueError("Cannot create a namespace without a current catalog")
    return catalog.create_namespace(identifier)

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

    if resolved := self._resolve_catalog(identifier):
        cat, identifier = resolved
        return cat.create_namespace_if_not_exists(identifier)

    if not (catalog := self.current_catalog()):
        raise ValueError("Cannot create a namespace without a current catalog")
    return catalog.create_namespace_if_not_exists(identifier)

def create_table(self, identifier: Identifier | str, source: Schema | DataFrame, **properties: Any) -> Table:
    """Creates a table in the current catalog.

    If no namespace is specified, the current namespace is used.

    Returns:
        Table: the newly created table instance.
    """
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    if resolved := self._resolve_catalog(identifier):
        cat, identifier = resolved
        return cat.create_table(identifier, source, properties)

    if not (catalog := self.current_catalog()):
        # TODO relax this constraint by joining with the catalog name
        raise ValueError("Cannot create a table without a current catalog")

    if len(identifier) == 1 and (ns := self.current_namespace()):
        identifier = ns + identifier

    return catalog.create_table(identifier, source, properties)

def create_table_if_not_exists(
    self,
    identifier: Identifier | str,
    source: Schema | DataFrame,
    **properties: Any,
) -> Table:
    """Creates a table in the current catalog if it does not already exist.

    If no namespace is specified, the current namespace is used.

    Returns:
        Table: the newly created instance, or the existing table instance.
    """
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    if resolved := self._resolve_catalog(identifier):
        cat, identifier = resolved
        return cat.create_table_if_not_exists(identifier, source, properties)

    if not (catalog := self.current_catalog()):
        # TODO relax this constraint by joining with the catalog name
        raise ValueError("Cannot create a table without a current catalog")

    if len(identifier) == 1 and (ns := self.current_namespace()):
        identifier = ns + identifier

    return catalog.create_table_if_not_exists(identifier, source, properties)

def create_temp_table(self, identifier: str, source: Schema | DataFrame) -> Table:
    """Creates a temp table scoped to this session's lifetime.

    Args:
        identifier (str): table identifier (name)
        source (TableSource|object): table source like a schema or dataframe

    Returns:
        Table: new table instance

    Examples:
        >>> import daft
        >>> from daft.session import Session
        >>> sess = Session()
        >>> sess.create_temp_table("T", daft.from_pydict({"x": [1, 2, 3]}))
        >>> sess.create_temp_table("S", daft.from_pydict({"y": [4, 5, 6]}))
        >>> sess.list_tables()
        [Identifier(''T''), Identifier(''S'')]

    Args:
        identifier (str): table identifier (name)
        source (Schema | DataFrame): table source is either a Schema or Dataframe

    Returns:
        Table: new table instance
    """
    py_source = self._to_py_table_source(source)
    return self._create_temp_table_with_source(identifier, py_source, replace=True)

def create_temp_view(self, identifier: str, view: DataFrame) -> Table:
    """Creates or replaces a non-materialized temporary view from a DataFrame."""
    py_source = self._to_py_table_source(view)
    return self._create_temp_table_with_source(identifier, py_source, replace=True)

def current_catalog(self) -> Catalog | None:
    """Get the session's current catalog or None.

    Returns:
        Catalog: current catalog or None if one is not set
    """
    return self._session.current_catalog()

def current_model(self) -> str | None:
    """Get the session's current model or None.

    Returns:
        str: the session's default model identifier
    """
    return self._session.current_model()

def current_namespace(self) -> Identifier | None:
    """Get the session's current namespace or None.

    Returns:
        Identifier: current namespace or none if one is not set
    """
    ident = self._session.current_namespace()
    return Identifier._from_pyidentifier(ident) if ident else None

def current_provider(self) -> Provider | None:
    """Get the session's current provider or None.

    Returns:
        str: the session's default provider identifier
    """
    return self._session.current_provider()

def detach_catalog(self, alias: str) -> None:
    """Detaches the catalog from this session or raises if the catalog does not exist.

    Args:
        alias (str): catalog alias to detach
    """
    return self._session.detach_catalog(alias)

def detach_function(self, alias: str) -> None:
    """Detaches a Python function as a UDF in the current session."""
    self._session.detach_function(alias)

def detach_provider(self, alias: str) -> None:
    """Detaches the provider from this session or raises if the provider does not exist.

    Args:
        alias (str): provider alias to detach
    """
    return self._session.detach_provider(alias)

def detach_table(self, alias: str) -> None:
    """Detaches the table from this session or raises if the table does not exist.

    Args:
        alias (str): catalog alias to detach
    """
    return self._session.detach_table(alias)

def drop_namespace(self, identifier: Identifier | str) -> None:
    """Drop the given namespace in the current catalog.

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

    if resolved := self._resolve_catalog(identifier):
        cat, identifier = resolved
        return cat.drop_namespace(identifier)

    if not (catalog := self.current_catalog()):
        raise ValueError("Cannot drop a namespace without a current catalog")
    return catalog.drop_namespace(identifier)

def drop_table(self, identifier: Identifier | str) -> None:
    """Drop the given table in the current catalog.

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

    if resolved := self._resolve_catalog(identifier):
        cat, identifier = resolved
        return cat.drop_table(identifier)

    if not (catalog := self.current_catalog()):
        raise ValueError("Cannot drop a table without a current catalog")

    if len(identifier) == 1 and (ns := self.current_namespace()):
        identifier = ns + identifier

    return catalog.drop_table(identifier)

def get_aggregate_function(self, name: str, *args: Expression) -> Expression:
    """Returns an aggregate function expression from the current session.

    Args:
        name (str): aggregate function name as registered by an extension
        *args: Expression arguments to pass to the aggregate function

    Returns:
        Expression: aggregate result expression
    """
    return Expression._from_pyexpr(self._session.get_aggregate_function(name, *[a._expr for a in args]))

def get_catalog(self, identifier: str) -> Catalog:
    """Returns the catalog or raises an exception if it does not exist.

    Args:
        identifier (str): catalog identifier (name)

    Returns:
        Catalog: The catalog object.

    Raises:
        ValueError: If the catalog does not exist.
    """
    return self._session.get_catalog(identifier)

def get_function(self, name: str, *args: Expression) -> Expression:
    """Returns the function from the current session or raises an exception if it does not exist.

    Args:
        name (str): function name as registered by an extension
        *args: Expression arguments to pass to the function

    Returns:
        Expression: result expression
    """
    return Expression._from_pyexpr(self._session.get_function(name, *[a._expr for a in args]))

def get_provider(self, identifier: str) -> Provider:
    """Returns the provider or raises an exception if it does not exist.

    Args:
        identifier (str): provider identifier e.g. "openai", "anthropic", "transformers"

    Returns:
        Provider: The provider object.

    Raises:
        ValueError: If the provider does not exist.
    """
    return self._session.get_provider(identifier)

def get_table(self, identifier: Identifier | str) -> Table:
    """Returns the table or raises an exception if it does not exist.

    Args:
        identifier (Identifier|str): table identifier or identifier string

    Returns:
        Table: The table object.

    Raises:
        ValueError: If the table does not exist.
    """
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)
    return self._session.get_table(identifier._ident)

def has_catalog(self, identifier: str) -> bool:
    """Returns true if a catalog with the given identifier exists."""
    return self._session.has_catalog(identifier)

def has_namespace(self, identifier: Identifier | str) -> bool:
    """Returns true if a namespace with the given identifier exists."""
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    if resolved := self._resolve_catalog(identifier):
        cat, identifier = resolved
        return cat.has_namespace(identifier)

    if not (catalog := self.current_catalog()):
        raise ValueError("Cannot call has_namespace without a current catalog")
    return catalog.has_namespace(identifier)

def has_provider(self, identifier: str) -> bool:
    """Returns true if a provider with the given identifier exists."""
    return self._session.has_provider(identifier)

def has_table(self, identifier: Identifier | str) -> bool:
    """Returns true if a table with the given identifier exists."""
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)
    return self._session.has_table(identifier._ident)

def list_catalogs(self, pattern: str | None = None) -> list[str]:
    """Returns a list of available catalogs matching the pattern.

    This API currently returns a list of catalog names for backwards compatibility.
    In 0.5.0 this API will return a list of Catalog objects.

    Args:
        pattern (str): catalog name pattern

    Returns:
        list[str]: list of available catalog names
    """
    return self._session.list_catalogs(pattern)

def list_namespaces(self, pattern: str | None = None) -> list[Identifier]:
    r"""Returns a list of available namespaces.

    Args:
        - pattern (str, optional): Pattern to match namespace names. Pattern syntax is catalog-dependent:
            - Native/Memory and Postgres catalogs: Use SQL LIKE syntax (`%`, `_`, `\`). Supports qualified patterns like `"cat1.%"` to dispatch to attached catalog `cat1`.
            - Other catalogs: Pattern behavior varies (e.g., prefix matching for Iceberg/S3 Tables, AWS Glue expressions for Glue).

    Returns:
        list[Identifier]: list of available namespaces

    Examples:
        >>> sess.list_namespaces()  # List all namespaces in current catalog
        >>> sess.list_namespaces("ns%")  # Namespaces starting with "ns" (native catalog)
        >>> sess.list_namespaces("cat1.%")  # All namespaces in attached catalog "cat1"
    """
    return [Identifier._from_pyidentifier(i) for i in self._session.list_namespaces(pattern)]

def list_tables(self, pattern: str | None = None) -> list[Identifier]:
    r"""Returns a list of available tables.

    Args:
        - pattern (str, optional): Pattern to match table names. Pattern syntax is catalog-dependent:
            - Native/Memory and Postgres catalogs: Use SQL LIKE syntax (`%`, `_`, `\`). Supports qualified patterns like `"ns1.table%"`.
            - Other catalogs: Pattern behavior varies (e.g., prefix matching for Iceberg/S3 Tables, AWS Glue expressions for Glue).

    Returns:
        list[Identifier]: list of available tables

    Examples:
        >>> sess.list_tables()  # List all tables
        >>> sess.list_tables("table%")  # Tables starting with "table" (native catalog)
        >>> sess.list_tables("ns1.%")  # All tables in namespace "ns1" (native catalog)
    """
    return [Identifier._from_pyidentifier(i) for i in self._session.list_tables(pattern)]

def load_extension(self, extension: str | types.ModuleType | Path) -> None:
    """Load a native extension by module symbol or an explicit file path.

    .. warning::
        This API is experimental and may change in future releases.

    Args:
        extension: A module with a native library or a direct file path to the shared library.
    """
    warnings.warn(
        "Native extensions are experimental and may change in future releases.",
        stacklevel=2,
    )
    if isinstance(extension, str):
        path = extension
    elif isinstance(extension, Path):
        path = str(extension)
    elif isinstance(extension, types.ModuleType):
        path = _get_shared_lib(extension)
    else:
        raise TypeError(f"Expected string, Path, or module, got {type(extension)}")

    # Load the shared library globally so that symbols are visible to other libraries.
    ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)

    # Load the extension into the session on the Rust side.
    self._session.load_extension(path)

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

    Args:
        identifier (Identifier|str): table identifier

    Returns:
        DataFrame:

    Raises:
        ValueError: If the tables does not exist.
    """
    return self.get_table(identifier).read(**options)

def set_catalog(self, identifier: str | None) -> None:
    """Set the given catalog as current_catalog or raises an err if it does not exist.

    Args:
        identifier (str): sets the current catalog

    Raises:
        ValueError: If the catalog does not exist.
    """
    self._session.set_catalog(identifier)

def set_model(self, identifier: str | None) -> None:
    """Set the default model type.

    Args:
        identifier (str | None): model identifier string.
    """
    self._session.set_model(identifier)

def set_namespace(self, identifier: Identifier | str | None) -> None:
    """Set the given namespace as current_namespace for table resolution.

    Args:
        identifier (Identifier | str): namespace identifier
    """
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)
    self._session.set_namespace(identifier._ident if identifier else None)

def set_provider(self, identifier: str | Provider | None, **options: Any) -> None:
    """Set the default model provider with associated options.

    Args:
        identifier (str | Provider | None): provider identifier string, Provider instance, or None.
            If a Provider instance is given, it is attached to the session before being set as current.
        **options (Any): provider specific options such as an API key or retry limit.

    Note:
        If there are no providers, and you give a known provider identifier
        like "openai", then we will create and attach this known provider.
        For example, `daft.set_provider("openai")` works.
    """
    # If a Provider instance is given, attach it and resolve to its name
    if isinstance(identifier, Provider):
        if options:
            raise TypeError(
                "Cannot pass keyword options when providing a Provider instance. Configure the Provider instance directly instead."
            )
        self.attach_provider(identifier)
        identifier = identifier.name

    # consider using @overload on known providers for better type hints
    if identifier is not None and not self._session.has_provider(identifier) and identifier in PROVIDERS:
        # upsert semantic for known providers e.g. daft.set_provider("openai")
        provider = load_provider(identifier, name=None, **options)
        self.attach_provider(provider)
    self._session.set_provider(identifier)

def sql(self, sql: str) -> DataFrame | None:
    """Executes the SQL statement using this session.

    Args:
        sql (str): input SQL statement

    Returns:
        DataFrame: dataframe instance if this was a data statement (DQL, DDL, DML).
    """
    py_sess = self._session
    py_config = get_context().daft_planning_config
    py_object = sql_exec(sql, py_sess, {}, py_config)
    if py_object is None:
        return None
    elif isinstance(py_object, PyBuilder):
        return DataFrame(LogicalPlanBuilder(py_object))
    else:
        raise ValueError(f"Unsupported return type from sql exec: {type(py_object)}")

def use(self, identifier: Identifier | str | None = None) -> None:
    """Use sets the current catalog and namespace."""
    if identifier is None:
        self.set_catalog(None)
        self.set_namespace(None)
        return
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)
    if len(identifier) == 1:
        self.set_catalog(str(identifier[0]))
    else:
        self.set_catalog(str(identifier[0]))
        self.set_namespace(identifier.drop(1))

def write_table(
    self,
    identifier: Identifier | str,
    df: DataFrame,
    mode: Literal["append", "overwrite"] = "append",
    **options: dict[str, Any],
) -> None:
    """Writes the DataFrame to the table specified by the identifier.

    Args:
        identifier (Identifier|str): table identifier
        df (DataFrame): dataframe to write
        mode ("append"|"overwrite"): write mode, defaults to "append"
        **options (dict[str,Any]): additional, format-specific write options
    """
    if isinstance(identifier, str):
        identifier = Identifier.from_str(identifier)

    self._session.get_table(identifier._ident).write(df, mode=mode, **options)