fix: use the cloud catalog when doing save_as_table in a cloud session

Author: german-typedef

src/fenic/_backends/cloud/execution.py

@@ -28,19 +28,24 @@
     SaveToFileExecutionRequest,
     ShowExecutionRequest,
     StartExecutionRequest,
-    TableIdentifier,
+)
+from fenic_cloud.protos.engine.v1.engine_pb2 import (
+    TableIdentifier as TableIdentifierProto,
 )
 from fenic_cloud.protos.engine.v1.engine_pb2_grpc import EngineServiceStub
 
 from fenic._backends.cloud.metrics import get_query_execution_metrics
 from fenic._backends.schema_serde import deserialize_schema, serialize_schema
+from fenic._backends.utils.catalog_utils import TableIdentifier
 from fenic.core._interfaces import BaseExecution
+from fenic.core._logical_plan.plans.sink import TableSink
 from fenic.core._logical_plan.serde import LogicalPlanSerde
 from fenic.core.error import (
     CloudExecutionError,
     CloudSessionError,
     ExecutionError,
     InternalError,
+    PlanError,
     ValidationError,
 )
 from fenic.core.metrics import LMMetrics, PhysicalPlanRepr, QueryMetrics, RMMetrics
@@ -164,16 +169,67 @@ def save_as_table(
     ) -> QueryMetrics:
         """Execute the logical plan and save the result as a table."""
         logger.debug(f"Saving plan {logical_plan} as table: {table_name}")
+
+        if isinstance(logical_plan, TableSink):
+            if not logical_plan.location:
+                raise ValidationError(
+                    f"Cannot save to table '{table_name}' - location is required. "
+                    f"Provide a location.")
+
+        table_identifier = TableIdentifier.from_string(table_name).enrich(
+            self.session_state.catalog.get_current_catalog(),
+            self.session_state.catalog.get_current_database(),
+        )
+
+        # If the table doesn't exist, create it, this has to be done in the user's context.
+        table_identifier_str = str(table_identifier)
+        table_exists = self.session_state.catalog.does_table_exist(table_identifier_str)
+        if table_exists:
+            if mode == "error":
+                raise PlanError(
+                    f"Cannot save to table '{table_name}' - it already exists and mode is 'error'. "
+                    f"Choose a different approach: "
+                    f"1) Use mode='overwrite' to replace the existing table, "
+                    f"2) Use mode='append' to add data to the existing table, "
+                    f"3) Use mode='ignore' to skip saving if table exists, "
+                    f"4) Use a different table name.")
+            if mode == "ignore":
+                logger.warning(f"Table {table_name} already exists, ignoring write.")
+                return QueryMetrics()
+            if mode == "append":
+                saved_schema = self.session_state.catalog.describe_table(table_identifier_str)
+                plan_schema = logical_plan.schema()
+                if saved_schema != plan_schema:
+                    raise PlanError(
+                        f"Cannot append to table '{table_name}' - schema mismatch detected. "
+                        f"The existing table has a different schema than your DataFrame. "
+                        f"Existing schema: {saved_schema} "
+                        f"Your DataFrame schema: {plan_schema} "
+                        f"To fix this: "
+                        f"1) Use mode='overwrite' to replace the table with your DataFrame's schema, "
+                        f"2) Modify your DataFrame to match the existing table's schema, "
+                        f"3) Use a different table name.")
+        else:
+            logger.debug(f"Creating table {table_identifier_str} with location: {logical_plan.location}")
+            # Create the table in the catalog.
+            result =self.session_state.catalog.create_table(
+                table_identifier_str,
+                logical_plan.schema(),
+                location=logical_plan.location,
+                ignore_if_exists=mode == "ignore",
+                file_format="PARQUET")
+            logger.debug(f"Table {table_identifier_str} created with result: {result}")
+
         # TODO (DY): check that current catalog and schema (if specified in table_name) match session state
