chore: import upstream snapshot with attribution

python/ray/data/_internal/datasource/sql_datasource.py

@@ -0,0 +1,211 @@
+import logging
+import math
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator, List, Optional
+
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource.datasource import Datasource, ReadTask
+
+Connection = Any  # A Python DB API2-compliant `Connection` object.
+Cursor = Any  # A Python DB API2-compliant `Cursor` object.
+
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+
+logger = logging.getLogger(__name__)
+
+
+def _cursor_to_block(cursor) -> Block:
+    import pyarrow as pa
+
+    rows = cursor.fetchall()
+    # Each `column_description` is a 7-element sequence. The first element is the column
+    # name. To learn more, read https://peps.python.org/pep-0249/#description.
+    columns = [column_description[0] for column_description in cursor.description]
+    pydict = {column: [row[i] for row in rows] for i, column in enumerate(columns)}
+    return pa.Table.from_pydict(pydict)
+
+
+def _check_connection_is_dbapi2_compliant(connection) -> None:
+    for attr in "close", "commit", "cursor":
+        if not hasattr(connection, attr):
+            raise ValueError(
+                "Your `connection_factory` created a `Connection` object without a "
+                f"{attr!r} method, but this method is required by the Python DB API2 "
+                "specification. Check that your database connector is DB API2-"
+                "compliant. To learn more, read https://peps.python.org/pep-0249/."
+            )
+
+
+def _check_cursor_is_dbapi2_compliant(cursor) -> None:
+    # These aren't all the methods required by the specification, but it's all the ones
+    # we care about.
+    for attr in "execute", "executemany", "fetchone", "fetchall", "description":
+        if not hasattr(cursor, attr):
+            raise ValueError(
+                "Your database connector created a `Cursor` object without a "
+                f"{attr!r} method, but this method is required by the Python DB API2 "
+                "specification. Check that your database connector is DB API2-"
+                "compliant. To learn more, read https://peps.python.org/pep-0249/."
+            )
+
+
+@contextmanager
+def _connect(connection_factory: Callable[[], Connection]) -> Iterator[Cursor]:
+    connection = connection_factory()
+    _check_connection_is_dbapi2_compliant(connection)
+
+    try:
+        cursor = connection.cursor()
+        _check_cursor_is_dbapi2_compliant(cursor)
+        yield cursor
+        connection.commit()
+    except Exception:
+        # `rollback` is optional since not all databases provide transaction support.
+        try:
+            connection.rollback()
+        except Exception as e:
+            # Each connector implements its own `NotSupportError` class, so we check
+            # the exception's name instead of using `isinstance`.
+            if (
+                isinstance(e, AttributeError)
+                or e.__class__.__name__ == "NotSupportedError"
+            ):
+                pass
+        raise
+    finally:
+        connection.close()
+
+
+def _execute(cursor: Cursor, sql: str, params: Optional[Any]) -> None:
+    if params is None:
+        cursor.execute(sql)
+    else:
+        cursor.execute(sql, params)
+
+
+class SQLDatasource(Datasource):
+    MIN_ROWS_PER_READ_TASK = 50
+
+    def __init__(
+        self,
+        sql: str,
+        connection_factory: Callable[[], Connection],
+        shard_hash_fn: str,
+        shard_keys: Optional[List[str]] = None,
+        sql_params: Optional[Any] = None,
+    ):
+        self.sql = sql
+        if shard_keys and len(shard_keys) > 1:
+            self.shard_keys = f"CONCAT({','.join(shard_keys)})"
+        elif shard_keys and len(shard_keys) == 1:
+            self.shard_keys = f"{shard_keys[0]}"
+        else:
+            self.shard_keys = None
+        self.shard_hash_fn = shard_hash_fn
+        self.connection_factory = connection_factory
+        self.sql_params = sql_params
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        return None
+
+    def supports_sharding(self, parallelism: int) -> bool:
+        """Check if database supports sharding with MOD/ABS/CONCAT operations.
+
+        Args:
+            parallelism: The number of shards to split the read into.
+
+        Returns:
+            bool: True if sharding is supported, False otherwise.
+        """
+        if parallelism <= 1 or self.shard_keys is None:
+            return False
+
+        # Test if database supports required operations (MOD, ABS, MD5, CONCAT)
+        # by executing a sample query
+        hash_fn = self.shard_hash_fn
+        query = (
+            f"SELECT COUNT(1) FROM ({self.sql}) as T"
+            f" WHERE MOD(ABS({hash_fn}({self.shard_keys})), {parallelism}) = 0"
+        )
+        try:
+            with _connect(self.connection_factory) as cursor:
+                _execute(cursor, query, self.sql_params)
+            return True
+        except Exception as e:
+            logger.info(f"Database does not support sharding: {str(e)}.")
+            return False
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        def fallback_read_fn() -> Iterable[Block]:
+            """Read all data in a single block when sharding is not supported."""
+            with _connect(self.connection_factory) as cursor:
+                _execute(cursor, self.sql, self.sql_params)
+                return [_cursor_to_block(cursor)]
+
+        # Check if sharding is supported by the database first
+        # If not, fall back to reading all data in a single task without counting rows
+        if not self.supports_sharding(parallelism):
+            logger.info(
+                "Sharding is not supported. "
+                "Falling back to reading all data in a single task."
+            )
+            metadata = BlockMetadata(None, None, None, None)
+            return [ReadTask(fallback_read_fn, metadata)]
+
+        # Only perform the expensive COUNT(*) query if sharding is supported
+        num_rows_total = self._get_num_rows()
+
+        if num_rows_total == 0:
+            return []
+
+        parallelism = min(
+            parallelism, math.ceil(num_rows_total / self.MIN_ROWS_PER_READ_TASK)
+        )
+        num_rows_per_block = num_rows_total // parallelism
+        num_blocks_with_extra_row = num_rows_total % parallelism
+
+        tasks = []
+        for i in range(parallelism):
+            num_rows = num_rows_per_block
+            if i < num_blocks_with_extra_row:
+                num_rows += 1
+            read_fn = self._create_parallel_read_fn(i, parallelism)
+            metadata = BlockMetadata(
+                num_rows=num_rows,
+                size_bytes=None,
+                input_files=None,
+                exec_stats=None,
+            )
+            tasks.append(
+                ReadTask(read_fn, metadata, per_task_row_limit=per_task_row_limit)
+            )
+
+        return tasks
+
+    def _get_num_rows(self) -> int:
+        with _connect(self.connection_factory) as cursor:
+            _execute(cursor, f"SELECT COUNT(*) FROM ({self.sql}) as T", self.sql_params)
+            return cursor.fetchone()[0]
+
+    def _create_parallel_read_fn(self, task_id: int, parallelism: int):
+        hash_fn = self.shard_hash_fn
+        query = (
+            f"SELECT * FROM ({self.sql}) as T "
+            f"WHERE MOD(ABS({hash_fn}({self.shard_keys})), {parallelism}) = {task_id}"
+        )
+
+        def read_fn() -> Iterable[Block]:
+            with _connect(self.connection_factory) as cursor:
+                _execute(cursor, query, self.sql_params)
+                block = _cursor_to_block(cursor)
+                return [block]
+
+        return read_fn