Snorkeling with Snowflake

  
from concurrent.futures import ThreadPoolExecutor

from snowflake.connector import connect

connect_args = {...}
query = "select * from SNOWFLAKE_SAMPLE_DATA.TPCH_SF1.LINEITEM"

def _compute(_sf_batch: ResultBatch):
    arrow_table = _sf_batch.to_arrow()
    # do more compute on this arrow table

with connect(**connect_args) as conn:
    with conn.cursor() as cur:
        cur.execute(query)
        batches = cur.get_result_batches()

with ThreadPoolExecutor() as _texec:
    for batch in batches:
        _texec.submit(_compute, batch)

_texec.shutdown()
  

  
import ray
import pyarrow as pa

from snowflake.connector import connect

connect_args = {...}
query = "select * from SNOWFLAKE_SAMPLE_DATA.TPCH_SF1.LINEITEM"

def _compute(arrow_table: pa.table):
    # do more compute on this arrow table
    pass

snowflake_datasource = SnowflakeDatasource(connect_args, query)
rds = ray.data.read_datasource(snowflake_datasource)
rds.map_batches(_compute, batch_format="pyarrow")
  

  
from ray.data.datasource import Reader
from ray.data.block import BlockMetadata
from ray.data import ReadTask

from snowflake.connector.result_batch import ResultBatch

# The reader performs the heavy lifting
class _SnowflakeDatasourceReader(Reader):
    def __init__(self, connection_args: dict, query: str):
        # connection info like snowflake account name & credentials.
        self._connection_args = connection_args
 
        # the query to execute.
        self._query = query

    # this method will be reused in both creating the read_tasks
    # as well as calc
    @cached_property
    def _result_batches(self):
        # connect with snowflake
        with connect(**self._connection_args) as conn:
            # get the cursor
            with conn.cursor() as cur:
                cur.execute(self._query)
                # Get the result as batches.
                # This API has a minimal memory footprint because
                # the ResultBatch doesn't have any data. It only
                # tells us how to pull the data and what size/schema
                # to expect from this data once it lands.
                # The driver hence won't have any memory footprint
                # and can safely do the work of creating relevant
                # Block (s) for ray.
                batches = cur.get_result_batches()
        return batches

    def estimate_inmemory_data_size(self) -> Optional[int]:
        sz = None

        for batch in self._result_batches:
            sz = (sz or 0) + (batch.uncompressed_size or 0)

        ray_data_logger.info("Estimating in-memory data size %s", sz)
        return sz

    def get_read_tasks(self, parallelism: int) -> list[ReadTask]:
            read_tasks = []

            for batch in self._result_batches:
                # Map the batch metadata to the ray block metadata.
                metadata = BlockMetadata(
                    num_rows=batch.rowcount,
                    size_bytes=batch.uncompressed_size,
                    schema=pa.schema(
                        [
                            pa.field(
                                s.name,
                                FIELD_TYPE_TO_PA_TYPE[
                                    s.type_code
                                ]
                            )
                            for s in batch.schema
                        ]
                    ),
                    input_files=None,
                    exec_stats=None
                )

                # create a lazy handler that will load up the
                # ResultBatch in the worker and do the actual
                # pull from snowflake.
                _r_task = LazyReadTask(
                    arrow_batch=batch,
                    metadata=metadata
                )

                read_tasks.append(_r_task)

            return read_tasks

# This read task is what executes in the worker(s) and pulls the data
# from snowflake and returns an PyArrow table.
class LazyReadTask(ReadTask):
    def __init__(self, arrow_batch: ResultBatch, metadata: BlockMetadata):
        self._arrow_batch = arrow_batch
        self._metadata = metadata

    def _read_fn(self) -> Iterable[pa.Table]:
        ray_data_logger.debug(
            "Reading %s rows from Snowflake", self._metadata.num_rows
        )
        return [self._arrow_batch.to_arrow()]
          

  
from ray.data.block import Block
from ray.data.datasource import Reader
from snowflake.connector import connect

class SnowflakeDatasource(Datasource):
    def __init__(self, connection_args: dict, query: str):
        self._connection_args = connection_args
        self._query = query

    def create_reader(self, **read_args) -> Reader:
        # Yesss! This is the Reader you had just implemented.
        return _SnowflakeDatasourceReader(
            connection_args=self._connection_args,
            query=self._query
        )
# This is it. You are not missing anything.
# To reaffirm. This is it. You are not missing anything.