Coverage for python / lsst / dax / apdb / cassandra / cassandra_utils.py: 20%
132 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-29 01:35 -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 "literal",
28 "quote_id",
29 "raw_data_factory",
30 "select_concurrent",
31]
33import logging
34import warnings
35from collections.abc import Collection, Iterable, Iterator, Sequence
36from datetime import datetime, timedelta
37from typing import Any
38from uuid import UUID
40import felis.datamodel
41import numpy as np
42import pandas
44# If cassandra-driver is not there the module can still be imported
45# but things will not work.
46try:
47 import cassandra.concurrent
48 from cassandra.cluster import EXEC_PROFILE_DEFAULT, Session
49 from cassandra.query import PreparedStatement
51 CASSANDRA_IMPORTED = True
52except ImportError:
53 CASSANDRA_IMPORTED = False
54 EXEC_PROFILE_DEFAULT = object()
56from .. import schema_model
57from ..apdbReplica import ApdbTableData
59_LOG = logging.getLogger(__name__)
62class ApdbCassandraTableData(ApdbTableData):
63 """Implementation of ApdbTableData that wraps Cassandra raw data."""
65 def __init__(self, columns: list[str], rows: list[tuple]):
66 self._columns = columns
67 self._rows = rows
68 self._column_types: dict[str, felis.datamodel.DataType] = {}
70 def set_column_types(self, types: dict[str, felis.datamodel.DataType]) -> None:
71 """Update column types.
73 Parameters
74 ----------
75 types : `dict`[`str`, `felis.datamodel.DataType`]
76 Mapping of column name its type.
78 Notes
79 -----
80 Due to the way how instances of this class are constructed it is
81 impossible to pass types of columns to the constructor, instead we will
82 need to make a call to this method after construction.
83 """
84 self._column_types = types
86 def column_names(self) -> Sequence[str]:
87 # docstring inherited
88 return self._columns
90 def column_defs(self) -> Sequence[tuple[str, felis.datamodel.DataType]]:
91 return tuple((column, self._column_types[column]) for column in self._columns)
93 def rows(self) -> Collection[tuple]:
94 # docstring inherited
95 return self._rows
97 def append(self, other: ApdbCassandraTableData) -> None:
98 """Extend rows in this table with rows in other table"""
99 if self._columns != other._columns:
100 raise ValueError(f"Different columns returned by queries: {self._columns} and {other._columns}")
101 self._rows.extend(other._rows)
103 def project(self, *, drop: Iterable[str] = set()) -> None:
104 """Modify data in place by droppiing some columns."""
105 drop_set = set(drop)
106 if not drop_set:
107 return
109 drop_idx = []
110 for idx, col_name in enumerate(self._columns):
111 if col_name in drop_set:
112 drop_idx.append(idx)
113 # Have to reverse it so deletion does not change index.
114 drop_idx.reverse()
116 for row_idx in range(len(self._rows)):
117 row = list(self._rows[row_idx])
118 for idx in drop_idx:
119 del row[idx]
120 self._rows[row_idx] = tuple(row)
122 for idx in drop_idx:
123 del self._columns[idx]
125 def to_pandas(self, table: schema_model.Table) -> pandas.DataFrame:
126 """Convert data to pandas DataFrame.
128 Parameters
129 ----------
130 table : `schema_model.Table`
131 Table schema matching the data in this instance.
133 Returns
134 -------
135 dataframe : `pandas.DataFrame`
136 Resulting DataFrame.
137 """
138 column_types = {column_def.name: column_def.pandas_type for column_def in table.columns}
140 # In rare cases there could be columns that are not in the configured
141 # schema, e.g. during schema migrations. Use object column type for
142 # them but also produce a warning.
143 extra_columns = [column for column in self._columns if column not in column_types]
144 if extra_columns:
145 warnings.warn(
146 f"Query result includes column(s) do not appear in schema for table {table.name}: "
147 f"{', '.join(extra_columns)}",
148 stacklevel=2,
149 )
151 if not self._rows:
152 column_data = {}
153 for column in self._columns:
154 column_data[column] = pandas.Series(dtype=column_types.get(column, object))
155 return pandas.DataFrame(column_data)
157 # To avoid nested loops convert everything to ndarray.
158 array = np.array(self._rows, dtype=object)
159 array = array.T
160 column_data = {}
161 for i, column in enumerate(self._columns):
162 column_data[column] = pandas.Series(array[i], dtype=column_types.get(column, object))
163 return pandas.DataFrame(column_data)
165 def __iter__(self) -> Iterator[tuple]:
166 """Make it look like a row iterator, needed for some odd logic."""
167 return iter(self._rows)
170class PreparedStatementCache:
171 """Cache for prepared Cassandra statements"""
173 def __init__(self, session: Session) -> None:
174 self._session = session
175 self._prepared_statements: dict[str, PreparedStatement] = {}
177 def prepare(self, query: str) -> PreparedStatement:
178 """Convert query string into prepared statement."""
179 stmt = self._prepared_statements.get(query)
180 if stmt is None:
181 stmt = self._session.prepare(query)
182 self._prepared_statements[query] = stmt
183 return stmt
186def raw_data_factory(colnames: list[str], rows: list[tuple]) -> ApdbCassandraTableData:
187 """Make 2-element tuple containing unmodified data: list of column names
188 and list of rows.
190 Parameters
191 ----------
192 colnames : `list` [ `str` ]
193 Names of the columns.
194 rows : `list` of `tuple`
195 Result rows.
197 Returns
198 -------
199 data : `ApdbCassandraTableData`
200 Input data wrapped into ApdbCassandraTableData.
202 Notes
203 -----
204 When using this method as row factory for Cassandra, the resulting
205 object should be accessed in a non-standard way using
206 `ResultSet._current_rows` attribute.
207 """
208 return ApdbCassandraTableData(colnames, rows)
211def execute_concurrent(
212 session: Session,
213 statements: list[tuple],
214 *,
215 execution_profile: object = EXEC_PROFILE_DEFAULT,
216 concurrency: int = 100,
217) -> None:
218 """Wrapp call to `cassandra.concurrent.execute_concurrent` to avoid
219 importing cassandra in other modules.
220 """
221 cassandra.concurrent.execute_concurrent(
222 session,
223 statements,
224 concurrency=concurrency,
225 execution_profile=execution_profile,
226 )
229def select_concurrent(
230 session: Session, statements: list[tuple], execution_profile: str, concurrency: int
231) -> pandas.DataFrame | ApdbCassandraTableData | list:
232 """Execute bunch of queries concurrently and merge their results into
233 a single result.
235 Parameters
236 ----------
237 statements : `list` [ `tuple` ]
238 List of statements and their parameters, passed directly to
239 ``execute_concurrent()``.
240 execution_profile : `str`
241 Execution profile name.
243 Returns
244 -------
245 result
246 Combined result of multiple statements, type of the result depends on
247 specific row factory defined in execution profile. If row factory is
248 `pandas_dataframe_factory` then pandas DataFrame is created from a
249 combined result. If row factory is `raw_data_factory` then
250 `ApdbCassandraTableData` is built from all records. Otherwise a list of
251 rows is returned, type of each row is determined by the row factory.
253 Notes
254 -----
255 This method can raise any exception that is raised by one of the provided
256 statements.
257 """
258 results = cassandra.concurrent.execute_concurrent(
259 session,
260 statements,
261 results_generator=True,
262 raise_on_first_error=False,
263 concurrency=concurrency,
264 execution_profile=execution_profile,
265 )
267 ep = session.get_execution_profile(execution_profile)
268 if ep.row_factory is raw_data_factory:
269 # Collect rows into a single list and build Dataframe out of that
270 _LOG.debug("making raw data out of rows/columns")
271 table_data: ApdbCassandraTableData | None = None
272 for success, result in results:
273 if success:
274 data = result._current_rows
275 assert isinstance(data, ApdbCassandraTableData)
276 if table_data is None:
277 table_data = data
278 else:
279 table_data.append(data)
280 else:
281 _LOG.error("error returned by query: %s", result)
282 raise result
283 if table_data is None:
284 table_data = ApdbCassandraTableData([], [])
285 return table_data
287 else:
288 # Just concatenate all rows into a single collection.
289 rows = []
290 for success, result in results:
291 if success:
292 rows.extend(result)
293 else:
294 _LOG.error("error returned by query: %s", result)
295 raise result
296 _LOG.debug("number of rows: %s", len(rows))
297 return rows
300def literal(v: Any) -> Any:
301 """Transform object into a value for the query."""
302 if v is None or v is pandas.NA:
303 v = None
304 elif isinstance(v, datetime):
305 v = int((v - datetime(1970, 1, 1)) / timedelta(seconds=1) * 1000)
306 elif isinstance(v, bytes | str | UUID | int):
307 pass
308 elif isinstance(v, np.bool_):
309 v = bool(v)
310 else:
311 try:
312 if not np.isfinite(v):
313 v = None
314 except TypeError:
315 pass
316 return v
319def quote_id(columnName: str) -> str:
320 """Smart quoting for column names. Lower-case names are not quoted."""
321 if not columnName.islower():
322 columnName = '"' + columnName + '"'
323 return columnName