sqlframe 3.23.0__py3-none-any.whl → 3.24.1__py3-none-any.whl

sqlframe/databricks/catalog.py

@@ -3,36 +3,49 @@
 from __future__ import annotations
 
 import fnmatch
-import json
 import typing as t
 
+import sqlglot as sg
 from sqlglot import exp, parse_one
 
-from sqlframe.base.catalog import Column, Function, _BaseCatalog
+from sqlframe.base.catalog import TABLE, Column, Function, _BaseCatalog
 from sqlframe.base.mixins.catalog_mixins import (
+    CreateTableFromFunctionMixin,
     GetCurrentCatalogFromFunctionMixin,
     GetCurrentDatabaseFromFunctionMixin,
     ListCatalogsFromInfoSchemaMixin,
     ListDatabasesFromInfoSchemaMixin,
     ListTablesFromInfoSchemaMixin,
-    SetCurrentCatalogFromUseMixin,
     SetCurrentDatabaseFromUseMixin,
 )
-from sqlframe.base.util import normalize_string, schema_, to_schema
+from sqlframe.base.types import StructType
+from sqlframe.base.util import (
+    get_column_mapping_from_schema_input,
+    normalize_string,
+    schema_,
+    to_csv,
+    to_schema,
+)
 
 if t.TYPE_CHECKING:
     from sqlframe.databricks.session import DatabricksSession  # noqa
     from sqlframe.databricks.dataframe import DatabricksDataFrame  # noqa
+    from sqlframe.databricks.table import DatabricksTable  # noqa
 
 
 class DatabricksCatalog(
-    GetCurrentCatalogFromFunctionMixin["DatabricksSession", "DatabricksDataFrame"],
-    GetCurrentDatabaseFromFunctionMixin["DatabricksSession", "DatabricksDataFrame"],
-    ListDatabasesFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame"],
-    ListCatalogsFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame"],
-    SetCurrentDatabaseFromUseMixin["DatabricksSession", "DatabricksDataFrame"],
-    ListTablesFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame"],
-    _BaseCatalog["DatabricksSession", "DatabricksDataFrame"],
+    GetCurrentCatalogFromFunctionMixin[
+        "DatabricksSession", "DatabricksDataFrame", "DatabricksTable"
+    ],
+    GetCurrentDatabaseFromFunctionMixin[
+        "DatabricksSession", "DatabricksDataFrame", "DatabricksTable"
+    ],
+    CreateTableFromFunctionMixin["DatabricksSession", "DatabricksDataFrame", "DatabricksTable"],
+    ListDatabasesFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame", "DatabricksTable"],
+    ListCatalogsFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame", "DatabricksTable"],
+    SetCurrentDatabaseFromUseMixin["DatabricksSession", "DatabricksDataFrame", "DatabricksTable"],
+    ListTablesFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame", "DatabricksTable"],
+    _BaseCatalog["DatabricksSession", "DatabricksDataFrame", "DatabricksTable"],
 ):
     CURRENT_CATALOG_EXPRESSION: exp.Expression = exp.func("current_catalog")
     UPPERCASE_INFO_SCHEMA = True
@@ -310,3 +323,164 @@ class DatabricksCatalog(
             )
 
         return columns
