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

snowflake/snowpark_connect/relation/read/metadata_utils.py

@@ -0,0 +1,170 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Utilities for handling internal metadata columns in file-based DataFrames.
+"""
+
+import os
+
+import pandas
+from pyspark.errors.exceptions.base import AnalysisException
+
+from snowflake import snowpark
+from snowflake.snowpark.column import METADATA_FILENAME
+from snowflake.snowpark.functions import col
+from snowflake.snowpark.types import StructField, StructType
+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
+
+# Constant for the metadata filename column name
+METADATA_FILENAME_COLUMN = "METADATA$FILENAME"
+
+
+def add_filename_metadata_to_reader(
+    reader: snowpark.DataFrameReader,
+    options: dict | None = None,
+) -> snowpark.DataFrameReader:
+    """
+    Add filename metadata to a DataFrameReader based on configuration.
+
+    Args:
+        reader: Snowpark DataFrameReader instance
+        options: Dictionary of options to check for metadata configuration
+
+    Returns:
+        DataFrameReader with filename metadata enabled if configured, otherwise unchanged
+    """
+    # NOTE: SNOWPARK_POPULATE_FILE_METADATA_DEFAULT is an internal environment variable
+    # used only for CI testing to verify no metadata columns leak in regular file operations.
+    # This environment variable should NOT be exposed to end users. Users should only use snowpark.populateFileMetadata
+    # to enable metadata population.
+    metadata_default = os.environ.get(
+        "SNOWPARK_POPULATE_FILE_METADATA_DEFAULT", "false"
+    )
+
+    populate_metadata = (
+        options.get("snowpark.populateFileMetadata", metadata_default)
+        if options
+        else metadata_default
+    ).lower() == "true"
+
+    if populate_metadata:
+        return reader.with_metadata(METADATA_FILENAME)
+    else:
+        return reader
+
+
+def get_non_metadata_fields(schema_fields: list[StructField]) -> list[StructField]:
+    """
+    Filter out METADATA$FILENAME fields from a list of schema fields.
+
+    Args:
+        schema_fields: List of StructField objects from a DataFrame schema
+
+    Returns:
+        List of StructField objects excluding METADATA$FILENAME
+    """
+    return [field for field in schema_fields if field.name != METADATA_FILENAME_COLUMN]
+
+
+def get_non_metadata_column_names(schema_fields: list[StructField]) -> list[str]:
+    """
+    Get column names from schema fields, excluding METADATA$FILENAME.
+
+    Args:
+        schema_fields: List of StructField objects from a DataFrame schema
+
+    Returns:
+        List of column names (strings) excluding METADATA$FILENAME
+    """
+    return [
+        field.name for field in schema_fields if field.name != METADATA_FILENAME_COLUMN
+    ]
+
+
+def filter_metadata_column_name(column_names: list[str]) -> list[str]:
+    """
+    Get column names from column_names, excluding METADATA$FILENAME.
+
+    Returns:
+        List of column names (strings) excluding METADATA$FILENAME
+    """
+    return [
+        col_name for col_name in column_names if col_name != METADATA_FILENAME_COLUMN
+    ]
+
+
+def without_internal_columns(
+    result_container: DataFrameContainer | pandas.DataFrame | None,
+) -> DataFrameContainer | pandas.DataFrame | None:
+    """
+    Filters internal columns like:
+     * METADATA$FILENAME from DataFrame container for execution and write operations
+     * hidden columns needed for outer joins implementation
+
+    Args:
+        result_container: DataFrameContainer or pandas DataFrame to filter
+
+    Returns:
+        Filtered container (callers can access dataframe via container.dataframe)
+    """
+    # Handle pandas DataFrame case - return as-is
+    if isinstance(result_container, pandas.DataFrame):
+        return result_container
+
+    if result_container is None:
+        return None
+
+    result_container = result_container.without_hidden_columns()
+    result_df = result_container.dataframe
+    if not isinstance(result_df, snowpark.DataFrame):
+        return result_container
+
+    df_columns = result_container.column_map.get_snowpark_columns()
+    has_metadata_filename = any(name == METADATA_FILENAME_COLUMN for name in df_columns)
+
+    if not has_metadata_filename:
+        return result_container
+
+    non_metadata_columns = filter_metadata_column_name(df_columns)
+
+    if len(non_metadata_columns) == 0:
+        # DataFrame contains only metadata columns (METADATA$FILENAME), no actual data columns remaining.
+        # We don't have a way to return an empty dataframe.
+        exception = AnalysisException(
+            "[DATAFRAME_MISSING_DATA_COLUMNS] Cannot perform operation on DataFrame that contains no data columns."
+        )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
+
+    filtered_df = result_df.select([col(name) for name in non_metadata_columns])
+
+    original_spark_columns = result_container.column_map.get_spark_columns()
+    original_snowpark_columns = result_container.column_map.get_snowpark_columns()
+
+    filtered_spark_columns = []
+    filtered_snowpark_columns = []
+
+    for i, colname in enumerate(df_columns):
+        if colname != METADATA_FILENAME_COLUMN:
+            filtered_spark_columns.append(original_spark_columns[i])
+            filtered_snowpark_columns.append(original_snowpark_columns[i])
+
+    new_container = DataFrameContainer.create_with_column_mapping(
+        dataframe=filtered_df,
+        spark_column_names=filtered_spark_columns,
+        snowpark_column_names=filtered_snowpark_columns,
+        column_metadata=result_container.column_map.column_metadata,
+        table_name=result_container.table_name,
+        alias=result_container.alias,
+        partition_hint=result_container.partition_hint,
+        # we don't want to evaluate `filtered_df` schema since it will always trigger a describe query
+        cached_schema_getter=lambda: StructType(
+            [f for f in result_df.schema if f.name != METADATA_FILENAME_COLUMN]
+        ),
+    )
+
+    return new_container

snowflake/snowpark_connect/relation/read/reader_config.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 from typing import Any
 
-from snowflake.snowpark_connect.config import str_to_bool
+from snowflake.snowpark_connect.config import global_config, str_to_bool
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 
 
@@ -126,6 +126,8 @@ CSV_READ_SUPPORTED_OPTIONS = lowercase_set(
         "compression",
         # "escapeQuotes",
         # "quoteAll",
+        "rowsToInferSchema",  # Snowflake specific option, number of rows to infer schema
+        "relaxTypesToInferSchema",  # Snowflake specific option, whether to relax types to infer schema
     }
 )
 
@@ -201,6 +203,21 @@ def csv_convert_to_snowpark_args(snowpark_config: dict[str, Any]) -> dict[str, A
     if snowpark_config["escape"] and snowpark_config["escape"] == "\\":
         snowpark_config["escape"] = "\\\\"
 
+    # Snowflake specific option, number of rows to infer schema for CSV files
+    if "rowstoinferschema" in snowpark_config:
+        rows_to_infer_schema = snowpark_config["rowstoinferschema"]
+        del snowpark_config["rowstoinferschema"]
+        relax_types_to_infer_schema = True
+        if "relaxtypestoinferschema" in snowpark_config:
+            relax_types_to_infer_schema = str_to_bool(
+                str(snowpark_config["relaxtypestoinferschema"])
+            )
+            del snowpark_config["relaxtypestoinferschema"]
+        snowpark_config["INFER_SCHEMA_OPTIONS"] = {
+            "MAX_RECORDS_PER_FILE": int(rows_to_infer_schema),
+            "USE_RELAXED_TYPES": relax_types_to_infer_schema,
+        }
+
     # Rename the keys to match the Snowpark configuration.
     for spark_arg, snowpark_arg in renamed_args.items():
         if spark_arg not in snowpark_config:
@@ -339,7 +356,7 @@ class JsonReaderConfig(ReaderWriterConfig):
                     "dropFieldIfAllNull",
                     "encoding",
                     # "locale",
-                    # "pathGlobFilter",
+                    "pathGlobFilter",
                     # "recursiveFileLookup",
                     # "modifiedBefore",
                     # "modifiedAfter",
@@ -366,6 +383,7 @@ class JsonReaderConfig(ReaderWriterConfig):
             "dateFormat": "DATE_FORMAT",
             "timestampFormat": "TIMESTAMP_FORMAT",
             "multiLine": "STRIP_OUTER_ARRAY",
+            "pathGlobFilter": "PATTERN",
         }
         renamed_args = lowercase_dict_keys(renamed_args)
         snowpark_config = super().convert_to_snowpark_args()
@@ -385,7 +403,7 @@ class ParquetReaderConfig(ReaderWriterConfig):
                 default_config={},
                 supported_options={
                     # "mergeSchema",
-                    # "pathGlobFilter",
+                    "pathGlobFilter",
                     # "recursiveFileLookup",
                     # "modifiedBefore",
                     # "modifiedAfter",
@@ -402,10 +420,25 @@ class ParquetReaderConfig(ReaderWriterConfig):
         )
 
     def convert_to_snowpark_args(self) -> dict[str, Any]:
+        renamed_args = {
+            "pathGlobFilter": "PATTERN",
+        }
+        renamed_args = lowercase_dict_keys(renamed_args)
         snowpark_args = super().convert_to_snowpark_args()
 
+        for spark_arg, snowpark_arg in renamed_args.items():
+            if spark_arg not in snowpark_args:
+                continue
+            snowpark_args[snowpark_arg] = snowpark_args[spark_arg]
+            del snowpark_args[spark_arg]
+
         # Should be determined by spark.sql.parquet.binaryAsString, but currently Snowpark Connect only supports
         # the default value (false). TODO: Add support for spark.sql.parquet.binaryAsString equal to "true".
         snowpark_args["BINARY_AS_TEXT"] = False
 
+        # Set USE_VECTORIZED_SCANNER from global config. This will become the default in a future BCR.
+        snowpark_args["USE_VECTORIZED_SCANNER"] = global_config._get_config_setting(
+            "snowpark.connect.parquet.useVectorizedScanner"
+        )
+
         return snowpark_args

snowflake/snowpark_connect/relation/read/utils.py

@@ -40,6 +40,47 @@ DATA_SOURCE_SQL_COMMENT = (
 INDEXED_COLUMN_NAME_PATTERN = re.compile(r"(^\"c)(\d+)(\"$)")
 
 
+def apply_metadata_exclusion_pattern(options: dict) -> None:
+    """
+    Exclude metadata and hidden files from reads, matching Spark's behavior.
+
+    Automatically filters out internal metadata files that should never be read as data:
+        - _SUCCESS, _metadata, _common_metadata (Spark/Parquet metadata)
+        - .crc (Hadoop checksum files)
+        - .DS_Store (macOS system files)
+        - Any file starting with _ or .
+
+    Pattern used: ".*/[^_.][^/]*$|^[^_.][^/]*$"
+        - Matches files where filename does NOT start with _ or .
+        - Works at any directory depth (flat or partitioned data)
+        - Allows files with or without extensions
+
+    Examples of excluded files:
+        ❌ _SUCCESS, _metadata, _common_metadata (Spark/Parquet metadata)
+        ❌ .crc, .DS_Store, .hidden (system/hidden files)
+        ❌ year=2024/_SUCCESS (metadata in partitioned directories)
+
+    Examples of allowed files:
+        ✅ part-00000.parquet, data.csv, output.json (data files)
+        ✅ success, myfile (files without extensions, don't start with _ or .)
+        ✅ year=2024/month=01/part-00000.parquet (partitioned data)
+
+    User pattern handling:
+        - No pattern or "*" or ".*" → Apply metadata exclusion
+        - Custom patterns → Default to user provided pattern.
+
+    Leak cases (user explicitly requests metadata files and are intentional):
+        ⚠️ "_*" → Matches _SUCCESS, _metadata (explicit underscore prefix)
+        ⚠️ "*SUCCESS*" → Matches _SUCCESS (broad wildcard side effect)
+        ⚠️ "[_.].*" → Matches _SUCCESS, .crc (character class includes _)
+
+    Args:
+        options: Dictionary of Snowpark read options (modified in place)
+    """
+    if "PATTERN" not in options or options["PATTERN"] in ("*", ".*"):
+        options["PATTERN"] = ".*/[^_.][^/]*$|^[^_.][^/]*$"
+
+
 def subtract_one(match: re.Match[str]) -> str:
     """Spark column names are 0 indexed, Snowpark is 1 indexed."""
     return f"_c{str(int(match.group(2)) - 1)}"
@@ -73,13 +114,17 @@ def rename_columns_as_snowflake_standard(
         return df, []
 
     new_columns = make_column_names_snowpark_compatible(df.columns, plan_id)
-    return (
-        df.select(
-            *(df.col(orig).alias(alias) for orig, alias in zip(df.columns, new_columns))
-        ),
-        new_columns,
+    result_df = df.select(
+        *(df.col(orig).alias(alias) for orig, alias in zip(df.columns, new_columns))
     )
 
+    # do not flatten initial rename when reading table
+    # TODO: remove once SNOW-2203826 is done
+    if result_df._select_statement is not None:
+        result_df._select_statement.flatten_disabled = True
+
+    return result_df, new_columns
+
 
 class Connection(Protocol):
     """External datasource connection created from user-input create_connection function."""

snowflake/snowpark_connect/relation/stage_locator.py

@@ -5,17 +5,20 @@
 import os
 
 from fsspec.core import url_to_fs
+from pyspark.errors.exceptions.base import AnalysisException
 from s3fs.core import S3FileSystem
 
 from snowflake import snowpark
 from snowflake.snowpark.session import Session
 from snowflake.snowpark_connect.config import sessions_config
+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.io_utils import (
     get_cloud_from_url,
     parse_azure_url,
 )
 from snowflake.snowpark_connect.relation.utils import random_string
-from snowflake.snowpark_connect.utils.context import get_session_id
+from snowflake.snowpark_connect.utils.context import get_spark_session_id
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 
 
@@ -33,37 +36,44 @@ def get_paths_from_stage(
 
     # TODO : What if GCP?
     # TODO: What if already stage path?
-    if get_cloud_from_url(paths[0]) == "azure":
-        rewrite_paths = []
-        for p in paths:
-            _, bucket_name, path = parse_azure_url(p)
-            rewrite_paths.append(f"{stage_name}/{path}")
-        paths = rewrite_paths
-    else:
-        filesystem, parsed_path = url_to_fs(paths[0])
-        if isinstance(filesystem, S3FileSystem):  # aws
-            # Remove bucket name from the path since the stage name will replace
-            # the bucket name in the path.
-            paths = [
-                f"{stage_name}/{'/'.join(url_to_fs(p)[1].split('/')[1:])}"
-                for p in paths
-            ]
-        else:  # local
-            # For local files, we need to preserve directory structure for partitioned data
-            # Instead of just using basename, we'll use the last few path components
-            new_paths = []
+    match get_cloud_from_url(paths[0]):
+        case "azure":
+            rewrite_paths = []
             for p in paths:
-                # Split the path and take the last 2-3 components to preserve structure
-                # but avoid very long paths
-                path_parts = p.split(os.sep)
-                if len(path_parts) >= 2:
-                    # Take last 2 components (e.g., "base_case/x=abc")
-                    relative_path = "/".join(path_parts[-2:])
-                else:
-                    # Single component, use basename
-                    relative_path = os.path.basename(p)
-                new_paths.append(f"{stage_name}/{relative_path}")
-            paths = new_paths
+                _, bucket_name, path = parse_azure_url(p)
+                rewrite_paths.append(f"{stage_name}/{path}")
+            paths = rewrite_paths
+        case "gcp":
+            exception = AnalysisException(
+                "You must configure an integration for Google Cloud Storage to perform I/O operations rather than accessing the URL directly. Reference: https://docs.snowflake.com/en/user-guide/data-load-gcs-config"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
+        case _:
+            filesystem, parsed_path = url_to_fs(paths[0])
+            if isinstance(filesystem, S3FileSystem):  # aws
+                # Remove bucket name from the path since the stage name will replace
+                # the bucket name in the path.
+                paths = [
+                    f"{stage_name}/{'/'.join(url_to_fs(p)[1].split('/')[1:])}"
+                    for p in paths
+                ]
+            else:  # local
+                # For local files, we need to preserve directory structure for partitioned data
+                # Instead of just using basename, we'll use the last few path components
+                new_paths = []
+                for p in paths:
+                    # Split the path and take the last 2-3 components to preserve structure
+                    # but avoid very long paths
+                    path_parts = p.split(os.sep)
+                    if len(path_parts) >= 2:
+                        # Take last 2 components (e.g., "base_case/x=abc")
+                        relative_path = "/".join(path_parts[-2:])
+                    else:
+                        # Single component, use basename
+                        relative_path = os.path.basename(p)
+                    new_paths.append(f"{stage_name}/{relative_path}")
+                paths = new_paths
 
     return paths
 
@@ -89,7 +99,7 @@ class StageLocator:
         self,
         url: str = "/",
     ) -> str:
-        spark_session_id = get_session_id()
+        spark_session_id = get_spark_session_id()
 
         match get_cloud_from_url(url):
             case "azure":
@@ -102,15 +112,21 @@ class StageLocator:
                 sql_query = f"CREATE OR REPLACE TEMP STAGE {stage_name[1:]} URL='azure://{account}.blob.core.windows.net/{bucket_name}'"
 
                 credential_session_key = (
-                    f"fs.azure.sas.{bucket_name}.{account}.blob.core.windows.net"
+                    f"fs.azure.sas.fixed.token.{account}.dfs.core.windows.net",
+                    f"fs.azure.sas.{bucket_name}.{account}.blob.core.windows.net",
                 )
                 credential = sessions_config.get(spark_session_id, None)
-                if (
-                    credential is not None
-                    and credential.get(credential_session_key) is not None
-                    and credential.get(credential_session_key).strip() != ""
-                ):
-                    sql_query += f" CREDENTIALS = (AZURE_SAS_TOKEN = '{credential.get(credential_session_key)}')"
+                sas_token = None
+                for session_key in credential_session_key:
+                    if (
+                        credential is not None
+                        and credential.get(session_key) is not None
+                        and credential.get(session_key).strip() != ""
+                    ):
+                        sas_token = credential.get(session_key)
+                        break
+                if sas_token is not None:
+                    sql_query += f" CREDENTIALS = (AZURE_SAS_TOKEN = '{sas_token}')"
 
                 logger.info(self.session.sql(sql_query).collect())
                 self.stages_for_azure[bucket_name] = stage_name
@@ -128,24 +144,44 @@ class StageLocator:
                     # but the rest of the time it's used, it does. We just drop it here.
                     sql_query = f"CREATE OR REPLACE TEMP STAGE {stage_name[1:]} URL='s3://{parsed_path.split('/')[0]}'"
                     credential = sessions_config.get(spark_session_id, None)
-                    if (
-                        credential is not None
-                        and credential.get("spark.hadoop.fs.s3a.access.key") is not None
-                        and credential.get("spark.hadoop.fs.s3a.secret.key") is not None
-                        and credential.get("spark.hadoop.fs.s3a.access.key").strip()
-                        != ""
-                        and credential.get("spark.hadoop.fs.s3a.secret.key").strip()
-                        != ""
-                    ):
-                        aws_keys = f" AWS_KEY_ID = '{credential.get('spark.hadoop.fs.s3a.access.key')}'"
-                        aws_keys += f" AWS_SECRET_KEY = '{credential.get('spark.hadoop.fs.s3a.secret.key')}'"
-                        if (
-                            credential.get("spark.hadoop.fs.s3a.session.token")
+                    if credential is not None:
+                        if (  # USE AWS KEYS to connect
+                            credential.get("spark.hadoop.fs.s3a.access.key") is not None
+                            and credential.get("spark.hadoop.fs.s3a.secret.key")
+                            is not None
+                            and credential.get("spark.hadoop.fs.s3a.access.key").strip()
+                            != ""
+                            and credential.get("spark.hadoop.fs.s3a.secret.key").strip()
+                            != ""
+                        ):
+                            aws_keys = f" AWS_KEY_ID = '{credential.get('spark.hadoop.fs.s3a.access.key')}'"
+                            aws_keys += f" AWS_SECRET_KEY = '{credential.get('spark.hadoop.fs.s3a.secret.key')}'"
+                            if (
+                                credential.get("spark.hadoop.fs.s3a.session.token")
+                                is not None
+                            ):
+                                aws_keys += f" AWS_TOKEN = '{credential.get('spark.hadoop.fs.s3a.session.token')}'"
+                            sql_query += f" CREDENTIALS = ({aws_keys})"
+                            sql_query += " ENCRYPTION = ( TYPE = 'AWS_SSE_S3' )"
+                        elif (  # USE AWS ROLE and KMS KEY to connect
+                            credential.get(
+                                "spark.hadoop.fs.s3a.server-side-encryption.key"
+                            )
+                            is not None
+                            and credential.get(
+                                "spark.hadoop.fs.s3a.server-side-encryption.key"
+                            ).strip()
+                            != ""
+                            and credential.get("spark.hadoop.fs.s3a.assumed.role.arn")
                             is not None
+                            and credential.get(
+                                "spark.hadoop.fs.s3a.assumed.role.arn"
+                            ).strip()
+                            != ""
                         ):
-                            aws_keys += f" AWS_TOKEN = '{credential.get('spark.hadoop.fs.s3a.session.token')}'"
-                        sql_query += f" CREDENTIALS = ({aws_keys})"
-                        sql_query += " ENCRYPTION = ( TYPE = 'AWS_SSE_S3' )"
+                            aws_role = f" AWS_ROLE = '{credential.get('spark.hadoop.fs.s3a.assumed.role.arn')}'"
+                            sql_query += f" CREDENTIALS = ({aws_role})"
+                            sql_query += f" ENCRYPTION = ( TYPE='AWS_SSE_KMS' KMS_KEY_ID = '{credential.get('spark.hadoop.fs.s3a.server-side-encryption.key')}' )"
 
                     logger.info(self.session.sql(sql_query).collect())
                     self.stages_for_aws[bucket_name] = stage_name