Coverage for python/lsst/dax/apdb/cassandra/apdbCassandra.py: 8%
924 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-07 09:42 +0000
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 json
28import logging
29import random
30import uuid
31import warnings
32from collections import Counter, defaultdict
33from collections.abc import Iterable, Mapping, Set
34from typing import TYPE_CHECKING, Any, cast
36import numpy as np
37import pandas
39# If cassandra-driver is not there the module can still be imported
40# but ApdbCassandra cannot be instantiated.
41try:
42 import cassandra
43 import cassandra.query
44 from cassandra.query import UNSET_VALUE
46 CASSANDRA_IMPORTED = True
47except ImportError:
48 CASSANDRA_IMPORTED = False
50import astropy.time
51import felis.datamodel
53from lsst import sphgeom
54from lsst.utils.iteration import chunk_iterable
56from ..apdb import Apdb, ApdbConfig
57from ..apdbConfigFreezer import ApdbConfigFreezer
58from ..apdbReplica import ApdbTableData, ReplicaChunk
59from ..apdbSchema import ApdbSchema, ApdbTables
60from ..apdbUpdateRecord import (
61 ApdbCloseDiaObjectValidityRecord,
62 ApdbReassignDiaSourceToDiaObjectRecord,
63 ApdbUpdateNDiaSourcesRecord,
64 ApdbWithdrawDiaForcedSourceRecord,
65 ApdbWithdrawDiaSourceRecord,
66)
67from ..monitor import MonAgent
68from ..recordIds import DiaForcedSourceId, DiaObjectId, DiaSourceId
69from ..schema_model import Table
70from ..timer import Timer
71from ..versionTuple import VersionTuple
72from .apdbCassandraAdmin import ApdbCassandraAdmin
73from .apdbCassandraReplica import ApdbCassandraReplica
74from .apdbCassandraSchema import ApdbCassandraSchema, CreateTableOptions, ExtraTables
75from .apdbMetadataCassandra import ApdbMetadataCassandra
76from .cassandra_utils import (
77 ApdbCassandraTableData,
78 execute_concurrent,
79 literal,
80 select_concurrent,
81)
82from .config import ApdbCassandraConfig, ApdbCassandraConnectionConfig, ApdbCassandraTimePartitionRange
83from .connectionContext import ConnectionContext, DbVersions
84from .exceptions import CassandraMissingError
85from .partitioner import Partitioner
86from .queries import Column as C # noqa: N817
87from .queries import ColumnExpr, Delete, Insert, QExpr, Select, Update
88from .sessionFactory import SessionContext, SessionFactory
90if TYPE_CHECKING:
91 from ..apdbMetadata import ApdbMetadata
92 from ..apdbUpdateRecord import ApdbUpdateRecord
94_LOG = logging.getLogger(__name__)
96_MON = MonAgent(__name__)
98VERSION = VersionTuple(1, 3, 0)
99"""Version for the code controlling non-replication tables. This needs to be
100updated following compatibility rules when schema produced by this code
101changes.
102"""
105class ApdbCassandra(Apdb):
106 """Implementation of APDB database with Apache Cassandra backend.
108 Parameters
109 ----------
110 config : `ApdbCassandraConfig`
111 Configuration object.
112 """
114 def __init__(self, config: ApdbCassandraConfig):
115 if not CASSANDRA_IMPORTED:
116 raise CassandraMissingError()
118 self._config = config
119 self._keyspace = config.keyspace
120 self._schema = ApdbSchema(config.schema_file, config.ss_schema_file)
122 self._session_factory = SessionFactory(config)
123 self._connection_context: ConnectionContext | None = None
125 def close(self) -> None:
126 # docstring is inherited from a base class
127 self._session_factory.close()
129 @property
130 def _context(self) -> ConnectionContext:
131 """Establish connection if not established and return context."""
132 if self._connection_context is None:
133 current_versions = DbVersions(
134 schema_version=self.schema.schemaVersion(),
135 code_version=self.apdbImplementationVersion(),
136 replica_version=ApdbCassandraReplica.apdbReplicaImplementationVersion(),
137 )
138 _LOG.debug("Current versions: %s", current_versions)
140 session = self._session_factory.session()
141 self._connection_context = ConnectionContext(
142 session, self._config, self.schema.tableSchemas, current_versions
143 )
145 if _LOG.isEnabledFor(logging.DEBUG):
146 _LOG.debug("ApdbCassandra Configuration: %s", self._connection_context.config.model_dump())
148 return self._connection_context
150 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
151 """Create `Timer` instance given its name."""
152 return Timer(name, _MON, tags=tags)
154 @classmethod
155 def apdbImplementationVersion(cls) -> VersionTuple:
156 """Return version number for current APDB implementation.
158 Returns
159 -------
160 version : `VersionTuple`
161 Version of the code defined in implementation class.
162 """
163 return VERSION
165 def getConfig(self) -> ApdbCassandraConfig:
166 # docstring is inherited from a base class
167 return self._context.config
169 def tableDef(self, table: ApdbTables) -> Table | None:
170 # docstring is inherited from a base class
171 return self.schema.tableSchemas.get(table)
173 @classmethod
174 def init_database(
175 cls,
176 hosts: tuple[str, ...],
177 keyspace: str,
178 *,
179 schema_file: str | None = None,
180 ss_schema_file: str | None = None,
181 read_sources_months: int | None = None,
182 read_forced_sources_months: int | None = None,
183 enable_replica: bool = False,
184 replica_skips_diaobjects: bool = False,
185 port: int | None = None,
186 username: str | None = None,
187 dbauth_alias: str | None = None,
188 prefix: str | None = None,
189 part_pixelization: str | None = None,
190 part_pix_level: int | None = None,
191 time_partition_tables: bool = True,
192 time_partition_start: str | None = None,
193 time_partition_end: str | None = None,
194 read_consistency: str | None = None,
195 write_consistency: str | None = None,
196 read_timeout: int | None = None,
197 write_timeout: int | None = None,
198 ra_dec_columns: tuple[str, str] | None = None,
199 replication_factor: int | None = None,
200 drop: bool = False,
201 table_options: CreateTableOptions | None = None,
202 ) -> ApdbCassandraConfig:
203 """Initialize new APDB instance and make configuration object for it.
205 Parameters
206 ----------
207 hosts : `tuple` [`str`, ...]
208 List of host names or IP addresses for Cassandra cluster.
209 keyspace : `str`
210 Name of the keyspace for APDB tables.
211 schema_file : `str`, optional
212 Location of (YAML) configuration file with APDB schema. If not
213 specified then default location will be used.
214 ss_schema_file : `str`, optional
215 Location of (YAML) configuration file with SSO schema. If not
216 specified then default location will be used.
217 read_sources_months : `int`, optional
218 Number of months of history to read from DiaSource.
219 read_forced_sources_months : `int`, optional
220 Number of months of history to read from DiaForcedSource.
221 enable_replica : `bool`, optional
222 If True, make additional tables used for replication to PPDB.
223 replica_skips_diaobjects : `bool`, optional
224 If `True` then do not fill regular ``DiaObject`` table when
225 ``enable_replica`` is `True`.
226 port : `int`, optional
227 Port number to use for Cassandra connections.
228 username : `str`, optional
229 User name for Cassandra connections.
230 dbauth_alias : `str`, optional
231 If specified then this string will be used to as a host name when
232 checking credentials in db-auth.yaml in addition to regular host
233 names in contact_points. For example if
234 dbauth_alias='pp_apdb_prod_cluster' then the entry
235 'cassandra://pp_apdb_prod_cluster/' will match. Port number should
236 not be used in that entry. Alias has higher priority than host
237 names.
238 prefix : `str`, optional
239 Optional prefix for all table names.
240 part_pixelization : `str`, optional
241 Name of the MOC pixelization used for partitioning.
242 part_pix_level : `int`, optional
243 Pixelization level.
244 time_partition_tables : `bool`, optional
245 Create per-partition tables.
246 time_partition_start : `str`, optional
247 Starting time for per-partition tables, in yyyy-mm-ddThh:mm:ss
248 format, in TAI.
249 time_partition_end : `str`, optional
250 Ending time for per-partition tables, in yyyy-mm-ddThh:mm:ss
251 format, in TAI.
252 read_consistency : `str`, optional
253 Name of the consistency level for read operations.
254 write_consistency : `str`, optional
255 Name of the consistency level for write operations.
256 read_timeout : `int`, optional
257 Read timeout in seconds.
258 write_timeout : `int`, optional
259 Write timeout in seconds.
260 ra_dec_columns : `tuple` [`str`, `str`], optional
261 Names of ra/dec columns in DiaObject table.
262 replication_factor : `int`, optional
263 Replication factor used when creating new keyspace, if keyspace
264 already exists its replication factor is not changed.
265 drop : `bool`, optional
266 If `True` then drop existing tables before re-creating the schema.
267 table_options : `CreateTableOptions`, optional
268 Options used when creating Cassandra tables.
270 Returns
271 -------
272 config : `ApdbCassandraConfig`
273 Resulting configuration object for a created APDB instance.
274 """
275 # Some non-standard defaults for connection parameters, these can be
276 # changed later in generated config. Check Cassandra driver
277 # documentation for what these parameters do. These parameters are not
278 # used during database initialization, but they will be saved with
279 # generated config.
280 connection_config = ApdbCassandraConnectionConfig(
281 extra_parameters={
282 "idle_heartbeat_interval": 0,
283 "idle_heartbeat_timeout": 30,
284 "control_connection_timeout": 100,
285 },
286 )
287 config = ApdbCassandraConfig(
288 contact_points=hosts,
289 keyspace=keyspace,
290 enable_replica=enable_replica,
291 replica_skips_diaobjects=replica_skips_diaobjects,
292 connection_config=connection_config,
293 )
294 config.partitioning.time_partition_tables = time_partition_tables
295 if schema_file is not None:
296 config.schema_file = schema_file
297 if ss_schema_file is not None:
298 config.ss_schema_file = ss_schema_file
299 if read_sources_months is not None:
300 config.read_sources_months = read_sources_months
301 if read_forced_sources_months is not None:
302 config.read_forced_sources_months = read_forced_sources_months
303 if port is not None:
304 config.connection_config.port = port
305 if username is not None:
306 config.connection_config.username = username
307 if dbauth_alias is not None:
308 config.connection_config.dbauth_alias = dbauth_alias
309 if prefix is not None:
310 config.prefix = prefix
311 if part_pixelization is not None:
312 config.partitioning.part_pixelization = part_pixelization
313 if part_pix_level is not None:
314 config.partitioning.part_pix_level = part_pix_level
315 if time_partition_start is not None:
316 config.partitioning.time_partition_start = time_partition_start
317 if time_partition_end is not None:
318 config.partitioning.time_partition_end = time_partition_end
319 if read_consistency is not None:
320 config.connection_config.read_consistency = read_consistency
321 if write_consistency is not None:
322 config.connection_config.write_consistency = write_consistency
323 if read_timeout is not None:
324 config.connection_config.read_timeout = read_timeout
325 if write_timeout is not None:
326 config.connection_config.write_timeout = write_timeout
327 if ra_dec_columns is not None:
328 config.ra_dec_columns = ra_dec_columns
330 cls._makeSchema(config, drop=drop, replication_factor=replication_factor, table_options=table_options)
332 return config
334 def get_replica(self) -> ApdbCassandraReplica:
335 """Return `ApdbReplica` instance for this database."""
336 # Note that this instance has to stay alive while replica exists, so
337 # we pass reference to self.
338 return ApdbCassandraReplica(self)
340 @classmethod
341 def _makeSchema(
342 cls,
343 config: ApdbConfig,
344 *,
345 drop: bool = False,
346 replication_factor: int | None = None,
347 table_options: CreateTableOptions | None = None,
348 ) -> None:
349 # docstring is inherited from a base class
351 if not isinstance(config, ApdbCassandraConfig):
352 raise TypeError(f"Unexpected type of configuration object: {type(config)}")
354 simple_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
356 with SessionContext(config) as session:
357 schema = ApdbCassandraSchema(
358 session=session,
359 keyspace=config.keyspace,
360 table_schemas=simple_schema.tableSchemas,
361 prefix=config.prefix,
362 time_partition_tables=config.partitioning.time_partition_tables,
363 enable_replica=config.enable_replica,
364 replica_skips_diaobjects=config.replica_skips_diaobjects,
365 )
367 # Ask schema to create all tables.
368 part_range_config: ApdbCassandraTimePartitionRange | None = None
369 if config.partitioning.time_partition_tables:
370 partitioner = Partitioner(config)
371 time_partition_start = astropy.time.Time(
372 config.partitioning.time_partition_start, format="isot", scale="tai"
373 )
374 time_partition_end = astropy.time.Time(
375 config.partitioning.time_partition_end, format="isot", scale="tai"
376 )
377 part_range_config = ApdbCassandraTimePartitionRange(
378 start=partitioner.time_partition(time_partition_start),
379 end=partitioner.time_partition(time_partition_end),
380 )
381 schema.makeSchema(
382 drop=drop,
383 part_range=part_range_config,
384 replication_factor=replication_factor,
385 table_options=table_options,
386 )
387 else:
388 schema.makeSchema(
389 drop=drop, replication_factor=replication_factor, table_options=table_options
390 )
392 meta_table_name = ApdbTables.metadata.table_name(config.prefix)
393 metadata = ApdbMetadataCassandra(
394 session, meta_table_name, config.keyspace, "read_tuples", "write"
395 )
397 # Fill version numbers, overrides if they existed before.
398 metadata.set(
399 ConnectionContext.metadataSchemaVersionKey, str(simple_schema.schemaVersion()), force=True
400 )
401 metadata.set(
402 ConnectionContext.metadataCodeVersionKey, str(cls.apdbImplementationVersion()), force=True
403 )
405 if config.enable_replica:
406 # Only store replica code version if replica is enabled.
407 metadata.set(
408 ConnectionContext.metadataReplicaVersionKey,
409 str(ApdbCassandraReplica.apdbReplicaImplementationVersion()),
410 force=True,
411 )
413 # Store frozen part of a configuration in metadata.
414 freezer = ApdbConfigFreezer[ApdbCassandraConfig](ConnectionContext.frozen_parameters)
415 metadata.set(ConnectionContext.metadataConfigKey, freezer.to_json(config), force=True)
417 # Store time partition range.
418 if part_range_config:
419 part_range_config.save_to_meta(metadata)
421 def getDiaObjects(self, region: sphgeom.Region) -> pandas.DataFrame:
422 # docstring is inherited from a base class
423 context = self._context
424 config = context.config
426 sp_where, num_sp_part = context.partitioner.spatial_where(region)
427 _LOG.debug("getDiaObjects: #partitions: %s", len(sp_where))
429 # We need to exclude extra partitioning columns from result.
430 column_names = context.schema.apdbColumnNames(ApdbTables.DiaObjectLast)
431 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
432 query = Select(self._keyspace, table_name, column_names)
433 statements: list[tuple] = []
434 for where_clause in sp_where:
435 full_query = query.where(where_clause)
436 statements.append(context.stmt_factory.with_params(full_query, prepare=True))
437 _LOG.debug("getDiaObjects: #queries: %s", len(statements))
439 with self._timer("select_time", tags={"table": "DiaObject", "method": "getDiaObjects"}) as timer:
440 raw_objects = cast(
441 ApdbCassandraTableData,
442 select_concurrent(
443 context.session,
444 statements,
445 "read_raw_multi",
446 config.connection_config.read_concurrency,
447 ),
448 )
449 objects = raw_objects.to_pandas(context.schema._table_schema(ApdbTables.DiaObjectLast))
450 timer.add_values(row_count=len(objects), num_sp_part=num_sp_part, num_queries=len(statements))
452 self._fix_null_columns(ApdbTables.DiaObjectLast, objects)
454 _LOG.debug("found %s DiaObjects", objects.shape[0])
455 return objects
457 def getDiaSources(
458 self,
459 region: sphgeom.Region,
460 object_ids: Iterable[int] | None,
461 visit_time: astropy.time.Time,
462 start_time: astropy.time.Time | None = None,
463 ) -> pandas.DataFrame | None:
464 # docstring is inherited from a base class
465 context = self._context
466 config = context.config
468 months = config.read_sources_months
469 if start_time is None and months == 0:
470 return None
472 mjd_end = float(visit_time.tai.mjd)
473 if start_time is None:
474 mjd_start = mjd_end - months * 30
475 else:
476 mjd_start = float(start_time.tai.mjd)
478 sources = self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaSource)
480 self._fix_null_columns(ApdbTables.DiaSource, sources)
482 return sources
484 def getDiaForcedSources(
485 self,
486 region: sphgeom.Region,
487 object_ids: Iterable[int] | None,
488 visit_time: astropy.time.Time,
489 start_time: astropy.time.Time | None = None,
490 ) -> pandas.DataFrame | None:
491 # docstring is inherited from a base class
492 context = self._context
493 config = context.config
495 months = config.read_forced_sources_months
496 if start_time is None and months == 0:
497 return None
499 mjd_end = float(visit_time.tai.mjd)
500 if start_time is None:
501 mjd_start = mjd_end - months * 30
502 else:
503 mjd_start = float(start_time.tai.mjd)
505 sources = self._getSources(region, object_ids, mjd_start, mjd_end, ApdbTables.DiaForcedSource)
507 self._fix_null_columns(ApdbTables.DiaForcedSource, sources)
509 return sources
511 def getDiaObjectsForDedup(self, since: astropy.time.Time | None = None) -> pandas.DataFrame:
512 # docstring is inherited from a base class
513 context = self._context
514 config = context.config
516 if not context.has_dedup_table:
517 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
519 if since is None:
520 # Read last deduplication time from metadata.
521 dedup_str = context.metadata.get(context.metadataDedupKey)
522 if dedup_str is not None:
523 dedup_state = json.loads(dedup_str)
524 dedup_time_str = dedup_state["dedup_time_iso_tai"]
525 since = astropy.time.Time(dedup_time_str, format="iso", scale="tai")
527 column_names = context.schema.apdbColumnNames(ExtraTables.DiaObjectDedup)
529 validity_start_column = self._timestamp_column_name("validityStart")
530 timestamp = None if since is None else self._timestamp_column_value(since)
532 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
533 query = Select(self._keyspace, table_name, column_names, extra_clause="ALLOW FILTERING")
534 query = query.where(C("dedup_part") == 0)
535 if since is not None:
536 query = query.where(C(validity_start_column) >= 0)
538 statement = context.stmt_factory(query, prepare=False)
540 num_part = config.partitioning.num_part_dedup
541 statements = []
542 for dedup_part in range(num_part):
543 params = (dedup_part,) if timestamp is None else (dedup_part, timestamp)
544 statements.append((statement, params))
546 with self._timer(
547 "select_time", tags={"table": "DiaObjectDedup", "method": "getDiaObjectsForDedup"}
548 ) as timer:
549 objects_raw = cast(
550 ApdbCassandraTableData,
551 select_concurrent(
552 context.session,
553 statements,
554 "read_raw_multi_dedup",
555 config.connection_config.read_concurrency,
556 ),
557 )
558 objects = objects_raw.to_pandas(context.schema._table_schema(ExtraTables.DiaObjectDedup))
559 timer.add_values(row_count=len(objects), num_queries=num_part)
561 _LOG.debug("found %s DiaObjectDedup records", objects.shape[0])
562 return objects
564 def getDiaSourcesForDiaObjects(
565 self, objects: list[DiaObjectId], start_time: astropy.time.Time, max_dist_arcsec: float = 1.0
566 ) -> pandas.DataFrame:
567 # docstring is inherited from a base class
568 context = self._context
569 config = context.config
571 # Which tables to query and temporal constraints.
572 end_time = self._current_time()
573 tables, temporal_where = context.partitioner.temporal_where(
574 ApdbTables.DiaSource,
575 start_time,
576 end_time,
577 partitons_range=context.time_partitions_range,
578 query_per_time_part=False,
579 )
580 if not tables:
581 warnings.warn(
582 f"Query time range ({start_time.isot} - {end_time.isot}) does not overlap database "
583 "time partitions."
584 )
586 # Group DiaObjects by partition.
587 partitioned_object_ids = self._group_dia_objects_by_partition(
588 context.partitioner, objects, max_dist_arcsec
589 )
591 # Columns to return.
592 column_names = context.schema.apdbColumnNames(ApdbTables.DiaSource)
594 # Make a bunch of queries.
595 statements = []
596 for apdb_part, diaObjectIds in partitioned_object_ids.items():
597 spatial_where = [C("apdb_part") == apdb_part]
598 for table in tables:
599 query = Select(self._keyspace, table, column_names, extra_clause="ALLOW FILTERING")
600 for id_chunk in chunk_iterable(diaObjectIds, 10_000):
601 id_where = C("diaObjectId").in_(id_chunk)
602 for clause in QExpr.combine(spatial_where, temporal_where, extra=id_where):
603 statements.append(
604 context.stmt_factory.with_params(query.where(clause), prepare=False)
605 )
607 _LOG.debug("getDiaSourcesForDiaObjects #queries: %s", len(statements))
609 with self._timer(
610 "select_time", tags={"table": "DiaSource", "method": "getDiaSourcesForDiaObjects"}
611 ) as timer:
612 table_data_raw = cast(
613 ApdbCassandraTableData,
614 select_concurrent(
615 context.session,
616 statements,
617 "read_raw_multi",
618 config.connection_config.read_concurrency,
619 ),
620 )
621 catalog = table_data_raw.to_pandas(context.schema._table_schema(ApdbTables.DiaSource))
622 timer.add_values(row_count_from_db=len(catalog), num_queries=len(statements))
624 # precise filtering on midpointMjdTai
625 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] >= start_time.tai.mjd])
627 timer.add_values(row_count=len(catalog))
629 _LOG.debug("found %d DiaSources", len(catalog))
630 return catalog
632 def containsVisitDetector(
633 self,
634 visit: int,
635 detector: int,
636 region: sphgeom.Region | None = None,
637 visit_time: astropy.time.Time | None = None,
638 ) -> bool:
639 # docstring is inherited from a base class
640 context = self._context
642 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
643 query = Select(self._keyspace, table_name, [ColumnExpr("count(*)")])
644 query = query.where((C("visit") == visit) & (C("detector") == detector))
645 stmt, params = context.stmt_factory.with_params(query, prepare=False)
647 with self._timer("contains_visit_detector_time", tags={"table": table_name}):
648 result = context.session.execute(stmt, params)
649 return bool(result.one()[0])
651 def store(
652 self,
653 visit_time: astropy.time.Time,
654 objects: pandas.DataFrame,
655 sources: pandas.DataFrame | None = None,
656 forced_sources: pandas.DataFrame | None = None,
657 ) -> None:
658 # docstring is inherited from a base class
659 context = self._context
660 config = context.config
662 # Store visit/detector in a special table, this has to be done
663 # before all other writes so if there is a failure at any point
664 # later we still have a record for attempted write.
665 visit_detector: set[tuple[int, int]] = set()
666 for df in sources, forced_sources:
667 if df is not None and not df.empty:
668 df = df[["visit", "detector"]]
669 for visit, detector in df.itertuples(index=False):
670 visit_detector.add((visit, detector))
672 if visit_detector:
673 # Typically there is only one entry, do not bother with
674 # concurrency.
675 table_name = context.schema.tableName(ExtraTables.ApdbVisitDetector)
676 query = Insert(self._keyspace, table_name, ("visit", "detector"))
677 stmt = context.stmt_factory(query)
678 for item in visit_detector:
679 context.session.execute(stmt, item, execution_profile="write")
681 objects = self._fix_input_timestamps(objects)
682 if sources is not None:
683 sources = self._fix_input_timestamps(sources)
684 if forced_sources is not None:
685 forced_sources = self._fix_input_timestamps(forced_sources)
687 replica_chunk: ReplicaChunk | None = None
688 if context.schema.replication_enabled:
689 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, config.replica_chunk_seconds)
690 self._storeReplicaChunk(replica_chunk)
692 # fill region partition column for DiaObjects
693 objects = self._add_apdb_part(objects)
694 self._storeDiaObjects(objects, visit_time, replica_chunk)
696 if sources is not None and len(sources) > 0:
697 # copy apdb_part column from DiaObjects to DiaSources
698 sources = self._add_apdb_part(sources)
699 subchunk = self._storeDiaSources(ApdbTables.DiaSource, sources, replica_chunk)
700 self._storeDiaSourcesPartitions(sources, visit_time, replica_chunk, subchunk)
702 if forced_sources is not None and len(forced_sources) > 0:
703 forced_sources = self._add_apdb_part(forced_sources)
704 self._storeDiaSources(ApdbTables.DiaForcedSource, forced_sources, replica_chunk)
706 def reassignDiaSourcesToDiaObjects(
707 self,
708 idMap: Mapping[DiaSourceId, int],
709 *,
710 increment_nDiaSources: bool = True,
711 decrement_nDiaSources: bool = True,
712 ) -> None:
713 # docstring is inherited from a base class
714 context = self._context
715 config = context.config
717 source_ids = {source_id.diaSourceId for source_id in idMap}
719 # Find all DiaSources.
720 found_sources = self._get_diasource_data(
721 idMap, "apdb_part", "diaObjectId", "ra", "dec", "midpointMjdTai"
722 )
724 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}):
725 raise LookupError(f"Some source IDs were not found in DiaSource table: {missing_ids}")
727 found_sources_by_id = {row.diaSourceId: row for row in found_sources}
729 # Make sure that all DiaObjects exist, we also want to know
730 # nDiaSources count for current and new records because we want to
731 # send updated values to replica.
732 current_object_ids = {
733 DiaObjectId(diaObjectId=row.diaObjectId, ra=row.ra, dec=row.dec) for row in found_sources
734 }
735 # Assume that DiaSource ra/dec are very close to re-assigned objects.
736 new_object_ids = {
737 DiaObjectId(diaObjectId=diaObjectId, ra=source_id.ra, dec=source_id.dec)
738 for source_id, diaObjectId in idMap.items()
739 }
740 all_object_ids = new_object_ids | current_object_ids
741 found_objects = self._get_diaobject_data(all_object_ids, "apdb_part", "ra", "dec", "nDiaSources")
743 if missing_ids := (
744 {row.diaObjectId for row in all_object_ids} - {row.diaObjectId for row in found_objects}
745 ):
746 raise LookupError(f"Some object IDs were not found in DiaObjectLast table: {missing_ids}")
748 update_records: list[ApdbUpdateRecord] = []
749 update_order = 0
750 current_time = self._current_time()
751 current_time_ns = int(current_time.unix_tai * 1e9)
753 # Update DiaSources.
754 statements: list[tuple] = []
755 for source_id, diaObjectId in idMap.items():
756 source_row = found_sources_by_id[source_id.diaSourceId]
757 apdb_part = source_row.apdb_part
758 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
760 if config.partitioning.time_partition_tables:
761 table_name = context.schema.tableName(ApdbTables.DiaSource, time_part)
762 update = (
763 Update(self._keyspace, table_name)
764 .values(C("diaObjectId").update(diaObjectId))
765 .where(C("apdb_part") == apdb_part)
766 .where(C("diaSourceId") == source_id.diaSourceId)
767 )
768 else:
769 table_name = context.schema.tableName(ApdbTables.DiaSource)
770 update = (
771 Update(self._keyspace, table_name)
772 .values(C("diaObjectId").update(diaObjectId))
773 .where(C("apdb_part") == apdb_part)
774 .where(C("apdb_time_part") == time_part)
775 .where(C("diaSourceId") == source_id.diaSourceId)
776 )
777 statements.append(context.stmt_factory.with_params(update, prepare=True))
779 if context.schema.replication_enabled:
780 update_records.append(
781 ApdbReassignDiaSourceToDiaObjectRecord(
782 diaSourceId=source_id.diaSourceId,
783 ra=source_id.ra,
784 dec=source_id.dec,
785 midpointMjdTai=source_id.midpointMjdTai,
786 diaObjectId=diaObjectId,
787 update_time_ns=current_time_ns,
788 update_order=update_order,
789 )
790 )
791 update_order += 1
793 with self._timer(
794 "update_time", tags={"table": "DiaSource", "method": "reassignDiaSourcesToDiaObjects"}
795 ) as timer:
796 execute_concurrent(context.session, statements, execution_profile="write")
797 timer.add_values(num_queries=len(statements))
799 # Update nDiaSources in DiaObjectLast. We do not update DiaObject table
800 # here because it may not even exist. PPDB updates DiaObject from
801 # update records.
802 if increment_nDiaSources or decrement_nDiaSources:
803 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
804 update = (
805 Update(self._keyspace, table_name)
806 .values(C("nDiaSources").update(-1))
807 .where(C("apdb_part") == -1)
808 .where(C("diaObjectId") == -1)
809 )
810 statement = context.stmt_factory(update, prepare=True)
811 statements = []
813 # Calculate increments/decrements for all affected DiaObjects.
814 increments: Counter = Counter()
815 if increment_nDiaSources:
816 increments.update(idMap.values())
817 if decrement_nDiaSources:
818 increments.subtract(row.diaObjectId for row in found_sources)
820 for row in found_objects:
821 if increments.get(row.diaObjectId):
822 nDiaSources = row.nDiaSources + increments[row.diaObjectId]
823 statements.append((statement, (nDiaSources, row.apdb_part, row.diaObjectId)))
825 # Also send updated values to replica.
826 if context.schema.replication_enabled:
827 update_records.append(
828 ApdbUpdateNDiaSourcesRecord(
829 diaObjectId=row.diaObjectId,
830 ra=row.ra,
831 dec=row.dec,
832 nDiaSources=nDiaSources,
833 update_time_ns=current_time_ns,
834 update_order=update_order,
835 )
836 )
837 update_order += 1
839 if statements:
840 with self._timer(
841 "update_time", tags={"table": table_name, "method": "reassignDiaSourcesToDiaObjects"}
842 ) as timer:
843 execute_concurrent(context.session, statements, execution_profile="write")
844 timer.add_values(num_queries=len(statements))
846 if update_records:
847 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
848 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
850 def setValidityEnd(
851 self, objects: list[DiaObjectId], validityEnd: astropy.time.Time, raise_on_missing_id: bool = False
852 ) -> int:
853 # docstring is inherited from a base class
854 if not objects:
855 return 0
857 context = self._context
858 config = context.config
860 pad_arcsec = 1.0
861 partitioned_object_ids = self._group_dia_objects_by_partition(
862 context.partitioner, objects, pad_arcsec
863 )
865 # Check that all objects exist.
866 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
867 statements: list[tuple] = []
868 for apdb_part, diaObjectIds in partitioned_object_ids.items():
869 query = Select(self._keyspace, table_name, ["apdb_part", "diaObjectId"])
870 query = query.where(C("apdb_part") == apdb_part)
871 query = query.where(C("diaObjectId").in_(diaObjectIds))
872 statements.append(context.stmt_factory.with_params(query, prepare=False))
874 with self._timer("select_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
875 records = cast(
876 list[tuple[int, int]],
877 select_concurrent(
878 context.session,
879 statements,
880 "read_tuples",
881 config.connection_config.read_concurrency,
882 ),
883 )
884 timer.add_values(row_count=len(objects))
886 requested_ids = {obj.diaObjectId for obj in objects}
887 found_ids = {rec[1] for rec in records}
888 if extra_ids := (found_ids - requested_ids):
889 raise RuntimeError(f"Consistency error - found duplicate records for object IDs: {extra_ids}")
890 if raise_on_missing_id:
891 if missing_ids := (requested_ids - found_ids):
892 raise LookupError(f"Some object IDs are missing from DiaObjectLast table: {missing_ids}")
894 # Filter existing records.
895 if len(objects) != len(found_ids):
896 objects = [obj for obj in objects if obj.diaObjectId in found_ids]
898 if not objects:
899 return 0
901 # Group by partitions again.
902 grouped_object_ids: dict[int, list[int]] = defaultdict(list)
903 for apdb_part, diaObjectId in records:
904 grouped_object_ids[apdb_part].append(diaObjectId)
906 # Remove all matching rows from DiaObjectLast.
907 statements = []
908 for apdb_part, diaObjectIds in grouped_object_ids.items():
909 delete = (
910 Delete(self._keyspace, table_name)
911 .where(C("apdb_part") == apdb_part)
912 .where(C("diaObjectId").in_(diaObjectIds))
913 )
914 statements.append(context.stmt_factory.with_params(delete))
916 # Also remove from DiaObjectLastToPartition.
917 reverse_table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
918 delete = Delete(self._keyspace, reverse_table_name).where(
919 C("diaObjectId").in_([rec[1] for rec in records])
920 )
921 statements.append(context.stmt_factory.with_params(delete))
923 with self._timer("delete_time", tags={"table": table_name, "method": "setValidityEnd"}) as timer:
924 execute_concurrent(context.session, statements, execution_profile="write")
925 timer.add_values(row_count=len(records))
927 # If repication is enabled then send all updates.
928 if context.schema.replication_enabled:
929 current_time = self._current_time()
930 current_time_ns = int(current_time.unix_tai * 1e9)
931 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
933 update_records = [
934 ApdbCloseDiaObjectValidityRecord(
935 diaObjectId=obj.diaObjectId,
936 ra=obj.ra,
937 dec=obj.dec,
938 update_time_ns=current_time_ns,
939 update_order=index,
940 validityEndMjdTai=float(validityEnd.tai.mjd),
941 nDiaSources=None,
942 )
943 for index, obj in enumerate(objects)
944 ]
946 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
948 return len(objects)
950 def resetDedup(self, dedup_time: astropy.time.Time | None = None) -> None:
951 # docstring is inherited from a base class
952 context = self._context
954 if not context.has_dedup_table:
955 raise TypeError("DiaObjectDedup table does not exist in this APDB instance.")
957 if dedup_time is None:
958 dedup_time = self._current_time()
960 validity_start_column = self._timestamp_column_name("validityStart")
962 # Find latest timestamp in deduplication table.
963 table_name = context.schema.tableName(ExtraTables.DiaObjectDedup)
964 query = Select(self._keyspace, table_name, [ColumnExpr(f'MAX("{validity_start_column}")')])
965 stmt = context.stmt_factory(query, prepare=False)
967 result = context.session.execute(stmt, execution_profile="read_tuples")
968 max_value = result.one()[0]
969 if self._schema.has_mjd_timestamps:
970 max_validity_start = astropy.time.Time(max_value, format="mjd", scale="tai")
971 else:
972 max_validity_start = astropy.time.Time(max_value, format="datetime", scale="tai")
974 # If max time is lower than dedup time we can do TRUNCATE.
975 if dedup_time >= max_validity_start:
976 query_str = f'TRUNCATE TABLE "{self._keyspace}"."{table_name}"'
977 context.session.execute(query_str, execution_profile="write")
978 else:
979 dedup_time_value = self._timestamp_column_value(dedup_time)
980 delete = Delete(self._keyspace, table_name).where(C(validity_start_column) < dedup_time_value)
981 stmt, params = context.stmt_factory.with_params(delete)
982 context.session.execute(stmt, params, execution_profile="write")
984 # Store dedup time.
985 data = {"dedup_time_iso_tai": dedup_time.tai.to_value("iso")}
986 data_json = json.dumps(data)
987 context.metadata.set(context.metadataDedupKey, data_json, force=True)
989 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
990 # docstring is inherited from a base class
991 context = self._context
992 config = context.config
994 now = self._current_time()
995 reassign_time_column = self._timestamp_column_name("ssObjectReassocTime")
996 reassignTime = self._timestamp_column_value(now)
998 # To update a record we need to know its exact primary key (including
999 # partition key) so we start by querying for diaSourceId to find the
1000 # primary keys.
1002 table_name = context.schema.tableName(ExtraTables.DiaSourceToPartition)
1003 # split it into 1k IDs per query
1004 selects: list[tuple] = []
1005 columns = ["diaSourceId", "apdb_part", "apdb_time_part", "apdb_replica_chunk"]
1006 query = Select(self._keyspace, table_name, columns)
1007 for ids in chunk_iterable(idMap.keys(), 1_000):
1008 full_query = query.where(C("diaSourceId").in_(ids))
1009 selects.append(context.stmt_factory.with_params(full_query, prepare=False))
1011 # No need for DataFrame here, read data as tuples.
1012 result = cast(
1013 list[tuple[int, int, int, int | None]],
1014 select_concurrent(
1015 context.session, selects, "read_tuples", config.connection_config.read_concurrency
1016 ),
1017 )
1019 # Make mapping from source ID to its partition.
1020 id2partitions: dict[int, tuple[int, int]] = {}
1021 id2chunk_id: dict[int, int] = {}
1022 for row in result:
1023 id2partitions[row[0]] = row[1:3]
1024 if row[3] is not None:
1025 id2chunk_id[row[0]] = row[3]
1027 # make sure we know partitions for each ID
1028 if set(id2partitions) != set(idMap):
1029 missing = ",".join(str(item) for item in set(idMap) - set(id2partitions))
1030 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
1032 # Reassign in standard tables
1033 queries: list[tuple[cassandra.query.PreparedStatement, tuple]] = []
1034 for diaSourceId, ssObjectId in idMap.items():
1035 apdb_part, apdb_time_part = id2partitions[diaSourceId]
1036 if config.partitioning.time_partition_tables:
1037 table_name = context.schema.tableName(ApdbTables.DiaSource, apdb_time_part)
1038 update = (
1039 Update(self._keyspace, table_name)
1040 .values(
1041 C("ssObjectId").update(ssObjectId),
1042 C("diaObjectId").update(None),
1043 C(reassign_time_column).update(reassignTime),
1044 )
1045 .where(C("apdb_part") == apdb_part)
1046 .where(C("diaSourceId") == diaSourceId)
1047 )
1048 else:
1049 table_name = context.schema.tableName(ApdbTables.DiaSource)
1050 update = (
1051 Update(self._keyspace, table_name)
1052 .values(
1053 C("ssObjectId").update(ssObjectId),
1054 C("diaObjectId").update(None),
1055 C(reassign_time_column).update(reassignTime),
1056 )
1057 .where(C("apdb_part") == apdb_part)
1058 .where(C("apdb_time_part") == apdb_time_part)
1059 .where(C("diaSourceId") == diaSourceId)
1060 )
1061 queries.append(context.stmt_factory.with_params(update, prepare=True))
1063 # TODO: (DM-50190) Replication for updated records is not implemented.
1064 if id2chunk_id:
1065 warnings.warn("Replication of reassigned DiaSource records is not implemented.", stacklevel=2)
1067 _LOG.debug("%s: will update %d records", table_name, len(idMap))
1068 with self._timer("source_reassign_time") as timer:
1069 execute_concurrent(context.session, queries, execution_profile="write")
1070 timer.add_values(source_count=len(idMap))
1072 def withdrawDiaSources(
1073 self,
1074 diaSourceIds: Iterable[DiaSourceId],
1075 *,
1076 timeWithdrawn: astropy.time.Time | None = None,
1077 ) -> None:
1078 # docstring is inherited from a base class
1079 context = self._context
1080 config = context.config
1082 if timeWithdrawn is None:
1083 timeWithdrawn = self._current_time()
1084 time_value = self._timestamp_column_value(timeWithdrawn)
1085 column_name = self._timestamp_column_name("time_withdrawn")
1087 diaSourceIds = list(diaSourceIds)
1088 source_ids = {source_id.diaSourceId for source_id in diaSourceIds}
1090 # Find all DiaSources.
1091 found_sources = self._get_diasource_data(
1092 diaSourceIds, "apdb_part", "diaObjectId", "ra", "dec", "midpointMjdTai", column_name
1093 )
1095 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}):
1096 raise LookupError(f"Some source IDs were not found in DiaSource table: {missing_ids}")
1098 found_sources_by_id = {row.diaSourceId: row for row in found_sources}
1100 update_records: list[ApdbUpdateRecord] = []
1101 update_order = 0
1102 current_time = self._current_time()
1103 current_time_ns = int(current_time.unix_tai * 1e9)
1105 # Update DiaSources.
1106 statements: list[tuple] = []
1107 for source_id in diaSourceIds:
1108 source_row = found_sources_by_id[source_id.diaSourceId]
1109 # Ignore sources already withdrawn.
1110 if getattr(source_row, column_name) is not None:
1111 continue
1113 apdb_part = source_row.apdb_part
1114 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
1116 if config.partitioning.time_partition_tables:
1117 table_name = context.schema.tableName(ApdbTables.DiaSource, time_part)
1118 update = (
1119 Update(self._keyspace, table_name)
1120 .values(C(column_name).update(time_value))
1121 .where(C("apdb_part") == apdb_part)
1122 .where(C("diaSourceId") == source_id.diaSourceId)
1123 )
1124 else:
1125 table_name = context.schema.tableName(ApdbTables.DiaSource)
1126 update = (
1127 Update(self._keyspace, table_name)
1128 .values(C(column_name).update(time_value))
1129 .where(C("apdb_part") == apdb_part)
1130 .where(C("apdb_time_part") == time_part)
1131 .where(C("diaSourceId") == source_id.diaSourceId)
1132 )
1133 statements.append(context.stmt_factory.with_params(update, prepare=True))
1135 if context.schema.replication_enabled:
1136 update_records.append(
1137 ApdbWithdrawDiaSourceRecord(
1138 diaSourceId=source_id.diaSourceId,
1139 ra=source_id.ra,
1140 dec=source_id.dec,
1141 midpointMjdTai=source_id.midpointMjdTai,
1142 update_time_ns=current_time_ns,
1143 update_order=update_order,
1144 timeWithdrawnMjdTai=float(timeWithdrawn.tai.mjd),
1145 )
1146 )
1147 update_order += 1
1149 with self._timer("update_time", tags={"table": "DiaSource", "method": "withdrawDiaSources"}) as timer:
1150 execute_concurrent(context.session, statements, execution_profile="write")
1151 timer.add_values(num_queries=len(statements))
1153 if update_records:
1154 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
1155 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
1157 def withdrawDiaForcedSources(
1158 self,
1159 diaForcedSourceIds: Iterable[DiaForcedSourceId],
1160 *,
1161 timeWithdrawn: astropy.time.Time | None = None,
1162 ) -> None:
1163 # docstring is inherited from a base class
1164 context = self._context
1165 config = context.config
1167 def _fsrc_id(fsource: Any) -> tuple[int, int, int]:
1168 return (fsource.diaObjectId, fsource.visit, fsource.detector)
1170 if timeWithdrawn is None:
1171 timeWithdrawn = self._current_time()
1172 time_value = self._timestamp_column_value(timeWithdrawn)
1173 column_name = self._timestamp_column_name("time_withdrawn")
1175 diaForcedSourceIds = list(diaForcedSourceIds)
1176 fsource_keys = {_fsrc_id(source) for source in diaForcedSourceIds}
1178 found_fsources = self._get_diaforcedsource_data(
1179 diaForcedSourceIds, "apdb_part", "ra", "dec", "midpointMjdTai", column_name
1180 )
1182 found_keys = {_fsrc_id(row) for row in found_fsources}
1183 if missing_ids := (fsource_keys - found_keys):
1184 raise LookupError(f"Some source IDs were not found in DiaForcedSource table: {missing_ids}")
1186 statements: list[tuple] = []
1187 update_records = []
1188 update_order = 0
1189 current_time = self._current_time()
1190 current_time_ns = int(current_time.unix_tai * 1e9)
1192 for source_row in found_fsources:
1193 # Ignore sources already withdrawn.
1194 if getattr(source_row, column_name) is not None:
1195 continue
1197 apdb_part = source_row.apdb_part
1198 time_part = context.partitioner.time_partition(source_row.midpointMjdTai)
1200 if config.partitioning.time_partition_tables:
1201 table_name = context.schema.tableName(ApdbTables.DiaForcedSource, time_part)
1202 update = (
1203 Update(self._keyspace, table_name)
1204 .values(C(column_name).update(time_value))
1205 .where(C("apdb_part") == apdb_part)
1206 .where(C("diaObjectId") == source_row.diaObjectId)
1207 .where(C("visit") == source_row.visit)
1208 .where(C("detector") == source_row.detector)
1209 )
1210 else:
1211 table_name = context.schema.tableName(ApdbTables.DiaForcedSource)
1212 update = (
1213 Update(self._keyspace, table_name)
1214 .values(C(column_name).update(time_value))
1215 .where(C("apdb_part") == apdb_part)
1216 .where(C("apdb_time_part") == time_part)
1217 .where(C("diaObjectId") == source_row.diaObjectId)
1218 .where(C("visit") == source_row.visit)
1219 .where(C("detector") == source_row.detector)
1220 )
1221 statements.append(context.stmt_factory.with_params(update, prepare=True))
1223 if context.schema.replication_enabled:
1224 update_records.append(
1225 ApdbWithdrawDiaForcedSourceRecord(
1226 diaObjectId=source_row.diaObjectId,
1227 visit=source_row.visit,
1228 detector=source_row.detector,
1229 ra=source_row.ra,
1230 dec=source_row.dec,
1231 midpointMjdTai=source_row.midpointMjdTai,
1232 update_time_ns=current_time_ns,
1233 update_order=update_order,
1234 timeWithdrawnMjdTai=float(timeWithdrawn.tai.mjd),
1235 )
1236 )
1237 update_order += 1
1239 if statements:
1240 with self._timer(
1241 "update_time", tags={"table": "DiaForcedSource", "method": "withdrawDiaForcedSources"}
1242 ) as timer:
1243 execute_concurrent(context.session, statements, execution_profile="write")
1244 timer.add_values(num_queries=len(statements))
1246 if update_records:
1247 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, config.replica_chunk_seconds)
1248 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
1250 def countUnassociatedObjects(self) -> int:
1251 # docstring is inherited from a base class
1253 # It's too inefficient to implement it for Cassandra in current schema.
1254 raise NotImplementedError()
1256 @property
1257 def schema(self) -> ApdbSchema:
1258 # docstring is inherited from a base class
1259 return self._schema
1261 @property
1262 def metadata(self) -> ApdbMetadata:
1263 # docstring is inherited from a base class
1264 context = self._context
1265 return context.metadata
1267 @property
1268 def admin(self) -> ApdbCassandraAdmin:
1269 # docstring is inherited from a base class
1270 return ApdbCassandraAdmin(self)
1272 def _getSources(
1273 self,
1274 region: sphgeom.Region,
1275 object_ids: Iterable[int] | None,
1276 mjd_start: float,
1277 mjd_end: float,
1278 table_name: ApdbTables,
1279 ) -> pandas.DataFrame:
1280 """Return catalog of DiaSource instances given set of DiaObject IDs.
1282 Parameters
1283 ----------
1284 region : `lsst.sphgeom.Region`
1285 Spherical region.
1286 object_ids :
1287 Collection of DiaObject IDs
1288 mjd_start : `float`
1289 Lower bound of time interval.
1290 mjd_end : `float`
1291 Upper bound of time interval.
1292 table_name : `ApdbTables`
1293 Name of the table.
1295 Returns
1296 -------
1297 catalog : `pandas.DataFrame`, or `None`
1298 Catalog containing DiaSource records. Empty catalog is returned if
1299 ``object_ids`` is empty.
1300 """
1301 context = self._context
1302 config = context.config
1304 object_id_set: Set[int] = set()
1305 if object_ids is not None:
1306 object_id_set = set(object_ids)
1307 if len(object_id_set) == 0:
1308 return self._make_empty_catalog(table_name)
1310 sp_where, num_sp_part = context.partitioner.spatial_where(region)
1311 tables, temporal_where = context.partitioner.temporal_where(
1312 table_name, mjd_start, mjd_end, partitons_range=context.time_partitions_range
1313 )
1314 if not tables:
1315 start = astropy.time.Time(mjd_start, format="mjd", scale="tai")
1316 end = astropy.time.Time(mjd_end, format="mjd", scale="tai")
1317 warnings.warn(
1318 f"Query time range ({start.isot} - {end.isot}) does not overlap database time partitions."
1319 )
1321 # We need to exclude extra partitioning columns from result.
1322 column_names = context.schema.apdbColumnNames(table_name)
1324 # Build all queries
1325 statements: list[tuple] = []
1326 for table in tables:
1327 query = Select(self._keyspace, table, column_names)
1328 for clause in QExpr.combine(sp_where, temporal_where):
1329 statements.append(context.stmt_factory.with_params(query.where(clause), prepare=True))
1330 _LOG.debug("_getSources %s: #queries: %s", table_name, len(statements))
1332 with self._timer("select_time", tags={"table": table_name.name, "method": "_getSources"}) as timer:
1333 table_data_raw = cast(
1334 ApdbCassandraTableData,
1335 select_concurrent(
1336 context.session,
1337 statements,
1338 "read_raw_multi",
1339 config.connection_config.read_concurrency,
1340 ),
1341 )
1342 catalog = table_data_raw.to_pandas(context.schema._table_schema(table_name))
1343 timer.add_values(
1344 row_count_from_db=len(catalog), num_sp_part=num_sp_part, num_queries=len(statements)
1345 )
1347 # filter by given object IDs
1348 if len(object_id_set) > 0:
1349 catalog = cast(pandas.DataFrame, catalog[catalog["diaObjectId"].isin(object_id_set)])
1351 # precise filtering on midpointMjdTai
1352 catalog = cast(pandas.DataFrame, catalog[catalog["midpointMjdTai"] > mjd_start])
1354 timer.add_values(row_count=len(catalog))
1356 _LOG.debug("found %d %ss", catalog.shape[0], table_name.name)
1357 return catalog
1359 def _storeReplicaChunk(self, replica_chunk: ReplicaChunk) -> None:
1360 context = self._context
1361 config = context.config
1363 # Cassandra timestamp uses milliseconds since epoch
1364 timestamp = int(replica_chunk.last_update_time.unix_tai * 1000)
1366 # everything goes into a single partition
1367 partition = 0
1369 table_name = context.schema.tableName(ExtraTables.ApdbReplicaChunks)
1371 columns = ["partition", "apdb_replica_chunk", "last_update_time", "unique_id"]
1372 values = [partition, replica_chunk.id, timestamp, replica_chunk.unique_id]
1373 if context.has_chunk_sub_partitions:
1374 columns.append("has_subchunks")
1375 values.append(True)
1377 query = Insert(self._keyspace, table_name, columns)
1378 stmt = context.stmt_factory(query)
1380 context.session.execute(
1381 stmt,
1382 values,
1383 timeout=config.connection_config.write_timeout,
1384 execution_profile="write",
1385 )
1387 def _queryDiaObjectLastPartitions(self, ids: Iterable[int]) -> Mapping[int, int]:
1388 """Return existing mapping of diaObjectId to its last partition."""
1389 context = self._context
1390 config = context.config
1392 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1393 queries = []
1394 object_count = 0
1395 for id_chunk in chunk_iterable(ids, 10_000):
1396 id_chunk_list = tuple(id_chunk)
1397 query = Select(self._keyspace, table_name, ("diaObjectId", "apdb_part"))
1398 query = query.where(C("diaObjectId").in_(id_chunk_list))
1399 queries.append(context.stmt_factory.with_params(query, prepare=False))
1400 object_count += len(id_chunk_list)
1402 with self._timer("query_object_last_partitions", tags={"table": table_name}) as timer:
1403 data = cast(
1404 ApdbTableData,
1405 select_concurrent(
1406 context.session,
1407 queries,
1408 "read_raw_multi",
1409 config.connection_config.read_concurrency,
1410 ),
1411 )
1412 timer.add_values(object_count=object_count, row_count=len(data.rows()))
1414 if data.column_names() != ["diaObjectId", "apdb_part"]:
1415 raise RuntimeError(f"Unexpected column names in query result: {data.column_names()}")
1417 return {row[0]: row[1] for row in data.rows()}
1419 def _deleteMovingObjects(self, objs: pandas.DataFrame) -> None:
1420 """Objects in DiaObjectsLast can move from one spatial partition to
1421 another. For those objects inserting new version does not replace old
1422 one, so we need to explicitly remove old versions before inserting new
1423 ones.
1424 """
1425 context = self._context
1427 # Extract all object IDs.
1428 new_partitions = dict(zip(objs["diaObjectId"], objs["apdb_part"]))
1429 old_partitions = self._queryDiaObjectLastPartitions(objs["diaObjectId"])
1431 moved_oids: dict[int, tuple[int, int]] = {}
1432 for oid, old_part in old_partitions.items():
1433 new_part = new_partitions.get(oid, old_part)
1434 if new_part != old_part:
1435 moved_oids[oid] = (old_part, new_part)
1436 _LOG.debug("DiaObject IDs that moved to new partition: %s", moved_oids)
1438 if moved_oids:
1439 # Delete old records from DiaObjectLast.
1440 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
1441 query = Delete(self._keyspace, table_name)
1442 query = query.where('apdb_part = {} AND "diaObjectId" = {}', (-1, -1))
1443 statement = context.stmt_factory(query, prepare=True)
1444 queries = []
1445 for oid, (old_part, _) in moved_oids.items():
1446 queries.append((statement, (old_part, oid)))
1447 with self._timer("delete_object_last", tags={"table": table_name}) as timer:
1448 execute_concurrent(context.session, queries, execution_profile="write")
1449 timer.add_values(row_count=len(moved_oids))
1451 # Add all new records to the map.
1452 table_name = context.schema.tableName(ExtraTables.DiaObjectLastToPartition)
1453 insert = Insert(self._keyspace, table_name, ("diaObjectId", "apdb_part"))
1454 statement = context.stmt_factory(insert, prepare=True)
1456 queries = []
1457 for oid, new_part in new_partitions.items():
1458 queries.append((statement, (oid, new_part)))
1460 with self._timer("update_object_last_partition", tags={"table": table_name}) as timer:
1461 execute_concurrent(context.session, queries, execution_profile="write")
1462 timer.add_values(row_count=len(queries))
1464 def _storeDiaObjects(
1465 self, objs: pandas.DataFrame, visit_time: astropy.time.Time, replica_chunk: ReplicaChunk | None
1466 ) -> None:
1467 """Store catalog of DiaObjects from current visit.
1469 Parameters
1470 ----------
1471 objs : `pandas.DataFrame`
1472 Catalog with DiaObject records
1473 visit_time : `astropy.time.Time`
1474 Time of the current visit.
1475 replica_chunk : `ReplicaChunk` or `None`
1476 Replica chunk identifier if replication is configured.
1477 """
1478 if len(objs) == 0:
1479 _LOG.debug("No objects to write to database.")
1480 return
1482 context = self._context
1483 config = context.config
1485 self._deleteMovingObjects(objs)
1487 validity_start_column = self._timestamp_column_name("validityStart")
1488 timestamp = self._timestamp_column_value(visit_time)
1490 # DiaObjectLast did not have this column in the past.
1491 extra_columns: dict[str, Any] = {}
1492 if context.schema.check_column(ApdbTables.DiaObjectLast, validity_start_column):
1493 extra_columns[validity_start_column] = timestamp
1495 self._storeObjectsPandas(objs, ApdbTables.DiaObjectLast, extra_columns=extra_columns)
1497 extra_columns[validity_start_column] = timestamp
1498 visit_time_part = context.partitioner.time_partition(visit_time)
1499 time_part: int | None = visit_time_part
1500 if (time_partitions_range := context.time_partitions_range) is not None:
1501 self._check_time_partitions([visit_time_part], time_partitions_range)
1502 if not config.partitioning.time_partition_tables:
1503 extra_columns["apdb_time_part"] = time_part
1504 time_part = None
1506 # Only store DiaObects if not doing replication or explicitly
1507 # configured to always store them.
1508 if replica_chunk is None or not config.replica_skips_diaobjects:
1509 self._storeObjectsPandas(
1510 objs, ApdbTables.DiaObject, extra_columns=extra_columns, time_part=time_part
1511 )
1513 if replica_chunk is not None:
1514 extra_columns = {"apdb_replica_chunk": replica_chunk.id, validity_start_column: timestamp}
1515 table = ExtraTables.DiaObjectChunks
1516 if context.has_chunk_sub_partitions:
1517 table = ExtraTables.DiaObjectChunks2
1518 # Use a random number for a second part of partitioning key so
1519 # that different clients could wrtite to different partitions.
1520 # This makes it not exactly reproducible.
1521 extra_columns["apdb_replica_subchunk"] = random.randrange(config.replica_sub_chunk_count)
1522 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1524 # Store copy of the records in dedup table.
1525 if context.has_dedup_table:
1526 table = ExtraTables.DiaObjectDedup
1527 extra_columns = {
1528 "dedup_part": random.randrange(config.partitioning.num_part_dedup),
1529 validity_start_column: timestamp,
1530 }
1531 self._storeObjectsPandas(objs, table, extra_columns=extra_columns)
1533 def _storeDiaSources(
1534 self,
1535 table_name: ApdbTables,
1536 sources: pandas.DataFrame,
1537 replica_chunk: ReplicaChunk | None,
1538 ) -> int | None:
1539 """Store catalog of DIASources or DIAForcedSources from current visit.
1541 Parameters
1542 ----------
1543 table_name : `ApdbTables`
1544 Table where to store the data.
1545 sources : `pandas.DataFrame`
1546 Catalog containing DiaSource records
1547 visit_time : `astropy.time.Time`
1548 Time of the current visit.
1549 replica_chunk : `ReplicaChunk` or `None`
1550 Replica chunk identifier if replication is configured.
1552 Returns
1553 -------
1554 subchunk : `int` or `None`
1555 Subchunk number for resulting replica data, `None` if relication is
1556 not enabled ot subchunking is not enabled.
1557 """
1558 context = self._context
1559 config = context.config
1561 # Time partitioning has to be based on midpointMjdTai, not visit_time
1562 # as visit_time is not really a visit time.
1563 tp_sources = sources.copy(deep=False)
1564 tp_sources["apdb_time_part"] = tp_sources["midpointMjdTai"].apply(context.partitioner.time_partition)
1565 if (time_partitions_range := context.time_partitions_range) is not None:
1566 self._check_time_partitions(tp_sources["apdb_time_part"], time_partitions_range)
1567 extra_columns: dict[str, Any] = {}
1568 if not config.partitioning.time_partition_tables:
1569 self._storeObjectsPandas(tp_sources, table_name)
1570 else:
1571 # Group by time partition
1572 partitions = set(tp_sources["apdb_time_part"])
1573 if len(partitions) == 1:
1574 # Single partition - just save the whole thing.
1575 time_part = partitions.pop()
1576 self._storeObjectsPandas(sources, table_name, time_part=time_part)
1577 else:
1578 # group by time partition.
1579 for time_part, sub_frame in tp_sources.groupby(by="apdb_time_part"):
1580 sub_frame.drop(columns="apdb_time_part", inplace=True)
1581 self._storeObjectsPandas(sub_frame, table_name, time_part=time_part)
1583 subchunk: int | None = None
1584 if replica_chunk is not None:
1585 extra_columns = {"apdb_replica_chunk": replica_chunk.id}
1586 if context.has_chunk_sub_partitions:
1587 subchunk = random.randrange(config.replica_sub_chunk_count)
1588 extra_columns["apdb_replica_subchunk"] = subchunk
1589 if table_name is ApdbTables.DiaSource:
1590 extra_table = ExtraTables.DiaSourceChunks2
1591 else:
1592 extra_table = ExtraTables.DiaForcedSourceChunks2
1593 else:
1594 if table_name is ApdbTables.DiaSource:
1595 extra_table = ExtraTables.DiaSourceChunks
1596 else:
1597 extra_table = ExtraTables.DiaForcedSourceChunks
1598 self._storeObjectsPandas(sources, extra_table, extra_columns=extra_columns)
1600 return subchunk
1602 def _check_time_partitions(
1603 self, partitions: Iterable[int], time_partitions_range: ApdbCassandraTimePartitionRange
1604 ) -> None:
1605 """Check that time partitons for new data actually exist.
1607 Parameters
1608 ----------
1609 partitions : `~collections.abc.Iterable` [`int`]
1610 Time partitions for new data.
1611 time_partitions_range : `ApdbCassandraTimePartitionRange`
1612 Currrent time partition range.
1613 """
1614 partitions = set(partitions)
1615 min_part = min(partitions)
1616 max_part = max(partitions)
1617 if min_part < time_partitions_range.start or max_part > time_partitions_range.end:
1618 raise ValueError(
1619 "Attempt to store data for time partitions that do not yet exist. "
1620 f"Partitons for new records: {min_part}-{max_part}. "
1621 f"Database partitons: {time_partitions_range.start}-{time_partitions_range.end}."
1622 )
1623 # Make a noise when writing to the last partition.
1624 if max_part == time_partitions_range.end:
1625 warnings.warn(
1626 "Writing into the last temporal partition. Partition range needs to be extended soon.",
1627 stacklevel=3,
1628 )
1630 def _storeDiaSourcesPartitions(
1631 self,
1632 sources: pandas.DataFrame,
1633 visit_time: astropy.time.Time,
1634 replica_chunk: ReplicaChunk | None,
1635 subchunk: int | None,
1636 ) -> None:
1637 """Store mapping of diaSourceId to its partitioning values.
1639 Parameters
1640 ----------
1641 sources : `pandas.DataFrame`
1642 Catalog containing DiaSource records
1643 visit_time : `astropy.time.Time`
1644 Time of the current visit.
1645 replica_chunk : `ReplicaChunk` or `None`
1646 Replication chunk, or `None` when replication is disabled.
1647 subchunk : `int` or `None`
1648 Replication sub-chunk, or `None` when replication is disabled or
1649 sub-chunking is not used.
1650 """
1651 context = self._context
1653 id_map = cast(pandas.DataFrame, sources[["diaSourceId", "apdb_part"]])
1654 extra_columns = {
1655 "apdb_time_part": context.partitioner.time_partition(visit_time),
1656 "apdb_replica_chunk": replica_chunk.id if replica_chunk is not None else None,
1657 }
1658 if context.has_chunk_sub_partitions:
1659 extra_columns["apdb_replica_subchunk"] = subchunk
1661 self._storeObjectsPandas(
1662 id_map, ExtraTables.DiaSourceToPartition, extra_columns=extra_columns, time_part=None
1663 )
1665 def _storeObjectsPandas(
1666 self,
1667 records: pandas.DataFrame,
1668 table_name: ApdbTables | ExtraTables,
1669 extra_columns: Mapping | None = None,
1670 time_part: int | None = None,
1671 ) -> None:
1672 """Store generic objects.
1674 Takes Pandas catalog and stores a bunch of records in a table.
1676 Parameters
1677 ----------
1678 records : `pandas.DataFrame`
1679 Catalog containing object records
1680 table_name : `ApdbTables`
1681 Name of the table as defined in APDB schema.
1682 extra_columns : `dict`, optional
1683 Mapping (column_name, column_value) which gives fixed values for
1684 columns in each row, overrides values in ``records`` if matching
1685 columns exist there.
1686 time_part : `int`, optional
1687 If not `None` then insert into a per-partition table.
1689 Notes
1690 -----
1691 If Pandas catalog contains additional columns not defined in table
1692 schema they are ignored. Catalog does not have to contain all columns
1693 defined in a table, but partition and clustering keys must be present
1694 in a catalog or ``extra_columns``.
1695 """
1696 context = self._context
1698 # use extra columns if specified
1699 if extra_columns is None:
1700 extra_columns = {}
1701 extra_fields = list(extra_columns.keys())
1703 # Fields that will come from dataframe.
1704 df_fields = [column for column in records.columns if column not in extra_fields]
1706 column_map = context.schema.getColumnMap(table_name)
1707 # list of columns (as in felis schema)
1708 fields = [column_map[field].name for field in df_fields if field in column_map]
1709 fields += extra_fields
1711 # check that all partitioning and clustering columns are defined
1712 partition_columns = context.schema.partitionColumns(table_name)
1713 required_columns = partition_columns + context.schema.clusteringColumns(table_name)
1714 missing_columns = [column for column in required_columns if column not in fields]
1715 if missing_columns:
1716 raise ValueError(f"Primary key columns are missing from catalog: {missing_columns}")
1718 batch_size = self._batch_size(table_name)
1720 with self._timer("insert_build_time", tags={"table": table_name.name}):
1721 # Multi-partition batches are problematic in general, so we want to
1722 # group records in a batch by their partition key.
1723 values_by_key: dict[tuple, list[list]] = defaultdict(list)
1724 for rec in records.itertuples(index=False):
1725 values = []
1726 partitioning_values: dict[str, Any] = {}
1727 for field in df_fields:
1728 if field not in column_map:
1729 continue
1730 value = getattr(rec, field)
1731 if column_map[field].datatype is felis.datamodel.DataType.timestamp:
1732 if isinstance(value, pandas.Timestamp):
1733 value = value.to_pydatetime()
1734 elif value is pandas.NaT:
1735 value = None
1736 else:
1737 # Assume it's seconds since epoch, Cassandra
1738 # datetime is in milliseconds
1739 value = int(value * 1000)
1740 value = literal(value)
1741 values.append(UNSET_VALUE if value is None else value)
1742 if field in partition_columns:
1743 partitioning_values[field] = value
1744 for field in extra_fields:
1745 value = literal(extra_columns[field])
1746 values.append(UNSET_VALUE if value is None else value)
1747 if field in partition_columns:
1748 partitioning_values[field] = value
1750 key = tuple(partitioning_values[field] for field in partition_columns)
1751 values_by_key[key].append(values)
1753 table = context.schema.tableName(table_name, time_part)
1755 query = Insert(self._keyspace, table, fields)
1756 statement = context.stmt_factory(query, prepare=True)
1757 # Cassandra has 64k limit on batch size, normally that should be
1758 # enough but some tests generate too many forced sources.
1759 queries = []
1760 for key_values in values_by_key.values():
1761 for values_chunk in chunk_iterable(key_values, batch_size):
1762 batch = cassandra.query.BatchStatement()
1763 for row_values in values_chunk:
1764 batch.add(statement, row_values)
1765 queries.append((batch, None))
1766 assert batch.routing_key is not None and batch.keyspace is not None
1768 _LOG.debug("%s: will store %d records", context.schema.tableName(table_name), records.shape[0])
1769 with self._timer(
1770 "insert_time", tags={"table": table_name.name, "method": "_storeObjectsPandas"}
1771 ) as timer:
1772 execute_concurrent(context.session, queries, execution_profile="write")
1773 timer.add_values(row_count=len(records), num_batches=len(queries))
1775 def _storeUpdateRecords(
1776 self, records: Iterable[ApdbUpdateRecord], chunk: ReplicaChunk, *, store_chunk: bool = False
1777 ) -> None:
1778 """Store ApdbUpdateRecords in the replica table for those records.
1780 Parameters
1781 ----------
1782 records : `list` [`ApdbUpdateRecord`]
1783 Records to store.
1784 chunk : `ReplicaChunk`
1785 Replica chunk for these records.
1786 store_chunk : `bool`
1787 If True then also store replica chunk.
1789 Raises
1790 ------
1791 TypeError
1792 Raised if replication is not enabled for this instance.
1793 """
1794 context = self._context
1795 config = context.config
1797 if not context.schema.replication_enabled:
1798 raise TypeError("Replication is not enabled for this APDB instance.")
1800 if store_chunk:
1801 self._storeReplicaChunk(chunk)
1803 apdb_replica_chunk = chunk.id
1804 # Do not use unique_if from ReplicaChunk as it could be reused in
1805 # multiple calls to this method.
1806 update_unique_id = uuid.uuid4()
1808 rows = []
1809 for record in records:
1810 rows.append(
1811 [
1812 apdb_replica_chunk,
1813 record.update_time_ns,
1814 record.update_order,
1815 update_unique_id,
1816 record.to_json(),
1817 ]
1818 )
1819 columns = [
1820 "apdb_replica_chunk",
1821 "update_time_ns",
1822 "update_order",
1823 "update_unique_id",
1824 "update_payload",
1825 ]
1826 if context.has_chunk_sub_partitions:
1827 subchunk = random.randrange(config.replica_sub_chunk_count)
1828 for row in rows:
1829 row.append(subchunk)
1830 columns.append("apdb_replica_subchunk")
1832 table_name = context.schema.tableName(ExtraTables.ApdbUpdateRecordChunks)
1833 query = Insert(self._keyspace, table_name, columns)
1834 stmt = context.stmt_factory(query)
1835 queries = [(stmt, row) for row in rows]
1837 with self._timer("store_update_record", tags={"table": table_name}) as timer:
1838 execute_concurrent(context.session, queries, execution_profile="write")
1839 timer.add_values(row_count=len(queries))
1841 def _add_apdb_part(self, df: pandas.DataFrame) -> pandas.DataFrame:
1842 """Calculate spatial partition for each record and add it to a
1843 DataFrame.
1845 Parameters
1846 ----------
1847 df : `pandas.DataFrame`
1848 DataFrame which has to contain ra/dec columns, names of these
1849 columns are defined by configuration ``ra_dec_columns`` field.
1851 Returns
1852 -------
1853 df : `pandas.DataFrame`
1854 DataFrame with ``apdb_part`` column which contains pixel index
1855 for ra/dec coordinates.
1857 Notes
1858 -----
1859 This overrides any existing column in a DataFrame with the same name
1860 (``apdb_part``). Original DataFrame is not changed, copy of a DataFrame
1861 is returned.
1862 """
1863 context = self._context
1864 config = context.config
1866 # Calculate pixelization index for every record.
1867 apdb_part = np.zeros(df.shape[0], dtype=np.int64)
1868 ra_col, dec_col = config.ra_dec_columns
1869 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1870 idx = context.partitioner.pixel(ra, dec)
1871 apdb_part[i] = idx
1872 df = df.copy()
1873 df["apdb_part"] = apdb_part
1874 return df
1876 def _make_empty_catalog(self, table_name: ApdbTables) -> pandas.DataFrame:
1877 """Make an empty catalog for a table with a given name.
1879 Parameters
1880 ----------
1881 table_name : `ApdbTables`
1882 Name of the table.
1884 Returns
1885 -------
1886 catalog : `pandas.DataFrame`
1887 An empty catalog.
1888 """
1889 table = self.schema.tableSchemas[table_name]
1891 data = {columnDef.name: pandas.Series(dtype=columnDef.pandas_type) for columnDef in table.columns}
1892 return pandas.DataFrame(data)
1894 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1895 """Update timestamp columns in input DataFrame to be naive datetime
1896 type.
1898 Clients may or may not generate aware timestamps, code in this class
1899 assumes that timestamps are naive, so we convert them to UTC and
1900 drop timezone.
1901 """
1902 # Find all columns with aware timestamps.
1903 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1904 for column in columns:
1905 # tz_convert(None) will convert to UTC and drop timezone.
1906 df[column] = df[column].dt.tz_convert(None)
1907 return df
1909 def _batch_size(self, table: ApdbTables | ExtraTables) -> int:
1910 """Calculate batch size based on config parameters."""
1911 context = self._context
1912 config = context.config
1914 # Cassandra limit on number of statements in a batch is 64k.
1915 batch_size = 65_535
1916 if 0 < config.batch_statement_limit < batch_size:
1917 batch_size = config.batch_statement_limit
1918 if config.batch_size_limit > 0:
1919 # The purpose of this limit is to try not to exceed batch size
1920 # threshold which is set on server side. Cassandra wire protocol
1921 # for prepared queries (and batches) only sends column values with
1922 # with an additional 4 bytes per value specifying size. Value is
1923 # not included for NULL or NOT_SET values, but the size is always
1924 # there. There is additional small per-query overhead, which we
1925 # ignore.
1926 row_size = context.schema.table_row_size(table)
1927 row_size += 4 * len(context.schema.getColumnMap(table))
1928 batch_size = min(batch_size, (config.batch_size_limit // row_size) + 1)
1929 return batch_size
1931 def _group_dia_objects_by_partition(
1932 self, partitioner: Partitioner, objects: list[DiaObjectId], pad_arcsec: float
1933 ) -> Mapping[int, list[int]]:
1934 """Group DiaObjects by partition.
1936 Parameters
1937 ----------
1938 partitioner : `Partitioner`
1939 Objects which knows how to partition things.
1940 objects : `list` [`DiaObjectId`]
1941 Collection of objects to partition.
1942 pad_arcsec : `float`
1943 Additional padding around object position.
1945 Returns
1946 -------
1947 grouped_objects
1948 Mapping of spatial patition ID to list ob object IDs that it
1949 contains. Some objects may belong to more than one partition.
1950 """
1951 partitioned_object_ids: dict[int, list[int]] = defaultdict(list)
1952 for obj_id in objects:
1953 partitions = partitioner.pixelization.circle_pixels(obj_id.ra, obj_id.dec, pad_arcsec)
1954 for pixel in partitions:
1955 partitioned_object_ids[pixel].append(obj_id.diaObjectId)
1956 return partitioned_object_ids
1958 def _timestamp_column_name(self, column: str) -> str:
1959 """Return column name before/after schema migration to MJD TAI."""
1960 return self._schema.timestamp_column_name(column)
1962 def _timestamp_column_value(self, time: astropy.time.Time) -> float | int:
1963 """Return column value before/after schema migration to MJD TAI."""
1964 if self._schema.has_mjd_timestamps:
1965 return float(time.tai.mjd)
1966 else:
1967 return int(time.datetime.astimezone(tz=datetime.UTC).timestamp() * 1000)
1969 def _get_diasource_data(self, source_ids: Iterable[DiaSourceId], *columns: str) -> list:
1970 """Select records from DiaSource table by diaSourceId and return all
1971 records as a list of named tuples.
1972 """
1973 context = self._context
1974 config = context.config
1975 partitioner = context.partitioner
1977 columns = ("diaSourceId",) + columns
1979 # Allow some uncertainty for coordinates and time when calculating
1980 # partitions.
1981 statements: list[tuple] = []
1982 pad_arcsec = 1.0
1983 pad_time_day = 10 / (24 * 3600)
1984 for source_id in source_ids:
1985 center = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(source_id.ra, source_id.dec))
1986 region = sphgeom.Circle(center, sphgeom.Angle.fromDegrees(pad_arcsec / 3600.0))
1987 spatial_where, _ = partitioner.spatial_where(region)
1989 tables, temporal_where = partitioner.temporal_where(
1990 ApdbTables.DiaSource,
1991 source_id.midpointMjdTai - pad_time_day,
1992 source_id.midpointMjdTai + pad_time_day,
1993 partitons_range=context.time_partitions_range,
1994 query_per_time_part=True,
1995 )
1997 id_where = QExpr('"diaSourceId" = {}', (source_id.diaSourceId,))
1999 for table in tables:
2000 query = Select(self._keyspace, table, columns)
2001 for clause in QExpr.combine(spatial_where, temporal_where, extra=id_where):
2002 statements.append(context.stmt_factory.with_params(query.where(clause), prepare=True))
2004 with self._timer(
2005 "select_time", tags={"table": "DiaSource", "method": "_get_diasource_data"}
2006 ) as timer:
2007 result = cast(
2008 list[tuple],
2009 select_concurrent(
2010 context.session,
2011 statements,
2012 "read_named_tuples",
2013 config.connection_config.read_concurrency,
2014 ),
2015 )
2016 timer.add_values(row_count=len(result), num_queries=len(statements))
2018 return result
2020 def _get_diaforcedsource_data(self, source_ids: Iterable[DiaForcedSourceId], *columns: str) -> list:
2021 """Select records from DiaForcedSource table by (diaObjectId, visit,
2022 detector) and return all records as a list of named tuples.
2023 """
2024 context = self._context
2025 config = context.config
2026 partitioner = context.partitioner
2028 columns = ("diaObjectId", "visit", "detector") + columns
2030 # Allow some uncertainty for coordinates and time when calculating
2031 # partitions.
2032 statements: list[tuple] = []
2033 pad_arcsec = 1.0
2034 pad_time_day = 10 / (24 * 3600)
2035 for source_id in source_ids:
2036 center = sphgeom.UnitVector3d(sphgeom.LonLat.fromDegrees(source_id.ra, source_id.dec))
2037 region = sphgeom.Circle(center, sphgeom.Angle.fromDegrees(pad_arcsec / 3600.0))
2038 spatial_where, _ = partitioner.spatial_where(region)
2040 tables, temporal_where = partitioner.temporal_where(
2041 ApdbTables.DiaForcedSource,
2042 source_id.midpointMjdTai - pad_time_day,
2043 source_id.midpointMjdTai + pad_time_day,
2044 partitons_range=context.time_partitions_range,
2045 query_per_time_part=True,
2046 )
2048 id_where = (
2049 (C("diaObjectId") == source_id.diaObjectId)
2050 & (C("visit") == source_id.visit)
2051 & (C("detector") == source_id.detector)
2052 )
2054 for table in tables:
2055 query = Select(self._keyspace, table, columns)
2056 for clause in QExpr.combine(spatial_where, temporal_where, extra=id_where):
2057 statements.append(context.stmt_factory.with_params(query.where(clause), prepare=True))
2059 with self._timer(
2060 "select_time", tags={"table": "DiaForcedSource", "method": "_get_diaforcedsource_data"}
2061 ) as timer:
2062 result = cast(
2063 list[tuple],
2064 select_concurrent(
2065 context.session,
2066 statements,
2067 "read_named_tuples",
2068 config.connection_config.read_concurrency,
2069 ),
2070 )
2071 timer.add_values(row_count=len(result), num_queries=len(statements))
2073 return result
2075 def _get_diaobject_data(self, object_ids: Iterable[DiaObjectId], *columns: str) -> list:
2076 """Select records from DiaObjectLast table by diaObjectId and return
2077 all records as a list of named tuples.
2078 """
2079 context = self._context
2080 config = context.config
2081 partitioner = context.partitioner
2083 table_name = context.schema.tableName(ApdbTables.DiaObjectLast)
2084 columns = ("diaObjectId",) + columns
2086 # Allow some uncertainty for coordinates when calculating partitions.
2087 pad_arcsec = 1.0
2088 ids_by_partition = defaultdict(list)
2089 for object_id in object_ids:
2090 pixels = partitioner.pixelization.circle_pixels(object_id.ra, object_id.dec, pad_arcsec)
2091 for pixel in pixels:
2092 ids_by_partition[pixel].append(object_id.diaObjectId)
2094 statements: list[tuple] = []
2095 for apdb_part, diaObjectIds in ids_by_partition.items():
2096 query = Select(self._keyspace, table_name, columns)
2097 query = query.where(C("apdb_part") == apdb_part)
2098 query = query.where(C("diaObjectId").in_(diaObjectIds))
2099 statements.append(context.stmt_factory.with_params(query, prepare=False))
2101 with self._timer(
2102 "select_time", tags={"table": "DiaObjectLast", "method": "_get_diaobject_data"}
2103 ) as timer:
2104 result = cast(
2105 list[tuple],
2106 select_concurrent(
2107 context.session,
2108 statements,
2109 "read_named_tuples",
2110 config.connection_config.read_concurrency,
2111 ),
2112 )
2113 timer.add_values(row_count=len(result), num_queries=len(statements))
2115 return result
2117 def _fix_null_columns(self, table: ApdbTables, df: pandas.DataFrame) -> None:
2118 """For non-null columns replace NULL values with column defaults.
2120 When new non-null columns are added to the schema we do not populate
2121 values for existing records. Instead when reading old data back we
2122 fill NULLs with the defaults. The dataflame is updated in-place.
2123 """
2124 # Find all non-null columns which also have default values defined.
2125 table_schema = self._schema.tableSchemas[table]
2126 defaults = {
2127 column.name: column.value
2128 for column in table_schema.columns
2129 if not column.nullable and column.value is not None
2130 }
2131 if defaults:
2132 df.fillna(defaults, inplace=True)