+
+    def createTable(
+        self,
+        tableName: str,
+        path: t.Optional[str] = None,
+        source: t.Optional[str] = None,
+        schema: t.Optional[StructType] = None,
+        description: t.Optional[str] = None,
+        **options: str,
+    ) -> DatabricksTable:
+        """Creates a table based on the dataset in a data source.
+
+        .. versionadded:: 2.2.0
+
+        Parameters
+        ----------
+        tableName : str
+            name of the table to create.
+
+            .. versionchanged:: 3.4.0
+               Allow ``tableName`` to be qualified with catalog name.
+
+        path : str, t.Optional
+            the path in which the data for this table exists.
+            When ``path`` is specified, an external table is
+            created from the data at the given path. Otherwise a managed table is created.
+        source : str, t.Optional
+            the source of this table such as 'parquet, 'orc', etc.
+            If ``source`` is not specified, the default data source configured by
+            ``spark.sql.sources.default`` will be used.
+        schema : class:`StructType`, t.Optional
+            the schema for this table.
+        description : str, t.Optional
+            the description of this table.
+
+            .. versionchanged:: 3.1.0
+                Added the ``description`` parameter.
+
+        **options : dict, t.Optional
+            extra options to specify in the table.
+
+        Returns
+        -------
+        :class:`DataFrame`
+            The DataFrame associated with the table.
+
+        Examples
+        --------
+        Creating a managed table.
+
+        >>> _ = spark.catalog.createTable("tbl1", schema=spark.range(1).schema, source='parquet')
+        >>> _ = spark.sql("DROP TABLE tbl1")
+
+        Creating an external table
+
+        >>> import tempfile
+        >>> with tempfile.TemporaryDirectory() as d:
+        ...     _ = spark.catalog.createTable(
+        ...         "tbl2", schema=spark.range(1).schema, path=d, source='parquet')
+        >>> _ = spark.sql("DROP TABLE tbl2")
+        """
+        if not isinstance(tableName, str):
+            raise TypeError("tableName must be a string")
+        if path is not None and not isinstance(path, str):
+            raise TypeError("path must be a string")
+        if source is not None and not isinstance(source, str):
+            raise TypeError("source must be a string")
+        if schema is not None and not isinstance(schema, StructType):
+            raise TypeError("schema must be a StructType")
+        if description is not None and not isinstance(description, str):
+            raise TypeError("description must be a string")
+
+        source = (source or "delta").lower()
+        replace: t.Union[str, bool, None] = options.pop("replace", None)
+        exists: t.Union[str, bool, None] = options.pop("exists", None)
+        table_properties: t.Union[str, t.Dict[str, str]] = options.pop("properties", {})
+        partitionBy: t.Union[t.List[str], str, None] = options.pop("partitionBy", None)
+        clusterBy: t.Union[t.List[str], str, None] = options.pop("clusterBy", None)
+
+        if isinstance(replace, str) and replace.lower() == "true":
+            replace = True
+        if isinstance(exists, str) and exists.lower() == "true":
+            exists = True
+
+        schema_expressions: t.List[exp.Expression] = []
+        if schema is not None and isinstance(schema, StructType):
+            column_mapping = get_column_mapping_from_schema_input(
+                schema, dialect=self.session.input_dialect
+            )
+            schema_expressions = [
+                exp.ColumnDef(
+                    this=exp.parse_identifier(k, dialect=self.session.input_dialect), kind=v
+                )
+                for k, v in column_mapping.items()
+            ]
+
+        name = normalize_string(tableName, from_dialect="input", is_table=True)
+        properties: t.List[exp.Expression] = []
+        if source is not None:
+            properties.append(exp.FileFormatProperty(this=exp.Var(this=source.upper())))
+        if path is not None:
+            properties.append(exp.LocationProperty(this=exp.convert(path)))
+            if replace and source != "delta":
+                replace = None
+                drop_expression = exp.Drop(
+                    this=exp.to_table(name, dialect=self.session.input_dialect),
+                    kind="TABLE",
+                    exists=True,
+                )
+                if self.session._has_connection:
+                    self.session._collect(drop_expression)
+        if description is not None:
+            properties.append(exp.SchemaCommentProperty(this=exp.convert(description)))
+        if partitionBy is not None:
+            if isinstance(partitionBy, str):
+                partition_by = [partitionBy]
+            else:
+                partition_by = partitionBy
+            properties.append(
+                exp.PartitionedByProperty(
+                    this=exp.Tuple(expressions=list(map(sg.to_identifier, partition_by)))
+                )
+            )
+        if clusterBy is not None:
+            if isinstance(clusterBy, str):
+                cluster_by = [clusterBy]
+            else:
+                cluster_by = clusterBy
+            properties.append(
+                exp.Cluster(
+                    expressions=[exp.Tuple(expressions=list(map(sg.to_identifier, cluster_by)))]
+                )
+            )
+
+        properties.extend(
+            exp.Property(this=sg.to_identifier(name=k), value=exp.convert(value=v))
+            for k, v in (table_properties if isinstance(table_properties, dict) else {}).items()
+        )
+
+        format_options: dict[str, t.Union[bool, float, int, str, None]] = {
+            key: f"'{val}'" for key, val in options.items() if val is not None
+        }
+        format_options_str = to_csv(format_options, " ")
+
+        output_expression_container = exp.Create(
+            this=exp.Schema(
+                this=exp.to_table(name, dialect=self.session.input_dialect),
+                expressions=schema_expressions,
+            ),
+            kind="TABLE",
+            exists=exists,
+            replace=replace,
+            properties=exp.Properties(expressions=properties),
+        )
+        if self.session._has_connection:
+            sql = self.session._to_sql(output_expression_container, quote_identifiers=True)
+            sql += f" OPTIONS ({format_options_str})" if format_options_str else ""
+            self.session._collect(sql)
+
+        df = self.session.table(name)
+        return df

