migrate data source partitioner to inheriant from DataSource

Author: sfc-gh-yuwang

src/snowflake/snowpark/__init__.py

@@ -41,6 +41,9 @@
     "AsyncJob",
     "StoredProcedureProfiler",
     "UDFProfiler",
+    "DataSource",
+    "DataSourceReader",
+    "InputPartition",
 ]
 
 
@@ -85,6 +88,7 @@
     WhenNotMatchedClause,
 )
 from snowflake.snowpark.window import Window, WindowSpec
+from snowflake.snowpark.data_source import DataSourceReader, DataSource, InputPartition
 
 _deprecation_warning_msg = (
     "Python Runtime 3.8 reached its End-Of-Life (EOL) on October 14, 2024, there will be no further bug fixes "

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

@@ -19,7 +19,10 @@
     DRIVER_MAP,
 )
 
-from snowflake.snowpark._internal.data_source.datasource_reader import DataSourceReader
+from snowflake.snowpark._internal.data_source.datasource_reader import (
+    DbapiDataSourceReader,
+)
+from snowflake.snowpark.data_source import DataSource
 from snowflake.snowpark._internal.type_utils import type_string_to_type_object
 from snowflake.snowpark._internal.data_source.datasource_typing import Connection
 from snowflake.snowpark._internal.utils import generate_random_alphanumeric
@@ -38,7 +41,7 @@
 logger = logging.getLogger(__name__)
 
 
-class DataSourcePartitioner:
+class DataSourcePartitioner(DataSource):
     def __init__(
         self,
         create_connection: Callable[..., "Connection"],
@@ -56,6 +59,7 @@ def __init__(
         fetch_merge_count: Optional[int] = 1,
         connection_parameters: Optional[dict] = None,
     ) -> None:
+        super().__init__()
         self.create_connection = create_connection
         self.table_or_query = table_or_query
         self.is_query = is_query
@@ -91,11 +95,11 @@ def __init__(
             else None
         )
 
-    def reader(self) -> DataSourceReader:
-        return DataSourceReader(
+    def reader(self, schema) -> DbapiDataSourceReader:
+        return DbapiDataSourceReader(
             self.driver_class,
             self.create_connection,
-            self.schema,
+            schema,
             self.dbms_type,
             self.fetch_size,
             self.query_timeout,

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

@@ -7,6 +7,7 @@
 
 from typing import List, Any, Iterator, Type, Callable, Optional
 
+from snowflake.snowpark.data_source import DataSourceReader
 from snowflake.snowpark._internal.data_source.datasource_typing import Connection
 from snowflake.snowpark._internal.data_source.drivers.base_driver import BaseDriver
 from snowflake.snowpark.exceptions import SnowparkDataframeReaderException
@@ -17,7 +18,7 @@
 logger = logging.getLogger(__name__)
 
 
-class DataSourceReader:
+class DbapiDataSourceReader(DataSourceReader):
     def __init__(
         self,
         driver_class: Type[BaseDriver],
@@ -30,6 +31,7 @@ def __init__(
         fetch_merge_count: Optional[int] = 1,
         connection_parameters: Optional[dict] = None,
     ) -> None:
+        super().__init__(schema)
         # we use cloudpickle to pickle the callback function so that local function and function defined in
         # __main__ can be pickled and unpickled in subprocess
         self.pickled_create_connection_callback = cloudpickle.dumps(

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

@@ -608,8 +608,9 @@ def track_data_source_statement_params(
 
 def local_ingestion(
     session: "snowflake.snowpark.Session",
-    partitioner: "snowflake.snowpark._internal.data_source.datasource_partitioner.DataSourcePartitioner",
+    data_source: "snowflake.snowpark.DataSource",
     partitioned_queries: List[str],
+    schema: StructType,
     max_workers: int,
     snowflake_stage_name: str,
     snowflake_table_name: str,
@@ -651,7 +652,7 @@ def local_ingestion(
                         partition_queue,
                         parquet_queue,
                         process_or_thread_error_indicator,
-                        partitioner.reader(),
+                        data_source.reader(schema),
                     ),
                 )
                 process.start()
@@ -667,7 +668,7 @@ def local_ingestion(
                     partition_queue,
                     parquet_queue,
                     process_or_thread_error_indicator,
-                    partitioner.reader(),
+                    data_source.reader(schema),
                     data_fetching_thread_stop_event,
                 )
                 for _worker_id in range(max_workers)

src/snowflake/snowpark/data_source.py

@@ -5,7 +5,10 @@
 from typing import Any, Tuple, Iterator, Union, List
 
 from snowflake.snowpark.types import StructType
-import snowflake
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import snowflake.snowpark
 
 
 class InputPartition:
@@ -45,5 +48,5 @@ def _partitions(self) -> List[InputPartition]:
             self._internal_partitions = self.reader(self.schema()).partitions()
         return self._internal_partitions
 
-    def udtf_ingestion(self) -> "snowflake.snowpark.DataFrame":
+    def _udtf_ingestion(self) -> "snowflake.snowpark.DataFrame":
         pass

src/snowflake/snowpark/dataframe_reader.py

@@ -2027,7 +2027,8 @@ def create_oracledb_connection():
 
         local_ingestion(
             session=self._session,
-            partitioner=partitioner,
+            schema=struct_schema,
+            data_source=partitioner,
             partitioned_queries=partitioned_queries,
             max_workers=max_workers,
             fetch_with_process=fetch_with_process,