snowpark-connect 0.27.0__py3-none-any.whl → 0.28.1__py3-none-any.whl

snowflake/snowpark_connect/relation/read/map_read_table.py

@@ -11,11 +11,17 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     unquote_if_quoted,
 )
 from snowflake.snowpark.exceptions import SnowparkSQLException
+from snowflake.snowpark.types import StructField, StructType
+from snowflake.snowpark_connect.column_name_handler import (
+    ColumnNameMap,
+    make_column_names_snowpark_compatible,
+)
 from snowflake.snowpark_connect.config import auto_uppercase_non_column_identifiers
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.utils import (
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.utils.context import get_processed_views
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
@@ -23,6 +29,7 @@ from snowflake.snowpark_connect.utils.session import _get_current_snowpark_sessi
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
+from snowflake.snowpark_connect.utils.temporary_view_cache import get_temp_view
 
 
 def post_process_df(
@@ -64,15 +71,66 @@ def post_process_df(
         raise
 
 
+def _get_temporary_view(
+    temp_view: DataFrameContainer, table_name: str, plan_id: int
+) -> DataFrameContainer:
+    fields_names = [field.name for field in temp_view.dataframe.schema.fields]
+    fields_types = [field.datatype for field in temp_view.dataframe.schema.fields]
+
+    snowpark_column_names = make_column_names_snowpark_compatible(fields_names, plan_id)
+    # Rename columns in dataframe to prevent conflicting names during joins
+    renamed_df = temp_view.dataframe.select(
+        *(
+            temp_view.dataframe.col(orig).alias(alias)
+            for orig, alias in zip(fields_names, snowpark_column_names)
+        )
+    )
+
+    new_column_map = ColumnNameMap(
+        spark_column_names=temp_view.column_map.get_spark_columns(),
+        snowpark_column_names=snowpark_column_names,
+        column_metadata=temp_view.column_map.column_metadata,
+        column_qualifiers=[split_fully_qualified_spark_name(table_name)]
+        * len(temp_view.column_map.get_spark_columns()),
+        parent_column_name_map=temp_view.column_map.get_parent_column_name_map(),
+    )
+
+    schema = StructType(
+        [
+            StructField(name, type, _is_column=False)
+            for name, type in zip(snowpark_column_names, fields_types)
+        ]
+    )
+    return DataFrameContainer(
+        dataframe=renamed_df,
+        column_map=new_column_map,
+        table_name=temp_view.table_name,
+        alias=temp_view.alias,
+        partition_hint=temp_view.partition_hint,
+        cached_schema_getter=lambda: schema,
+    )
+
+
 def get_table_from_name(
     table_name: str, session: snowpark.Session, plan_id: int
 ) -> DataFrameContainer:
     """Get table from name returning a container."""
+
+    # Verify if recursive view read is not attempted
+    if table_name in get_processed_views():
+        raise AnalysisException(
+            f"[RECURSIVE_VIEW] Recursive view `{table_name}` detected (cycle: `{table_name}` -> `{table_name}`)"
+        )
+
     snowpark_name = ".".join(
         quote_name_without_upper_casing(part)
         for part in split_fully_qualified_spark_name(table_name)
     )
 
+    temp_view = get_temp_view(snowpark_name)
+    if temp_view:
+        return _get_temporary_view(temp_view, table_name, plan_id)
+
     if auto_uppercase_non_column_identifiers():
         snowpark_name = snowpark_name.upper()
 

snowflake/snowpark_connect/relation/write/map_write.py

@@ -50,6 +50,7 @@ from snowflake.snowpark_connect.utils.identifiers import (
     spark_to_sf_single_id,
     split_fully_qualified_spark_name,
 )
+from snowflake.snowpark_connect.utils.io_utils import get_table_type
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
@@ -217,8 +218,9 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 },
                 "overwrite": overwrite,
             }
-            # By default, download from the same prefix we wrote to.
-            download_stage_path = temp_file_prefix_on_stage
+            # Download from the base write path to ensure we fetch whatever Snowflake produced.
+            # Using the base avoids coupling to exact filenames/prefixes.
+            download_stage_path = write_path
 
             # Check for partition hint early to determine precedence over single option
             partition_hint = result.partition_hint
@@ -237,13 +239,19 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     raise SnowparkConnectNotImplementedError(
                         "Partitioning is only supported for parquet format"
                     )
-                partitioning_columns = [f'"{c}"' for c in write_op.partitioning_columns]
-                if len(partitioning_columns) > 1:
-                    raise SnowparkConnectNotImplementedError(
-                        "Multiple partitioning columns are not yet supported"
-                    )
-                else:
-                    parameters["partition_by"] = partitioning_columns[0]
+                # Build Spark-style directory structure: col1=value1/col2=value2/...
+                # Example produced expression (Snowflake SQL):
+                #   'department=' || TO_VARCHAR("department") || '/' || 'region=' || TO_VARCHAR("region")
+                partitioning_column_names = list(write_op.partitioning_columns)
+                partition_expr_parts: list[str] = []
+                for col_name in partitioning_column_names:
+                    quoted = f'"{col_name}"'
+                    segment = f"'{col_name}=' || COALESCE(TO_VARCHAR({quoted}), '__HIVE_DEFAULT_PARTITION__')"
+                    partition_expr_parts.append(segment)
+                parameters["partition_by"] = " || '/' || ".join(partition_expr_parts)
+                # When using PARTITION BY, Snowflake writes into subdirectories under the base path.
+                # Download from the base write path to preserve partition directories locally.
+                download_stage_path = write_path
 
             # If a partition hint is present (from DataFrame.repartition(n)), optionally split the
             # write into n COPY INTO calls by assigning a synthetic partition id. Controlled by config.
@@ -311,7 +319,10 @@ def map_write(request: proto_base.ExecutePlanRequest):
 
             match write_mode:
                 case None | "error" | "errorifexists":
-                    if check_snowflake_table_existence(snowpark_table_name, session):
+                    table_schema_or_error = _get_table_schema_or_error(
+                        snowpark_table_name, session
+                    )
+                    if isinstance(table_schema_or_error, DataType):  # Table exists
                         raise AnalysisException(
                             f"Table {snowpark_table_name} already exists"
                         )
@@ -322,29 +333,45 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         snowpark_session=session,
                     )
                     _validate_schema_and_get_writer(
-                        input_df, "append", snowpark_table_name
+                        input_df, "append", snowpark_table_name, table_schema_or_error
                     ).saveAsTable(
                         table_name=snowpark_table_name,
                         mode="append",
                         column_order=_column_order_for_write,
                     )
                 case "append":