-        table_identifier = TableIdentifier(
-            catalog=self.session_state.catalog,
-            schema=self.session_state.schema,
+        table_identifier_proto = TableIdentifierProto(
+            catalog=table_identifier.catalog,
+            schema=table_identifier.db,
             table=table_name,
         )
         request = StartExecutionRequest(
             save_as_table=SaveAsTableExecutionRequest(
                 serialized_plan=LogicalPlanSerde.serialize(logical_plan),
-                table_identifier=table_identifier,
+                table_identifier=table_identifier_proto,
                 mode=mode,
             )
         )

src/fenic/_backends/local/physical_plan/sink.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, List, Literal, Tuple
+from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
 
 if TYPE_CHECKING:
     from fenic._backends.local.session_state import LocalSessionState
@@ -80,13 +80,15 @@ def __init__(
         cache_info: CacheInfo,
         session_state: LocalSessionState,
         schema: Schema,
+        location: Optional[str] = None,
     ):
         super().__init__(
             children=[child], cache_info=cache_info, session_state=session_state
         )
         self.table_name = table_name
         self.mode = mode
         self.schema = schema
+        self.location = location
 
     def _execute(self, child_dfs: List[pl.DataFrame]) -> pl.DataFrame:
         if len(child_dfs) != 1:

src/fenic/_backends/local/transpiler/plan_converter.py

@@ -368,6 +368,7 @@ def convert(
                 cache_info=logical.cache_info,
                 session_state=self.session_state,
                 schema=logical.schema(),
+                location=logical.location,
             )
 
         elif isinstance(logical, SQL):

src/fenic/_backends/utils/catalog_utils.py

@@ -103,6 +103,13 @@ def enrich(self, catalog_name: str, db_name: str) -> "TableIdentifier":
             table=self.table,
         )
 
+    def __str__(self) -> str:
+        str_identifier = self.table
+        if self.db:
+            str_identifier = f"{self.db}.{str_identifier}"
+        if self.catalog:
+            str_identifier = f"{self.catalog}.{str_identifier}"
+        return str_identifier
 
 @dataclass(frozen=True)
 class DBIdentifier(BaseIdentifier):
@@ -132,6 +139,12 @@ def enrich(self, catalog_name: str) -> "DBIdentifier":
             return self
         return DBIdentifier(catalog=catalog_name, db=self.db)
 
+    def __str__(self) -> str:
+        str_identifier = self.db
+        if self.catalog:
+            str_identifier = f"{self.catalog}.{str_identifier}"
+        return str_identifier
+
 
 def compare_object_names(object_name_1: str, object_name_2: str) -> bool:
     """Compare two object names, ignoring case."""

src/fenic/api/io/writer.py

@@ -4,7 +4,7 @@
 
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING, Literal, Union
+from typing import TYPE_CHECKING, Literal, Optional, Union
 
 if TYPE_CHECKING:
     from fenic.api.dataframe import DataFrame
@@ -36,6 +36,7 @@ def save_as_table(
         self,
         table_name: str,
         mode: Literal["error", "append", "overwrite", "ignore"] = "error",
+        location: Optional[str] = None,
     ) -> QueryMetrics:
         """Saves the content of the DataFrame as the specified table.
 
@@ -46,6 +47,7 @@ def save_as_table(
                  - append: Appends data to table if it exists
                  - overwrite: Overwrites existing table
                  - ignore: Silently ignores operation if table exists
+            location: location where the table will be saved (only available for cloud execution)
 
         Returns:
             QueryMetrics: The query metrics
@@ -66,7 +68,7 @@ def save_as_table(
             ```
         """
         sink_plan = TableSink(
-            child=self._dataframe._logical_plan, table_name=table_name, mode=mode
+            child=self._dataframe._logical_plan, table_name=table_name, mode=mode, location=location
         )
 
         metrics = self._dataframe._logical_plan.session_state.execution.save_as_table(

src/fenic/core/_logical_plan/plans/sink.py

@@ -1,4 +1,4 @@
-from typing import List, Literal
+from typing import List, Literal, Optional
 
 from fenic.core._logical_plan.plans.base import LogicalPlan
 from fenic.core.error import InternalError
@@ -80,6 +80,7 @@ def __init__(
         child: LogicalPlan,
         table_name: str,
         mode: Literal["error", "append", "overwrite", "ignore"] = "error",
+        location: Optional[str] = None,
     ):
         """Initialize a table sink node.
 
@@ -91,10 +92,12 @@ def __init__(
                  - append: Appends data to table if it exists
                  - overwrite: Overwrites existing table
                  - ignore: Silently ignores operation if table exists
+            location: location where the table will be saved (only available for cloud execution)
         """
         self.child = child
         self.table_name = table_name
         self.mode = mode
+        self.location = location
         super().__init__(self.child.session_state)
 
     def children(self) -> List[LogicalPlan]:
@@ -107,7 +110,7 @@ def _build_schema(self) -> Schema:
 
     def _repr(self) -> str:
         """Return the string representation for this table sink plan."""
-        return f"TableSink(table_name='{self.table_name}', mode='{self.mode}')"
+        return f"TableSink(table_name='{self.table_name}', mode='{self.mode}', location='{self.location}')"
 
     def with_children(self, children: List[LogicalPlan]) -> LogicalPlan:
         """Create a new table sink with the same properties but different children.
@@ -129,4 +132,5 @@ def with_children(self, children: List[LogicalPlan]) -> LogicalPlan:
             child=children[0],
             table_name=self.table_name,
             mode=self.mode,
+            location=self.location,
         )