Coverage for bbconf/modules/bedfiles.py: 48%
703 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 22:05 +0000
1import datetime
2import os
3from logging import getLogger
5import numpy as np
6from geniml.bbclient import BBClient
7from geniml.search.backends import QdrantBackend
8from gtars.models import RegionSet as GRegionSet
9from pydantic import BaseModel
10from qdrant_client import models
11from qdrant_client.http.exceptions import UnexpectedResponse
12from qdrant_client.http.models import PointStruct
13from qdrant_client.models import PointIdsList
14from sqlalchemy import and_, cast, delete, func, or_, select
15from sqlalchemy.dialects import postgresql
16from sqlalchemy.exc import IntegrityError
17from sqlalchemy.orm import Session, aliased, selectinload
18from sqlalchemy.orm.attributes import flag_modified
19from tqdm import tqdm
21from bbconf.config_parser.bedbaseconfig import BedBaseConfig
22from bbconf.const import (
23 DEFAULT_LICENSE,
24 DEFAULT_QDRANT_GENOME_DIGESTS,
25 PKG_NAME,
26 ZARR_TOKENIZED_FOLDER,
27)
28from bbconf.db_utils import (
29 Bed,
30 BedMetadata,
31 BedStats,
32 Files,
33 GenomeRefStats,
34 TokenizedBed,
35 Universes,
36)
37from bbconf.exceptions import (
38 BedBaseConfError,
39 BedFIleExistsError,
40 BEDFileNotFoundError,
41 QdrantInstanceNotInitializedError,
42 TokenizeFileExistsError,
43 TokenizeFileNotExistError,
44 UniverseNotFoundError,
45)
46from bbconf.models.bed_models import (
47 BedClassification,
48 BedEmbeddingResult,
49 BedFiles,
50 BedListResult,
51 BedListSearchResult,
52 BedMetadataAll,
53 BedMetadataBasic,
54 BedPlots,
55 BedSetMinimal,
56 BedStatsModel,
57 FileModel,
58 QdrantSearchResult,
59 RefGenValidModel,
60 RefGenValidReturnModel,
61 StandardMeta,
62 TokenizedBedResponse,
63 TokenizedPathResponse,
64 UniverseMetadata,
65 VectorMetadata,
66)
68_LOGGER = getLogger(PKG_NAME)
70# QDRANT_GENOME = "hg38"
73def _bed_id_from_point_id(point_id) -> str:
74 """
75 Convert a Qdrant point id to a plain 32-character BED id.
77 Qdrant stores BED ids as UUIDs and returns them with dashes.
79 Args:
80 point_id: Qdrant point id (dashed UUID string or UUID).
82 Returns:
83 BED id without dashes.
84 """
85 return str(point_id).replace("-", "")
88class BedAgentBedFile:
89 """
90 Class that represents a BED file in the Database.
92 Provides methods to add, delete, get BED files and metadata from the database.
93 """
95 def __init__(self, config: BedBaseConfig, bbagent_obj=None):
96 """
97 Initialize BedAgentBedFile.
99 Args:
100 config: Config object with database and qdrant engine and credentials.
101 bbagent_obj: BedBaseAgent object (Parent object).
102 """
103 self._sa_engine = config.db_engine.engine
104 self._db_engine = config.db_engine
105 self._boto3_client = config.boto3_client
106 self.config = config
107 self.bb_agent = bbagent_obj
109 def get(self, identifier: str, full: bool = False) -> BedMetadataAll:
110 """
111 Get file metadata by identifier.
113 Args:
114 identifier: Bed file identifier.
115 full: If True, return full metadata, including statistics and files.
117 Returns:
118 BED file metadata.
119 """
120 statement = select(Bed).where(and_(Bed.id == identifier))
122 with Session(self._sa_engine) as session:
123 bed_object = session.scalar(statement)
124 if not bed_object:
125 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
127 return self._build_metadata(bed_object, full=full)
129 def _build_metadata(self, bed_object: Bed, full: bool = False) -> BedMetadataAll:
130 """
131 Build a BedMetadataAll model from a Bed ORM object.
133 For ``full=True`` this assembles plots, files, stats, bedsets, and
134 universe metadata (which lazy-load relationships, so the caller must
135 keep the SQLAlchemy session open). For ``full=False`` only scalar
136 columns and the (joined-loaded) annotations are accessed, so the
137 Bed object may be detached from its session.
139 Args:
140 bed_object: Bed ORM object to build metadata from.
141 full: If True, return full metadata, including statistics and files.
143 Returns:
144 BED file metadata.
145 """
146 identifier = bed_object.id
148 bed_plots = BedPlots()
149 bed_files = BedFiles()
151 if full:
152 for result in bed_object.files:
153 # PLOTS
154 if result.name in BedPlots.model_fields:
155 setattr(
156 bed_plots,
157 result.name,
158 FileModel(
159 **result.__dict__,
160 object_id=f"bed.{identifier}.{result.name}",
161 access_methods=self.config.construct_access_method_list(
162 result.path
163 ),
164 ),
165 )
166 # FILES
167 elif result.name in BedFiles.model_fields:
168 (
169 setattr(
170 bed_files,
171 result.name,
172 FileModel(
173 **result.__dict__,
174 object_id=f"bed.{identifier}.{result.name}",
175 access_methods=self.config.construct_access_method_list(
176 result.path
177 ),
178 ),
179 ),
180 )
182 else:
183 _LOGGER.error(
184 f"Unknown file type: {result.name}. And is not in the model fields. Skipping.."
185 )
186 bed_stats = BedStatsModel(**bed_object.stats.__dict__)
187 bed_bedsets = []
188 for relation in bed_object.bedsets:
189 bed_bedsets.append(
190 BedSetMinimal(
191 id=relation.bedset.id,
192 description=relation.bedset.description,
193 name=relation.bedset.name,
194 bedfile_count=relation.bedset.bedfile_count,
195 )
196 )
198 if bed_object.universe:
199 universe_meta = UniverseMetadata(**bed_object.universe.__dict__)
200 else:
201 universe_meta = UniverseMetadata()
202 else:
203 bed_plots = None
204 bed_files = None
205 bed_stats = None
206 universe_meta = None
207 bed_bedsets = []
209 return BedMetadataAll(
210 id=bed_object.id,
211 name=bed_object.name,
212 stats=bed_stats,
213 plots=bed_plots,
214 files=bed_files,
215 description=bed_object.description,
216 submission_date=bed_object.submission_date,
217 last_update_date=bed_object.last_update_date,
218 genome_alias=bed_object.genome_alias,
219 genome_digest=bed_object.genome_digest,
220 bed_compliance=bed_object.bed_compliance,
221 data_format=bed_object.data_format,
222 is_universe=bed_object.is_universe,
223 compliant_columns=bed_object.compliant_columns,
224 non_compliant_columns=bed_object.non_compliant_columns,
225 license_id=bed_object.license_id or DEFAULT_LICENSE,
226 universe_metadata=universe_meta,
227 bedsets=bed_bedsets,
228 processed=bed_object.processed,
229 annotation=StandardMeta(
230 **bed_object.annotations.__dict__ if bed_object.annotations else {}
231 ),
232 )
234 def get_stats(self, identifier: str) -> BedStatsModel:
235 """
236 Get file statistics by identifier.
238 Args:
239 identifier: Bed file identifier.
241 Returns:
242 Project statistics as BedStats object.
243 """
244 statement = select(BedStats).where(and_(BedStats.id == identifier))
246 with Session(self._sa_engine) as session:
247 bed_object = session.scalar(statement)
248 if not bed_object:
249 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
250 bed_stats = BedStatsModel(**bed_object.__dict__)
252 return bed_stats
254 def get_plots(self, identifier: str) -> BedPlots:
255 """
256 Get file plots by identifier.
258 Args:
259 identifier: Bed file identifier.
261 Returns:
262 Project plots.
263 """
264 statement = select(Bed).where(and_(Bed.id == identifier))
266 with Session(self._sa_engine) as session:
267 bed_object = session.scalar(statement)
268 if not bed_object:
269 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
270 bed_plots = BedPlots()
271 for result in bed_object.files:
272 if result.name in BedPlots.model_fields:
273 setattr(
274 bed_plots,
275 result.name,
276 FileModel(
277 **result.__dict__,
278 object_id=f"bed.{identifier}.{result.name}",
279 access_methods=self.config.construct_access_method_list(
280 result.path
281 ),
282 ),
283 )
284 return bed_plots
286 def get_neighbours(
287 self, identifier: str, limit: int = 10, offset: int = 0
288 ) -> BedListSearchResult:
289 """
290 Get nearest neighbours of bed file from qdrant.
292 Args:
293 identifier: Bed file identifier.
294 limit: Number of results to return.
295 offset: Offset to start from.
297 Returns:
298 List of nearest neighbours.
299 """
300 if not self.exists(identifier):
301 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
302 s = identifier
303 try:
304 results = self.config.qdrant_file_backend.qd_client.query_points(
305 collection_name=self.config.config.qdrant.file_collection,
306 query="-".join([s[:8], s[8:12], s[12:16], s[16:20], s[20:]]),
307 limit=limit,
308 offset=offset,
309 )
310 # Hydrate all neighbours with a single batched query instead of one
311 # SELECT per neighbour (was an N+1). annotations is joined-loaded,
312 # but selectinload keeps that explicit for this detached-object path.
313 ids = [_bed_id_from_point_id(result.id) for result in results.points]
314 with Session(self._sa_engine) as session:
315 beds = {
316 bed.id: bed
317 for bed in session.scalars(
318 select(Bed)
319 .where(Bed.id.in_(ids))
320 .options(selectinload(Bed.annotations))
321 ).all()
322 }
323 result_list = [
324 QdrantSearchResult(
325 id=bed_id,
326 payload=result.payload,
327 score=result.score,
328 metadata=self._build_metadata(beds[bed_id], full=False),
329 )
330 for bed_id, result in zip(ids, results.points)
331 # skip stale Qdrant points that no longer exist in the database
332 if bed_id in beds
333 ]
334 except UnexpectedResponse as err:
335 _LOGGER.error(
336 f"Qdrant request failed. Error: {err}. Returning empty result set."
337 )
338 result_list = []
340 return BedListSearchResult(
341 count=self.bb_agent.get_stats().bedfiles_number,
342 limit=limit,
343 offset=offset,
344 results=result_list,
345 )
347 def get_files(self, identifier: str) -> BedFiles:
348 """
349 Get file files by identifier.
351 Args:
352 identifier: Bed file identifier.
354 Returns:
355 Project files.
356 """
357 statement = select(Bed).where(and_(Bed.id == identifier))
359 with Session(self._sa_engine) as session:
360 bed_object = session.scalar(statement)
361 if not bed_object:
362 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
363 bed_files = BedFiles()
364 for result in bed_object.files:
365 if result.name in BedFiles.model_fields:
366 name = result.name
367 setattr(
368 bed_files,
369 name,
370 FileModel(
371 **result.__dict__,
372 object_id=f"bed.{identifier}.{result.name}",
373 access_methods=self.config.construct_access_method_list(
374 result.path
375 ),
376 ),
377 )
378 return bed_files
380 def get_classification(self, identifier: str) -> BedClassification:
381 """
382 Get file classification by identifier.
384 Args:
385 identifier: Bed file identifier.
387 Returns:
388 Project classification.
389 """
390 statement = select(Bed).where(and_(Bed.id == identifier))
392 with Session(self._sa_engine) as session:
393 bed_object = session.scalar(statement)
394 if not bed_object:
395 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
396 bed_classification = BedClassification(**bed_object.__dict__)
398 return bed_classification
400 def get_objects(self, identifier: str) -> dict[str, FileModel]:
401 """
402 Get all object related to bedfile.
404 Args:
405 identifier: Bed file identifier.
407 Returns:
408 Project objects dict.
409 """
410 statement = select(Bed).where(and_(Bed.id == identifier))
411 return_dict = {}
413 with Session(self._sa_engine) as session:
414 bed_object = session.scalar(statement)
415 if not bed_object:
416 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
417 for result in bed_object.files:
418 return_dict[result.name] = FileModel(**result.__dict__)
420 return return_dict
422 def get_embedding(self, identifier: str) -> BedEmbeddingResult:
423 """
424 Get bed file embedding of bed file from qdrant.
426 Args:
427 identifier: Bed file identifier.
429 Returns:
430 Bed file embedding.
431 """
432 if not self.exists(identifier):
433 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
434 result = self.config.qdrant_file_backend.qd_client.retrieve(
435 collection_name=self.config.config.qdrant.file_collection,
436 ids=[identifier],
437 with_vectors=True,
438 with_payload=True,
439 )
440 if not result:
441 raise BEDFileNotFoundError(
442 f"Bed file with id: {identifier} not found in qdrant database."
443 )
444 return BedEmbeddingResult(
445 identifier=identifier, embedding=result[0].vector, payload=result[0].payload
446 )
448 def get_ids_list(
449 self,
450 limit: int = 100,
451 offset: int = 0,
452 genome: str = None,
453 bed_compliance: str = None,
454 after: str | None = None,
455 ) -> BedListResult:
456 """
457 Get list of bed file identifiers.
459 To enumerate everything, page with ``after=<last id>`` (keyset paging).
460 Offset paging costs grow with the offset, since Postgres reads and
461 discards ``offset`` rows per page; keyset paging is constant cost.
463 Args:
464 limit: Number of results to return.
465 offset: Offset to start from. Ignored when ``after`` is set.
466 genome: Filter by genome.
467 bed_compliance: Filter by bed type. e.g. 'bed6+4'.
468 after: Return records with id greater than this. Keyset paging;
469 ignores ``offset``.
471 Returns:
472 List of bed file identifiers. ``count`` is the total number of
473 matching records on the first page (``offset == 0`` and no
474 ``after``), and ``None`` on every other page, since the count query
475 is skipped there.
476 """
477 statement = select(Bed)
478 count_statement = select(func.count(Bed.id))
480 # TODO: make it generic, like in PEPhub
481 if genome:
482 statement = statement.where(and_(Bed.genome_alias == genome))
483 count_statement = count_statement.where(and_(Bed.genome_alias == genome))
485 if bed_compliance:
486 statement = statement.where(and_(Bed.bed_compliance == bed_compliance))
487 count_statement = count_statement.where(
488 and_(Bed.bed_compliance == bed_compliance)
489 )
491 if after is not None:
492 statement = statement.where(Bed.id > after)
493 offset = 0
495 statement = statement.order_by(Bed.id).limit(limit).offset(offset)
497 result_list = []
498 count = None
499 with Session(self._sa_engine) as session:
500 bed_ids = session.scalars(statement)
501 if offset == 0 and after is None:
502 count = session.execute(count_statement).one()[0]
504 for result in bed_ids:
505 annotation = StandardMeta(
506 **result.annotations.__dict__ if result.annotations else {}
507 )
508 result_list.append(
509 BedMetadataBasic(**result.__dict__, annotation=annotation)
510 )
512 return BedListResult(
513 count=count,
514 limit=limit,
515 offset=offset,
516 results=result_list,
517 )
519 def get_reference_validation(self, identifier: str) -> RefGenValidReturnModel:
520 """
521 Get results of reference genome validation for the bed file.
523 Args:
524 identifier: Bed file identifier.
526 Returns:
527 Reference genome validation results.
528 """
530 if not self.exists(identifier):
531 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
533 with Session(self._sa_engine) as session:
534 statement = select(GenomeRefStats).where(
535 GenomeRefStats.bed_id == identifier
536 )
538 results = session.scalars(statement)
540 result_list = []
542 results = [item for item in results]
544 if not results:
545 return RefGenValidReturnModel(
546 id=identifier,
547 provided_genome=None,
548 compared_genome=[],
549 )
551 for result in results:
552 result_list.append(
553 RefGenValidModel(
554 provided_genome=result.provided_genome,
555 compared_genome=(
556 result.genome_object.alias if result.genome_object else None
557 ),
558 genome_digest=result.genome_digest,
559 xs=result.xs,
560 oobr=result.oobr,
561 sequence_fit=result.sequence_fit,
562 assigned_points=result.assigned_points,
563 tier_ranking=result.tier_ranking,
564 )
565 )
567 return RefGenValidReturnModel(
568 id=identifier,
569 provided_genome=result.provided_genome,
570 compared_genome=result_list,
571 )
573 def add(
574 self,
575 identifier: str,
576 stats: dict,
577 metadata: dict = None,
578 plots: dict = None,
579 files: dict = None,
580 classification: dict = None,
581 ref_validation: dict[str, BaseModel] | None = None,
582 license_id: str = DEFAULT_LICENSE,
583 upload_qdrant: bool = False,
584 upload_s3: bool = False,
585 local_path: str = None,
586 overwrite: bool = False,
587 nofail: bool = False,
588 processed: bool = True,
589 ) -> None:
590 """
591 Add bed file to the database.
593 Args:
594 identifier: Bed file identifier.
595 stats: Bed file results {statistics, plots, files, metadata}.
596 metadata: Bed file metadata.
597 plots: Bed file plots.
598 files: Bed file files.
599 classification: Bed file classification.
600 ref_validation: Reference validation data. RefGenValidModel.
601 license_id: Bed file license id (default: 'DUO:0000042'). Full list of licenses:
602 https://raw.githubusercontent.com/EBISPOT/DUO/master/duo.csv
603 upload_qdrant: Add bed file to qdrant indexes.
604 upload_s3: Upload files to s3.
605 local_path: Local path to the output files.
606 overwrite: Overwrite bed file if it already exists.
607 nofail: Do not raise an error for error in s3/qdrant or record exists and not overwrite.
608 processed: True if bedfile was processed and statistics and plots were calculated.
610 Returns:
611 None.
612 """
614 _LOGGER.info(f"Adding bed file to database. bed_id: {identifier}")
616 if self.exists(identifier):
617 _LOGGER.warning(f"Bed file with id: {identifier} exists in the database.")
618 if not overwrite:
619 bed_metadata = StandardMeta(**metadata)
621 ## UPDATE ONLY sources:
622 self._update_sources(
623 identifier=identifier,
624 global_sample_id=bed_metadata.global_sample_id,
625 global_experiment_id=bed_metadata.global_experiment_id,
626 )
628 ## IF we want to update all metadata on add if exists: (DO not remove this code)
629 # with Session(self._sa_engine) as session:
630 # statement = select(Bed).where(Bed.id == identifier)
631 # bed_object = session.scalar(statement)
632 # self._update_metadata(
633 # sa_session=session,
634 # bed_object=bed_object,
635 # bed_metadata=bed_metadata,
636 # )
637 if not nofail:
638 raise BedFIleExistsError(
639 f"Bed file with id: {identifier} already exists in the database."
640 )
641 _LOGGER.warning("Overwrite set to False. Skipping..")
642 return None
643 else:
644 self.delete(identifier)
646 if license_id not in self.bb_agent.list_of_licenses:
647 raise BedBaseConfError(
648 f"License: {license_id} is not in the list of licenses. Please provide a valid license."
649 f"List of licenses: {self.bb_agent.list_of_licenses}"
650 )
652 stats = BedStatsModel(**stats)
653 # TODO: we should not check for specific keys, of the plots!
654 plots = BedPlots(**plots)
655 files = BedFiles(**files)
656 bed_metadata = StandardMeta(**metadata)
658 classification = BedClassification(**classification)
660 if upload_qdrant:
661 if classification.genome_alias == "hg38":
662 _LOGGER.info(f"Uploading bed file to qdrant.. [{identifier}]")
663 self.upload_file_qdrant(
664 identifier,
665 files.bed_file.path,
666 bed_metadata.model_dump(exclude_none=False),
667 )
668 _LOGGER.info(f"File uploaded to qdrant. {identifier}")
669 else:
670 _LOGGER.warning(
671 f"Could not upload to qdrant. Genome: {classification.genome_alias} is not supported."
672 )
673 else:
674 _LOGGER.info("upload_qdrant set to false. Skipping qdrant..")
676 # Upload files to s3
677 if upload_s3:
678 if files:
679 files = self.config.upload_files_s3(
680 identifier, files=files, base_path=local_path, type="files"
681 )
683 if plots:
684 plots = self.config.upload_files_s3(
685 identifier, files=plots, base_path=local_path, type="plots"
686 )
687 with Session(self._sa_engine) as session:
688 new_bed = Bed(
689 id=identifier,
690 **classification.model_dump(),
691 description=bed_metadata.description,
692 license_id=license_id,
693 indexed=upload_qdrant,
694 processed=processed,
695 )
696 session.add(new_bed)
697 if upload_s3:
698 for k, v in files:
699 if v:
700 new_file = Files(
701 **v.model_dump(
702 exclude_none=True,
703 exclude_unset=True,
704 exclude={"object_id", "access_methods"},
705 ),
706 bedfile_id=identifier,
707 type="file",
708 )
709 session.add(new_file)
710 for k, v in plots:
711 if v:
712 new_plot = Files(
713 **v.model_dump(
714 exclude_none=True,
715 exclude_unset=True,
716 exclude={"object_id", "access_methods"},
717 ),
718 bedfile_id=identifier,
719 type="plot",
720 )
721 session.add(new_plot)
723 new_bedstat = BedStats(**stats.model_dump(), id=identifier)
724 new_metadata = BedMetadata(
725 **bed_metadata.model_dump(exclude={"description"}), id=identifier
726 )
728 session.add(new_bedstat)
729 session.add(new_metadata)
731 if ref_validation:
732 new_gen_refs = self._create_ref_validation_models(
733 ref_validation=ref_validation,
734 bed_id=identifier,
735 provided_genome=classification.genome_alias,
736 )
737 session.add_all(new_gen_refs)
739 # for ref_gen_check, data in ref_validation.items():
740 # new_gen_ref = GenomeRefStats(
741 # **RefGenValidModel(
742 # **data.model_dump(),
743 # provided_genome=classification.genome_alias,
744 # compared_genome=ref_gen_check,
745 # genome_digest=ref_gen_check,
746 # ).model_dump(),
747 # bed_id=identifier,
748 # )
749 # session.add(new_gen_ref)
750 session.commit()
752 return None
754 def update(
755 self,
756 identifier: str,
757 stats: dict | None = None,
758 metadata: dict | None = None,
759 plots: dict | None = None,
760 files: dict | None = None,
761 classification: dict | None = None,
762 ref_validation: dict[str, BaseModel] | None = None,
763 license_id: str = DEFAULT_LICENSE,
764 upload_qdrant: bool = False,
765 upload_s3: bool = True,
766 local_path: str = None,
767 overwrite: bool = False,
768 nofail: bool = False,
769 processed: bool = False,
770 ) -> None:
771 """
772 Update bed file to the database.
774 Args:
775 identifier: Bed file identifier.
776 stats: Bed file results {statistics, plots, files, metadata}.
777 metadata: Bed file metadata.
778 plots: Bed file plots.
779 files: Bed file files.
780 classification: Bed file classification.
781 ref_validation: Reference validation data. RefGenValidModel.
782 license_id: Bed file license id (default: 'DUO:0000042').
783 upload_qdrant: Add bed file to qdrant indexes.
784 upload_s3: Upload files to s3.
785 local_path: Local path to the output files.
786 overwrite: Overwrite bed file if it already exists.
787 nofail: Do not raise an error for error in s3/qdrant or record exists and not overwrite.
788 processed: True if bedfile was processed and statistics and plots were calculated.
790 Returns:
791 None.
792 """
793 if not self.exists(identifier):
794 raise BEDFileNotFoundError(
795 f"Bed file with id: {identifier} not found. Cannot update."
796 )
797 _LOGGER.info(f"Updating bed file: '{identifier}'")
799 if license_id not in self.bb_agent.list_of_licenses and not license_id:
800 raise BedBaseConfError(
801 f"License: {license_id} is not in the list of licenses. Please provide a valid license."
802 f"List of licenses: {self.bb_agent.list_of_licenses}"
803 )
805 stats = BedStatsModel(**stats if stats else {})
806 plots = BedPlots(**plots if plots else {})
807 files = BedFiles(**files if files else {})
808 bed_metadata = StandardMeta(**metadata if metadata else {})
809 classification = BedClassification(**classification if classification else {})
811 if upload_qdrant:
812 if classification.genome_alias == "hg38":
813 _LOGGER.info(f"Uploading bed file to qdrant.. [{identifier}]")
814 self.upload_file_qdrant(
815 identifier,
816 files.bed_file.path,
817 bed_metadata.model_dump(exclude_none=False),
818 )
819 _LOGGER.info(f"File uploaded to qdrant. {identifier}")
820 else:
821 _LOGGER.warning(
822 f"Could not upload to qdrant. Genome: {classification.genome_alias} is not supported."
823 )
825 with Session(self._sa_engine) as session:
826 bed_statement = select(Bed).where(and_(Bed.id == identifier))
827 bed_object = session.scalar(bed_statement)
829 self._update_classification(
830 sa_session=session, bed_object=bed_object, classification=classification
831 )
833 self._update_metadata(
834 sa_session=session,
835 bed_object=bed_object,
836 bed_metadata=bed_metadata,
837 )
838 self._update_stats(sa_session=session, bed_object=bed_object, stats=stats)
840 if upload_s3:
841 self._update_plots(
842 sa_session=session,
843 bed_object=bed_object,
844 plots=plots,
845 local_path=local_path,
846 )
847 self._update_files(
848 sa_session=session,
849 bed_object=bed_object,
850 files=files,
851 local_path=local_path,
852 )
854 self._update_ref_validation(
855 sa_session=session,
856 bed_id=identifier,
857 ref_validation=ref_validation,
858 provided_genome=bed_object.genome_alias or "",
859 )
861 bed_object.processed = processed
862 bed_object.indexed = upload_qdrant
863 bed_object.last_update_date = datetime.datetime.now(datetime.timezone.utc)
864 if bed_metadata.description:
865 bed_object.description = bed_metadata.description
867 session.commit()
869 return None
871 @staticmethod
872 def _update_classification(
873 sa_session: Session, bed_object: Bed, classification: BedClassification
874 ) -> None:
875 """
876 Update bed file classification.
878 Args:
879 sa_session: Sqlalchemy session.
880 bed_object: Bed sqlalchemy object.
881 classification: Bed file classification as BedClassification object.
883 Returns:
884 None.
885 """
886 classification_dict = classification.model_dump(
887 exclude_defaults=True, exclude_none=True, exclude_unset=True
888 )
889 for k, v in classification_dict.items():
890 setattr(bed_object, k, v)
892 sa_session.commit()
894 @staticmethod
895 def _update_stats(
896 sa_session: Session, bed_object: Bed, stats: BedStatsModel
897 ) -> None:
898 """
899 Update bed file statistics.
901 Args:
902 sa_session: Sqlalchemy session.
903 bed_object: Bed sqlalchemy object.
904 stats: Bed file statistics as BedStatsModel object.
906 Returns:
907 None.
908 """
910 stats_dict = stats.model_dump(
911 exclude_defaults=True, exclude_none=True, exclude_unset=True
912 )
913 if not bed_object.stats:
914 new_bedstat = BedStats(**stats.model_dump(), id=bed_object.id)
915 sa_session.add(new_bedstat)
916 else:
917 for k, v in stats_dict.items():
918 setattr(bed_object.stats, k, v)
920 sa_session.commit()
922 def _update_metadata(
923 self, sa_session: Session, bed_object: Bed, bed_metadata: StandardMeta
924 ) -> None:
925 """
926 Update bed file metadata.
928 Args:
929 sa_session: Sqlalchemy session.
930 bed_object: Bed sqlalchemy object.
931 bed_metadata: Bed file metadata as StandardMeta object.
933 Returns:
934 None.
935 """
937 self._update_sources(
938 identifier=bed_object.id,
939 global_sample_id=bed_metadata.global_sample_id,
940 global_experiment_id=bed_metadata.global_experiment_id,
941 )
943 bed_metadata.global_experiment_id = None
944 bed_metadata.global_sample_id = None
946 metadata_dict = bed_metadata.model_dump(
947 exclude_defaults=True, exclude_none=True, exclude_unset=True
948 )
949 if not bed_object.annotations:
950 new_metadata = BedMetadata(
951 **bed_metadata.model_dump(exclude={"description"}), id=bed_object.id
952 )
953 sa_session.add(new_metadata)
954 else:
955 for k, v in metadata_dict.items():
956 setattr(bed_object.annotations, k, v)
958 sa_session.commit()
960 def _update_plots(
961 self,
962 sa_session: Session,
963 bed_object: Bed,
964 plots: BedPlots,
965 local_path: str = None,
966 ) -> None:
967 """
968 Update bed file plots.
970 Args:
971 sa_session: Sqlalchemy session.
972 bed_object: Bed sqlalchemy object.
973 plots: Bed file plots.
974 local_path: Local path to the output files.
975 """
977 _LOGGER.info("Updating bed file plots..")
978 if plots:
979 plots = self.config.upload_files_s3(
980 bed_object.id, files=plots, base_path=local_path, type="plots"
981 )
982 plots_dict = plots.model_dump(
983 exclude_defaults=True, exclude_none=True, exclude_unset=True
984 )
985 if not plots_dict:
986 return None
988 for k, v in plots:
989 if v:
990 new_plot = Files(
991 **v.model_dump(
992 exclude_none=True,
993 exclude_unset=True,
994 exclude={"object_id", "access_methods"},
995 ),
996 bedfile_id=bed_object.id,
997 type="plot",
998 )
999 try:
1000 sa_session.add(new_plot)
1001 sa_session.commit()
1002 except IntegrityError as _:
1003 sa_session.rollback()
1004 _LOGGER.debug(
1005 f"Plot with name: {v.name} already exists. Updating.."
1006 )
1008 return None
1010 def _update_files(
1011 self,
1012 sa_session: Session,
1013 bed_object: Bed,
1014 files: BedFiles,
1015 local_path: str = None,
1016 ) -> None:
1017 """
1018 Update bed files.
1020 Args:
1021 sa_session: Sqlalchemy session.
1022 bed_object: Bed sqlalchemy object.
1023 files: Bed file files.
1024 """
1026 _LOGGER.info("Updating bed files..")
1027 if files:
1028 files = self.config.upload_files_s3(
1029 bed_object.id, files=files, base_path=local_path, type="files"
1030 )
1032 files_dict = files.model_dump(
1033 exclude_defaults=True, exclude_none=True, exclude_unset=True
1034 )
1035 if not files_dict:
1036 return None
1038 for k, v in files:
1039 if v:
1040 new_file = Files(
1041 **v.model_dump(
1042 exclude_none=True,
1043 exclude_unset=True,
1044 exclude={"object_id", "access_methods"},
1045 ),
1046 bedfile_id=bed_object.id,
1047 type="file",
1048 )
1050 try:
1051 sa_session.add(new_file)
1052 sa_session.commit()
1053 except IntegrityError as _:
1054 sa_session.rollback()
1055 _LOGGER.debug(
1056 f"File with name: {v.name} already exists. Updating.."
1057 )
1059 def _update_ref_validation(
1060 self,
1061 sa_session: Session,
1062 bed_id: str,
1063 ref_validation: dict[str, BaseModel],
1064 provided_genome: str = "",
1065 ) -> None:
1066 """
1067 Update reference validation data.
1069 This function won't update the reference validation data, if it exists, it will skip it.
1071 Args:
1072 sa_session: Sqlalchemy session.
1073 bed_id: Bed sqlalchemy object.
1074 ref_validation: Bed file metadata.
1075 provided_genome: Genome reference that was provided by user.
1076 """
1078 if not ref_validation:
1079 return None
1081 sa_session.execute(
1082 delete(GenomeRefStats).where(GenomeRefStats.bed_id == bed_id)
1083 )
1085 new_gen_refs = self._create_ref_validation_models(
1086 ref_validation=ref_validation,
1087 bed_id=bed_id,
1088 provided_genome=provided_genome,
1089 )
1090 sa_session.add_all(new_gen_refs)
1092 # One commit for both deletes and inserts
1093 sa_session.commit()
1095 return None
1097 def _create_ref_validation_models(
1098 self,
1099 ref_validation: dict[str, BaseModel],
1100 bed_id: str,
1101 provided_genome: str = None,
1102 ) -> list[GenomeRefStats]:
1104 compatibility = {}
1106 for k, v in ref_validation.items():
1107 if v.tier_ranking < 4:
1108 compatibility[k] = v
1110 # Add all new GenomeRefStats objects in one go
1111 new_gen_refs: list[GenomeRefStats] = [
1112 GenomeRefStats(
1113 **RefGenValidModel(
1114 **data.model_dump(),
1115 provided_genome=provided_genome,
1116 compared_genome=ref_gen_check,
1117 genome_digest=ref_gen_check,
1118 ).model_dump(),
1119 bed_id=bed_id,
1120 )
1121 for ref_gen_check, data in compatibility.items()
1122 ]
1123 return new_gen_refs
1125 def delete(self, identifier: str) -> None:
1126 """
1127 Delete bed file from the database.
1129 Args:
1130 identifier: Bed file identifier.
1132 Returns:
1133 None.
1134 """
1135 _LOGGER.info(f"Deleting bed file from database. bed_id: {identifier}")
1136 if not self.exists(identifier):
1137 raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.")
1139 with Session(self._sa_engine) as session:
1140 statement = select(Bed).where(and_(Bed.id == identifier))
1141 bed_object = session.scalar(statement)
1143 files = [FileModel(**k.__dict__) for k in bed_object.files]
1144 delete_qdrant = bed_object.indexed
1146 session.delete(bed_object)
1147 session.commit()
1149 if delete_qdrant:
1150 self.delete_qdrant_point(identifier)
1151 self.config.delete_files_s3(files)
1153 def upload_file_qdrant(
1154 self,
1155 bed_id: str,
1156 bed_file: str | GRegionSet,
1157 payload: dict = None,
1158 ) -> None:
1159 """
1160 Convert bed file to vector and add it to qdrant database.
1162 Warning: only hg38 genome can be added to qdrant!
1164 Args:
1165 bed_id: Bed file id.
1166 bed_file: Path to the bed file, or RegionSet object.
1167 payload: Additional metadata to store alongside vectors.
1169 Returns:
1170 None.
1171 """
1173 _LOGGER.debug(f"Adding bed file to qdrant. bed_id: {bed_id}")
1175 if not isinstance(self.config.qdrant_file_backend, QdrantBackend):
1176 raise QdrantInstanceNotInitializedError("Could not upload file.")
1178 bed_embedding = self._embed_file(bed_file)
1180 self.config.qdrant_file_backend.load(
1181 ids=[bed_id],
1182 vectors=bed_embedding,
1183 payloads=[{**payload}],
1184 )
1185 return None
1187 def _embed_file(self, bed_file: str | GRegionSet) -> np.ndarray:
1188 """
1189 Create embedding for bed file.
1191 Args:
1192 bed_file: Path to the bed file, or RegionSet object.
1194 Returns:
1195 Numpy array of embeddings.
1196 """
1197 if self.config.qdrant_file_backend is None:
1198 raise QdrantInstanceNotInitializedError
1199 if not self.config.r2v_encoder:
1200 raise BedBaseConfError(
1201 "Could not add region to qdrant. Invalid type, or path. "
1202 )
1204 if isinstance(bed_file, str):
1205 # Use try if file is corrupted. In Python RegionSet we have functionality to tackle this problem
1206 try:
1207 bed_region_set = GRegionSet(bed_file)
1208 except RuntimeError as _:
1209 bed_region_set = GRegionSet(bed_file)
1210 elif isinstance(bed_file, GRegionSet) or isinstance(bed_file, GRegionSet):
1211 bed_region_set = bed_file
1212 else:
1213 raise BedBaseConfError(
1214 "Could not add region to qdrant. Invalid type, or path. "
1215 )
1216 bed_embedding = np.mean(self.config.r2v_encoder.encode(bed_region_set), axis=0)
1217 vec_dim = bed_embedding.shape[0]
1218 return bed_embedding.reshape(1, vec_dim)
1220 def _get_umap_file(self, bed_file: str | GRegionSet) -> np.ndarray:
1221 """
1222 Create UMAP for bed file.
1224 Args:
1225 bed_file: Bed file path or region set.
1226 """
1228 if self.config.umap_encoder is None:
1229 raise BedBaseConfError("UMAP model is not initialized.")
1231 bed_embedding = self._embed_file(bed_file)
1232 bed_umap = self.config.umap_encoder.transform(bed_embedding)
1233 return bed_umap
1235 def text_to_bed_search(
1236 self,
1237 query: str,
1238 limit: int = 10,
1239 offset: int = 0,
1240 with_metadata: bool = True,
1241 ) -> BedListSearchResult:
1242 """
1243 Search for bed files by text query in qdrant database.
1244 This is bivec_search.
1246 Args:
1247 query: Text query.
1248 limit: Number of results to return.
1249 offset: Offset to start from.
1250 with_metadata: If True, will return metadata for each result.
1252 Returns:
1253 List of bed file metadata.
1254 """
1255 _LOGGER.info(f"Looking for: {query}")
1257 results = self.config.bivec_search_interface.query_search(
1258 query, limit=limit, offset=offset
1259 )
1260 results_list = []
1261 for result in results:
1262 result_id = _bed_id_from_point_id(result["id"])
1263 try:
1264 if with_metadata:
1265 result_meta = self.get(result_id)
1266 else:
1267 result_meta = None
1268 except BEDFileNotFoundError as e:
1269 _LOGGER.warning(
1270 f"Could not retrieve metadata for bed file: {result_id}. Error: {e}"
1271 )
1272 continue
1273 results_list.append(
1274 QdrantSearchResult(
1275 **{**result, "id": result_id},
1276 metadata=result_meta if with_metadata else None,
1277 )
1278 )
1280 if with_metadata:
1281 count = self.config.qdrant_client.get_collection(
1282 collection_name=self.config.config.qdrant.file_collection
1283 ).points_count
1284 else:
1285 count = 0
1286 return BedListSearchResult(
1287 count=count,
1288 limit=limit,
1289 offset=offset,
1290 results=results_list,
1291 )
1293 def bed_to_bed_search(
1294 self,
1295 region_set: GRegionSet,
1296 limit: int = 10,
1297 offset: int = 0,
1298 ) -> BedListSearchResult:
1299 """
1300 Search for bed files by using region set in qdrant database.
1302 Args:
1303 region_set: RegionSet object to search for (bed file).
1304 limit: Number of results to return.
1305 offset: Offset to start from.
1307 Returns:
1308 BedListSetResults.
1309 """
1310 results = self.config.b2b_search_interface.query_search(
1311 region_set, limit=limit, offset=offset
1312 )
1313 results_list = []
1314 for result in results:
1315 result_id = _bed_id_from_point_id(result["id"])
1316 try:
1317 result_meta = self.get(result_id)
1318 except BEDFileNotFoundError as e:
1319 _LOGGER.warning(
1320 f"Could not retrieve metadata for bed file: {result_id}. Error: {e}"
1321 )
1322 continue
1323 if result_meta:
1324 results_list.append(
1325 QdrantSearchResult(
1326 **{**result, "id": result_id}, metadata=result_meta
1327 )
1328 )
1330 # Count of the searchable pool (indexed bed vectors), not the total number
1331 # of bed files in the database (which overcounts unindexed genomes).
1332 count = self.config.qdrant_client.count(
1333 collection_name=self.config.config.qdrant.file_collection,
1334 exact=True,
1335 ).count
1336 return BedListSearchResult(
1337 count=count,
1338 limit=limit,
1339 offset=offset,
1340 results=results_list,
1341 )
1343 def sql_search(
1344 self,
1345 query: str,
1346 genome: str = None,
1347 assay: str = None,
1348 limit: int = 10,
1349 offset: int = 0,
1350 ) -> BedListSearchResult:
1351 """
1352 Search for bed files by using sql exact search.
1353 This search will search files by id, name, and description.
1355 Args:
1356 query: Text query.
1357 genome: Genome alias to filter results.
1358 assay: Filter by assay type.
1359 limit: Number of results to return.
1360 offset: Offset to start from.
1362 Returns:
1363 List of bed file metadata.
1364 """
1366 _LOGGER.debug(f"Looking for: {query}")
1368 statement = select(Bed).join(BedMetadata)
1370 sql_search_str = f"%{query}%"
1371 or_statement = or_(
1372 Bed.id.ilike(sql_search_str),
1373 Bed.name.ilike(sql_search_str),
1374 Bed.description.ilike(sql_search_str),
1375 BedMetadata.cell_line.ilike(sql_search_str),
1376 BedMetadata.tissue.ilike(sql_search_str),
1377 BedMetadata.cell_type.ilike(sql_search_str),
1378 )
1380 condition_statement = or_statement
1382 if genome_alias := genome:
1383 _LOGGER.debug(f"Filtering by genome: {genome_alias}")
1385 condition_statement = and_(
1386 Bed.genome_alias == genome_alias, condition_statement
1387 )
1389 if assay:
1390 _LOGGER.debug(f"Filtering by assay: {assay}")
1392 condition_statement = and_(BedMetadata.assay == assay, condition_statement)
1394 statement = statement.where(condition_statement)
1396 with Session(self._sa_engine) as session:
1397 statement = statement.limit(limit).offset(offset)
1398 bed_objects = session.scalars(statement)
1399 results = [
1400 BedMetadataBasic(
1401 **bedfile_obj.__dict__,
1402 annotation=StandardMeta(
1403 **(
1404 bedfile_obj.annotations.__dict__
1405 if bedfile_obj.annotations
1406 else {}
1407 )
1408 ),
1409 )
1410 for bedfile_obj in bed_objects
1411 ]
1412 result_list = [
1413 QdrantSearchResult(id=result.id, score=1, metadata=result)
1414 for result in results
1415 ]
1417 return BedListSearchResult(
1418 count=self._sql_search_count(condition_statement),
1419 limit=limit,
1420 offset=offset,
1421 results=result_list,
1422 )
1424 def _sql_search_count(self, condition_statement) -> int:
1425 """
1426 Get number of total found files in the database.
1428 Args:
1429 condition_statement: Sql alchemy condition statement to filter results.
1431 Returns:
1432 Number of found files.
1433 """
1435 with Session(self._sa_engine) as session:
1436 statement = (
1437 select(func.count())
1438 .select_from(Bed)
1439 .join(BedMetadata)
1440 .where(condition_statement)
1441 )
1442 count = session.execute(statement).one()
1443 return count[0]
1445 def reindex_qdrant(self, batch: int = 100, purge: bool = False) -> None:
1446 """
1447 Re-upload all files to qdrant.
1448 Warning: only hg38 genome can be added to qdrant!
1450 If you want to fully reindex/reupload to qdrant, first delete collection and create new one.
1451 Upload all files to qdrant.
1453 Args:
1454 batch: Number of files to upload in one batch.
1455 """
1456 bb_client = BBClient()
1458 with Session(self._sa_engine) as session:
1459 if purge:
1460 _LOGGER.info("Purging indexed files in the database ...")
1461 session.query(Bed).update({Bed.file_indexed: False})
1462 session.commit()
1463 _LOGGER.info("Purged indexed files in the database successfully!")
1465 statement = (
1466 select(Bed)
1467 .join(BedMetadata, Bed.id == BedMetadata.id)
1468 .where(
1469 and_(
1470 Bed.file_indexed.is_(False),
1471 # Bed.genome_alias == QDRANT_GENOME,
1472 # BedMetadata.global_experiment_id.contains(['encode']) # If we want only encode data
1473 Bed.genome_digest.in_(DEFAULT_QDRANT_GENOME_DIGESTS),
1474 )
1475 )
1476 .limit(150000)
1477 )
1479 annotation_results = session.scalars(statement)
1481 results: list[Bed] = [result for result in annotation_results]
1482 if not results:
1483 _LOGGER.info("No files to reindex in qdrant.")
1484 return None
1486 with tqdm(total=len(results), position=0, leave=True) as pbar:
1487 points_list = []
1488 processed_number = 0
1489 for record in results:
1490 try:
1491 bed_region_set_obj = GRegionSet(bb_client.seek(record.id))
1492 except FileNotFoundError:
1493 bed_region_set_obj = bb_client.load_bed(record.id)
1494 except RuntimeError:
1495 _LOGGER.warning(f"Unable to open file: {record.id}")
1496 continue
1497 except Exception:
1498 _LOGGER.warning(
1499 f"File not found or corrupted. Skipping.. {record.id}"
1500 )
1501 continue
1503 pbar.set_description(f"Processing file: {record.id}")
1505 file_embedding = self._embed_file(bed_region_set_obj)
1507 bed_metadata = VectorMetadata(
1508 id=record.id,
1509 name=record.name,
1510 description=record.description,
1511 genome_alias=record.genome_alias,
1512 genome_digest=record.genome_digest,
1513 cell_line=record.annotations.cell_line,
1514 cell_type=record.annotations.cell_type,
1515 tissue=record.annotations.tissue,
1516 target=record.annotations.target,
1517 treatment=record.annotations.treatment,
1518 assay=record.annotations.assay,
1519 species_name=record.annotations.species_name,
1520 )
1522 record.file_indexed = True
1524 points_list.append(
1525 PointStruct(
1526 id=record.id,
1527 vector=file_embedding.tolist()[0],
1528 payload=(bed_metadata.model_dump()),
1529 )
1530 )
1531 processed_number += 1
1533 if processed_number % batch == 0:
1534 pbar.set_description(
1535 "Uploading points to qdrant using batch..."
1536 )
1537 operation_info = self.config.qdrant_file_backend.qd_client.upsert(
1538 collection_name=self.config.config.qdrant.file_collection,
1539 points=points_list,
1540 )
1541 pbar.write("Uploaded batch to qdrant.")
1542 points_list = []
1543 assert operation_info.status == "completed"
1545 session.commit()
1547 pbar.write(f"File: {record.id} successfully indexed.")
1548 pbar.update(1)
1550 _LOGGER.info("Uploading points to qdrant using batches...")
1551 operation_info = self.config.qdrant_file_backend.qd_client.upsert(
1552 collection_name=self.config.config.qdrant.file_collection,
1553 points=points_list,
1554 )
1555 assert operation_info.status == "completed"
1556 return None
1558 def delete_qdrant_point(self, identifier: str) -> None:
1559 """
1560 Delete bed file from qdrant.
1562 Args:
1563 identifier: Bed file identifier.
1565 Returns:
1566 None.
1567 """
1569 result = self.config.qdrant_file_backend.qd_client.delete(
1570 collection_name=self.config.config.qdrant.file_collection,
1571 points_selector=PointIdsList(
1572 points=[identifier],
1573 ),
1574 )
1575 if result.status == "completed":
1576 _LOGGER.info(
1577 f"File with id: {identifier} successfully deleted from qdrant."
1578 )
1579 return None
1581 def exists(self, identifier: str) -> bool:
1582 """
1583 Check if bed file exists in the database.
1585 Args:
1586 identifier: Bed file identifier.
1588 Returns:
1589 True if bed file exists, False otherwise.
1590 """
1591 statement = select(Bed).where(and_(Bed.id == identifier))
1593 with Session(self._sa_engine) as session:
1594 bed_object = session.scalar(statement)
1595 if not bed_object:
1596 return False
1597 return True
1599 def exists_universe(self, identifier: str) -> bool:
1600 """
1601 Check if universe exists in the database.
1603 Args:
1604 identifier: Universe identifier.
1606 Returns:
1607 True if universe exists, False otherwise.
1608 """
1609 statement = select(Universes).where(and_(Universes.id == identifier))
1611 with Session(self._sa_engine) as session:
1612 bed_object = session.scalar(statement)
1613 if not bed_object:
1614 return False
1615 return True
1617 def add_universe(
1618 self, bedfile_id: str, bedset_id: str = None, construct_method: str = None
1619 ) -> str:
1620 """
1621 Add universe to the database.
1623 Args:
1624 bedfile_id: Bed file identifier.
1625 bedset_id: Bedset identifier.
1626 construct_method: Method used to construct the universe.
1628 Returns:
1629 Universe identifier.
1630 """
1632 if not self.exists(bedfile_id):
1633 raise BEDFileNotFoundError
1634 with Session(self._sa_engine) as session:
1635 new_univ = Universes(
1636 id=bedfile_id, bedset_id=bedset_id, method=construct_method
1637 )
1638 session.add(new_univ)
1639 session.commit()
1641 _LOGGER.info(f"Universe added to the database successfully. id: {bedfile_id}")
1642 return bedfile_id
1644 def delete_universe(self, identifier: str) -> None:
1645 """
1646 Delete universe from the database.
1648 Args:
1649 identifier: Universe identifier.
1651 Returns:
1652 None.
1653 """
1654 if not self.exists_universe(identifier):
1655 raise UniverseNotFoundError(f"Universe not found. id: {identifier}")
1657 with Session(self._sa_engine) as session:
1658 statement = delete(Universes).where(and_(Universes.id == identifier))
1659 session.execute(statement)
1660 session.commit()
1662 def add_tokenized(
1663 self, bed_id: str, universe_id: str, token_vector: list, overwrite: bool = False
1664 ) -> str:
1665 """
1666 Add tokenized bed file to the database.
1668 Args:
1669 bed_id: Bed file identifier.
1670 universe_id: Universe identifier.
1671 token_vector: List of tokens.
1672 overwrite: Overwrite tokenized file if it already exists.
1674 Returns:
1675 Token path.
1676 """
1678 with Session(self._sa_engine) as session:
1679 if not self.exists_universe(universe_id):
1680 raise UniverseNotFoundError(
1681 f"Universe not found in the database. id: {universe_id}"
1682 f"Please add universe first."
1683 )
1685 if self.exist_tokenized(bed_id, universe_id):
1686 if not overwrite:
1687 raise TokenizeFileExistsError(
1688 "Tokenized file already exists in the database. "
1689 "Set overwrite to True to overwrite it."
1690 )
1691 else:
1692 self.delete_tokenized(bed_id, universe_id)
1694 path = self._add_zarr_s3(
1695 bed_id=bed_id,
1696 universe_id=universe_id,
1697 tokenized_vector=token_vector,
1698 overwrite=overwrite,
1699 )
1700 path = os.path.join(f"s3://{self.config.config.s3.bucket}", path)
1701 new_token = TokenizedBed(bed_id=bed_id, universe_id=universe_id, path=path)
1703 session.add(new_token)
1704 session.commit()
1705 return path
1707 def _add_zarr_s3(
1708 self,
1709 universe_id: str,
1710 bed_id: str,
1711 tokenized_vector: list,
1712 overwrite: bool = False,
1713 ) -> str:
1714 """
1715 Add zarr file to the database.
1717 Args:
1718 universe_id: Universe identifier.
1719 bed_id: Bed file identifier.
1720 tokenized_vector: Tokenized vector.
1722 Returns:
1723 Zarr path.
1724 """
1725 univers_group = self.config.zarr_root.require_group(universe_id)
1727 if not univers_group.get(bed_id):
1728 _LOGGER.info("Saving tokenized vector to s3")
1729 path = univers_group.create_dataset(bed_id, data=tokenized_vector).path
1730 elif overwrite:
1731 _LOGGER.info("Overwriting tokenized vector in s3")
1732 path = univers_group.create_dataset(
1733 bed_id, data=tokenized_vector, overwrite=True
1734 ).path
1735 else:
1736 raise TokenizeFileExistsError(
1737 "Tokenized file already exists in the database. "
1738 "Set overwrite to True to overwrite it."
1739 )
1741 return str(os.path.join(ZARR_TOKENIZED_FOLDER, path))
1743 def get_tokenized(self, bed_id: str, universe_id: str) -> TokenizedBedResponse:
1744 """
1745 Get zarr file from the database.
1747 Args:
1748 bed_id: Bed file identifier.
1749 universe_id: Universe identifier.
1751 Returns:
1752 Zarr path.
1753 """
1755 if not self.exist_tokenized(bed_id, universe_id):
1756 raise TokenizeFileNotExistError("Tokenized file not found in the database.")
1757 univers_group = self.config.zarr_root.require_group(universe_id)
1759 return TokenizedBedResponse(
1760 universe_id=universe_id,
1761 bed_id=bed_id,
1762 tokenized_bed=list(univers_group[bed_id]),
1763 )
1765 def delete_tokenized(self, bed_id: str, universe_id: str) -> None:
1766 """
1767 Delete tokenized bed file from the database.
1769 Args:
1770 bed_id: Bed file identifier.
1771 universe_id: Universe identifier.
1773 Returns:
1774 None.
1775 """
1776 if not self.exist_tokenized(bed_id, universe_id):
1777 raise TokenizeFileNotExistError("Tokenized file not found in the database.")
1778 univers_group = self.config.zarr_root.require_group(universe_id)
1780 del univers_group[bed_id]
1782 with Session(self._sa_engine) as session:
1783 statement = delete(TokenizedBed).where(
1784 and_(
1785 TokenizedBed.bed_id == bed_id,
1786 TokenizedBed.universe_id == universe_id,
1787 )
1788 )
1789 session.execute(statement)
1790 session.commit()
1792 return None
1794 def _get_tokenized_path(self, bed_id: str, universe_id: str) -> str:
1795 """
1796 Get tokenized path to tokenized file.
1798 Args:
1799 bed_id: Bed file identifier.
1800 universe_id: Universe identifier.
1802 Returns:
1803 Token path.
1804 """
1805 if not self.exist_tokenized(bed_id, universe_id):
1806 raise TokenizeFileNotExistError("Tokenized file not found in the database.")
1808 with Session(self._sa_engine) as session:
1809 statement = select(TokenizedBed).where(
1810 and_(
1811 TokenizedBed.bed_id == bed_id,
1812 TokenizedBed.universe_id == universe_id,
1813 ),
1814 )
1815 tokenized_object = session.scalar(statement)
1816 return str(tokenized_object.path)
1818 def exist_tokenized(self, bed_id: str, universe_id: str) -> bool:
1819 """
1820 Check if tokenized bed file exists in the database.
1822 Args:
1823 bed_id: Bed file identifier.
1824 universe_id: Universe identifier.
1826 Returns:
1827 True if tokenized bed file exists, False otherwise.
1828 """
1829 with Session(self._sa_engine) as session:
1830 statement = select(TokenizedBed).where(
1831 and_(
1832 TokenizedBed.bed_id == bed_id,
1833 TokenizedBed.universe_id == universe_id,
1834 )
1835 )
1836 tokenized_object = session.scalar(statement)
1837 if not tokenized_object:
1838 return False
1839 return True
1841 def get_tokenized_link(
1842 self, bed_id: str, universe_id: str
1843 ) -> TokenizedPathResponse:
1844 """
1845 Get tokenized link to tokenized file.
1847 Args:
1848 bed_id: Bed file identifier.
1849 universe_id: Universe identifier.
1851 Returns:
1852 Token link.
1854 Raises:
1855 TokenizeFileNotExistError: If the tokenized file does not exist.
1856 """
1857 file_path = self._get_tokenized_path(bed_id, universe_id)
1859 return TokenizedPathResponse(
1860 endpoint_url=self.config.config.s3.endpoint_url,
1861 file_path=file_path,
1862 bed_id=bed_id,
1863 universe_id=universe_id,
1864 )
1866 def get_missing_plots(
1867 self, plot_name: str, limit: int = 1000, offset: int = 0
1868 ) -> list[str]:
1869 """
1870 Get list of bed files that are missing plot.
1872 Args:
1873 plot_name: Plot name.
1874 limit: Number of results to return.
1875 offset: Offset to start from.
1877 Returns:
1878 List of bed file identifiers.
1879 """
1880 if plot_name not in list(BedPlots.model_fields.keys()):
1881 raise BedBaseConfError(
1882 f"Plot name: {plot_name} is not valid. Valid names: {list(BedPlots.model_fields.keys())}"
1883 )
1885 with Session(self._sa_engine) as session:
1886 # Alias for subquery
1887 t2_alias = aliased(Files)
1889 # Define the subquery
1890 subquery = (
1891 select(t2_alias).where(and_(t2_alias.name == plot_name)).subquery()
1892 )
1894 query = (
1895 select(Bed.id)
1896 .outerjoin(subquery, Bed.id == subquery.c.bedfile_id)
1897 .where(subquery.c.bedfile_id.is_(None))
1898 .limit(limit)
1899 .offset(offset)
1900 )
1902 results = session.scalars(query)
1904 results = [result for result in results]
1906 return results
1908 def get_missing_stats(self, limit: int = 1000, offset: int = 0) -> list[str]:
1909 """
1910 Get list of bed files that are missing statistics.
1912 Args:
1913 limit: Number of results to return.
1914 offset: Offset to start from.
1916 Returns:
1917 List of bed file identifiers.
1918 """
1920 with Session(self._sa_engine) as session:
1921 query = (
1922 select(BedStats)
1923 .where(BedStats.number_of_regions.is_(None))
1924 .limit(limit)
1925 .offset(offset)
1926 )
1928 results = session.scalars(query)
1930 results = [result.id for result in results]
1932 return results
1934 def get_missing_files(self, limit: int = 1000, offset: int = 0) -> list[str]:
1935 """
1936 Get list of bed files that are missing files (bigBed files).
1938 Args:
1939 limit: Number of results to return.
1940 offset: Offset to start from.
1942 Returns:
1943 List of bed file identifiers.
1944 """
1946 with Session(self._sa_engine) as session:
1947 # Alias for subquery
1948 t2_alias = aliased(Files)
1950 # Define the subquery
1951 subquery = select(t2_alias).where(t2_alias.name == "bigbed_file").subquery()
1953 query = (
1954 select(Bed.id)
1955 .outerjoin(subquery, Bed.id == subquery.c.bedfile_id)
1956 .where(subquery.c.bedfile_id.is_(None))
1957 .limit(limit)
1958 .offset(offset)
1959 )
1961 results = session.scalars(query)
1963 results = [result for result in results]
1965 return results
1967 def get_unprocessed(
1968 self, limit: int = 1000, offset: int = 0, genome: str | list | None = None
1969 ) -> BedListResult:
1970 """
1971 Get bed files that are not processed.
1973 Args:
1974 limit: Number of results to return.
1975 offset: Offset to start from.
1976 genome: Genome alias or list of genome aliases to filter by. e.g. "hg38" or
1977 ["hg38", "mm10"]. By default None, which means no filtering by genome.
1979 Returns:
1980 List of bed file identifiers.
1981 """
1983 if isinstance(genome, str):
1984 genome = [genome]
1986 with Session(self._sa_engine) as session:
1987 query = select(Bed).where(Bed.processed.is_(False))
1989 count_query = select(func.count()).where(Bed.processed.is_(False))
1991 if genome:
1992 or_statement = or_(*[Bed.genome_alias == k for k in genome])
1993 query = query.where(or_statement)
1994 count_query = count_query.where(or_statement)
1996 query = query.limit(limit).offset(offset)
1998 count = session.execute(count_query).one()[0]
2000 bed_results = session.scalars(query)
2002 results = []
2003 for bed_object in bed_results:
2004 results.append(
2005 BedMetadataBasic(
2006 id=bed_object.id,
2007 name=bed_object.name,
2008 genome_alias=bed_object.genome_alias,
2009 genome_digest=bed_object.genome_digest,
2010 bed_compliance=bed_object.bed_compliance,
2011 data_format=bed_object.data_format,
2012 description=bed_object.description,
2013 annotation=StandardMeta(
2014 **(
2015 bed_object.annotations.__dict__
2016 if bed_object.annotations
2017 else {}
2018 )
2019 ),
2020 last_update_date=bed_object.last_update_date,
2021 submission_date=bed_object.submission_date,
2022 is_universe=bed_object.is_universe,
2023 license_id=bed_object.license_id,
2024 )
2025 )
2027 return BedListResult(
2028 count=count,
2029 limit=limit,
2030 offset=offset,
2031 results=results,
2032 )
2034 def _update_sources(
2035 self,
2036 identifier: str,
2037 global_sample_id: list[str] | None = None,
2038 global_experiment_id: list[str] | None = None,
2039 ) -> None:
2040 """
2041 Add global sample and experiment ids to the bed file if they are missing.
2043 Args:
2044 identifier: Bed file identifier.
2045 global_sample_id: List of global sample ids.
2046 global_experiment_id: List of global experiment ids.
2048 Returns:
2049 None.
2050 """
2051 _LOGGER.info(f"Updating sources for bed file: {identifier}")
2053 with Session(self._sa_engine) as session:
2054 bed_statement = select(BedMetadata).where(
2055 and_(BedMetadata.id == identifier)
2056 )
2058 bedmetadata_object = session.scalar(bed_statement)
2060 if (
2061 global_sample_id
2062 and global_sample_id[0] not in bedmetadata_object.global_sample_id
2063 ):
2064 bedmetadata_object.global_sample_id.append(global_sample_id[0])
2066 flag_modified(bedmetadata_object, "global_sample_id")
2068 if (
2069 global_experiment_id
2070 and global_experiment_id[0]
2071 not in bedmetadata_object.global_experiment_id
2072 ):
2073 bedmetadata_object.global_experiment_id.append(global_experiment_id[0])
2074 flag_modified(bedmetadata_object, "global_experiment_id")
2076 session.commit()
2078 def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None:
2079 """
2080 Reindex all bed files for semantic database.
2082 Args:
2083 batch: Number of files to upload in one batch.
2084 purge: Resets indexed in database for all files to False.
2086 Returns:
2087 None.
2088 """
2090 # Add column that will indicate if this file is indexed or not
2091 statement = (
2092 select(Bed)
2093 .join(BedMetadata, Bed.id == BedMetadata.id)
2094 .where(Bed.indexed.is_(False))
2095 )
2097 with Session(self._sa_engine) as session:
2098 if purge:
2099 _LOGGER.info("Purging indexed files in the database ...")
2100 session.query(Bed).update({Bed.indexed: False})
2101 session.commit()
2102 _LOGGER.info("Purged indexed files in the database successfully!")
2104 _LOGGER.info("Fetching data from the database ...")
2105 results = session.scalars(statement)
2107 _LOGGER.info("Fetch data successfully!")
2109 points = []
2110 results = [result for result in results]
2112 with tqdm(total=len(results), position=0, leave=True) as pbar:
2113 processed_number = 0
2114 for result in results:
2115 text = (
2116 f"biosample is {result.annotations.cell_line} / {result.annotations.cell_type} / "
2117 f"{result.annotations.tissue} with target {result.annotations.target} "
2118 f"assay {result.annotations.assay}."
2119 f"File name {result.name} with summary {result.description}"
2120 )
2122 embeddings_list = list(self.config.dense_encoder.embed(text))
2124 if self.config.sparse_encoder:
2125 sparse_result = self.config.sparse_encoder.encode(
2126 text
2127 ).coalesce()
2129 sparse_embeddings = models.SparseVector(
2130 indices=sparse_result.indices().tolist()[0],
2131 values=sparse_result.values().tolist(),
2132 )
2134 point_vectors = {
2135 "dense": list(embeddings_list[0]),
2136 "sparse": sparse_embeddings,
2137 }
2138 else:
2139 point_vectors = {
2140 "dense": list(embeddings_list[0]),
2141 }
2142 # result_list.append(
2143 data = VectorMetadata(
2144 id=result.id,
2145 name=result.name,
2146 description=result.description,
2147 genome_alias=result.genome_alias,
2148 genome_digest=result.genome_digest,
2149 cell_line=result.annotations.cell_line,
2150 cell_type=result.annotations.cell_type,
2151 tissue=result.annotations.tissue,
2152 target=result.annotations.target,
2153 treatment=result.annotations.treatment,
2154 assay=result.annotations.assay,
2155 species_name=result.annotations.species_name,
2156 )
2158 points.append(
2159 PointStruct(
2160 id=result.id,
2161 vector=point_vectors,
2162 payload=data.model_dump(),
2163 )
2164 )
2165 processed_number += 1
2166 result.indexed = True
2168 if processed_number % batch == 0:
2169 pbar.set_description(
2170 "Uploading points to qdrant using batch..."
2171 )
2172 operation_info = self.config.qdrant_client.upsert(
2173 collection_name=self.config.config.qdrant.hybrid_collection,
2174 points=points,
2175 wait=False,
2176 )
2177 pbar.write("Uploaded batch to qdrant.")
2178 points = []
2179 _LOGGER.info(operation_info.status)
2180 assert (
2181 operation_info.status == "completed"
2182 or operation_info.status == "acknowledged"
2183 )
2184 session.commit()
2186 pbar.write(f"File: {result.id} successfully indexed.")
2187 pbar.update(1)
2189 if points:
2190 operation_info = self.config.qdrant_client.upsert(
2191 collection_name=self.config.config.qdrant.hybrid_collection,
2192 points=points,
2193 )
2194 assert (
2195 operation_info.status == "completed"
2196 or operation_info.status == "acknowledged"
2197 )
2198 session.commit()
2200 return None
2202 def hybrid_search(
2203 self,
2204 query: str = "liver",
2205 genome_alias: str = "",
2206 assay: str = "",
2207 limit: int = 100,
2208 offset: int = 0,
2209 with_metadata: bool = True,
2210 ) -> BedListSearchResult:
2211 """
2212 Run semantic search for bed files using qdrant.
2213 This is not bivec search, but usual qdrant search with sparse and dense embeddings.
2215 Args:
2216 query: Text query to search for.
2217 genome_alias: Genome alias to filter results.
2218 assay: Filter by assay type.
2219 limit: Number of results to return.
2220 offset: Offset to start from.
2221 with_metadata: If True, metadata will be returned in the results. Default is True.
2223 Returns:
2224 List of bed file metadata.
2225 """
2227 must_statement = []
2229 if genome_alias:
2230 must_statement.append(
2231 models.FieldCondition(
2232 key="genome_alias",
2233 match=models.MatchValue(value=genome_alias),
2234 )
2235 )
2236 if assay:
2237 must_statement.append(
2238 models.FieldCondition(
2239 key="assay",
2240 match=models.MatchValue(value=assay),
2241 )
2242 )
2244 dense_query = list(list(self.config.dense_encoder.embed(query))[0])
2245 if self.config.sparse_encoder:
2246 sparse_result = self.config.sparse_encoder.encode(query).coalesce()
2247 sparse_embeddings = models.SparseVector(
2248 indices=sparse_result.indices().tolist()[0],
2249 values=sparse_result.values().tolist(),
2250 )
2252 hybrid_query = [
2253 # Dense retrieval: semantic understanding
2254 models.Prefetch(
2255 query=dense_query, using="dense", limit=limit + offset + 100
2256 ),
2257 # Sparse retrieval: exact technical term matching
2258 models.Prefetch(
2259 query=sparse_embeddings, using="sparse", limit=limit + offset + 100
2260 ),
2261 ]
2262 else:
2263 hybrid_query = [
2264 # Dense retrieval: semantic understanding
2265 models.Prefetch(
2266 query=dense_query, using="dense", limit=limit + offset + 100
2267 ),
2268 ]
2270 results = self.config.qdrant_client.query_points(
2271 collection_name=self.config.config.qdrant.hybrid_collection,
2272 limit=limit,
2273 offset=offset,
2274 prefetch=hybrid_query,
2275 query=models.FusionQuery(fusion=models.Fusion.RRF),
2276 with_payload=True,
2277 with_vectors=True,
2278 search_params=models.SearchParams(
2279 exact=True,
2280 ),
2281 query_filter=(
2282 models.Filter(must=must_statement) if must_statement else None
2283 ),
2284 )
2286 result_list = []
2287 for result in results.points:
2288 result_id = _bed_id_from_point_id(result.id)
2290 if with_metadata:
2291 metadata = self.get(result_id, full=False)
2292 else:
2293 metadata = None
2295 result_list.append(
2296 QdrantSearchResult(
2297 id=result_id,
2298 payload=result.payload,
2299 score=result.score,
2300 metadata=metadata,
2301 )
2302 )
2304 if with_metadata:
2305 count = self.bb_agent.get_stats().bedfiles_number
2306 else:
2307 count = 0
2309 return BedListSearchResult(
2310 count=count,
2311 limit=limit,
2312 offset=offset,
2313 results=result_list,
2314 )
2316 def search_external_file(self, source: str, accession: str) -> BedListSearchResult:
2317 """
2318 Search for bed files by external source and accession number.
2319 e.g. source='geo', accession='GSE12345'.
2321 Args:
2322 source: External source, e.g. 'geo' or 'encode'.
2323 accession: Accession number, e.g. 'GSE12345' or 'ENCSR12345'.
2325 Returns:
2326 List of bed file metadata.
2327 """
2329 if source not in ["geo", "encode"]:
2330 raise BedBaseConfError(
2331 f"Source {source} is not supported. Supported sources are: 'geo', 'encode'."
2332 )
2334 if source == "geo" and accession.upper().startswith("GSE"):
2335 statement = (
2336 select(Bed)
2337 .join(BedMetadata, Bed.id == BedMetadata.id)
2338 .where(
2339 BedMetadata.global_experiment_id.contains(
2340 cast(
2341 [f"{source}:{accession}"],
2342 postgresql.ARRAY(postgresql.VARCHAR),
2343 )
2344 )
2345 )
2346 )
2347 elif source == "geo" and accession.upper().startswith("GSM"):
2348 statement = (
2349 select(Bed)
2350 .join(BedMetadata, Bed.id == BedMetadata.id)
2351 .where(
2352 BedMetadata.global_sample_id.contains(
2353 cast(
2354 [f"{source}:{accession}"],
2355 postgresql.ARRAY(postgresql.VARCHAR),
2356 )
2357 )
2358 )
2359 )
2360 elif source == "encode" and accession.upper().startswith("ENCSR"):
2361 accession = accession.upper()
2362 statement = (
2363 select(Bed)
2364 .join(BedMetadata, Bed.id == BedMetadata.id)
2365 .where(
2366 BedMetadata.global_experiment_id.contains(
2367 cast(
2368 [f"{source}:{accession}"],
2369 postgresql.ARRAY(postgresql.VARCHAR),
2370 )
2371 )
2372 )
2373 )
2374 print(statement)
2375 else:
2376 statement = (
2377 select(Bed)
2378 .join(BedMetadata, Bed.id == BedMetadata.id)
2379 .where(
2380 BedMetadata.global_sample_id.contains(
2381 cast(
2382 [f"{source}:{accession}"],
2383 postgresql.ARRAY(postgresql.VARCHAR),
2384 )
2385 )
2386 )
2387 )
2389 with Session(self._sa_engine) as session:
2390 bed_objects = session.scalars(statement)
2391 results = [
2392 BedMetadataBasic(
2393 **bedfile_obj.__dict__,
2394 annotation=StandardMeta(
2395 **(
2396 bedfile_obj.annotations.__dict__
2397 if bedfile_obj.annotations
2398 else {}
2399 )
2400 ),
2401 )
2402 for bedfile_obj in bed_objects
2403 ]
2404 result_list = [
2405 QdrantSearchResult(id=result.id, score=1, metadata=result)
2406 for result in results
2407 ]
2409 return BedListSearchResult(
2410 count=len(result_list),
2411 limit=len(result_list),
2412 offset=0,
2413 results=result_list,
2414 )