Coverage for python/lsst/dax/apdb/cassandra/cassandra_utils.py: 22%
147 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-07 01:39 -0800
1# This file is part of dax_apdb.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22from __future__ import annotations
24__all__ = [
25 "ApdbCassandraTableData",
26 "PreparedStatementCache",
27 "StatementFactory",
28 "literal",
29 "quote_id",
30 "raw_data_factory",
31 "select_concurrent",
32]
34import logging
35import warnings
36from collections.abc import Collection, Iterable, Iterator, Sequence
37from datetime import datetime, timedelta
38from typing import Any
39from uuid import UUID
41import felis.datamodel
42import numpy as np
43import pandas
45# If cassandra-driver is not there the module can still be imported
46# but things will not work.
47try:
48 import cassandra.concurrent
49 from cassandra.cluster import EXEC_PROFILE_DEFAULT, Session
50 from cassandra.query import PreparedStatement, SimpleStatement
52 CASSANDRA_IMPORTED = True
53except ImportError:
54 CASSANDRA_IMPORTED = False
55 EXEC_PROFILE_DEFAULT = object()
57from .. import schema_model
58from ..apdbReplica import ApdbTableData
59from .queries import Query
61_LOG = logging.getLogger(__name__)
64class ApdbCassandraTableData(ApdbTableData):
65 """Implementation of ApdbTableData that wraps Cassandra raw data."""
67 def __init__(self, columns: list[str], rows: list[tuple]):
68 self._columns = columns
69 self._rows = rows
70 self._column_types: dict[str, felis.datamodel.DataType] = {}
72 def set_column_types(self, types: dict[str, felis.datamodel.DataType]) -> None:
73 """Update column types.
75 Parameters
76 ----------
77 types : `dict`[`str`, `felis.datamodel.DataType`]
78 Mapping of column name its type.
80 Notes
81 -----
82 Due to the way how instances of this class are constructed it is
83 impossible to pass types of columns to the constructor, instead we will
84 need to make a call to this method after construction.
85 """
86 self._column_types = types
88 def column_names(self) -> Sequence[str]:
89 # docstring inherited
90 return self._columns
92 def column_index(self, name: str) -> int:
93 # docstring inherited
94 return self._columns.index(name)
96 def column_defs(self) -> Sequence[tuple[str, felis.datamodel.DataType]]:
97 return tuple((column, self._column_types[column]) for column in self._columns)
99 def rows(self) -> Collection[tuple]:
100 # docstring inherited
101 return self._rows
103 def append(self, other: ApdbCassandraTableData) -> None:
104 """Extend rows in this table with rows in other table"""
105 if self._columns != other._columns:
106 raise ValueError(f"Different columns returned by queries: {self._columns} and {other._columns}")
107 self._rows.extend(other._rows)
109 def project(self, *, drop: Iterable[str] = set()) -> None:
110 """Modify data in place by droppiing some columns."""
111 drop_set = set(drop)
112 if not drop_set:
113 return
115 drop_idx = []
116 for idx, col_name in enumerate(self._columns):
117 if col_name in drop_set:
118 drop_idx.append(idx)
119 # Have to reverse it so deletion does not change index.
120 drop_idx.reverse()
122 for row_idx in range(len(self._rows)):
123 row = list(self._rows[row_idx])
124 for idx in drop_idx:
125 del row[idx]
126 self._rows[row_idx] = tuple(row)
128 for idx in drop_idx:
129 del self._columns[idx]
131 def to_pandas(self, table: schema_model.Table) -> pandas.DataFrame:
132 """Convert data to pandas DataFrame.
134 Parameters
135 ----------
136 table : `schema_model.Table`
137 Table schema matching the data in this instance.
139 Returns
140 -------
141 dataframe : `pandas.DataFrame`
142 Resulting DataFrame.
143 """
144 column_types = {column_def.name: column_def.pandas_type for column_def in table.columns}
146 # In rare cases there could be columns that are not in the configured
147 # schema, e.g. during schema migrations. Use object column type for
148 # them but also produce a warning.
149 extra_columns = [column for column in self._columns if column not in column_types]
150 if extra_columns:
151 warnings.warn(
152 f"Query result includes column(s) do not appear in schema for table {table.name}: "
153 f"{', '.join(extra_columns)}",
154 stacklevel=2,
155 )
157 if not self._rows:
158 column_data = {}
159 for column in self._columns:
160 column_data[column] = pandas.Series(dtype=column_types.get(column, object))
161 return pandas.DataFrame(column_data)
163 # To avoid nested loops convert everything to ndarray.
164 array = np.array(self._rows, dtype=object)
165 array = array.T
166 column_data = {}
167 for i, column in enumerate(self._columns):
168 column_data[column] = pandas.Series(array[i], dtype=column_types.get(column, object))
169 return pandas.DataFrame(column_data)
171 def __iter__(self) -> Iterator[tuple]:
172 """Make it look like a row iterator, needed for some odd logic."""
173 return iter(self._rows)
176class PreparedStatementCache:
177 """Cache for prepared Cassandra statements"""
179 def __init__(self, session: Session) -> None:
180 self._session = session
181 self._prepared_statements: dict[str, PreparedStatement] = {}
183 def prepare(self, query: str) -> PreparedStatement:
184 """Convert query string into prepared statement."""
185 stmt = self._prepared_statements.get(query)
186 if stmt is None:
187 stmt = self._session.prepare(query)
188 self._prepared_statements[query] = stmt
189 return stmt
192class StatementFactory:
193 """Class that builds Cassandra statements from Query objects."""
195 def __init__(self, session: Session, cache: PreparedStatementCache | None = None) -> None:
196 self._session = session
197 self._prepared_cache = cache
199 def __call__(self, query: Query, prepare: bool = False) -> PreparedStatement | SimpleStatement:
200 """Generate Cassandra statement from Query.
202 Parameters
203 ----------
204 query : `Query`
205 Query to convert to Cassandra statement.
206 prepare : `bool`, optional
207 if `True` then generate prepared statement (and only if
208 ``query.can_prepare`` is True).
210 Returns
211 -------
212 statement : `PreparedStatement` or `SimpleStatement`
213 Statement to execute.
214 """
215 if prepare and query.can_prepare and self._prepared_cache is not None:
216 stmt = self._prepared_cache.prepare(query.render("?"))
217 else:
218 stmt = SimpleStatement(query.render("%s"))
219 return stmt
221 def with_params(
222 self, query: Query, prepare: bool = False
223 ) -> tuple[PreparedStatement | SimpleStatement, tuple]:
224 """Generate Cassandra statement and its parameters from Query.
226 Parameters
227 ----------
228 query : `Query`
229 Query to convert to Cassandra statement.
230 prepare : `bool`, optional
231 if `True` then generate prepared statement (and only if
232 ``query.can_prepare`` is True).
234 Returns
235 -------
236 statement : `PreparedStatement` or `SimpleStatement`
237 Statement to execute.
238 parameters : `tuple`
239 Parameters for this statement.
240 """
241 stmt = self(query, prepare)
242 return stmt, query.parameters
245def raw_data_factory(colnames: list[str], rows: list[tuple]) -> ApdbCassandraTableData:
246 """Make 2-element tuple containing unmodified data: list of column names
247 and list of rows.
249 Parameters
250 ----------
251 colnames : `list` [ `str` ]
252 Names of the columns.
253 rows : `list` of `tuple`
254 Result rows.
256 Returns
257 -------
258 data : `ApdbCassandraTableData`
259 Input data wrapped into ApdbCassandraTableData.
261 Notes
262 -----
263 When using this method as row factory for Cassandra, the resulting
264 object should be accessed in a non-standard way using
265 `ResultSet._current_rows` attribute.
266 """
267 return ApdbCassandraTableData(colnames, rows)
270def execute_concurrent(
271 session: Session,
272 statements: list[tuple],
273 *,
274 execution_profile: object = EXEC_PROFILE_DEFAULT,
275 concurrency: int = 100,
276) -> None:
277 """Wrapp call to `cassandra.concurrent.execute_concurrent` to avoid
278 importing cassandra in other modules.
279 """
280 cassandra.concurrent.execute_concurrent(
281 session,
282 statements,
283 concurrency=concurrency,
284 execution_profile=execution_profile,
285 )
288def select_concurrent(
289 session: Session, statements: list[tuple], execution_profile: str, concurrency: int
290) -> pandas.DataFrame | ApdbCassandraTableData | list:
291 """Execute bunch of queries concurrently and merge their results into
292 a single result.
294 Parameters
295 ----------
296 statements : `list` [ `tuple` ]
297 List of statements and their parameters, passed directly to
298 ``execute_concurrent()``.
299 execution_profile : `str`
300 Execution profile name.
302 Returns
303 -------
304 result
305 Combined result of multiple statements, type of the result depends on
306 specific row factory defined in execution profile. If row factory is
307 `pandas_dataframe_factory` then pandas DataFrame is created from a
308 combined result. If row factory is `raw_data_factory` then
309 `ApdbCassandraTableData` is built from all records. Otherwise a list of
310 rows is returned, type of each row is determined by the row factory.
312 Notes
313 -----
314 This method can raise any exception that is raised by one of the provided
315 statements.
316 """
317 results = cassandra.concurrent.execute_concurrent(
318 session,
319 statements,
320 results_generator=True,
321 raise_on_first_error=False,
322 concurrency=concurrency,
323 execution_profile=execution_profile,
324 )
326 ep = session.get_execution_profile(execution_profile)
327 if ep.row_factory is raw_data_factory:
328 # Collect rows into a single list and build Dataframe out of that
329 _LOG.debug("making raw data out of rows/columns")
330 table_data: ApdbCassandraTableData | None = None
331 for success, result in results:
332 if success:
333 data = result._current_rows
334 assert isinstance(data, ApdbCassandraTableData)
335 if table_data is None:
336 table_data = data
337 else:
338 table_data.append(data)
339 else:
340 _LOG.error("error returned by query: %s", result)
341 raise result
342 if table_data is None:
343 table_data = ApdbCassandraTableData([], [])
344 return table_data
346 else:
347 # Just concatenate all rows into a single collection.
348 rows = []
349 for success, result in results:
350 if success:
351 rows.extend(result)
352 else:
353 _LOG.error("error returned by query: %s", result)
354 raise result
355 _LOG.debug("number of rows: %s", len(rows))
356 return rows
359def literal(v: Any) -> Any:
360 """Transform object into a value for the query."""
361 if v is None or v is pandas.NA:
362 v = None
363 elif isinstance(v, datetime):
364 v = int((v - datetime(1970, 1, 1)) / timedelta(seconds=1) * 1000)
365 elif isinstance(v, bytes | str | UUID | int):
366 pass
367 elif isinstance(v, np.bool_):
368 v = bool(v)
369 else:
370 try:
371 if not np.isfinite(v):
372 v = None
373 except TypeError:
374 pass
375 return v
378def quote_id(columnName: str) -> str:
379 """Smart quoting for column names. Lower-case names are not quoted."""
380 if not columnName.islower():
381 columnName = '"' + columnName + '"'
382 return columnName