-                    # TODO: SNOW-2299414 Fix the implementation of table type check
-                    # if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                    #     raise AnalysisException(
-                    #         f"Table {snowpark_table_name} is not an iceberg table"
-                    #     )
+                    table_schema_or_error = _get_table_schema_or_error(
+                        snowpark_table_name, session
+                    )
+                    if isinstance(table_schema_or_error, DataType):  # Table exists
+                        if get_table_type(snowpark_table_name, session) not in (
+                            "ICEBERG",
+                            "TABLE",
+                        ):
+                            raise AnalysisException(
+                                f"Table {snowpark_table_name} is not an iceberg table"
+                            )
+                    else:
+                        create_iceberg_table(
+                            snowpark_table_name=snowpark_table_name,
+                            location=write_op.options.get("location", None),
+                            schema=input_df.schema,
+                            snowpark_session=session,
+                        )
                     _validate_schema_and_get_writer(
-                        input_df, "append", snowpark_table_name
+                        input_df, "append", snowpark_table_name, table_schema_or_error
                     ).saveAsTable(
                         table_name=snowpark_table_name,
                         mode="append",
                         column_order=_column_order_for_write,
                     )
                 case "ignore":
-                    if not check_snowflake_table_existence(
+                    table_schema_or_error = _get_table_schema_or_error(
                         snowpark_table_name, session
-                    ):
+                    )
+                    if not isinstance(
+                        table_schema_or_error, DataType
+                    ):  # Table not exists
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
                             location=write_op.options.get("location", None),
@@ -359,13 +386,17 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             column_order=_column_order_for_write,
                         )
                 case "overwrite":
-                    if check_snowflake_table_existence(snowpark_table_name, session):
-                        # TODO: SNOW-2299414 Fix the implementation of table type check
-                        # if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                        #     raise AnalysisException(
-                        #         f"Table {snowpark_table_name} is not an iceberg table"
-                        #     )
-                        pass
+                    table_schema_or_error = _get_table_schema_or_error(
+                        snowpark_table_name, session
+                    )
+                    if isinstance(table_schema_or_error, DataType):  # Table exists
+                        if get_table_type(snowpark_table_name, session) not in (
+                            "ICEBERG",
+                            "TABLE",
+                        ):
+                            raise AnalysisException(
+                                f"Table {snowpark_table_name} is not an iceberg table"
+                            )
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -374,7 +405,7 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             snowpark_session=session,
                         )
                     _validate_schema_and_get_writer(
-                        input_df, "truncate", snowpark_table_name
+                        input_df, "truncate", snowpark_table_name, table_schema_or_error
                     ).saveAsTable(
                         table_name=snowpark_table_name,
                         mode="truncate",
@@ -393,33 +424,49 @@ def map_write(request: proto_base.ExecutePlanRequest):
             ):
                 match write_mode:
                     case "overwrite":
