|
50 | 50 | from ord_schema import message_helpers, validations |
51 | 51 | from ord_schema.logging import get_logger |
52 | 52 | from ord_schema.proto import reaction_pb2 |
53 | | -from psycopg import AsyncCursor |
| 53 | +from psycopg import AsyncCursor, sql |
54 | 54 | from pydantic import BaseModel |
55 | 55 | from rdkit import Chem |
56 | 56 | from rdkit.Chem import rdChemReactions |
@@ -435,47 +435,76 @@ class StatsResult(BaseModel): |
435 | 435 | times_appearing: int |
436 | 436 |
|
437 | 437 |
|
438 | | -async def fetch_dataset_most_used_smiles_for_inputs( |
439 | | - cursor: DictCursor, dataset_id: str, limit: int = 30 |
| 438 | +async def _fetch_dataset_most_used_smiles( |
| 439 | + cursor: DictCursor, |
| 440 | + dataset_id: str, |
| 441 | + *, |
| 442 | + compound_table: str, |
| 443 | + join_table: str, |
| 444 | + foreign_key: str, |
| 445 | + limit: int, |
440 | 446 | ) -> list[StatsResult]: |
441 | | - """Fetches the top K most used SMILES molecules in terms of reaction inputs for a given dataset.""" |
442 | | - query = """ |
443 | | - SELECT smiles, COUNT(*) as times_appearing |
444 | | - FROM ord.compound |
445 | | - JOIN ord.reaction_input ON ord.compound.reaction_input_id = ord.reaction_input.id |
446 | | - JOIN ord.reaction ON ord.reaction_input.reaction_id = ord.reaction.id |
| 447 | + """Fetches the top K most used SMILES for a dataset, joined through the given tables. |
| 448 | +
|
| 449 | + Args: |
| 450 | + cursor: Database cursor. |
| 451 | + dataset_id: Dataset to aggregate over. |
| 452 | + compound_table: Table holding SMILES, e.g. "compound" or "product_compound". |
| 453 | + join_table: Table linking compounds to reactions, e.g. "reaction_input". |
| 454 | + foreign_key: Column on ``compound_table`` referencing ``join_table``. |
| 455 | + limit: Maximum number of rows to return. |
| 456 | +
|
| 457 | + Returns: |
| 458 | + Most frequently appearing SMILES, in descending order of frequency. |
| 459 | + """ |
| 460 | + # Compose schema-qualified identifiers safely rather than interpolating raw |
| 461 | + # strings into the SQL text. |
| 462 | + compound = sql.Identifier("ord", compound_table) |
| 463 | + join = sql.Identifier("ord", join_table) |
| 464 | + query = sql.SQL( |
| 465 | + """ |
| 466 | + SELECT smiles, COUNT(*) AS times_appearing |
| 467 | + FROM {compound} |
| 468 | + JOIN {join} ON {compound}.{foreign_key} = {join}.id |
| 469 | + JOIN ord.reaction ON {join}.reaction_id = ord.reaction.id |
447 | 470 | JOIN ord.dataset ON ord.reaction.dataset_id = ord.dataset.id |
448 | 471 | WHERE ord.dataset.dataset_id = %s |
449 | 472 | AND smiles IS NOT NULL |
450 | 473 | GROUP BY smiles |
451 | 474 | ORDER BY times_appearing DESC |
452 | 475 | LIMIT %s |
453 | 476 | """ |
| 477 | + ).format(compound=compound, join=join, foreign_key=sql.Identifier(foreign_key)) |
454 | 478 | await cursor.execute(query, (dataset_id, limit)) |
455 | 479 | results = [] |
456 | 480 | async for row in cursor: |
457 | 481 | results.append(StatsResult(**row)) |
458 | 482 | return results |
459 | 483 |
|
460 | 484 |
|
| 485 | +async def fetch_dataset_most_used_smiles_for_inputs( |
| 486 | + cursor: DictCursor, dataset_id: str, limit: int = 30 |
| 487 | +) -> list[StatsResult]: |
| 488 | + """Fetches the top K most used SMILES molecules in terms of reaction inputs for a given dataset.""" |
| 489 | + return await _fetch_dataset_most_used_smiles( |
| 490 | + cursor, |
| 491 | + dataset_id, |
| 492 | + compound_table="compound", |
| 493 | + join_table="reaction_input", |
| 494 | + foreign_key="reaction_input_id", |
| 495 | + limit=limit, |
| 496 | + ) |
| 497 | + |
| 498 | + |
461 | 499 | async def fetch_dataset_most_used_smiles_for_products( |
462 | 500 | cursor: DictCursor, dataset_id: str, limit: int = 30 |
463 | 501 | ) -> list[StatsResult]: |
464 | 502 | """Fetches the top K most used SMILES molecules in terms of reaction products for a given dataset.""" |
465 | | - query = """ |
466 | | - SELECT smiles, COUNT(*) as times_appearing |
467 | | - FROM ord.product_compound |
468 | | - JOIN ord.reaction_outcome ON ord.product_compound.reaction_outcome_id = ord.reaction_outcome.id |
469 | | - JOIN ord.reaction ON ord.reaction_outcome.reaction_id = ord.reaction.id |
470 | | - JOIN ord.dataset ON ord.reaction.dataset_id = ord.dataset.id |
471 | | - WHERE ord.dataset.dataset_id = %s |
472 | | - AND smiles IS NOT NULL |
473 | | - GROUP BY smiles |
474 | | - ORDER BY times_appearing DESC |
475 | | - LIMIT %s |
476 | | - """ |
477 | | - await cursor.execute(query, (dataset_id, limit)) |
478 | | - results = [] |
479 | | - async for row in cursor: |
480 | | - results.append(StatsResult(**row)) |
481 | | - return results |
| 503 | + return await _fetch_dataset_most_used_smiles( |
| 504 | + cursor, |
| 505 | + dataset_id, |
| 506 | + compound_table="product_compound", |
| 507 | + join_table="reaction_outcome", |
| 508 | + foreign_key="reaction_outcome_id", |
| 509 | + limit=limit, |
| 510 | + ) |
0 commit comments