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

snowflake/snowpark_connect/relation/read/map_read_parquet.py

@@ -5,6 +5,7 @@
 import collections
 import re
 from collections.abc import Callable
+from typing import Any
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
@@ -20,12 +21,31 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
 )
 from snowflake.snowpark.column import METADATA_FILENAME
-from snowflake.snowpark.types import DataType, DoubleType, IntegerType, StringType
+from snowflake.snowpark.types import (
+    DataType,
+    DoubleType,
+    IntegerType,
+    StringType,
+    StructType,
+)
+from snowflake.snowpark_connect.config import external_table_location
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+from snowflake.snowpark_connect.relation.read.map_read_partitioned_parquet import (
+    read_partitioned_parquet_from_external_table,
+    use_external_table,
+)
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    add_filename_metadata_to_reader,
+)
 from snowflake.snowpark_connect.relation.read.reader_config import ReaderWriterConfig
 from snowflake.snowpark_connect.relation.read.utils import (
+    apply_metadata_exclusion_pattern,
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.type_support import emulate_integral_types
+from snowflake.snowpark_connect.utils.io_utils import cached_file_format
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -33,7 +53,7 @@ from snowflake.snowpark_connect.utils.telemetry import (
 
 def map_read_parquet(
     rel: relation_proto.Relation,
-    schema: str | None,
+    schema: StructType | None,
     session: snowpark.Session,
     paths: list[str],
     options: ReaderWriterConfig,
@@ -41,28 +61,62 @@ def map_read_parquet(
     """Read a Parquet file into a Snowpark DataFrame."""
 
     if rel.read.is_streaming is True:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Streaming is not supported for Parquet files."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
-    snowpark_options = options.convert_to_snowpark_args()
-    assert schema is None, "Read PARQUET does not support user schema"
+    converted_snowpark_options = options.convert_to_snowpark_args()
+    file_format_options = _parse_parquet_snowpark_options(converted_snowpark_options)
+    raw_options = rel.read.data_source.options
     assert len(paths) > 0, "Read PARQUET expects at least one path"
 
-    reader = session.read.options(snowpark_options)
+    snowpark_options = {
+        # Setting these two options prevents a significant number of additional CREATE TEMPORARY
+        # FILE FORMAT and DROP FILE FORMAT queries. If FORMAT_NAME is not set, the Snowpark DF reader
+        # will eagerly issue a CREATE TEMPORARY FILE FORMAT when inferring the schema of the result;
+        # if ENFORCE_EXISTING_FILE_FORMAT is not set, an additional CREATE ... command will be
+        # issued when the lazy DF is materialized by a cache_result call.
+        "FORMAT_NAME": converted_snowpark_options.get(
+            "FORMAT_NAME",
+            cached_file_format(session, "parquet", file_format_options),
+        ),
+        "ENFORCE_EXISTING_FILE_FORMAT": True,
+    }
+
+    if "PATTERN" in converted_snowpark_options:
+        snowpark_options["PATTERN"] = converted_snowpark_options.get("PATTERN")
+
+    apply_metadata_exclusion_pattern(snowpark_options)
+
+    reader = add_filename_metadata_to_reader(
+        session.read.options(snowpark_options), raw_options
+    )
 
     if len(paths) == 1:
-        df = _read_parquet_with_partitions(session, reader, paths[0])
+        df, read_using_external_table = _read_parquet_with_partitions(
+            session, reader, paths[0], schema, snowpark_options
+        )
+        can_be_cached = not read_using_external_table
     else:
         is_merge_schema = options.config.get("mergeschema")
-        df = _read_parquet_with_partitions(session, reader, paths[0])
+        df, read_using_external_table = _read_parquet_with_partitions(
+            session, reader, paths[0], schema, snowpark_options
+        )
+        can_be_cached = not read_using_external_table
         schema_cols = df.columns
         for p in paths[1:]:
             reader._user_schema = None
+            partition_df, read_using_external_table = _read_parquet_with_partitions(
+                session, reader, p, schema, snowpark_options
+            )
             df = df.union_all_by_name(
-                _read_parquet_with_partitions(session, reader, p),
+                partition_df,
                 allow_missing_columns=True,
             )
+            can_be_cached = can_be_cached and not read_using_external_table
+
         if not is_merge_schema:
             df = df.select(*schema_cols)
 
@@ -73,34 +127,92 @@ def map_read_parquet(
         dataframe=renamed_df,
         spark_column_names=[analyzer_utils.unquote_if_quoted(c) for c in df.columns],
         snowpark_column_names=snowpark_column_names,
-        snowpark_column_types=[f.datatype for f in df.schema.fields],
+        snowpark_column_types=[
+            emulate_integral_types(f.datatype) for f in df.schema.fields
+        ],
+        can_be_cached=can_be_cached,
     )
 
 
 def _read_parquet_with_partitions(
-    session: Session, reader: DataFrameReader, path: str
-) -> DataFrame:
-    """Reads parquet files and adds partition columns from subdirectories."""
+    session: Session,
+    reader: DataFrameReader,
+    path: str,
+    schema: StructType | None,
+    snowpark_options: dict[str, Any],
+) -> tuple[DataFrame, bool]:
+    """
+    Reads parquet files and adds partition columns from subdirectories.
+    Returns a tuple of read DataFrame and a boolean indicating if DataFrame was read from external table.
+    """
 
     partition_columns, inferred_types = _discover_partition_columns(session, path)
-    df = reader.with_metadata(METADATA_FILENAME).parquet(path)
 
-    if not partition_columns:
-        return df.drop(METADATA_FILENAME)
+    def _get_df() -> DataFrame:
+        if not partition_columns:
+            return reader.parquet(path)
+        else:
+            # In case of too big overhead we can always optimize by using option: MAX_FILE_COUNT and allow user to define how many files should be scanned
+            df = reader.with_metadata(METADATA_FILENAME).parquet(path)
+
+            for col_name in partition_columns:
+                quoted_col_name = quote_name_without_upper_casing(col_name)
+                escaped_col_name = re.escape(col_name)
+                regex_pattern = rf"{escaped_col_name}=([^/]+)"
+
+                raw_value = snowpark_fn.regexp_extract(
+                    METADATA_FILENAME, regex_pattern, 1
+                )
+                value_or_null = snowpark_fn.when(raw_value == "", None).otherwise(
+                    raw_value
+                )
+
+                df = df.with_column(
+                    quoted_col_name,
+                    snowpark_fn.cast(value_or_null, inferred_types[col_name]),
+                )
+            return df.drop(METADATA_FILENAME)
+
+    if use_external_table(session, path):
+        if schema is None:
+            schema = _get_df().schema
+        return (
+            read_partitioned_parquet_from_external_table(
+                session,
+                schema,
+                external_table_location(),
+                path[1:-1],
+                partition_columns,
+                inferred_types,
+                snowpark_options,
+            ),
+            True,
+        )
+    else:
+        # TODO: SNOW-2736756 support user schema
+        assert schema is None, "Read PARQUET does not support user schema"
+        return _get_df(), False
 
-    for col_name in partition_columns:
-        quoted_col_name = quote_name_without_upper_casing(col_name)
-        escaped_col_name = re.escape(col_name)
-        regex_pattern = rf"{escaped_col_name}=([^/]+)"
 
-        raw_value = snowpark_fn.regexp_extract(METADATA_FILENAME, regex_pattern, 1)
-        value_or_null = snowpark_fn.when(raw_value == "", None).otherwise(raw_value)
+_parquet_file_format_allowed_options = {
+    "COMPRESSION",
+    "SNAPPY_COMPRESSION",
+    "BINARY_AS_TEXT",
+    "TRIM_SPACE",
+    "USE_LOGICAL_TYPE",
+    "USE_VECTORIZED_SCANNER",
+    "REPLACE_INVALID_CHARACTERS",
+    "NULL_IF",
+}
 
-        df = df.with_column(
-            quoted_col_name, snowpark_fn.cast(value_or_null, inferred_types[col_name])
-        )
 
-    return df.drop(METADATA_FILENAME)
+def _parse_parquet_snowpark_options(snowpark_options: dict[str, Any]) -> dict[str, Any]:
+    file_format_options = dict()
+    for key, value in snowpark_options.items():
+        upper_key = key.upper()
+        if upper_key in _parquet_file_format_allowed_options:
+            file_format_options[upper_key] = value
+    return file_format_options
 
 
 def _extract_partitions_from_path(path: str) -> dict[str, str]:
@@ -149,10 +261,14 @@ def _discover_partition_columns(
                     if i not in dir_level_to_column_name:
                         dir_level_to_column_name[i] = key
                     elif dir_level_to_column_name[i] != key:
-                        raise ValueError(
+                        exception = ValueError(
                             f"Conflicting partition column names detected: '{dir_level_to_column_name[i]}' and '{key}' "
                             f"at the same directory level"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_OPERATION
+                        )
+                        raise exception
 
                     partition_columns_values[key].add(value)
 
@@ -160,10 +276,12 @@ def _discover_partition_columns(
     for level in sorted(dir_level_to_column_name.keys()):
         col_name = dir_level_to_column_name[level]
         if col_name in seen_columns:
-            raise ValueError(
+            exception = ValueError(
                 f"Found partition column '{col_name}' at multiple directory levels. "
                 f"A partition column can only appear at a single level."
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
         seen_columns.add(col_name)
 
     ordered_columns = [

snowflake/snowpark_connect/relation/read/map_read_partitioned_parquet.py

@@ -0,0 +1,142 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import re
+from copy import deepcopy
+from typing import Any
+
+from snowflake import snowpark
+from snowflake.snowpark import Session
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+    unquote_if_quoted,
+)
+from snowflake.snowpark.functions import col, lit
+from snowflake.snowpark.types import ArrayType, DataType, MapType, StructType
+from snowflake.snowpark_connect.config import external_table_location
+from snowflake.snowpark_connect.utils.context import (
+    get_spark_session_id,
+    register_request_external_table,
+)
+from snowflake.snowpark_connect.utils.io_utils import cached_file_format
+from snowflake.snowpark_connect.utils.scala_udf_utils import map_type_to_snowflake_type
+
+STRUCTURED_TYPE_PATTERN = re.compile(r"\([^)]*\)")
+
+
+def use_external_table(session: Session, path: str) -> bool:
+    external_table_path = external_table_location()
+    stripped_path = path[1:-1]
+
+    is_external_table_path_defined = external_table_path is not None
+    is_stage = stripped_path.startswith("@")
+
+    return (
+        is_external_table_path_defined
+        and is_stage
+        and _is_external_stage(session, stripped_path)
+    )
+
+
+def _is_external_stage(session: Session, path: str) -> bool:
+    try:
+        stage_description = (
+            session.sql(f"DESCRIBE STAGE {path.split('/')[0][1:]}")
+            .filter(col('"property"') == lit("URL"))
+            .collect()
+        )
+        return stage_description[0]["property_value"] != ""
+    except Exception:
+        return False
+
+
+def _get_count_of_non_partition_path_parts(path: str) -> int:
+    count = 0
+    # First element of a path is a stage identifier we need to ignore it to count relative path parts
+    for element in path.split("/")[1:]:
+        if "=" in element:
+            break
+        count += 1
+    return count
+
+
+def read_partitioned_parquet_from_external_table(
+    session: Session,
+    schema: StructType,
+    external_table_path: str,
+    path: str,
+    partition_columns: list[str],
+    inferred_types: dict[str, DataType],
+    snowpark_options: dict[str, Any],
+) -> snowpark.DataFrame:
+    skip_path_parts = _get_count_of_non_partition_path_parts(path)
+    snowpark_partition_columns = ", ".join(
+        [quote_name_without_upper_casing(col) for col in partition_columns]
+    )
+    snowpark_typed_partition_columns = ", ".join(
+        [
+            f"{quote_name_without_upper_casing(col)} {map_type_to_snowflake_type(inferred_types[col])} as (split_part(split_part(METADATA$FILENAME, '/', {i + skip_path_parts}), '=', 2)::{map_type_to_snowflake_type(inferred_types[col])})"
+            for col, i in zip(partition_columns, range(len(partition_columns)))
+        ]
+    )
+    snowpark_schema_columns = ",".join(
+        [
+            f"{field.name} {_map_snowpark_type_to_simplified_snowflake_type(field.datatype)} as (value:{field.name}::{_map_snowpark_type_to_simplified_snowflake_type(field.datatype)})"
+            for field in schema.fields
+            if unquote_if_quoted(field.name) not in snowpark_partition_columns
+        ]
+    )
+
+    table_name = f"{external_table_path}.{quote_name_without_upper_casing(path + get_spark_session_id())}"
+    snowpark_options_copy = deepcopy(snowpark_options)
+    # These options are only used in the Snowpark Python reader, but not the actual emitted SQL.
+    snowpark_options_copy.pop("PATTERN")
+    snowpark_options_copy.pop("FORMAT_NAME")
+    snowpark_options_copy.pop("ENFORCE_EXISTING_FILE_FORMAT")
+    file_format_name = cached_file_format(session, "parquet", snowpark_options_copy)
+    session.sql(
+        f"""
+        CREATE OR REPLACE EXTERNAL TABLE {table_name} (
+            {snowpark_typed_partition_columns},
+            {snowpark_schema_columns}
+        )
+        PARTITION BY ({snowpark_partition_columns})
+        WITH LOCATION = {path}
+        FILE_FORMAT = {file_format_name}
+        PATTERN = '{snowpark_options.get('PATTERN', '.*')}'
+        AUTO_REFRESH = false
+        """
+    ).collect()
+    register_request_external_table(table_name)
+    map_fields = ", ".join(
+        [
+            f"{field.name}::{_map_snowpark_type_to_snowflake(field.datatype)} as {field.name}"
+            if isinstance(field.datatype, (StructType, MapType, ArrayType))
+            else field.name
+            for field in schema.fields
+        ]
+    )
+    return session.sql(f"SELECT {map_fields} FROM {table_name}")
+
+
+def _map_snowpark_type_to_simplified_snowflake_type(datatype: DataType) -> str:
+    if isinstance(datatype, StructType):
+        return "OBJECT"
+    elif isinstance(datatype, MapType):
+        return "VARIANT"
+    else:
+        return STRUCTURED_TYPE_PATTERN.sub("", map_type_to_snowflake_type(datatype))
+
+
+def _map_snowpark_type_to_snowflake(datatype: DataType) -> str:
+    if isinstance(datatype, StructType):
+        object_fields = ", ".join(
+            [
+                f"{field.name} { _map_snowpark_type_to_snowflake(field.datatype)}"
+                for field in datatype.fields
+            ]
+        )
+        return f"OBJECT({object_fields})"
+    else:
+        return map_type_to_snowflake_type(datatype)

snowflake/snowpark_connect/relation/read/map_read_socket.py

@@ -9,6 +9,9 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+from snowflake.snowpark_connect.type_support import emulate_integral_types
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -30,7 +33,9 @@ def map_read_socket(
         host = options.get("host", None)
         port = options.get("port", None)
         if not host or not port:
-            raise ValueError("Host and port must be provided in options.")
+            exception = ValueError("Host and port must be provided in options.")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+            raise exception
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
             try:
                 s.connect((host, int(port)))
@@ -54,10 +59,17 @@ def map_read_socket(
                     dataframe=df,
                     spark_column_names=[spark_cname],
                     snowpark_column_names=[snowpark_cname],
+                    snowpark_column_types=[
+                        emulate_integral_types(f.datatype) for f in df.schema.fields
+                    ],
                 )
             except OSError as e:
-                raise Exception(f"Error connecting to {host}:{port} - {e}")
+                exception = Exception(f"Error connecting to {host}:{port} - {e}")
+                attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                raise exception
     else:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Socket reads are only supported in streaming mode."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception

snowflake/snowpark_connect/relation/read/map_read_table.py

@@ -11,11 +11,21 @@ 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.column_qualifier import ColumnQualifier
 from snowflake.snowpark_connect.config import auto_uppercase_non_column_identifiers
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.relation.read.utils import (
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.type_support import emulate_integral_types
+from snowflake.snowpark_connect.utils.context import get_processed_views
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
@@ -23,6 +33,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_helper import get_temp_view
 
 
 def post_process_df(
@@ -48,8 +59,10 @@ def post_process_df(
             dataframe=renamed_df,
             spark_column_names=true_names,
             snowpark_column_names=snowpark_column_names,
-            snowpark_column_types=[f.datatype for f in df.schema.fields],
-            column_qualifiers=[name_parts] * len(true_names)
+            snowpark_column_types=[
+                emulate_integral_types(f.datatype) for f in df.schema.fields
+            ],
+            column_qualifiers=[{ColumnQualifier(tuple(name_parts))} for _ in true_names]
             if source_table_name
             else None,
         )
@@ -57,22 +70,85 @@ def post_process_df(
         # Check if this is a table/view not found error
         # Snowflake error codes: 002003 (42S02) - Object does not exist or not authorized
         if hasattr(e, "sql_error_code") and e.sql_error_code == 2003:
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"[TABLE_OR_VIEW_NOT_FOUND] The table or view cannot be found. {source_table_name}"
-            ) from None  # Suppress original exception to reduce message size
+            )
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            raise exception from None  # Suppress original exception to reduce message size
         # Re-raise if it's not a table not found error
         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(
+        temp_view.column_map.get_spark_columns(), 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)
+        )
+    )
+    # do not flatten initial rename when reading table
+    # TODO: remove once SNOW-2203826 is done
+    if renamed_df._select_statement is not None:
+        renamed_df._select_statement.flatten_disabled = True
+
+    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=[
+            {ColumnQualifier(tuple(split_fully_qualified_spark_name(table_name)))}
+            for _ in range(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():
+        exception = AnalysisException(
+            f"[RECURSIVE_VIEW] Recursive view `{table_name}` detected (cycle: `{table_name}` -> `{table_name}`)"
+        )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
+
     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()
 
@@ -101,10 +177,14 @@ def map_read_table(
         and rel.read.data_source.format.lower() == "iceberg"
     ):
         if len(rel.read.data_source.paths) != 1:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unexpected paths: {rel.read.data_source.paths}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
         table_identifier = rel.read.data_source.paths[0]
     else:
-        raise ValueError("The relation must have a table identifier.")
+        exception = ValueError("The relation must have a table identifier.")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
     return get_table_from_name(table_identifier, session, rel.common.plan_id)

snowflake/snowpark_connect/relation/read/map_read_text.py

@@ -8,10 +8,13 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
     rename_columns_as_snowflake_standard,
 )
+from snowflake.snowpark_connect.type_support import emulate_integral_types
 from snowflake.snowpark_connect.utils.io_utils import file_format
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -24,11 +27,17 @@ def get_file_paths_from_stage(
 ) -> typing.List[str]:
     files_paths = []
     for listed_path_row in session.sql(f"LIST {path}").collect():
+        # Skip _SUCCESS marker files
+        if listed_path_row[0].endswith("_SUCCESS"):
+            continue
+
         listed_path = listed_path_row[0].split("/")
         if listed_path_row[0].startswith("s3://") or listed_path_row[0].startswith(
             "s3a://"
         ):
             listed_path = listed_path[3:]
+        elif listed_path_row[0].startswith("azure://"):
+            listed_path = listed_path[4:]
         else:
             listed_path = listed_path[1:]
         files_paths.append("/".join(listed_path))
@@ -43,7 +52,12 @@ def read_text(
 ) -> snowpark.DataFrame:
     # TODO: handle stage name with double quotes
     files_paths = get_file_paths_from_stage(path, session)
-    stage_name = path.split("/")[0]
+    # Remove matching quotes from both ends of the path to get the stage name, if present.
+    if path and len(path) > 1 and path[0] == path[-1] and path[0] in ('"', "'"):
+        unquoted_path = path[1:-1]
+    else:
+        unquoted_path = path
+    stage_name = unquoted_path.split("/")[0]
     line_sep = options.get("lineSep") or "\n"
     column_name = (
         schema[0].name if schema is not None and len(schema.fields) > 0 else '"value"'
@@ -59,7 +73,7 @@ def read_text(
     )
     for fp in files_paths:
         content = session.sql(
-            f"SELECT T.$1 AS {default_column_name} FROM {stage_name}/{fp} (FILE_FORMAT => {text_file_format}) AS T"
+            f"SELECT T.$1 AS {default_column_name} FROM '{stage_name}/{fp}' (FILE_FORMAT => {text_file_format}) AS T"
         ).collect()
         for row in content:
             result.append(row[0])
@@ -77,9 +91,11 @@ def map_read_text(
     """
     if rel.read.is_streaming is True:
         # TODO: Structured streaming implementation.
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Streaming is not supported for CSV files."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     df = read_text(paths[0], schema, session, rel.read.data_source.options)
     if len(paths) > 1:
@@ -102,5 +118,7 @@ def map_read_text(
         dataframe=renamed_df,
         spark_column_names=spark_column_names,
         snowpark_column_names=snowpark_column_names,
-        snowpark_column_types=[f.datatype for f in df.schema.fields],
+        snowpark_column_types=[
+            emulate_integral_types(f.datatype) for f in df.schema.fields
+        ],
     )