-                        if check_snowflake_table_existence(
+                        table_schema_or_error = _get_table_schema_or_error(
                             snowpark_table_name, session
-                        ):
-                            # TODO: SNOW-2299414 Fix the implementation of table type check
-                            # if (
-                            #     check_table_type(snowpark_table_name, session)
-                            #     != "TABLE"
-                            # ):
-                            #     raise AnalysisException(
-                            #         f"Table {snowpark_table_name} is not a FDN table"
-                            #     )
+                        )
+                        if isinstance(table_schema_or_error, DataType):  # Table exists
+                            if get_table_type(snowpark_table_name, session) not in (
+                                "NORMAL",
+                                "TABLE",
+                            ):
+                                raise AnalysisException(
+                                    f"Table {snowpark_table_name} is not a FDN table"
+                                )
                             write_mode = "truncate"
                         _validate_schema_and_get_writer(
-                            input_df, write_mode, snowpark_table_name
+                            input_df,
+                            write_mode,
+                            snowpark_table_name,
+                            table_schema_or_error,
                         ).saveAsTable(
                             table_name=snowpark_table_name,
                             mode=write_mode,
                             column_order=_column_order_for_write,
                         )
                     case "append":
-                        # TODO: SNOW-2299414 Fix the implementation of table type check
-                        # if check_table_type(snowpark_table_name, session) != "TABLE":
-                        #     raise AnalysisException(
-                        #         f"Table {snowpark_table_name} is not a FDN table"
-                        #     )
+                        table_schema_or_error = _get_table_schema_or_error(
+                            snowpark_table_name, session
+                        )
+                        if isinstance(
+                            table_schema_or_error, DataType
+                        ) and get_table_type(  # Table exists
+                            snowpark_table_name, session
+                        ) not in (
+                            "NORMAL",
+                            "TABLE",
+                        ):
+                            raise AnalysisException(
+                                f"Table {snowpark_table_name} is not a FDN table"
+                            )
+
                         _validate_schema_and_get_writer(
-                            input_df, write_mode, snowpark_table_name
+                            input_df,
+                            write_mode,
+                            snowpark_table_name,
+                            table_schema_or_error,
                         ).saveAsTable(
                             table_name=snowpark_table_name,
                             mode=write_mode,
@@ -466,7 +513,10 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
     if write_op.provider.lower() == "iceberg":
         match write_op.mode:
             case commands_proto.WriteOperationV2.MODE_CREATE:
-                if check_snowflake_table_existence(snowpark_table_name, session):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if isinstance(table_schema_or_error, DataType):  # Table exists
                     raise AnalysisException(
                         f"Table {snowpark_table_name} already exists"
                     )
@@ -477,24 +527,29 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     snowpark_session=session,
                 )
                 _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name
+                    input_df, "append", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="append",
                     column_order=_column_order_for_write,
                 )
             case commands_proto.WriteOperationV2.MODE_APPEND:
-                if not check_snowflake_table_existence(snowpark_table_name, session):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if not isinstance(table_schema_or_error, DataType):  # Table not exists
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
                     )
-                # TODO: SNOW-2299414 Fix the implementation of table type check
-                # if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                #     raise AnalysisException(
-                #         f"Table {snowpark_table_name} is not an iceberg table"
-                #     )
+                if get_table_type(snowpark_table_name, session) not in (
+                    "ICEBERG",
+                    "TABLE",
+                ):
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} is not an iceberg table"
+                    )
                 _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name
+                    input_df, "append", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="append",
@@ -502,26 +557,33 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                 )
             case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
                 # TODO: handle the filter condition for MODE_OVERWRITE
-                if check_snowflake_table_existence(snowpark_table_name, session):
-                    # TODO: SNOW-2299414 Fix the implementation of table type check
-                    # if check_table_type(snowpark_table_name, session) != "ICEBERG":
-                    #     raise AnalysisException(
-                    #         f"Table {snowpark_table_name} is not an iceberg table"
-                    #     )
-                    pass
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if isinstance(table_schema_or_error, DataType):  # Table exists
+                    if get_table_type(snowpark_table_name, session) not in (
+                        "ICEBERG",
+                        "TABLE",
+                    ):
+                        raise AnalysisException(
+                            f"Table {snowpark_table_name} is not an iceberg table"
+                        )
                 else:
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
                     )
                 _validate_schema_and_get_writer(
-                    input_df, "truncate", snowpark_table_name
+                    input_df, "truncate", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="truncate",
                     column_order=_column_order_for_write,
                 )
             case commands_proto.WriteOperationV2.MODE_REPLACE:
