Coverage for python/lsst/dax/apdb/cassandra/apdbCassandra.py: 10%
605 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:35 -0700
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__ = ["ApdbCassandra"]
26import datetime
27import logging
28import random
29import warnings
30from collections import defaultdict
31from collections.abc import Iterable, Iterator, Mapping, Set
32from typing import TYPE_CHECKING, Any, cast
34import numpy as np
35import pandas
37# If cassandra-driver is not there the module can still be imported
38# but ApdbCassandra cannot be instantiated.
39try:
40 import cassandra
41 import cassandra.query
42 from cassandra.query import UNSET_VALUE
44 CASSANDRA_IMPORTED = True
45except ImportError:
46 CASSANDRA_IMPORTED = False
48import astropy.time
49import felis.datamodel
51from lsst import sphgeom
52from lsst.utils.iteration import chunk_iterable
54from ..apdb import Apdb, ApdbConfig
55from ..apdbConfigFreezer import ApdbConfigFreezer
56from ..apdbReplica import ApdbTableData, ReplicaChunk
57from ..apdbSchema import ApdbSchema, ApdbTables
58from ..monitor import MonAgent
59from ..schema_model import Table
60from ..timer import Timer
61from ..versionTuple import IncompatibleVersionError, VersionTuple
62from .apdbCassandraAdmin import ApdbCassandraAdmin
63from .apdbCassandraReplica import ApdbCassandraReplica
64from .apdbCassandraSchema import ApdbCassandraSchema, CreateTableOptions, ExtraTables
65from .apdbMetadataCassandra import ApdbMetadataCassandra
66from .cassandra_utils import (
67 execute_concurrent,
68 literal,
69 quote_id,
70 select_concurrent,
71)
72from .config import ApdbCassandraConfig, ApdbCassandraConnectionConfig, ApdbCassandraTimePartitionRange
73from .connectionContext import ConnectionContext, DbVersions
74from .exceptions import CassandraMissingError
75from .partitioner import Partitioner
76from .sessionFactory import SessionContext, SessionFactory
78if TYPE_CHECKING:
79 from ..apdbMetadata import ApdbMetadata
81_LOG = logging.getLogger(__name__)
83_MON = MonAgent(__name__)
85VERSION = VersionTuple(1, 1, 0)
86"""Version for the code controlling non-replication tables. This needs to be
87updated following compatibility rules when schema produced by this code
88changes.
89"""
92class ApdbCassandra(Apdb):
93 """Implementation of APDB database with Apache Cassandra backend.
95 Parameters
96 ----------
97 config : `ApdbCassandraConfig`
98 Configuration object.
99 """
101 def __init__(self, config: ApdbCassandraConfig):
102 if not CASSANDRA_IMPORTED:
103 raise CassandraMissingError()
105 self._config = config
106 self._keyspace = config.keyspace
107 self._schema = ApdbSchema(config.schema_file, config.schema_name)
109 self._session_factory = SessionFactory(config)
110 self._connection_context: ConnectionContext | None = None
112 @property
113 def _context(self) -> ConnectionContext:
114 """Establish connection if not established and return context."""
115 if self._connection_context is None:
116 session = self._session_factory.session()
117 self._connection_context = ConnectionContext(session, self._config, self._schema.tableSchemas)
119 # Check version compatibility
120 current_versions = DbVersions(
121 schema_version=self._schema.schemaVersion(),
122 code_version=self.apdbImplementationVersion(),
123 replica_version=(
124 ApdbCassandraReplica.apdbReplicaImplementationVersion()
125 if self._connection_context.config.enable_replica
126 else None
127 ),
128 )
129 _LOG.debug("Current versions: %s", current_versions)
130 self._versionCheck(current_versions, self._connection_context.db_versions)
132 if _LOG.isEnabledFor(logging.DEBUG):
133 _LOG.debug("ApdbCassandra Configuration: %s", self._connection_context.config.model_dump())
135 return self._connection_context
137 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
138 """Create `Timer` instance given its name."""
139 return Timer(name, _MON, tags=tags)
141 def _versionCheck(self, current_versions: DbVersions, db_versions: DbVersions) -> None:
142 """Check schema version compatibility."""
143 if not current_versions.schema_version.checkCompatibility(db_versions.schema_version):
144 raise IncompatibleVersionError(
145 f"Configured schema version {current_versions.schema_version} "
146 f"is not compatible with database version {db_versions.schema_version}"
147 )
148 if not current_versions.code_version.checkCompatibility(db_versions.code_version):
149 raise IncompatibleVersionError(
150 f"Current code version {current_versions.code_version} "
151 f"is not compatible with database version {db_versions.code_version}"
152 )
154 # Check replica code version only if replica is enabled.
155 match current_versions.replica_version, db_versions.replica_version:
156 case None, None:
157 pass
158 case VersionTuple() as current, VersionTuple() as stored:
159 if not current.checkCompatibility(stored):
160 raise IncompatibleVersionError(
161 f"Current replication code version {current} "
162 f"is not compatible with database version {stored}"
163 )
164 case _:
165 raise IncompatibleVersionError(
166 f"Current replication code version {current_versions.replica_version} "
167 f"is not compatible with database version {db_versions.replica_version}"
168 )
170 @classmethod
171 def apdbImplementationVersion(cls) -> VersionTuple:
172 """Return version number for current APDB implementation.
174 Returns
175 -------
176 version : `VersionTuple`
177 Version of the code defined in implementation class.
178 """
179 return VERSION
181 def getConfig(self) -> ApdbCassandraConfig:
182 # docstring is inherited from a base class
183 return self._context.config
185 def tableDef(self, table: ApdbTables) -> Table | None:
186 # docstring is inherited from a base class
187 return self._schema.tableSchemas.get(table)
189 @classmethod
190 def init_database(
191 cls,
192 hosts: tuple[str, ...],
193 keyspace: str,
194 *,
195 schema_file: str | None = None,
196 schema_name: str | None = None,
197 read_sources_months: int | None = None,
198 read_forced_sources_months: int | None = None,
199 enable_replica: bool = False,
200 replica_skips_diaobjects: bool = False,
201 port: int | None = None,
202 username: str | None = None,
203 prefix: str | None = None,
204 part_pixelization: str | None = None,
205 part_pix_level: int | None = None,
206 time_partition_tables: bool = True,
207 time_partition_start: str | None = None,
208 time_partition_end: str | None = None,
209 read_consistency: str | None = None,
210 write_consistency: str | None = None,
211 read_timeout: int | None = None,
212 write_timeout: int | None = None,
213 ra_dec_columns: tuple[str, str] | None = None,
214 replication_factor: int | None = None,
215 drop: bool = False,
216 table_options: CreateTableOptions | None = None,
217 ) -> ApdbCassandraConfig:
218 """Initialize new APDB instance and make configuration object for it.
220 Parameters
221 ----------
222 hosts : `tuple` [`str`, ...]
223 List of host names or IP addresses for Cassandra cluster.
224 keyspace : `str`
225 Name of the keyspace for APDB tables.
226 schema_file : `str`, optional
227 Location of (YAML) configuration file with APDB schema. If not
228 specified then default location will be used.
229 schema_name : `str`, optional
230 Name of the schema in YAML configuration file. If not specified
231 then default name will be used.
232 read_sources_months : `int`, optional
233 Number of months of history to read from DiaSource.
234 read_forced_sources_months : `int`, optional
235 Number of months of history to read from DiaForcedSource.
236 enable_replica : `bool`, optional
237 If True, make additional tables used for replication to PPDB.
238 replica_skips_diaobjects : `bool`, optional
239 If `True` then do not fill regular ``DiaObject`` table when
240 ``enable_replica`` is `True`.
241 port : `int`, optional
242 Port number to use for Cassandra connections.
243 username : `str`, optional
244 User name for Cassandra connections.
245 prefix : `str`, optional
246 Optional prefix for all table names.
247 part_pixelization : `str`, optional
248 Name of the MOC pixelization used for partitioning.
249 part_pix_level : `int`, optional
250 Pixelization level.
251 time_partition_tables : `bool`, optional
252 Create per-partition tables.
253 time_partition_start : `str`, optional
254 Starting time for per-partition tables, in yyyy-mm-ddThh:mm:ss
255 format, in TAI.
256 time_partition_end : `str`, optional
257 Ending time for per-partition tables, in yyyy-mm-ddThh:mm:ss
258 format, in TAI.
259 read_consistency : `str`, optional
260 Name of the consistency level for read operations.
261 write_consistency : `str`, optional
262 Name of the consistency level for write operations.
263 read_timeout : `int`, optional
264 Read timeout in seconds.
265 write_timeout : `int`, optional
266 Write timeout in seconds.
267 ra_dec_columns : `tuple` [`str`, `str`], optional
268 Names of ra/dec columns in DiaObject table.
269 replication_factor : `int`, optional
270 Replication factor used when creating new keyspace, if keyspace
271 already exists its replication factor is not changed.
272 drop : `bool`, optional
273 If `True` then drop existing tables before re-creating the schema.
274 table_options : `CreateTableOptions`, optional
275 Options used when creating Cassandra tables.
277 Returns
278 -------
279 config : `ApdbCassandraConfig`
280 Resulting configuration object for a created APDB instance.
281 """
282 # Some non-standard defaults for connection parameters, these can be
283 # changed later in generated config. Check Cassandra driver
284 # documentation for what these parameters do. These parameters are not
285 # used during database initialization, but they will be saved with
286 # generated config.
287 connection_config = ApdbCassandraConnectionConfig(
288 extra_parameters={
289 "idle_heartbeat_interval": 0,
290 "idle_heartbeat_timeout": 30,
291 "control_connection_timeout": 100,
292 },
293 )
294 config = ApdbCassandraConfig(
295 contact_points=hosts,
296 keyspace=keyspace,
297 enable_replica=enable_replica,
298 replica_skips_diaobjects=replica_skips_diaobjects,
299 connection_config=connection_config,
300 )
301 config.partitioning.time_partition_tables = time_partition_tables
302 if schema_file is not None:
303 config.schema_file = schema_file
304 if schema_name is not None:
305 config.schema_name = schema_name
306 if read_sources_months is not None:
307 config.read_sources_months = read_sources_months
308 if read_forced_sources_months is not None:
309 config.read_forced_sources_months = read_forced_sources_months
310 if port is not None:
311 config.connection_config.port = port
312 if username is not None:
313 config.connection_config.username = username
314 if prefix is not None:
315 config.prefix = prefix
316 if part_pixelization is not None:
317 config.partitioning.part_pixelization = part_pixelization
318 if part_pix_level is not None:
319 config.partitioning.part_pix_level = part_pix_level
320 if time_partition_start is not None:
321 config.partitioning.time_partition_start = time_partition_start
322 if time_partition_end is not None:
323 config.partitioning.time_partition_end = time_partition_end
324 if read_consistency is not None:
325 config.connection_config.read_consistency = read_consistency
326 if write_consistency is not None:
327 config.connection_config.write_consistency = write_consistency
328 if read_timeout is not None:
329 config.connection_config.read_timeout = read_timeout
330 if write_timeout is not None:
331 config.connection_config.write_timeout = write_timeout
332 if ra_dec_columns is not None:
333 config.ra_dec_columns = ra_dec_columns
335 cls._makeSchema(config, drop=drop, replication_factor=replication_factor, table_options=table_options)
337 return config
339 def get_replica(self) -> ApdbCassandraReplica:
340 """Return `ApdbReplica` instance for this database."""
341 # Note that this instance has to stay alive while replica exists, so
342 # we pass reference to self.
343 return ApdbCassandraReplica(self)
345 @classmethod
346 def _makeSchema(
347 cls,
348 config: ApdbConfig,
349 *,
350 drop: bool = False,
351 replication_factor: int | None = None,
352 table_options: CreateTableOptions | None = None,
353 ) -> None:
354 # docstring is inherited from a base class
356 if not isinstance(config, ApdbCassandraConfig):
357 raise TypeError(f"Unexpected type of configuration object: {type(config)}")
359 simple_schema = ApdbSchema(config.schema_file, config.schema_name)
361 with SessionContext(config) as session:
362 schema = ApdbCassandraSchema(
363 session=session,
364 keyspace=config.keyspace,
365 table_schemas=simple_schema.tableSchemas,
366 prefix=config.prefix,
367 time_partition_tables=config.partitioning.time_partition_tables,
368 enable_replica=config.enable_replica,
369 replica_skips_diaobjects=config.replica_skips_diaobjects,
370 )
372 # Ask schema to create all tables.
373 part_range_config: ApdbCassandraTimePartitionRange | None = None
374 if config.partitioning.time_partition_tables:
375 partitioner = Partitioner(config)
376 time_partition_start = astropy.time.Time(
377 config.partitioning.time_partition_start, format="isot", scale="tai"
378 )
379 time_partition_end = astropy.time.Time(
380 config.partitioning.time_partition_end, format="isot", scale="tai"
381 )
382 part_range_config = ApdbCassandraTimePartitionRange(
383 start=partitioner.time_partition(time_partition_start),
384 end=partitioner.time_partition(time_partition_end),
385 )
386 schema.makeSchema(
387 drop=drop,
388 part_range=part_range_config,
389 replication_factor=replication_factor,
390 table_options=table_options,
391 )
392 else:
393 schema.makeSchema(
394 drop=drop, replication_factor=replication_factor, table_options=table_options
395 )
397 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
398 metadata = ApdbMetadataCassandra(
399 session, meta_table_name, config.keyspace, "read_tuples", "write"
400 )
402 # Fill version numbers, overrides if they existed before.
403 metadata.set(
404 ConnectionContext.metadataSchemaVersionKey, str(simple_schema.schemaVersion()), force=True
405 )
406 metadata.set(
407 ConnectionContext.metadataCodeVersionKey, str(cls.apdbImplementationVersion()), force=True
408 )
410 if config.enable_replica:
411 # Only store replica code version if replica is enabled.
412 metadata.set(
413 ConnectionContext.metadataReplicaVersionKey,
414 str(ApdbCassandraReplica.apdbReplicaImplementationVersion()),
415 force=True,
416 )
418 # Store frozen part of a configuration in metadata.
419 freezer = ApdbConfigFreezer[ApdbCassandraConfig](ConnectionContext.frozen_parameters)
420 metadata.set(ConnectionContext.metadataConfigKey, freezer.to_json(config), force=True)
422 # Store time partition range.
423 if part_range_config:
424 part_range_config.save_to_meta(metadata)
426 def getDiaObjects(self, region: sphgeom.Region) -> pandas.DataFrame:
427 # docstring is inherited from a base class
428 context = self._context
429 config = context.config
431 sp_where, num_sp_part = context.partitioner.spatial_where(region, for_prepare=True)
432 _LOG.debug("getDiaObjects: #partitions: %s", len(sp_where))
434 # We need to exclude extra partitioning columns from result.
435 column_names = context.schema.apdbColumnNames(ApdbTables.DiaObjectLast)
436 what = ",".join(quote_id(column) for column in column_names)
438 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
439 query = f'SELECT {what} from "{self._keyspace}"."{table_name}"'
440 statements: list[tuple] = []
441 for where, params in sp_where:
442 full_query = f"{query} WHERE {where}"
443 if params:
444 statement = context.preparer.prepare(full_query)
445 else:
446 # If there are no params then it is likely that query has a
447 # bunch of literals rendered already, no point trying to
448 # prepare it because it's not reusable.
449 statement = cassandra.query.SimpleStatement(full_query)
450 statements.append((statement, params))
451 _LOG.debug("getDiaObjects: #queries: %s", len(statements))
453 with _MON.context_tags({"table": "DiaObject"}):
454 _MON.add_record(
455 "select_query_stats", values={"num_sp_part": num_sp_part, "num_queries": len(statements)}
456 )
457 with self._timer("select_time") as timer:
458 objects = cast(
459 pandas.DataFrame,
460 select_concurrent(
461 context.session,
462 statements,
463 "read_pandas_multi",
464 config.connection_config.read_concurrency,
465 ),
466 )
467 timer.add_values(row_count=len(objects))
469 _LOG.debug("found %s DiaObjects", objects.shape[0])
470 return objects
472 def getDiaSources(
473 self,
474 region: sphgeom.Region,
475 object_ids: Iterable[int] | None,
476 visit_time: astropy.time.Time,
477 start_time: astropy.time.Time | None = None,
478 ) -> pandas.DataFrame | None:
479 # docstring is inherited from a base class
480 context = self._context
481 config = context.config
483 months = config.read_sources_months
484 if start_time is None and months == 0:
485 return None
487 mjd_end = float(visit_time.tai.mjd)
488 if start_time is None:
489 mjd_start = mjd_end - months * 30
490 else:
491 mjd_start = float(start_time.tai.mjd)
493 return self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaSource)
495 def getDiaForcedSources(
496 self,
497 region: sphgeom.Region,
498 object_ids: Iterable[int] | None,
499 visit_time: astropy.time.Time,
500 start_time: astropy.time.Time | None = None,
501 ) -> pandas.DataFrame | None:
502 # docstring is inherited from a base class
503 context = self._context
504 config = context.config
506 months = config.read_forced_sources_months
507 if start_time is None and months == 0:
508 return None
510 mjd_end = float(visit_time.tai.mjd)
511 if start_time is None:
512 mjd_start = mjd_end - months * 30
513 else:
514 mjd_start = float(start_time.tai.mjd)
516 return self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaForcedSource)
518 def containsVisitDetector(
519 self,
520 visit: int,
521 detector: int,
522 region: sphgeom.Region,
523 visit_time: astropy.time.Time,
524 ) -> bool:
525 # docstring is inherited from a base class
526 context = self._context
527 config = context.config
529 # If ApdbDetectorVisit table exists just check it.
530 if context.has_visit_detector_table:
531 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
532 query = (
533 f'SELECT count(*) FROM "{self._keyspace}"."{table_name}" WHERE visit = %s AND detector = %s'
534 )
535 with self._timer("contains_visit_detector_time"):
536 result = context.session.execute(query, (visit, detector))
537 return bool(result.one()[0])
539 # The order of checks corresponds to order in store(), on potential
540 # store failure earlier tables have higher probability containing
541 # stored records. With per-partition tables there will be many tables
542 # in the list, but it is unlikely that we'll use that setup in
543 # production.
544 sp_where, _ = context.partitioner.spatial_where(region, use_ranges=True, for_prepare=True)
545 visit_detector_where = ("visit = ? AND detector = ?", (visit, detector))
547 # Sources are partitioned on their midPointMjdTai. To avoid precision
548 # issues add some fuzzines to visit time.
549 mjd_start = float(visit_time.tai.mjd) - 1.0 / 24
550 mjd_end = float(visit_time.tai.mjd) + 1.0 / 24
552 statements: list[tuple] = []
553 for table_type in ApdbTables.DiaSource, ApdbTables.DiaForcedSource:
554 tables, temporal_where = context.partitioner.temporal_where(
555 table_type, mjd_start, mjd_end, query_per_time_part=True, for_prepare=True
556 )
557 for table in tables:
558 prefix = f'SELECT apdb_part FROM "{self._keyspace}"."{table}"'
559 # Needs ALLOW FILTERING as there is no PK constraint.
560 suffix = "PER PARTITION LIMIT 1 LIMIT 1 ALLOW FILTERING"
561 statements += list(
562 self._combine_where(prefix, sp_where, temporal_where, visit_detector_where, suffix)
563 )
565 with self._timer("contains_visit_detector_time"):
566 result = cast(
567 list[tuple[int] | None],
568 select_concurrent(
569 context.session,
570 statements,
571 "read_tuples",
572 config.connection_config.read_concurrency,
573 ),
574 )
575 return bool(result)
577 def getSSObjects(self) -> pandas.DataFrame:
578 # docstring is inherited from a base class
579 context = self._context
581 tableName = context.schema.tableName(ApdbTables.SSObject)
582 query = f'SELECT * from "{self._keyspace}"."{tableName}"'
584 objects = None
585 with self._timer("select_time", tags={"table": "SSObject"}) as timer:
586 result = context.session.execute(query, execution_profile="read_pandas")
587 objects = result._current_rows
588 timer.add_values(row_count=len(objects))
590 _LOG.debug("found %s SSObjects", objects.shape[0])
591 return objects
593 def store(
594 self,
595 visit_time: astropy.time.Time,
596 objects: pandas.DataFrame,
597 sources: pandas.DataFrame | None = None,
598 forced_sources: pandas.DataFrame | None = None,
599 ) -> None:
600 # docstring is inherited from a base class
601 context = self._context
602 config = context.config
604 if context.has_visit_detector_table:
605 # Store visit/detector in a special table, this has to be done
606 # before all other writes so if there is a failure at any point
607 # later we still have a record for attempted write.
608 visit_detector: set[tuple[int, int]] = set()
609 for df in sources, forced_sources:
610 if df is not None and not df.empty:
611 df = df[["visit", "detector"]]
612 for visit, detector in df.itertuples(index=False):
613 visit_detector.add((visit, detector))
615 if visit_detector:
616 # Typically there is only one entry, do not bother with
617 # concurrency.
618 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
619 query = f'INSERT INTO "{self._keyspace}"."{table_name}" (visit, detector) VALUES (%s, %s)'
620 for item in visit_detector:
621 context.session.execute(query, item, execution_profile="write")
623 objects = self._fix_input_timestamps(objects)
624 if sources is not None:
625 sources = self._fix_input_timestamps(sources)
626 if forced_sources is not None:
627 forced_sources = self._fix_input_timestamps(forced_sources)
629 replica_chunk: ReplicaChunk | None = None
630 if context.schema.replication_enabled:
631 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, config.replica_chunk_seconds)
632 self._storeReplicaChunk(replica_chunk, visit_time)
634 # fill region partition column for DiaObjects
635 objects = self._add_apdb_part(objects)
636 self._storeDiaObjects(objects, visit_time, replica_chunk)
638 if sources is not None and len(sources) > 0:
639 # copy apdb_part column from DiaObjects to DiaSources
640 sources = self._add_apdb_part(sources)
641 subchunk = self._storeDiaSources(ApdbTables.DiaSource, sources, replica_chunk)
642 self._storeDiaSourcesPartitions(sources, visit_time, replica_chunk, subchunk)
644 if forced_sources is not None and len(forced_sources) > 0:
645 forced_sources = self._add_apdb_part(forced_sources)
646 self._storeDiaSources(ApdbTables.DiaForcedSource, forced_sources, replica_chunk)
648 def storeSSObjects(self, objects: pandas.DataFrame) -> None:
649 # docstring is inherited from a base class
650 objects = self._fix_input_timestamps(objects)
651 self._storeObjectsPandas(objects, ApdbTables.SSObject)
653 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
654 # docstring is inherited from a base class
655 context = self._context
656 config = context.config
658 if self._schema.has_mjd_timestamps:
659 reassign_time_column = "ssObjectReassocTimeMjdTai"
660 reassignTime = float(astropy.time.Time.now().tai.mjd)
661 else:
662 reassign_time_column = "ssObjectReassocTime"
663 # Current time as milliseconds since epoch.
664 reassignTime = int(datetime.datetime.now(tz=datetime.UTC).timestamp() * 1000)
666 # To update a record we need to know its exact primary key (including
667 # partition key) so we start by querying for diaSourceId to find the
668 # primary keys.
670 table_name = context.schema.tableName(ExtraTables.DiaSourceToPartition)
671 # split it into 1k IDs per query
672 selects: list[tuple] = []
673 for ids in chunk_iterable(idMap.keys(), 1_000):
674 ids_str = ",".join(str(item) for item in ids)
675 selects.append(
676 (
677 (
678 'SELECT "diaSourceId", "apdb_part", "apdb_time_part", "apdb_replica_chunk" '
679 f'FROM "{self._keyspace}"."{table_name}" WHERE "diaSourceId" IN ({ids_str})'
680 ),
681 {},
682 )
683 )
685 # No need for DataFrame here, read data as tuples.
686 result = cast(
687 list[tuple[int, int, int, int | None]],
688 select_concurrent(
689 context.session, selects, "read_tuples", config.connection_config.read_concurrency
690 ),
691 )
693 # Make mapping from source ID to its partition.
694 id2partitions: dict[int, tuple[int, int]] = {}
695 id2chunk_id: dict[int, int] = {}
696 for row in result:
697 id2partitions[row[0]] = row[1:3]
698 if row[3] is not None:
699 id2chunk_id[row[0]] = row[3]
701 # make sure we know partitions for each ID
702 if set(id2partitions) != set(idMap):
703 missing = ",".join(str(item) for item in set(idMap) - set(id2partitions))
704 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
706 # Reassign in standard tables
707 queries: list[tuple[cassandra.query.PreparedStatement, tuple]] = []
708 for diaSourceId, ssObjectId in idMap.items():
709 apdb_part, apdb_time_part = id2partitions[diaSourceId]
710 values: tuple
711 if config.partitioning.time_partition_tables:
712 table_name = context.schema.tableName(ApdbTables.DiaSource, apdb_time_part)
713 query = (
714 f'UPDATE "{self._keyspace}"."{table_name}"'
715 f' SET "ssObjectId" = ?, "diaObjectId" = NULL, "{reassign_time_column}" = ?'
716 ' WHERE "apdb_part" = ? AND "diaSourceId" = ?'
717 )
718 values = (ssObjectId, reassignTime, apdb_part, diaSourceId)
719 else:
720 table_name = context.schema.tableName(ApdbTables.DiaSource)
721 query = (
722 f'UPDATE "{self._keyspace}"."{table_name}"'
723 f' SET "ssObjectId" = ?, "diaObjectId" = NULL, "{reassign_time_column}" = ?'
724 ' WHERE "apdb_part" = ? AND "apdb_time_part" = ? AND "diaSourceId" = ?'
725 )
726 values = (ssObjectId, reassignTime, apdb_part, apdb_time_part, diaSourceId)
727 queries.append((context.preparer.prepare(query), values))
729 # TODO: (DM-50190) Replication for updated records is not implemented.
730 if id2chunk_id:
731 warnings.warn("Replication of reassigned DiaSource records is not implemented.", stacklevel=2)
733 _LOG.debug("%s: will update %d records", table_name, len(idMap))
734 with self._timer("source_reassign_time") as timer:
735 execute_concurrent(context.session, queries, execution_profile="write")
736 timer.add_values(source_count=len(idMap))
738 def dailyJob(self) -> None:
739 # docstring is inherited from a base class
740 pass
742 def countUnassociatedObjects(self) -> int:
743 # docstring is inherited from a base class
745 # It's too inefficient to implement it for Cassandra in current schema.
746 raise NotImplementedError()
748 @property
749 def metadata(self) -> ApdbMetadata:
750 # docstring is inherited from a base class
751 context = self._context
752 return context.metadata
754 @property
755 def admin(self) -> ApdbCassandraAdmin:
756 # docstring is inherited from a base class
757 return ApdbCassandraAdmin(self)
759 def _getSources(
760 self,
761 region: sphgeom.Region,
762 object_ids: Iterable[int] | None,
763 mjd_start: float,
764 mjd_end: float,
765 table_name: ApdbTables,
766 ) -> pandas.DataFrame:
767 """Return catalog of DiaSource instances given set of DiaObject IDs.
769 Parameters
770 ----------
771 region : `lsst.sphgeom.Region`
772 Spherical region.
773 object_ids :
774 Collection of DiaObject IDs
775 mjd_start : `float`
776 Lower bound of time interval.
777 mjd_end : `float`
778 Upper bound of time interval.
779 table_name : `ApdbTables`
780 Name of the table.
782 Returns
783 -------
784 catalog : `pandas.DataFrame`, or `None`
785 Catalog containing DiaSource records. Empty catalog is returned if
786 ``object_ids`` is empty.
787 """
788 context = self._context
789 config = context.config
791 object_id_set: Set[int] = set()
792 if object_ids is not None:
793 object_id_set = set(object_ids)
794 if len(object_id_set) == 0:
795 return self._make_empty_catalog(table_name)
797 sp_where, num_sp_part = context.partitioner.spatial_where(region, for_prepare=True)
798 tables, temporal_where = context.partitioner.temporal_where(
799 table_name, mjd_start, mjd_end, for_prepare=True, partitons_range=context.time_partitions_range
800 )
801 if not tables:
802 start = astropy.time.Time(mjd_start, format="mjd", scale="tai")
803 end = astropy.time.Time(mjd_end, format="mjd", scale="tai")
804 warnings.warn(
805 f"Query time range ({start.isot} - {end.isot}) does not overlap database time partitions."
806 )
808 # We need to exclude extra partitioning columns from result.
809 column_names = context.schema.apdbColumnNames(table_name)
810 what = ",".join(quote_id(column) for column in column_names)
812 # Build all queries
813 statements: list[tuple] = []
814 for table in tables:
815 prefix = f'SELECT {what} from "{self._keyspace}"."{table}"'
816 statements += list(self._combine_where(prefix, sp_where, temporal_where))
817 _LOG.debug("_getSources %s: #queries: %s", table_name, len(statements))
819 with _MON.context_tags({"table": table_name.name}):
820 _MON.add_record(
821 "select_query_stats", values={"num_sp_part": num_sp_part, "num_queries": len(statements)}
822 )
823 with self._timer("select_time") as timer:
824 catalog = cast(
825 pandas.DataFrame,
826 select_concurrent(
827 context.session,
828 statements,
829 "read_pandas_multi",
830 config.connection_config.read_concurrency,
831 ),
832 )
833 timer.add_values(row_count_from_db=len(catalog))
835 # filter by given object IDs
836 if len(object_id_set) > 0:
837 catalog = cast(pandas.DataFrame, catalog[catalog["diaObjectId"].isin(object_id_set)])
839 # precise filtering on midpointMjdTai
840 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] > mjd_start])
842 timer.add_values(row_count=len(catalog))
844 _LOG.debug("found %d %ss", catalog.shape[0], table_name.name)
845 return catalog
847 def _storeReplicaChunk(self, replica_chunk: ReplicaChunk, visit_time: astropy.time.Time) -> None:
848 context = self._context
849 config = context.config
851 # Cassandra timestamp uses milliseconds since epoch
852 timestamp = int(replica_chunk.last_update_time.unix_tai * 1000)
854 # everything goes into a single partition
855 partition = 0
857 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)
859 columns = ["partition", "apdb_replica_chunk", "last_update_time", "unique_id"]
860 values = [partition, replica_chunk.id, timestamp, replica_chunk.unique_id]
861 if context.has_chunk_sub_partitions:
862 columns.append("has_subchunks")
863 values.append(True)
865 column_list = ", ".join(columns)
866 placeholders = ",".join(["%s"] * len(columns))
867 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ({column_list}) VALUES ({placeholders})'
869 context.session.execute(
870 query,
871 values,
872 timeout=config.connection_config.write_timeout,
873 execution_profile="write",
874 )
876 def _queryDiaObjectLastPartitions(self, ids: Iterable[int]) -> Mapping[int, int]:
877 """Return existing mapping of diaObjectId to its last partition."""
878 context = self._context
879 config = context.config
881 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
882 queries = []
883 object_count = 0
884 for id_chunk in chunk_iterable(ids, 10_000):
885 id_chunk_list = list(id_chunk)
886 query = (
887 f'SELECT "diaObjectId", apdb_part FROM "{self._keyspace}"."{table_name}" '
888 f'WHERE "diaObjectId" in ({",".join(str(oid) for oid in id_chunk_list)})'
889 )
890 queries.append((query, ()))
891 object_count += len(id_chunk_list)
893 with self._timer("query_object_last_partitions") as timer:
894 data = cast(
895 ApdbTableData,
896 select_concurrent(
897 context.session,
898 queries,
899 "read_raw_multi",
900 config.connection_config.read_concurrency,
901 ),
902 )
903 timer.add_values(object_count=object_count, row_count=len(data.rows()))
905 if data.column_names() != ["diaObjectId", "apdb_part"]:
906 raise RuntimeError(f"Unexpected column names in query result: {data.column_names()}")
908 return {row[0]: row[1] for row in data.rows()}
910 def _deleteMovingObjects(self, objs: pandas.DataFrame) -> None:
911 """Objects in DiaObjectsLast can move from one spatial partition to
912 another. For those objects inserting new version does not replace old
913 one, so we need to explicitly remove old versions before inserting new
914 ones.
915 """
916 context = self._context
918 # Extract all object IDs.
919 new_partitions = dict(zip(objs["diaObjectId"], objs["apdb_part"]))
920 old_partitions = self._queryDiaObjectLastPartitions(objs["diaObjectId"])
922 moved_oids: dict[int, tuple[int, int]] = {}
923 for oid, old_part in old_partitions.items():
924 new_part = new_partitions.get(oid, old_part)
925 if new_part != old_part:
926 moved_oids[oid] = (old_part, new_part)
927 _LOG.debug("DiaObject IDs that moved to new partition: %s", moved_oids)
929 if moved_oids:
930 # Delete old records from DiaObjectLast.
931 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
932 query = f'DELETE FROM "{self._keyspace}"."{table_name}" WHERE apdb_part = ? AND "diaObjectId" = ?'
933 statement = context.preparer.prepare(query)
934 queries = []
935 for oid, (old_part, _) in moved_oids.items():
936 queries.append((statement, (old_part, oid)))
937 with self._timer("delete_object_last") as timer:
938 execute_concurrent(context.session, queries, execution_profile="write")
939 timer.add_values(row_count=len(moved_oids))
941 # Add all new records to the map.
942 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
943 query = f'INSERT INTO "{self._keyspace}"."{table_name}" ("diaObjectId", apdb_part) VALUES (?,?)'
944 statement = context.preparer.prepare(query)
946 queries = []
947 for oid, new_part in new_partitions.items():
948 queries.append((statement, (oid, new_part)))
950 with self._timer("update_object_last_partition") as timer:
951 execute_concurrent(context.session, queries, execution_profile="write")
952 timer.add_values(row_count=len(queries))
954 def _storeDiaObjects(
955 self, objs: pandas.DataFrame, visit_time: astropy.time.Time, replica_chunk: ReplicaChunk | None
956 ) -> None:
957 """Store catalog of DiaObjects from current visit.
959 Parameters
960 ----------
961 objs : `pandas.DataFrame`
962 Catalog with DiaObject records
963 visit_time : `astropy.time.Time`
964 Time of the current visit.
965 replica_chunk : `ReplicaChunk` or `None`
966 Replica chunk identifier if replication is configured.
967 """
968 if len(objs) == 0:
969 _LOG.debug("No objects to write to database.")
970 return
972 context = self._context
973 config = context.config
975 if context.has_dia_object_last_to_partition:
976 self._deleteMovingObjects(objs)
978 timestamp: float | datetime.datetime
979 if self._schema.has_mjd_timestamps:
980 validity_start_column = "validityStartMjdTai"
981 timestamp = float(visit_time.tai.mjd)
982 else:
983 validity_start_column = "validityStart"
984 timestamp = visit_time.datetime
986 self._storeObjectsPandas(objs, ApdbTables.DiaObjectLast)
988 extra_columns: dict[str, Any] = {validity_start_column: timestamp}
989 visit_time_part = context.partitioner.time_partition(visit_time)
990 time_part: int | None = visit_time_part
991 if (time_partitions_range := context.time_partitions_range) is not None:
992 self._check_time_partitions([visit_time_part], time_partitions_range)
993 if not config.partitioning.time_partition_tables:
994 extra_columns["apdb_time_part"] = time_part
995 time_part = None
997 # Only store DiaObects if not doing replication or explicitly
998 # configured to always store them.
999 if replica_chunk is None or not config.replica_skips_diaobjects:
1000 self._storeObjectsPandas(
1001 objs, ApdbTables.DiaObject, extra_columns=extra_columns, time_part=time_part
1002 )
1004 if replica_chunk is not None:
1005 extra_columns = {"apdb_replica_chunk": replica_chunk.id, validity_start_column: timestamp}
1006 table = ExtraTables.DiaObjectChunks
1007 if context.has_chunk_sub_partitions:
1008 table = ExtraTables.DiaObjectChunks2
1009 # Use a random number for a second part of partitioning key so
1010 # that different clients could wrtite to different partitions.
1011 # This makes it not exactly reproducible.
1012 extra_columns["apdb_replica_subchunk"] = random.randrange(config.replica_sub_chunk_count)
1013 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1015 def _storeDiaSources(
1016 self,
1017 table_name: ApdbTables,
1018 sources: pandas.DataFrame,
1019 replica_chunk: ReplicaChunk | None,
1020 ) -> int | None:
1021 """Store catalog of DIASources or DIAForcedSources from current visit.
1023 Parameters
1024 ----------
1025 table_name : `ApdbTables`
1026 Table where to store the data.
1027 sources : `pandas.DataFrame`
1028 Catalog containing DiaSource records
1029 visit_time : `astropy.time.Time`
1030 Time of the current visit.
1031 replica_chunk : `ReplicaChunk` or `None`
1032 Replica chunk identifier if replication is configured.
1034 Returns
1035 -------
1036 subchunk : `int` or `None`
1037 Subchunk number for resulting replica data, `None` if relication is
1038 not enabled ot subchunking is not enabled.
1039 """
1040 context = self._context
1041 config = context.config
1043 # Time partitioning has to be based on midpointMjdTai, not visit_time
1044 # as visit_time is not really a visit time.
1045 tp_sources = sources.copy(deep=False)
1046 tp_sources["apdb_time_part"] = tp_sources["midpointMjdTai"].apply(context.partitioner.time_partition)
1047 if (time_partitions_range := context.time_partitions_range) is not None:
1048 self._check_time_partitions(tp_sources["apdb_time_part"], time_partitions_range)
1049 extra_columns: dict[str, Any] = {}
1050 if not config.partitioning.time_partition_tables:
1051 self._storeObjectsPandas(tp_sources, table_name)
1052 else:
1053 # Group by time partition
1054 partitions = set(tp_sources["apdb_time_part"])
1055 if len(partitions) == 1:
1056 # Single partition - just save the whole thing.
1057 time_part = partitions.pop()
1058 self._storeObjectsPandas(sources, table_name, time_part=time_part)
1059 else:
1060 # group by time partition.
1061 for time_part, sub_frame in tp_sources.groupby(by="apdb_time_part"):
1062 sub_frame.drop(columns="apdb_time_part", inplace=True)
1063 self._storeObjectsPandas(sub_frame, table_name, time_part=time_part)
1065 subchunk: int | None = None
1066 if replica_chunk is not None:
1067 extra_columns = {"apdb_replica_chunk": replica_chunk.id}
1068 if context.has_chunk_sub_partitions:
1069 subchunk = random.randrange(config.replica_sub_chunk_count)
1070 extra_columns["apdb_replica_subchunk"] = subchunk
1071 if table_name is ApdbTables.DiaSource:
1072 extra_table = ExtraTables.DiaSourceChunks2
1073 else:
1074 extra_table = ExtraTables.DiaForcedSourceChunks2
1075 else:
1076 if table_name is ApdbTables.DiaSource:
1077 extra_table = ExtraTables.DiaSourceChunks
1078 else:
1079 extra_table = ExtraTables.DiaForcedSourceChunks
1080 self._storeObjectsPandas(sources, extra_table, extra_columns=extra_columns)
1082 return subchunk
1084 def _check_time_partitions(
1085 self, partitions: Iterable[int], time_partitions_range: ApdbCassandraTimePartitionRange
1086 ) -> None:
1087 """Check that time partitons for new data actually exist.
1089 Parameters
1090 ----------
1091 partitions : `~collections.abc.Iterable` [`int`]
1092 Time partitions for new data.
1093 time_partitions_range : `ApdbCassandraTimePartitionRange`
1094 Currrent time partition range.
1095 """
1096 partitions = set(partitions)
1097 min_part = min(partitions)
1098 max_part = max(partitions)
1099 if min_part < time_partitions_range.start or max_part > time_partitions_range.end:
1100 raise ValueError(
1101 "Attempt to store data for time partitions that do not yet exist. "
1102 f"Partitons for new records: {min_part}-{max_part}. "
1103 f"Database partitons: {time_partitions_range.start}-{time_partitions_range.end}."
1104 )
1105 # Make a noise when writing to the last partition.
1106 if max_part == time_partitions_range.end:
1107 warnings.warn(
1108 "Writing into the last temporal partition. Partition range needs to be extended soon.",
1109 stacklevel=3,
1110 )
1112 def _storeDiaSourcesPartitions(
1113 self,
1114 sources: pandas.DataFrame,
1115 visit_time: astropy.time.Time,
1116 replica_chunk: ReplicaChunk | None,
1117 subchunk: int | None,
1118 ) -> None:
1119 """Store mapping of diaSourceId to its partitioning values.
1121 Parameters
1122 ----------
1123 sources : `pandas.DataFrame`
1124 Catalog containing DiaSource records
1125 visit_time : `astropy.time.Time`
1126 Time of the current visit.
1127 replica_chunk : `ReplicaChunk` or `None`
1128 Replication chunk, or `None` when replication is disabled.
1129 subchunk : `int` or `None`
1130 Replication sub-chunk, or `None` when replication is disabled or
1131 sub-chunking is not used.
1132 """
1133 context = self._context
1135 id_map = cast(pandas.DataFrame, sources[["diaSourceId", "apdb_part"]])
1136 extra_columns = {
1137 "apdb_time_part": context.partitioner.time_partition(visit_time),
1138 "apdb_replica_chunk": replica_chunk.id if replica_chunk is not None else None,
1139 }
1140 if context.has_chunk_sub_partitions:
1141 extra_columns["apdb_replica_subchunk"] = subchunk
1143 self._storeObjectsPandas(
1144 id_map, ExtraTables.DiaSourceToPartition, extra_columns=extra_columns, time_part=None
1145 )
1147 def _storeObjectsPandas(
1148 self,
1149 records: pandas.DataFrame,
1150 table_name: ApdbTables | ExtraTables,
1151 extra_columns: Mapping | None = None,
1152 time_part: int | None = None,
1153 ) -> None:
1154 """Store generic objects.
1156 Takes Pandas catalog and stores a bunch of records in a table.
1158 Parameters
1159 ----------
1160 records : `pandas.DataFrame`
1161 Catalog containing object records
1162 table_name : `ApdbTables`
1163 Name of the table as defined in APDB schema.
1164 extra_columns : `dict`, optional
1165 Mapping (column_name, column_value) which gives fixed values for
1166 columns in each row, overrides values in ``records`` if matching
1167 columns exist there.
1168 time_part : `int`, optional
1169 If not `None` then insert into a per-partition table.
1171 Notes
1172 -----
1173 If Pandas catalog contains additional columns not defined in table
1174 schema they are ignored. Catalog does not have to contain all columns
1175 defined in a table, but partition and clustering keys must be present
1176 in a catalog or ``extra_columns``.
1177 """
1178 context = self._context
1180 # use extra columns if specified
1181 if extra_columns is None:
1182 extra_columns = {}
1183 extra_fields = list(extra_columns.keys())
1185 # Fields that will come from dataframe.
1186 df_fields = [column for column in records.columns if column not in extra_fields]
1188 column_map = context.schema.getColumnMap(table_name)
1189 # list of columns (as in felis schema)
1190 fields = [column_map[field].name for field in df_fields if field in column_map]
1191 fields += extra_fields
1193 # check that all partitioning and clustering columns are defined
1194 partition_columns = context.schema.partitionColumns(table_name)
1195 required_columns = partition_columns + context.schema.clusteringColumns(table_name)
1196 missing_columns = [column for column in required_columns if column not in fields]
1197 if missing_columns:
1198 raise ValueError(f"Primary key columns are missing from catalog: {missing_columns}")
1200 qfields = [quote_id(field) for field in fields]
1201 qfields_str = ",".join(qfields)
1203 batch_size = self._batch_size(table_name)
1205 with self._timer("insert_build_time", tags={"table": table_name.name}):
1206 # Multi-partition batches are problematic in general, so we want to
1207 # group records in a batch by their partition key.
1208 values_by_key: dict[tuple, list[list]] = defaultdict(list)
1209 for rec in records.itertuples(index=False):
1210 values = []
1211 partitioning_values: dict[str, Any] = {}
1212 for field in df_fields:
1213 if field not in column_map:
1214 continue
1215 value = getattr(rec, field)
1216 if column_map[field].datatype is felis.datamodel.DataType.timestamp:
1217 if isinstance(value, pandas.Timestamp):
1218 value = value.to_pydatetime()
1219 elif value is pandas.NaT:
1220 value = None
1221 else:
1222 # Assume it's seconds since epoch, Cassandra
1223 # datetime is in milliseconds
1224 value = int(value * 1000)
1225 value = literal(value)
1226 values.append(UNSET_VALUE if value is None else value)
1227 if field in partition_columns:
1228 partitioning_values[field] = value
1229 for field in extra_fields:
1230 value = literal(extra_columns[field])
1231 values.append(UNSET_VALUE if value is None else value)
1232 if field in partition_columns:
1233 partitioning_values[field] = value
1235 key = tuple(partitioning_values[field] for field in partition_columns)
1236 values_by_key[key].append(values)
1238 table = context.schema.tableName(table_name, time_part)
1240 holders = ",".join(["?"] * len(qfields))
1241 query = f'INSERT INTO "{self._keyspace}"."{table}" ({qfields_str}) VALUES ({holders})'
1242 statement = context.preparer.prepare(query)
1243 # Cassandra has 64k limit on batch size, normally that should be
1244 # enough but some tests generate too many forced sources.
1245 queries = []
1246 for key_values in values_by_key.values():
1247 for values_chunk in chunk_iterable(key_values, batch_size):
1248 batch = cassandra.query.BatchStatement()
1249 for row_values in values_chunk:
1250 batch.add(statement, row_values)
1251 queries.append((batch, None))
1252 assert batch.routing_key is not None and batch.keyspace is not None
1254 _LOG.debug("%s: will store %d records", context.schema.tableName(table_name), records.shape[0])
1255 with self._timer("insert_time", tags={"table": table_name.name}) as timer:
1256 execute_concurrent(context.session, queries, execution_profile="write")
1257 timer.add_values(row_count=len(records), num_batches=len(queries))
1259 def _add_apdb_part(self, df: pandas.DataFrame) -> pandas.DataFrame:
1260 """Calculate spatial partition for each record and add it to a
1261 DataFrame.
1263 Parameters
1264 ----------
1265 df : `pandas.DataFrame`
1266 DataFrame which has to contain ra/dec columns, names of these
1267 columns are defined by configuration ``ra_dec_columns`` field.
1269 Returns
1270 -------
1271 df : `pandas.DataFrame`
1272 DataFrame with ``apdb_part`` column which contains pixel index
1273 for ra/dec coordinates.
1275 Notes
1276 -----
1277 This overrides any existing column in a DataFrame with the same name
1278 (``apdb_part``). Original DataFrame is not changed, copy of a DataFrame
1279 is returned.
1280 """
1281 context = self._context
1282 config = context.config
1284 # Calculate pixelization index for every record.
1285 apdb_part = np.zeros(df.shape[0], dtype=np.int64)
1286 ra_col, dec_col = config.ra_dec_columns
1287 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1288 uv3d = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(ra, dec))
1289 idx = context.partitioner.pixel(uv3d)
1290 apdb_part[i] = idx
1291 df = df.copy()
1292 df["apdb_part"] = apdb_part
1293 return df
1295 def _make_empty_catalog(self, table_name: ApdbTables) -> pandas.DataFrame:
1296 """Make an empty catalog for a table with a given name.
1298 Parameters
1299 ----------
1300 table_name : `ApdbTables`
1301 Name of the table.
1303 Returns
1304 -------
1305 catalog : `pandas.DataFrame`
1306 An empty catalog.
1307 """
1308 table = self._schema.tableSchemas[table_name]
1310 data = {
1311 columnDef.name: pandas.Series(dtype=self._schema.column_dtype(columnDef.datatype))
1312 for columnDef in table.columns
1313 }
1314 return pandas.DataFrame(data)
1316 def _combine_where(
1317 self,
1318 prefix: str,
1319 where1: list[tuple[str, tuple]],
1320 where2: list[tuple[str, tuple]],
1321 where3: tuple[str, tuple] | None = None,
1322 suffix: str | None = None,
1323 ) -> Iterator[tuple[cassandra.query.Statement, tuple]]:
1324 """Make cartesian product of two parts of WHERE clause into a series
1325 of statements to execute.
1327 Parameters
1328 ----------
1329 prefix : `str`
1330 Initial statement prefix that comes before WHERE clause, e.g.
1331 "SELECT * from Table"
1332 """
1333 context = self._context
1335 # If lists are empty use special sentinels.
1336 if not where1:
1337 where1 = [("", ())]
1338 if not where2:
1339 where2 = [("", ())]
1341 for expr1, params1 in where1:
1342 for expr2, params2 in where2:
1343 full_query = prefix
1344 wheres = []
1345 params = params1 + params2
1346 if expr1:
1347 wheres.append(expr1)
1348 if expr2:
1349 wheres.append(expr2)
1350 if where3:
1351 wheres.append(where3[0])
1352 params += where3[1]
1353 if wheres:
1354 full_query += " WHERE " + " AND ".join(wheres)
1355 if suffix:
1356 full_query += " " + suffix
1357 if params:
1358 statement = context.preparer.prepare(full_query)
1359 else:
1360 # If there are no params then it is likely that query
1361 # has a bunch of literals rendered already, no point
1362 # trying to prepare it.
1363 statement = cassandra.query.SimpleStatement(full_query)
1364 yield (statement, params)
1366 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1367 """Update timestamp columns in input DataFrame to be naive datetime
1368 type.
1370 Clients may or may not generate aware timestamps, code in this class
1371 assumes that timestamps are naive, so we convert them to UTC and
1372 drop timezone.
1373 """
1374 # Find all columns with aware timestamps.
1375 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1376 for column in columns:
1377 # tz_convert(None) will convert to UTC and drop timezone.
1378 df[column] = df[column].dt.tz_convert(None)
1379 return df
1381 def _batch_size(self, table: ApdbTables | ExtraTables) -> int:
1382 """Calculate batch size based on config parameters."""
1383 context = self._context
1384 config = context.config
1386 # Cassandra limit on number of statements in a batch is 64k.
1387 batch_size = 65_535
1388 if 0 < config.batch_statement_limit < batch_size:
1389 batch_size = config.batch_statement_limit
1390 if config.batch_size_limit > 0:
1391 # The purpose of this limit is to try not to exceed batch size
1392 # threshold which is set on server side. Cassandra wire protocol
1393 # for prepared queries (and batches) only sends column values with
1394 # with an additional 4 bytes per value specifying size. Value is
1395 # not included for NULL or NOT_SET values, but the size is always
1396 # there. There is additional small per-query overhead, which we
1397 # ignore.
1398 row_size = context.schema.table_row_size(table)
1399 row_size += 4 * len(context.schema.getColumnMap(table))
1400 batch_size = min(batch_size, (config.batch_size_limit // row_size) + 1)
1401 return batch_size