sqlframe/databricks/readwriter.py

@@ -6,7 +6,9 @@ import sys
 import typing as t
 
 import sqlglot as sg
+from databricks.sql import ServerOperationError
 from sqlglot import exp
+from sqlglot.helper import ensure_list
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -17,12 +19,21 @@ from sqlframe.base.mixins.readwriter_mixins import PandasLoaderMixin, PandasWrit
 from sqlframe.base.readerwriter import (
     _BaseDataFrameReader,
     _BaseDataFrameWriter,
+    _infer_format,
+)
+from sqlframe.base.util import (
+    ensure_column_mapping,
+    generate_random_identifier,
+    normalize_string,
+    split_filepath,
+    to_csv,
 )
-from sqlframe.base.util import normalize_string
 
 if t.TYPE_CHECKING:
-    from sqlframe.databricks.session import DatabricksSession  # noqa
+    from sqlframe.base._typing import OptionalPrimitiveType, PathOrPaths
+    from sqlframe.base.types import StructType
     from sqlframe.databricks.dataframe import DatabricksDataFrame  # noqa
+    from sqlframe.databricks.session import DatabricksSession  # noqa
     from sqlframe.databricks.table import DatabricksTable  # noqa
 
 
@@ -30,13 +41,232 @@ class DatabricksDataFrameReader(
     PandasLoaderMixin["DatabricksSession", "DatabricksDataFrame"],
     _BaseDataFrameReader["DatabricksSession", "DatabricksDataFrame", "DatabricksTable"],
 ):
-    pass
+    def load(
+        self,
+        path: t.Optional[PathOrPaths] = None,
+        format: t.Optional[str] = None,
+        schema: t.Optional[t.Union[StructType, str]] = None,
+        **options: OptionalPrimitiveType,
+    ) -> DatabricksDataFrame:
+        """Loads data from a data source and returns it as a :class:`DataFrame`.
+
+        .. versionadded:: 1.4.0
+
+        .. versionchanged:: 3.4.0
+            Supports Spark Connect.
+
+        Parameters
+        ----------
+        path : str or list, t.Optional
+            t.Optional string or a list of string for file-system backed data sources.
+        format : str, t.Optional
+            t.Optional string for format of the data source. Default to 'parquet'.
+        schema : :class:`pyspark.sql.types.StructType` or str, t.Optional
+            t.Optional :class:`pyspark.sql.types.StructType` for the input schema
+            or a DDL-formatted string (For example ``col0 INT, col1 DOUBLE``).
+        **options : dict
+            all other string options
+
+        Examples
+        --------
+        Load a CSV file with format, schema and options specified.
+
+        >>> import tempfile
+        >>> with tempfile.TemporaryDirectory() as d:
+        ...     # Write a DataFrame into a CSV file with a header
+        ...     df = spark.createDataFrame([{"age": 100, "name": "Hyukjin Kwon"}])
+        ...     df.write.option("header", True).mode("overwrite").format("csv").save(d)
+        ...
+        ...     # Read the CSV file as a DataFrame with 'nullValue' option set to 'Hyukjin Kwon',
+        ...     # and 'header' option set to `True`.
+        ...     df = spark.read.load(
+        ...         d, schema=df.schema, format="csv", nullValue="Hyukjin Kwon", header=True)
+        ...     df.printSchema()
+        ...     df.show()
+        root
+         |-- age: long (nullable = true)
+         |-- name: string (nullable = true)
+        +---+----+
+        |age|name|
+        +---+----+
+        |100|NULL|
+        +---+----+
+        """
+        assert path is not None, "path is required"
+        assert isinstance(path, str), "path must be a string"
+        format = format or self.state_format_to_read or _infer_format(path)
+        fs_prefix, filepath = split_filepath(path)
+
+        if fs_prefix == "":
+            return super().load(path, format, schema, **options)
+
+        if schema:
+            column_mapping = ensure_column_mapping(schema)
+            select_column_mapping = column_mapping.copy()
+            select_columns = [x.expression for x in self._to_casted_columns(select_column_mapping)]
+
+            if hasattr(schema, "simpleString"):
+                schema = schema.simpleString()
+        else:
+            select_columns = [exp.Star()]
+
+        if format == "delta":
+            from_clause = f"delta.`{fs_prefix + filepath}`"
+        elif format:
+            paths = ",".join([f"{path}" for path in ensure_list(path)])
+
+            format_options: dict[str, OptionalPrimitiveType] = {
+                k: v for k, v in options.items() if v is not None
+            }
+            format_options["format"] = format
+            format_options["schemaEvolutionMode"] = "none"
+            if schema:
+                format_options["schema"] = f"{schema}"
+            if "inferSchema" in format_options:
+                format_options["inferColumnTypes"] = format_options.pop("inferSchema")
+
+            format_options = {key: f"'{val}'" for key, val in format_options.items()}
+            format_options_str = to_csv(format_options, " => ")
+
+            from_clause = f"read_files('{paths}', {format_options_str})"
+        else:
+            from_clause = f"'{path}'"
+
+        df = self.session.sql(
+            exp.select(*select_columns).from_(from_clause, dialect=self.session.input_dialect),
+            qualify=False,
+        )
+        if select_columns == [exp.Star()] and df.schema:
+            return self.load(path=path, format=format, schema=df.schema, **options)
+        self.session._last_loaded_file = path  # type: ignore
+        return df
 
 
 class DatabricksDataFrameWriter(
     PandasWriterMixin["DatabricksSession", "DatabricksDataFrame"],
     _BaseDataFrameWriter["DatabricksSession", "DatabricksDataFrame"],
 ):