-                if check_snowflake_table_existence(snowpark_table_name, session):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if isinstance(table_schema_or_error, DataType):  # Table exists
                     create_iceberg_table(
                         snowpark_table_name=snowpark_table_name,
                         location=write_op.table_properties.get("location"),
@@ -534,7 +596,7 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                         f"Table {snowpark_table_name} does not exist"
                     )
                 _validate_schema_and_get_writer(
-                    input_df, "replace", snowpark_table_name
+                    input_df, "replace", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="append",
@@ -570,17 +632,22 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                     column_order=_column_order_for_write,
                 )
             case commands_proto.WriteOperationV2.MODE_APPEND:
-                if not check_snowflake_table_existence(snowpark_table_name, session):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if not isinstance(table_schema_or_error, DataType):  # Table not exists
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
                     )
-                # TODO: SNOW-2299414 Fix the implementation of table type check
-                # if check_table_type(snowpark_table_name, session) != "TABLE":
-                #     raise AnalysisException(
-                #         f"Table {snowpark_table_name} is not a FDN table"
-                #     )
+                if get_table_type(snowpark_table_name, session) not in (
+                    "NORMAL",
+                    "TABLE",
+                ):
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} is not a FDN table"
+                    )
                 _validate_schema_and_get_writer(
-                    input_df, "append", snowpark_table_name
+                    input_df, "append", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="append",
@@ -588,31 +655,38 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                 )
             case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
                 # TODO: handle the filter condition for MODE_OVERWRITE
-                if check_snowflake_table_existence(snowpark_table_name, session):
-                    # TODO: SNOW-2299414 Fix the implementation of table type check
-                    # if check_table_type(snowpark_table_name, session) != "TABLE":
-                    #     raise AnalysisException(
-                    #         f"Table {snowpark_table_name} is not a FDN table"
-                    #     )
-                    pass
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if isinstance(table_schema_or_error, DataType):  # Table exists
+                    if get_table_type(snowpark_table_name, session) not in (
+                        "NORMAL",
+                        "TABLE",
+                    ):
+                        raise AnalysisException(
+                            f"Table {snowpark_table_name} is not a FDN table"
+                        )
                 else:
                     raise AnalysisException(
                         f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
                     )
                 _validate_schema_and_get_writer(
-                    input_df, "truncate", snowpark_table_name
+                    input_df, "truncate", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="truncate",
                     column_order=_column_order_for_write,
                 )
             case commands_proto.WriteOperationV2.MODE_REPLACE:
-                if not check_snowflake_table_existence(snowpark_table_name, session):
+                table_schema_or_error = _get_table_schema_or_error(
+                    snowpark_table_name, session
+                )
+                if not isinstance(table_schema_or_error, DataType):  # Table not exists
                     raise AnalysisException(
                         f"Table {snowpark_table_name} does not exist"
                     )
                 _validate_schema_and_get_writer(
-                    input_df, "replace", snowpark_table_name
+                    input_df, "replace", snowpark_table_name, table_schema_or_error
                 ).saveAsTable(
                     table_name=snowpark_table_name,
                     mode="overwrite",
@@ -632,8 +706,20 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
                 )
 
 
