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