+    def save(
+        self,
+        path: str,
+        mode: t.Optional[str] = None,
+        format: t.Optional[str] = None,
+        partitionBy: t.Optional[t.Union[str, t.List[str]]] = None,
+        **options,
+    ):
+        format = str(format or self._state_format_to_write)
+        self._write(path, mode, format, partitionBy=partitionBy, **options)
+
+    def _write(self, path: str, mode: t.Optional[str], format: str, **options):  # type: ignore
+        fs_prefix, filepath = split_filepath(path)
+        if fs_prefix == "":
+            super()._write(filepath, mode, format, **options)
+        elif format == "delta":
+            self.saveAsTable(f"delta.`{fs_prefix + filepath}`", format, mode, **options)
+        else:
+            mode = str(mode or self._mode or "error")
+            partition_by = options.pop("partitionBy", None)
+            tmp_table = f"_{generate_random_identifier()}_tmp"
+            drop_expr = exp.Drop(
+                this=exp.to_table(tmp_table, dialect=self._session.input_dialect),
+                kind="TABLE",
+                exists=True,
+            )
+            if mode == "append" or mode == "default":
+                try:
+                    self._session.catalog.createTable(
+                        tmp_table,
+                        path=fs_prefix + filepath,
+                        source=format,
+                        **options,
+                    )
+                    self.byName.insertInto(tmp_table)
+                except ServerOperationError as e:
+                    if "UNABLE_TO_INFER_SCHEMA" in str(e):
+                        self.saveAsTable(
+                            tmp_table,
+                            format=format,
+                            mode="error",
+                            path=fs_prefix + filepath,
+                            **options,
+                        )
+                    else:
+                        raise e
+                finally:
+                    self._df.session._collect(drop_expr)
+            elif mode == "error" or mode == "errorifexists":
+                try:
+                    self._session.catalog.createTable(
+                        tmp_table,
+                        path=fs_prefix + filepath,
+                        source=format,
+                        **options,
+                    )
+                    raise FileExistsError(f"Path already exists: {fs_prefix + filepath}")
+                except ServerOperationError as e:
+                    if "UNABLE_TO_INFER_SCHEMA" in str(e):
+                        self.saveAsTable(
+                            tmp_table,
+                            format=format,
+                            mode=mode,
+                            path=fs_prefix + filepath,
+                            **options,
+                        )
+                finally:
+                    self._df.session._collect(drop_expr)
+            elif mode == "overwrite":
+                try:
+                    self.saveAsTable(
+                        tmp_table,
+                        format=format,
+                        mode=mode,
+                        path=fs_prefix + filepath,
+                        partitionBy=partition_by,
+                        **options,
+                    )
+                finally:
+                    self._df.session._collect(drop_expr)
+            elif mode == "ignore":
+                pass
+            else:
+                raise RuntimeError(f"Unssuported mode: {mode}")
+
+    def insertInto(
+        self,
+        tableName: str,
+        overwrite: t.Optional[bool] = None,
+        replaceWhere: t.Optional[str] = None,
+    ) -> Self:
+        from sqlframe.base.session import _BaseSession
+
+        tableName = normalize_string(tableName, from_dialect="input", is_table=True)
+        output_expression_container = exp.Insert(
+            **{
+                **{
+                    "this": exp.to_table(tableName, dialect=_BaseSession().input_dialect),
+                    "overwrite": overwrite,
+                },
+                **(
+                    {
+                        "by_name": self._by_name,
+                    }
+                    if self._by_name
+                    else {}
+                ),
+                **({"where": sg.parse_one(replaceWhere)} if replaceWhere else {}),
+            }
+        )
+        df = self._df.copy(output_expression_container=output_expression_container)
+        if self._by_name:
+            columns = self._session.catalog._schema.column_names(
+                tableName, only_visible=True, dialect=_BaseSession().input_dialect
+            )
+            df = df._convert_leaf_to_cte().select(*columns)
+
+        if self._session._has_connection:
+            df.collect()
+        return self.copy(_df=df)
+
     def saveAsTable(
         self,
         name: str,
@@ -44,17 +274,39 @@ class DatabricksDataFrameWriter(
         mode: t.Optional[str] = None,
         partitionBy: t.Optional[t.Union[str, t.List[str]]] = None,
         clusterBy: t.Optional[t.Union[str, t.List[str]]] = None,
-        **options,
-    ) -> Self:
-        if format is not None:
-            raise NotImplementedError("Providing Format in the save as table is not supported")
-        exists, replace, mode = None, None, mode or str(self._mode)
+        **options: OptionalPrimitiveType,
+    ):
+        format = (format or self._state_format_to_write or "delta").lower()
+        table_properties: t.Union[OptionalPrimitiveType, t.Dict[str, OptionalPrimitiveType]] = (
+            options.pop("properties", {})
+        )
+        path: OptionalPrimitiveType = options.pop("path", None)
+        if path is not None and not isinstance(path, str):
+            raise ValueError("path must be a string")
+
+        replace_where: OptionalPrimitiveType = options.pop("replaceWhere", None)
+        if replace_where is not None and not isinstance(replace_where, str):
+            raise ValueError("replaceWhere must be a string")
+
+        exists, replace, mode = None, None, str(mode or self._mode or "error")
         if mode == "append":
