custom data source place hoolder

Author: sfc-gh-yuwang

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

@@ -162,7 +162,7 @@ def detect_dbms_pyodbc(dbapi2_conn):
 
 def _task_fetch_data_from_source(
     worker: DataSourceReader,
-    partition: str,
+    partition: Any,
     partition_idx: int,
     parquet_queue: Union[mp.Queue, queue.Queue],
     stop_event: threading.Event = None,
@@ -202,7 +202,7 @@ def convert_to_parquet_bytesio(fetched_data, fetch_idx):
 
 def _task_fetch_data_from_source_with_retry(
     worker: DataSourceReader,
-    partition: str,
+    partition: Any,
     partition_idx: int,
     parquet_queue: Union[mp.Queue, queue.Queue],
     stop_event: threading.Event = None,
@@ -609,7 +609,7 @@ def track_data_source_statement_params(
 def local_ingestion(
     session: "snowflake.snowpark.Session",
     data_source: "snowflake.snowpark.DataSource",
-    partitioned_queries: List[str],
+    partitioned_queries: List[Any],
     schema: StructType,
     max_workers: int,
     snowflake_stage_name: str,
@@ -732,3 +732,30 @@ def local_ingestion(
             data_fetching_thread_pool_executor.shutdown(wait=True)
 
     logger.debug("All data has been successfully loaded into the Snowflake table.")
+
+
+def custom_data_source_udtf_class_builder(
+    data_source_class: type,
+    schema: StructType,
+):
+    import cloudpickle
+    import pickle
+
+    pickled_schema = cloudpickle.dumps(schema, protocol=pickle.HIGHEST_PROTOCOL)
+
+    class UDTFIngestion:
+        def process(self, pickled_partition: bytearray):
+            import cloudpickle
+
+            data_source_instance = data_source_class()
+            unpickled_schema = cloudpickle.loads(pickled_schema)
+            partition = cloudpickle.loads(pickled_partition)
+            reader = data_source_instance.reader(unpickled_schema)
+            for result in reader.read(partition):
+                if isinstance(result, list):
+                    yield from result
+                else:
+                    yield from list(reader.read(partition))
+                    break
+
+    return UDTFIngestion

src/snowflake/snowpark/data_source.py

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

src/snowflake/snowpark/dataframe_reader.py

@@ -39,6 +39,7 @@
     DATA_SOURCE_DBAPI_SIGNATURE,
     create_data_source_table_and_stage,
     local_ingestion,
+    custom_data_source_udtf_class_builder,
 )
 from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
 from snowflake.snowpark._internal.telemetry import set_api_call_source
@@ -2154,3 +2155,97 @@ def directory(self, stage_name: str, _emit_ast: bool = True) -> DataFrame:
     def register_custom_data_source(self, data_source: DataSource):
         self._data_source_format.append(data_source.name())
         self._custom_data_source_format[data_source.name()] = data_source
+
+    def _custom_data_source(
+        self,
+        data_source_name: str,
+        max_workers: Optional[int] = None,
+        udtf_configs: Optional[dict] = None,
+        fetch_with_process: bool = False,
+        _emit_ast: bool = True,
+    ) -> DataFrame:
+
+        data_source_class = self._custom_data_source_format[data_source_name]
+        data_source_instance = data_source_class()
+        partitions = data_source_instance._partitions()
+        if partitions is None:
+            partitions = [None]
+        schema = DataSourcePartitioner.formatting_custom_schema(
+            data_source_instance.schema()
+        )
+        statements_params_for_telemetry = {STATEMENT_PARAMS_DATA_SOURCE: "1"}
+        telemetry_json_string = {}
+
+        # udtf ingestion
+
+        if udtf_configs is not None:
+            import cloudpickle
+            import pickle
+
+            pickled_partitions = [
+                cloudpickle.dumps(par, protocol=pickle.HIGHEST_PROTOCOL)
+                for par in partitions
+            ]
+            partitions_table = random_name_for_temp_object(TempObjectType.TABLE)
+            self._session.create_dataframe(
+                [[par] for par in pickled_partitions], schema=["partition"]
+            ).write.save_as_table(partitions_table, table_type="temp")
+
+            udtf_name = random_name_for_temp_object(TempObjectType.TABLE_FUNCTION)
+            self._session.udtf.register(
+                custom_data_source_udtf_class_builder(data_source_class, schema),
+                name=udtf_name,
+                output_schema=StructType(
+                    [
+                        StructField(field.name, VariantType(), field.nullable)
+                        for field in schema.fields
+                    ]
+                ),
+                external_access_integrations=[
+                    udtf_configs["external_access_integrations"]
+                ],
+                packages=udtf_configs["packages"],
+                imports=udtf_configs["imports"],
+            )
+            call_udtf_sql = f"""
+                            select * from {partitions_table}, table({udtf_name}(partition))
+                            """
+            res = self._session.sql(call_udtf_sql, _emit_ast=_emit_ast)
+            cols = [
+                res[field.name].cast(field.datatype).alias(field.name)
+                for field in schema.fields
+            ]
+            return res.select(cols, _emit_ast=True)
+
+        # parquet ingestion
+
+        snowflake_table_name = random_name_for_temp_object(TempObjectType.TABLE)
+        snowflake_stage_name = random_name_for_temp_object(TempObjectType.STAGE)
+
+        create_data_source_table_and_stage(
+            session=self._session,
+            schema=schema,
+            snowflake_table_name=snowflake_table_name,
+            snowflake_stage_name=snowflake_stage_name,
+            statements_params_for_telemetry=statements_params_for_telemetry,
+        )
+
+        local_ingestion(
+            session=self._session,
+            schema=schema,
+            data_source=data_source_instance,
+            partitioned_queries=partitions,
+            max_workers=max_workers,
+            fetch_with_process=fetch_with_process,
+            snowflake_stage_name=snowflake_stage_name,
+            snowflake_table_name=snowflake_table_name,
+            statements_params_for_telemetry=statements_params_for_telemetry,
+            telemetry_json_string=telemetry_json_string,
+            _emit_ast=_emit_ast,
+        )
+
+        res_df = self._session.table(snowflake_table_name, _emit_ast=_emit_ast)
+
+        set_api_call_source(res_df, DATA_SOURCE_DBAPI_SIGNATURE)
+
+        return res_df

tests/integ/test_data_source_api.py

@@ -25,7 +25,9 @@
 from snowflake.snowpark._internal.data_source.datasource_partitioner import (
     DataSourcePartitioner,
 )
-from snowflake.snowpark._internal.data_source.datasource_reader import DataSourceReader
+from snowflake.snowpark._internal.data_source.datasource_reader import (
+    DbapiDataSourceReader,
+)
 from snowflake.snowpark._internal.data_source.drivers import (
     PyodbcDriver,
     SqliteDriver,
@@ -177,7 +179,7 @@ def test_dbapi_retry(session, fetch_with_process):
             SnowparkDataframeReaderException, match="\\[RuntimeError\\] Test error"
         ):
             _task_fetch_data_from_source_with_retry(
-                worker=DataSourceReader(
+                worker=DbapiDataSourceReader(
                     PyodbcDriver,
                     sql_server_create_connection,
                     StructType([StructField("col1", IntegerType(), False)]),
@@ -220,7 +222,7 @@ def test_dbapi_non_retryable_error(session, fetch_with_process):
         parquet_queue = multiprocessing.Queue() if fetch_with_process else queue.Queue()
         with pytest.raises(SnowparkDataframeReaderException, match="mock error"):
             _task_fetch_data_from_source_with_retry(
-                worker=DataSourceReader(
+                worker=DbapiDataSourceReader(
                     PyodbcDriver,
                     sql_server_create_connection,
                     StructType([StructField("col1", IntegerType(), False)]),
@@ -712,7 +714,7 @@ def test_task_fetch_from_data_source_with_fetch_size(
     parquet_queue = multiprocessing.Queue() if fetch_with_process else queue.Queue()
 
     params = {
-        "worker": DataSourceReader(
+        "worker": DbapiDataSourceReader(
             PyodbcDriver,
             sql_server_create_connection_small_data,
             schema=schema,
@@ -1147,7 +1149,7 @@ def test_fetch_merge_count_unit(fetch_size, fetch_merge_count, expected_batch_cn
     with tempfile.TemporaryDirectory() as temp_dir:
         dbpath = os.path.join(temp_dir, "testsqlite3.db")
         table_name, columns, example_data, _ = sqlite3_db(dbpath)
-        reader = DataSourceReader(
+        reader = DbapiDataSourceReader(
             SqliteDriver,
             functools.partial(create_connection_to_sqlite3_db, dbpath),
             schema=SQLITE3_DB_CUSTOM_SCHEMA_STRUCT_TYPE,
@@ -1268,8 +1270,8 @@ def test_worker_process_unit(fetch_with_process):
         dbpath = os.path.join(temp_dir, "testsqlite3.db")
         table_name, columns, example_data, _ = sqlite3_db(dbpath)
 
-        # Create DataSourceReader for sqlite3
-        reader = DataSourceReader(
+        # Create DbapiDataSourceReader for sqlite3
+        reader = DbapiDataSourceReader(
             SqliteDriver,
             functools.partial(create_connection_to_sqlite3_db, dbpath),
             schema=SQLITE3_DB_CUSTOM_SCHEMA_STRUCT_TYPE,

tests/unit/test_data_source.py

@@ -6,7 +6,9 @@
 import re
 from unittest.mock import Mock, patch
 from snowflake.snowpark._internal.data_source.drivers.base_driver import BaseDriver
-from snowflake.snowpark._internal.data_source.datasource_reader import DataSourceReader
+from snowflake.snowpark._internal.data_source.datasource_reader import (
+    DbapiDataSourceReader,
+)
 from snowflake.snowpark._internal.data_source.utils import DBMS_TYPE
 from snowflake.snowpark.types import StructType, StructField, StringType
 
@@ -67,7 +69,7 @@ def test_close_error_handling(cursor_fails, conn_fails):
     ],
 )
 def test_datasource_reader_close_error_handling(cursor_fails, conn_fails):
-    """Test that DataSourceReader handles cursor/connection close errors gracefully."""
+    """Test that DbapiDataSourceReader handles cursor/connection close errors gracefully."""
     # Setup mocks
     mock_create_connection = Mock()
     expected_schema = StructType([StructField("test_col", StringType())])
@@ -94,7 +96,7 @@ def test_datasource_reader_close_error_handling(cursor_fails, conn_fails):
     mock_driver_class = Mock(return_value=mock_driver)
 
     # Create reader with the mock driver class
-    reader = DataSourceReader(
+    reader = DbapiDataSourceReader(
         driver_class=mock_driver_class,
         create_connection=mock_create_connection,
         schema=expected_schema,