SNOW-2430625:encapsulate local ingestion and udtf ingestion, get rid of driver reference on top level (#3897)

Author: sfc-gh-yuwang

src/snowflake/snowpark/_internal/data_source/datasource_partitioner.py

@@ -6,7 +6,7 @@
 import decimal
 from collections import defaultdict
 from functools import cached_property
-from typing import Optional, Union, List, Callable
+from typing import Optional, Union, List, Callable, Dict
 import logging
 import pytz
 from dateutil import parser
@@ -30,6 +30,10 @@
     DateType,
     DataType,
 )
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import snowflake.snowpark
 
 logger = logging.getLogger(__name__)
 
@@ -171,6 +175,34 @@ def partitions(self) -> List[str]:
             self.num_partitions,
         )
 
+    def _udtf_ingestion(
+        self,
+        session: "snowflake.snowpark.Session",
+        schema: StructType,
+        partition_table: str,
+        external_access_integrations: str,
+        fetch_size: int = 1000,
+        imports: Optional[List[str]] = None,
+        packages: Optional[List[str]] = None,
+        session_init_statement: Optional[List[str]] = None,
+        query_timeout: Optional[int] = 0,
+        statement_params: Optional[Dict[str, str]] = None,
+        _emit_ast: bool = True,
+    ) -> "snowflake.snowpark.DataFrame":
+        return self.driver.udtf_ingestion(
+            session,
+            schema,
+            partition_table,
+            external_access_integrations,
+            fetch_size,
+            imports,
+            packages,
+            session_init_statement,
+            query_timeout,
+            statement_params,
+            _emit_ast,
+        )
+
     @staticmethod
     def generate_partitions(
         select_query: str,

src/snowflake/snowpark/_internal/data_source/utils.py

@@ -32,8 +32,10 @@
 )
 import snowflake
 from snowflake.snowpark._internal.data_source import DataSourceReader
+from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
+from snowflake.snowpark._internal.utils import get_temp_type_for_object
 from snowflake.snowpark.exceptions import SnowparkDataframeReaderException
-
+from snowflake.snowpark.types import StructType
 
 logger = logging.getLogger(__name__)
 
@@ -551,6 +553,37 @@ def process_parquet_queue_with_threads(
     return fetch_to_local_end_time, upload_to_sf_start_time, upload_to_sf_end_time
 
 
+def create_data_source_table_and_stage(
+    session: "snowflake.snowpark.Session",
+    schema: StructType,
+    snowflake_table_name: str,
+    snowflake_stage_name: str,
+    statements_params_for_telemetry: dict,
+) -> None:
+    snowflake_table_type = "TEMPORARY"
+    create_table_sql = (
+        "CREATE "
+        f"{snowflake_table_type} "
+        "TABLE "
+        f"identifier(?) "
+        f"""({" , ".join([f'{field.name} {convert_sp_to_sf_type(field.datatype)} {"NOT NULL" if not field.nullable else ""}' for field in schema.fields])})"""
+        f"""{DATA_SOURCE_SQL_COMMENT}"""
+    )
+    params = (snowflake_table_name,)
+    logger.debug(f"Creating temporary Snowflake table: {snowflake_table_name}")
+    session.sql(create_table_sql, params=params, _emit_ast=False).collect(
+        statement_params=statements_params_for_telemetry, _emit_ast=False
+    )
+    # create temp stage
+    sql_create_temp_stage = (
+        f"create {get_temp_type_for_object(session._use_scoped_temp_objects, True)} stage"
+        f" if not exists {snowflake_stage_name} {DATA_SOURCE_SQL_COMMENT}"
+    )
+    session.sql(sql_create_temp_stage, _emit_ast=False).collect(
+        statement_params=statements_params_for_telemetry, _emit_ast=False
+    )
+
+
 def track_data_source_statement_params(
     dataframe, statement_params: Optional[Dict] = None
 ) -> Optional[Dict]:

src/snowflake/snowpark/dataframe_reader.py

@@ -42,9 +42,9 @@
     worker_process,
     process_parquet_queue_with_threads,
     STATEMENT_PARAMS_DATA_SOURCE,
-    DATA_SOURCE_SQL_COMMENT,
     DATA_SOURCE_DBAPI_SIGNATURE,
     _MAX_WORKER_SCALE,
+    create_data_source_table_and_stage,
 )
 from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
 from snowflake.snowpark._internal.telemetry import set_api_call_source
@@ -69,7 +69,6 @@
     get_aliased_option_name,
     get_copy_into_table_options,
     get_stage_parts,
-    get_temp_type_for_object,
     is_in_stored_procedure,
     parse_positional_args_to_list_variadic,
     private_preview,
@@ -1992,7 +1991,7 @@ def create_oracledb_connection():
                 table_type="temp",
                 statement_params=statements_params_for_telemetry,
             )
-            df = partitioner.driver.udtf_ingestion(
+            df = partitioner._udtf_ingestion(
                 self._session,
                 struct_schema,
                 partitions_table,
@@ -2018,29 +2017,14 @@ def create_oracledb_connection():
             return df
 
         # parquet ingestion
-        snowflake_table_type = "TEMPORARY"
         snowflake_table_name = random_name_for_temp_object(TempObjectType.TABLE)
-        create_table_sql = (
-            "CREATE "
-            f"{snowflake_table_type} "
-            "TABLE "
-            f"identifier(?) "
-            f"""({" , ".join([f'{field.name} {convert_sp_to_sf_type(field.datatype)} {"NOT NULL" if not field.nullable else ""}' for field in struct_schema.fields])})"""
-            f"""{DATA_SOURCE_SQL_COMMENT}"""
-        )
-        params = (snowflake_table_name,)
-        logger.debug(f"Creating temporary Snowflake table: {snowflake_table_name}")
-        self._session.sql(create_table_sql, params=params, _emit_ast=False).collect(
-            statement_params=statements_params_for_telemetry, _emit_ast=False
-        )
-        # create temp stage
         snowflake_stage_name = random_name_for_temp_object(TempObjectType.STAGE)
-        sql_create_temp_stage = (
-            f"create {get_temp_type_for_object(self._session._use_scoped_temp_objects, True)} stage"
-            f" if not exists {snowflake_stage_name} {DATA_SOURCE_SQL_COMMENT}"
-        )
-        self._session.sql(sql_create_temp_stage, _emit_ast=False).collect(
-            statement_params=statements_params_for_telemetry, _emit_ast=False
+        create_data_source_table_and_stage(
+            session=self._session,
+            schema=struct_schema,
+            snowflake_table_name=snowflake_table_name,
+            snowflake_stage_name=snowflake_stage_name,
+            statements_params_for_telemetry=statements_params_for_telemetry,
         )
 
         data_fetching_thread_pool_executor = None