-            return self.insertInto(name)
+            self._session.catalog.createTable(
+                name,
+                path=path,
+                source=format,
+                schema=self._df.schema,
+                partitionBy=partitionBy,
+                clusterBy=clusterBy,
+                exists="true",
+                **options,
+            )
+            self.insertInto(name, replaceWhere=replace_where)
+            return
         if mode == "ignore":
             exists = True
         if mode == "overwrite":
             replace = True
+
         name = normalize_string(name, from_dialect="input", is_table=True)
 
         properties: t.List[exp.Expression] = []
@@ -79,9 +331,31 @@ class DatabricksDataFrameWriter(
                 )
             )
 
+        format_options_str = ""
+        if format is not None:
+            properties.append(exp.FileFormatProperty(this=exp.Var(this=format.upper())))
+            format_options: dict[str, OptionalPrimitiveType] = {
+                key: f"'{val}'" for key, val in options.items() if val is not None
+            }
+            format_options_str = to_csv(format_options, " ")
+
+        if path is not None and isinstance(path, str):
+            properties.append(exp.LocationProperty(this=exp.convert(path)))
+            if replace and format != "delta":
+                replace = None
+                drop_expression = exp.Drop(
+                    this=exp.to_table(name, dialect=self._session.input_dialect),
+                    kind="TABLE",
+                    exists=True,
+                )
+                if self._session._has_connection:
+                    self._session._collect(drop_expression)
+
         properties.extend(
             exp.Property(this=sg.to_identifier(name), value=exp.convert(value))
-            for name, value in (options or {}).items()
+            for name, value in (
+                (table_properties if isinstance(table_properties, dict) else {}).items()
+            )
         )
 
         output_expression_container = exp.Create(
@@ -91,7 +365,13 @@ class DatabricksDataFrameWriter(
             replace=replace,
             properties=exp.Properties(expressions=properties),
         )
-        df = self._df.copy(output_expression_container=output_expression_container)
         if self._session._has_connection:
-            df.collect()
-        return self.copy(_df=df)
+            create_sql = self._session._to_sql(output_expression_container, quote_identifiers=True)
+            df_sql = self._df.sql(self._session.execution_dialect, False, False)
+            sql = (
+                create_sql
+                + (f" OPTIONS ({format_options_str})" if format_options_str else "")
+                + " AS "
+                + df_sql
+            )
+            self._session._collect(sql)

sqlframe/duckdb/catalog.py

@@ -9,6 +9,7 @@ from sqlglot import exp
 
 from sqlframe.base.catalog import Function, _BaseCatalog
 from sqlframe.base.mixins.catalog_mixins import (
+    CreateTableFromFunctionMixin,
     GetCurrentCatalogFromFunctionMixin,
     GetCurrentDatabaseFromFunctionMixin,
     ListCatalogsFromInfoSchemaMixin,
@@ -23,18 +24,20 @@ from sqlframe.base.util import normalize_string, schema_, to_schema
 if t.TYPE_CHECKING:
     from sqlframe.duckdb.session import DuckDBSession  # noqa
     from sqlframe.duckdb.dataframe import DuckDBDataFrame  # noqa
+    from sqlframe.duckdb.table import DuckDBTable  # noqa
 
 
 class DuckDBCatalog(
-    GetCurrentCatalogFromFunctionMixin["DuckDBSession", "DuckDBDataFrame"],
-    SetCurrentCatalogFromUseMixin["DuckDBSession", "DuckDBDataFrame"],
-    GetCurrentDatabaseFromFunctionMixin["DuckDBSession", "DuckDBDataFrame"],
-    ListDatabasesFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame"],
-    ListCatalogsFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame"],
-    SetCurrentDatabaseFromUseMixin["DuckDBSession", "DuckDBDataFrame"],
-    ListTablesFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame"],
-    ListColumnsFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame"],
-    _BaseCatalog["DuckDBSession", "DuckDBDataFrame"],
+    GetCurrentCatalogFromFunctionMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    SetCurrentCatalogFromUseMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    GetCurrentDatabaseFromFunctionMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    CreateTableFromFunctionMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    ListDatabasesFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    ListCatalogsFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    SetCurrentDatabaseFromUseMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    ListTablesFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    ListColumnsFromInfoSchemaMixin["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
+    _BaseCatalog["DuckDBSession", "DuckDBDataFrame", "DuckDBTable"],
 ):
     TEMP_CATALOG_FILTER = exp.column("table_catalog").eq("temp")
 

sqlframe/postgres/catalog.py

@@ -9,6 +9,7 @@ from sqlglot import exp, parse_one
 
 from sqlframe.base.catalog import Column, Function, _BaseCatalog
 from sqlframe.base.mixins.catalog_mixins import (
+    CreateTableFromFunctionMixin,
     GetCurrentCatalogFromFunctionMixin,
     GetCurrentDatabaseFromFunctionMixin,
     ListCatalogsFromInfoSchemaMixin,
@@ -21,16 +22,18 @@ from sqlframe.base.util import normalize_string, to_schema
 if t.TYPE_CHECKING:
     from sqlframe.postgres.session import PostgresSession  # noqa
     from sqlframe.postgres.dataframe import PostgresDataFrame  # noqa
+    from sqlframe.postgres.table import PostgresTable  # noqa
 
 
 class PostgresCatalog(
-    GetCurrentCatalogFromFunctionMixin["PostgresSession", "PostgresDataFrame"],
-    GetCurrentDatabaseFromFunctionMixin["PostgresSession", "PostgresDataFrame"],
-    ListDatabasesFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame"],
-    ListCatalogsFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame"],
-    SetCurrentDatabaseFromSearchPathMixin["PostgresSession", "PostgresDataFrame"],
-    ListTablesFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame"],
-    _BaseCatalog["PostgresSession", "PostgresDataFrame"],
+    GetCurrentCatalogFromFunctionMixin["PostgresSession", "PostgresDataFrame", "PostgresTable"],
+    GetCurrentDatabaseFromFunctionMixin["PostgresSession", "PostgresDataFrame", "PostgresTable"],
+    CreateTableFromFunctionMixin["PostgresSession", "PostgresDataFrame", "PostgresTable"],
+    ListDatabasesFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame", "PostgresTable"],
+    ListCatalogsFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame", "PostgresTable"],
+    SetCurrentDatabaseFromSearchPathMixin["PostgresSession", "PostgresDataFrame", "PostgresTable"],
+    ListTablesFromInfoSchemaMixin["PostgresSession", "PostgresDataFrame", "PostgresTable"],
+    _BaseCatalog["PostgresSession", "PostgresDataFrame", "PostgresTable"],
 ):
     CURRENT_CATALOG_EXPRESSION: exp.Expression = exp.column("current_catalog")
     TEMP_SCHEMA_FILTER = exp.column("table_schema").like("pg_temp_%")

sqlframe/py.typed

@@ -0,0 +1 @@
+