chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,363 @@
|
||||
import logging
|
||||
import math
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
from ray.data._internal.util import _check_import
|
||||
from ray.data.block import Block, BlockAccessor, BlockMetadata
|
||||
from ray.data.datasource.datasource import Datasource, ReadTask
|
||||
from ray.util.annotations import DeveloperAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.context import DataContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _is_filter_string_safe(filter_str: str) -> bool:
|
||||
in_string = False
|
||||
escape_next = False
|
||||
for c in filter_str:
|
||||
if in_string:
|
||||
# If we're inside a string, check if we're closing it.
|
||||
if c == "'" and not escape_next:
|
||||
in_string = False
|
||||
escape_next = (c == "\\") and not escape_next
|
||||
else:
|
||||
# If we're not in a string, entering one if we see a single quote
|
||||
if c == "'":
|
||||
in_string = True
|
||||
escape_next = False
|
||||
# Disallow semicolon if we're not in a string
|
||||
elif c == ";":
|
||||
return False
|
||||
else:
|
||||
escape_next = False
|
||||
# If we end inside a string, it's suspicious, but let's allow
|
||||
# it to be further validated by the DB. Just return True here.
|
||||
return True
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class ClickHouseDatasource(Datasource):
|
||||
"""
|
||||
A Ray datasource for reading from ClickHouse.
|
||||
|
||||
Args:
|
||||
table: Fully qualified table or view identifier (e.g.,
|
||||
"default.table_name").
|
||||
dsn: A string in DSN (Data Source Name) HTTP format (e.g.,
|
||||
"clickhouse+http://username:password@host:8124/default").
|
||||
For more information, see `ClickHouse Connection String doc
|
||||
<https://clickhouse.com/docs/en/integrations/sql-clients/cli#connection_string>`_.
|
||||
columns: Optional List of columns to select from the data source.
|
||||
If no columns are specified, all columns will be selected by default.
|
||||
filter: Optional SQL filter string that will be used in the
|
||||
WHERE statement (e.g., "label = 2 AND text IS NOT NULL").
|
||||
The filter must be valid for use in a ClickHouse SQL WHERE clause.
|
||||
Note: Parallel reads are not currently supported when a filter is set.
|
||||
Specifying a filter forces the parallelism to 1 to ensure deterministic
|
||||
and consistent results. For more information, see
|
||||
`ClickHouse SQL WHERE Clause doc
|
||||
<https://clickhouse.com/docs/en/sql-reference/statements/select/where>`_.
|
||||
order_by: Optional Tuple containing a list of columns to order by
|
||||
and a boolean indicating the order. Note: order_by is required to
|
||||
support parallelism.
|
||||
client_settings: Optional ClickHouse server settings to be used with the
|
||||
session/every request. For more information, see
|
||||
`ClickHouse Client Settings doc
|
||||
<https://clickhouse.com/docs/en/integrations/python#settings-argument>`_.
|
||||
client_kwargs: Optional Additional keyword arguments to pass to the
|
||||
ClickHouse client. For more information,
|
||||
see `ClickHouse Core Settings doc
|
||||
<https://clickhouse.com/docs/en/integrations/python#additional-options>`_.
|
||||
"""
|
||||
|
||||
NUM_SAMPLE_ROWS = 100
|
||||
MIN_ROWS_PER_READ_TASK = 50
|
||||
_BASE_QUERY = "SELECT {select_clause} FROM {table}"
|
||||
_EXPLAIN_FILTERS_QUERY = "EXPLAIN SELECT 1 FROM {table} WHERE {filter_clause}"
|
||||
_SIZE_ESTIMATE_QUERY = "SELECT SUM(byteSize(*)) AS estimate FROM ({query})"
|
||||
_COUNT_ESTIMATE_QUERY = "SELECT COUNT(*) AS estimate FROM ({query})"
|
||||
_SAMPLE_BLOCK_QUERY = "{query} LIMIT {limit_row_count}"
|
||||
_FIRST_BLOCK_QUERY = """
|
||||
{query}
|
||||
FETCH FIRST {fetch_row_count} {fetch_row_or_rows} ONLY
|
||||
"""
|
||||
_NEXT_BLOCK_QUERY = """
|
||||
{query}
|
||||
OFFSET {offset_row_count} {offset_row_or_rows}
|
||||
FETCH NEXT {fetch_row_count} {fetch_row_or_rows} ONLY
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
table: str,
|
||||
dsn: str,
|
||||
columns: Optional[List[str]] = None,
|
||||
filter: Optional[str] = None,
|
||||
order_by: Optional[Tuple[List[str], bool]] = None,
|
||||
client_settings: Optional[Dict[str, Any]] = None,
|
||||
client_kwargs: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
self._table = table
|
||||
self._dsn = dsn
|
||||
self._columns = columns
|
||||
self._filter = filter
|
||||
self._order_by = order_by
|
||||
self._client_settings = client_settings or {}
|
||||
self._client_kwargs = client_kwargs or {}
|
||||
self._query = self._generate_query()
|
||||
|
||||
def _init_client(self):
|
||||
_check_import(self, module="clickhouse_connect", package="clickhouse-connect")
|
||||
import clickhouse_connect
|
||||
|
||||
return clickhouse_connect.get_client(
|
||||
dsn=self._dsn,
|
||||
settings=self._client_settings or {},
|
||||
**self._client_kwargs or {},
|
||||
)
|
||||
|
||||
def _validate_filter(self):
|
||||
if not self._filter:
|
||||
return
|
||||
# Minimal lexical check (regex or manual approach for semicolons, etc.).
|
||||
if not _is_filter_string_safe(self._filter):
|
||||
raise ValueError(
|
||||
f"Invalid characters outside of "
|
||||
f"string literals in filter: {self._filter}"
|
||||
)
|
||||
# Test "EXPLAIN" query to confirm parse-ability and catch expression errors.
|
||||
client = self._init_client()
|
||||
try:
|
||||
test_query = self._EXPLAIN_FILTERS_QUERY.format(
|
||||
table=self._table,
|
||||
filter_clause=self._filter,
|
||||
)
|
||||
client.query(test_query)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Invalid filter expression: {self._filter}. Error: {e}",
|
||||
)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
def _generate_query(self) -> str:
|
||||
query = self._BASE_QUERY.format(
|
||||
select_clause=", ".join(self._columns) if self._columns else "*",
|
||||
table=self._table,
|
||||
)
|
||||
if self._filter:
|
||||
self._validate_filter()
|
||||
query += f" WHERE {self._filter}"
|
||||
if self._order_by:
|
||||
columns, desc = self._order_by
|
||||
direction = " DESC" if desc else ""
|
||||
if len(columns) == 1:
|
||||
query += f" ORDER BY {columns[0]}{direction}"
|
||||
elif len(columns) > 1:
|
||||
columns_clause = ", ".join(columns)
|
||||
query += f" ORDER BY ({columns_clause}){direction}"
|
||||
return query
|
||||
|
||||
def _build_block_query(self, limit_row_count: int, offset_row_count: int) -> str:
|
||||
if offset_row_count == 0:
|
||||
# The first block query is optimized to use FETCH FIRST clause
|
||||
# with an OFFSET specified.
|
||||
return self._FIRST_BLOCK_QUERY.format(
|
||||
query=self._query,
|
||||
fetch_row_count=limit_row_count,
|
||||
fetch_row_or_rows="ROWS" if limit_row_count > 1 else "ROW",
|
||||
)
|
||||
# Subsequent block queries use OFFSET and FETCH NEXT clauses to read the
|
||||
# next block of data.
|
||||
return self._NEXT_BLOCK_QUERY.format(
|
||||
query=self._query,
|
||||
offset_row_count=offset_row_count,
|
||||
offset_row_or_rows="ROWS" if offset_row_count > 1 else "ROW",
|
||||
fetch_row_count=limit_row_count,
|
||||
fetch_row_or_rows="ROWS" if limit_row_count > 1 else "ROW",
|
||||
)
|
||||
|
||||
def _create_read_fn(
|
||||
self,
|
||||
query: str,
|
||||
) -> Callable[[], Iterable[Block]]:
|
||||
def read_fn() -> Iterable[Block]:
|
||||
return [self._execute_block_query(query)]
|
||||
|
||||
return read_fn
|
||||
|
||||
def _get_sampled_estimates(self):
|
||||
if self._order_by is not None:
|
||||
# If the query is ordered, we can use a FETCH clause to get a sample.
|
||||
# This reduces the CPU overhead on ClickHouse and speeds up the
|
||||
# estimation query.
|
||||
query = self._FIRST_BLOCK_QUERY.format(
|
||||
query=self._query,
|
||||
fetch_row_count=self.NUM_SAMPLE_ROWS,
|
||||
fetch_row_or_rows="ROWS" if self.NUM_SAMPLE_ROWS > 1 else "ROW",
|
||||
)
|
||||
else:
|
||||
# If the query is not ordered, we need to use a LIMIT clause to
|
||||
# get a sample.
|
||||
query = self._SAMPLE_BLOCK_QUERY.format(
|
||||
query=self._query,
|
||||
limit_row_count=self.NUM_SAMPLE_ROWS,
|
||||
)
|
||||
sample_block_accessor = BlockAccessor.for_block(
|
||||
self._execute_block_query(query)
|
||||
)
|
||||
estimated_size_bytes_per_row = math.ceil(
|
||||
sample_block_accessor.size_bytes() / sample_block_accessor.num_rows()
|
||||
)
|
||||
sample_block_schema = sample_block_accessor.schema()
|
||||
return estimated_size_bytes_per_row, sample_block_schema
|
||||
|
||||
def _get_estimate_count(self) -> Optional[int]:
|
||||
return self._execute_estimate_query(self._COUNT_ESTIMATE_QUERY)
|
||||
|
||||
def _get_estimate_size(self) -> Optional[int]:
|
||||
return self._execute_estimate_query(self._SIZE_ESTIMATE_QUERY)
|
||||
|
||||
def _execute_estimate_query(self, estimate_query: str) -> Optional[int]:
|
||||
client = self._init_client()
|
||||
try:
|
||||
# Estimate queries wrap around the primary query, self._query.
|
||||
# This allows us to use self._query as a sub-query to efficiently
|
||||
# and accurately estimate the size or count of the result set.
|
||||
query = estimate_query.format(query=self._query)
|
||||
result = client.query(query)
|
||||
if result and len(result.result_rows) > 0:
|
||||
estimate = result.result_rows[0][0]
|
||||
return int(estimate) if estimate is not None else None
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to execute estimate query: {e}")
|
||||
finally:
|
||||
client.close()
|
||||
return None
|
||||
|
||||
def _execute_block_query(self, query: str) -> Block:
|
||||
import pyarrow as pa
|
||||
|
||||
client = self._init_client()
|
||||
try:
|
||||
with client.query_arrow_stream(query) as stream:
|
||||
record_batches = list(stream) # Collect all record batches
|
||||
return pa.Table.from_batches(record_batches)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to execute block query: {e}")
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
def estimate_inmemory_data_size(self) -> Optional[int]:
|
||||
"""
|
||||
Estimate the in-memory data size for the query.
|
||||
|
||||
Returns:
|
||||
Estimated in-memory data size in bytes, or
|
||||
None if the estimation cannot be performed.
|
||||
"""
|
||||
return self._get_estimate_size()
|
||||
|
||||
def get_read_tasks(
|
||||
self,
|
||||
parallelism: int,
|
||||
per_task_row_limit: Optional[int] = None,
|
||||
data_context: Optional["DataContext"] = None,
|
||||
) -> List[ReadTask]:
|
||||
"""
|
||||
Create read tasks for the ClickHouse query.
|
||||
|
||||
Args:
|
||||
parallelism: The desired number of partitions to read the data into.
|
||||
- If ``order_by`` is not set, parallelism will be forced to 1.
|
||||
- If ``filter`` is set, parallelism will also be forced to 1
|
||||
to ensure deterministic results.
|
||||
per_task_row_limit: Maximum number of rows allowed in each emitted
|
||||
task. Blocks larger than this limit will be sliced before
|
||||
being yielded downstream.
|
||||
data_context: The data context to use to get read tasks. Not used by this
|
||||
datasource.
|
||||
|
||||
Returns:
|
||||
A list of read tasks to be executed.
|
||||
"""
|
||||
num_rows_total = self._get_estimate_count()
|
||||
if num_rows_total == 0 or num_rows_total is None:
|
||||
return []
|
||||
parallelism = min(
|
||||
parallelism, math.ceil(num_rows_total / self.MIN_ROWS_PER_READ_TASK)
|
||||
)
|
||||
# To ensure consistent order of query results, self._order_by
|
||||
# must be specified and self.filter must not be specified
|
||||
# in order to support parallelism.
|
||||
if self._filter is not None and parallelism > 1:
|
||||
logger.warning(
|
||||
"ClickHouse datasource does not currently support parallel reads "
|
||||
"when a filter is set; falling back to parallelism of 1."
|
||||
)
|
||||
# When filter is specified and parallelism is greater than 1,
|
||||
# we need to reduce parallelism to 1 to ensure consistent results.
|
||||
parallelism = 1
|
||||
# To ensure consistent order of query results, self._order_by
|
||||
# must be specified in order to support parallelism.
|
||||
if self._order_by is None and parallelism > 1:
|
||||
logger.warning(
|
||||
"ClickHouse datasource requires dataset to be explicitly ordered "
|
||||
"to support parallelism; falling back to parallelism of 1."
|
||||
)
|
||||
# When order_by is not specified and parallelism is greater than 1,
|
||||
# we need to reduce parallelism to 1 to ensure consistent results.
|
||||
parallelism = 1
|
||||
# By reducing parallelism to 1 when either of the conditions above are met,
|
||||
# we ensure the downstream process is treated exactly as a non-parallelized
|
||||
# (single block) process would be, thus ensuring output consistency.
|
||||
num_rows_per_block = num_rows_total // parallelism
|
||||
num_blocks_with_extra_row = num_rows_total % parallelism
|
||||
(
|
||||
estimated_size_bytes_per_row,
|
||||
sample_block_schema,
|
||||
) = self._get_sampled_estimates()
|
||||
|
||||
def _get_read_task(
|
||||
block_rows: int, offset_rows: int, parallelized: bool
|
||||
) -> ReadTask:
|
||||
if parallelized:
|
||||
# When parallelized, we need to build a block query with OFFSET
|
||||
# and FETCH clauses.
|
||||
query = self._build_block_query(block_rows, offset_rows)
|
||||
else:
|
||||
# When not parallelized, we can use the original query without
|
||||
# OFFSET and FETCH clauses.
|
||||
query = self._query
|
||||
return ReadTask(
|
||||
self._create_read_fn(query),
|
||||
BlockMetadata(
|
||||
num_rows=block_rows,
|
||||
size_bytes=estimated_size_bytes_per_row * block_rows,
|
||||
input_files=None,
|
||||
exec_stats=None,
|
||||
),
|
||||
schema=sample_block_schema,
|
||||
per_task_row_limit=per_task_row_limit,
|
||||
)
|
||||
|
||||
if parallelism == 1:
|
||||
# When parallelism is 1, we can read the entire dataset in a single task.
|
||||
# We then optimize this scenario by using self._query directly without
|
||||
# unnecessary OFFSET and FETCH clauses.
|
||||
return [_get_read_task(num_rows_total, 0, False)]
|
||||
|
||||
# Otherwise we need to split the dataset into multiple tasks.
|
||||
# Each task will include OFFSET and FETCH clauses to efficiently
|
||||
# read a subset of the dataset.
|
||||
read_tasks = []
|
||||
offset = 0
|
||||
for i in range(parallelism):
|
||||
this_block_size = num_rows_per_block
|
||||
if i < num_blocks_with_extra_row:
|
||||
this_block_size += 1
|
||||
read_tasks.append(_get_read_task(this_block_size, offset, True))
|
||||
offset += this_block_size
|
||||
return read_tasks
|
||||
Reference in New Issue
Block a user