+def _get_table_schema_or_error(
+    snowpark_table_name: str, snowpark_session: snowpark.Session
+) -> DataType | SnowparkSQLException:
+    try:
+        return snowpark_session.table(snowpark_table_name).schema
+    except SnowparkSQLException as e:
+        return e
+
+
 def _validate_schema_and_get_writer(
-    input_df: snowpark.DataFrame, write_mode: str, snowpark_table_name: str
+    input_df: snowpark.DataFrame,
+    write_mode: str,
+    snowpark_table_name: str,
+    table_schema_or_error: DataType | SnowparkSQLException | None = None,
 ) -> snowpark.DataFrameWriter:
     if write_mode is not None and write_mode.lower() in (
         "replace",
@@ -642,16 +728,26 @@ def _validate_schema_and_get_writer(
         return input_df.write
 
     table_schema = None
-    try:
-        table_schema = (
-            get_or_create_snowpark_session().table(snowpark_table_name).schema
-        )
-    except SnowparkSQLException as e:
-        msg = e.message
-        if "SQL compilation error" in msg and "does not exist" in msg:
-            pass
-        else:
-            raise e
+    if table_schema_or_error is not None:
+        if isinstance(table_schema_or_error, SnowparkSQLException):
+            msg = table_schema_or_error.message
+            if "SQL compilation error" in msg and "does not exist" in msg:
+                pass
+            else:
+                raise table_schema_or_error
+        elif isinstance(table_schema_or_error, DataType):
+            table_schema = table_schema_or_error
+    else:
+        try:
+            table_schema = (
+                get_or_create_snowpark_session().table(snowpark_table_name).schema
+            )
+        except SnowparkSQLException as e:
+            msg = e.message
+            if "SQL compilation error" in msg and "does not exist" in msg:
+                pass
+            else:
+                raise e
 
     if table_schema is None:
         # If table does not exist, we can skip the schema validation
@@ -889,7 +985,47 @@ def store_files_locally(
     )
     if overwrite and os.path.isdir(target_path):
         _truncate_directory(real_path)
-    snowpark.file_operation.FileOperation(session).get(stage_path, str(real_path))
+    # Per Snowflake docs: "The command does not preserve stage directory structure when transferring files to your client machine"
+    # https://docs.snowflake.com/en/sql-reference/sql/get
+    # Preserve directory structure under stage_path by listing files and
+    # downloading each into its corresponding local subdirectory when partition subdirs exist.
+    # Otherwise, fall back to a direct GET which flattens.
+
+    # TODO(SNOW-2326973): This can be parallelized further. Its not done here because it only affects
+    # write to local storage.
+
+    ls_dataframe = session.sql(f"LS {stage_path}")
+    ls_iterator = ls_dataframe.toLocalIterator()
+
+    # Build a normalized base prefix from stage_path to compute relatives
+    # Example: stage_path='@MY_STAGE/prefix' -> base_prefix='my_stage/prefix/'
+    base_prefix = stage_path.lstrip("@").rstrip("/") + "/"
+    base_prefix_lower = base_prefix.lower()
+
+    # Group by parent directory under the base prefix, then issue a GET per directory.
+    # This gives a small parallelism advantage if we have many files per partition directory.
+    parent_dirs: set[str] = set()
+    for row in ls_iterator:
+        name: str = row[0]
+        name_lower = name.lower()
+        rel_start = name_lower.find(base_prefix_lower)
+        relative = name[rel_start + len(base_prefix) :] if rel_start != -1 else name
+        parent_dir = os.path.dirname(relative)
+        if parent_dir and parent_dir != ".":
+            parent_dirs.add(parent_dir)
+
+    # If no parent directories were discovered (non-partitioned unload prefix), use direct GET.
+    if not parent_dirs:
+        snowpark.file_operation.FileOperation(session).get(stage_path, str(real_path))
+        return
+
+    file_op = snowpark.file_operation.FileOperation(session)
+    for parent_dir in sorted(parent_dirs):
+        local_dir = real_path / parent_dir
+        os.makedirs(local_dir, exist_ok=True)
+
+        src_dir = f"@{base_prefix}{parent_dir}"
+        file_op.get(src_dir, str(local_dir))
 
 
 def _truncate_directory(directory_path: Path) -> None:
@@ -904,31 +1040,3 @@ def _truncate_directory(directory_path: Path) -> None:
             file.unlink()
         elif file.is_dir():
             shutil.rmtree(file)
-
-
-def check_snowflake_table_existence(
-    snowpark_table_name: str,
-    snowpark_session: snowpark.Session,
-):
-    try:
-        snowpark_session.sql(f"SELECT 1 FROM {snowpark_table_name} LIMIT 1").collect()
-        return True
-    except Exception:
-        return False
-
-
-# TODO: SNOW-2299414 Fix the implementation of table type check
-# def check_table_type(
-#     snowpark_table_name: str,
-#     snowpark_session: snowpark.Session,
-# ) -> str:
-#     # currently we only support iceberg table and FDN table
-#     metadata = snowpark_session.sql(
-#         f"SHOW TABLES LIKE '{unquote_if_quoted(snowpark_table_name)}';"
-#     ).collect()
-#     if metadata is None or len(metadata) == 0:
-#         raise AnalysisException(f"Table {snowpark_table_name} does not exist")
-#     metadata = metadata[0]
-#     if metadata.as_dict().get("is_iceberg") == "Y":
-#         return "ICEBERG"
-#     return "TABLE"