Coverage for python/lsst/daf/relation/sql/_engine.py: 12%

296 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-01-10 02:26 -0800

1# This file is part of daf_relation. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://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 <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ("Engine",) 

25 

26import dataclasses 

27from collections.abc import Callable, Iterable, Mapping, Sequence, Set 

28from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast 

29 

30import sqlalchemy 

31 

32from .._binary_operation import BinaryOperation, IgnoreOne 

33from .._columns import ( 

34 ColumnExpression, 

35 ColumnExpressionSequence, 

36 ColumnFunction, 

37 ColumnInContainer, 

38 ColumnLiteral, 

39 ColumnRangeLiteral, 

40 ColumnReference, 

41 ColumnTag, 

42 LogicalAnd, 

43 LogicalNot, 

44 LogicalOr, 

45 Predicate, 

46 PredicateFunction, 

47 PredicateLiteral, 

48 PredicateReference, 

49 flatten_logical_and, 

50) 

51from .._engine import GenericConcreteEngine 

52from .._exceptions import EngineError, RelationalAlgebraError 

53from .._leaf_relation import LeafRelation 

54from .._materialization import Materialization 

55from .._operation_relations import BinaryOperationRelation, UnaryOperationRelation 

56from .._operations import ( 

57 Calculation, 

58 Chain, 

59 Deduplication, 

60 Join, 

61 PartialJoin, 

62 Projection, 

63 Selection, 

64 Slice, 

65 Sort, 

66 SortTerm, 

67) 

68from .._transfer import Transfer 

69from .._unary_operation import Identity, UnaryOperation 

70from ._payload import Payload 

71from ._select import Select 

72 

73if TYPE_CHECKING: 73 ↛ 74line 73 didn't jump to line 74, because the condition on line 73 was never true

74 from .._relation import Relation 

75 

76 

77_L = TypeVar("_L") 

78 

79 

80@dataclasses.dataclass(repr=False, eq=False, kw_only=True) 

81class Engine( 

82 GenericConcreteEngine[Callable[..., sqlalchemy.sql.ColumnElement]], 

83 Generic[_L], 

84): 

85 """A concrete engine class for relations backed by a SQL database. 

86 

87 See the `.sql` module documentation for details. 

88 """ 

89 

90 name: str = "sql" 

91 

92 EMPTY_COLUMNS_NAME: ClassVar[str] = "IGNORED" 

93 """Name of the column added to a SQL ``SELECT`` query in order to represent 

94 relations that have no real columns. 

95 """ 

96 

97 EMPTY_COLUMNS_TYPE: ClassVar[type] = sqlalchemy.Boolean 

98 """Type of the column added to a SQL ``SELECT`` query in order to represent 

99 relations that have no real columns. 

100 """ 

101 

102 def __str__(self) -> str: 

103 return self.name 

104 

105 def __repr__(self) -> str: 

106 return f"lsst.daf.relation.sql.Engine({self.name!r})@{id(self):0x}" 

107 

108 def make_leaf( 

109 self, 

110 columns: Set[ColumnTag], 

111 payload: Payload[_L], 

112 *, 

113 min_rows: int = 0, 

114 max_rows: int | None = None, 

115 name: str = "", 

116 messages: Sequence[str] = (), 

117 name_prefix: str = "leaf", 

118 parameters: Any = None, 

119 ) -> Relation: 

120 """Create a nontrivial leaf relation in this engine. 

121 

122 This is a convenience method that simply forwards all arguments to 

123 the `.LeafRelation` constructor, and then wraps the result in a 

124 `Select`; see `LeafRelation` for details. 

125 """ 

126 return Select.apply_skip( 

127 LeafRelation( 

128 self, 

129 frozenset(columns), 

130 payload, 

131 min_rows=min_rows, 

132 max_rows=max_rows, 

133 messages=messages, 

134 name=name, 

135 name_prefix=name_prefix, 

136 parameters=parameters, 

137 ) 

138 ) 

139 

140 def conform(self, relation: Relation) -> Select: 

141 # Docstring inherited. 

142 match relation: 

143 case Select(): 

144 return relation 

145 case UnaryOperationRelation(operation=operation, target=target): 

146 conformed_target = self.conform(target) 

147 return self._append_unary_to_select(operation, conformed_target) 

148 case BinaryOperationRelation(operation=operation, lhs=lhs, rhs=rhs): 

149 conformed_lhs = self.conform(lhs) 

150 conformed_rhs = self.conform(rhs) 

151 return self._append_binary_to_select(operation, conformed_lhs, conformed_rhs) 

152 case Transfer() | Materialization() | LeafRelation(): 

153 # We always conform upstream of transfers and materializations, 

154 # so no need to recurse. 

155 return Select.apply_skip(relation) 

156 case _: 

157 # Other marker relations. 

158 conformed_target = self.conform(target) 

159 return Select.apply_skip(relation) 

160 raise AssertionError("Match should be exhaustive and all branches return.") 

161 

162 def materialize( 

163 self, target: Relation, name: str | None = None, name_prefix: str = "materialization_" 

164 ) -> Select: 

165 # Docstring inherited. 

166 conformed_target = self.conform(target) 

167 if conformed_target.has_sort and not conformed_target.has_slice: 

168 raise RelationalAlgebraError( 

169 f"Materializing relation {conformed_target} will not preserve row order." 

170 ) 

171 return Select.apply_skip(super().materialize(conformed_target, name, name_prefix)) 

172 

173 def transfer(self, target: Relation, payload: Any | None = None) -> Select: 

174 # Docstring inherited. 

175 return Select.apply_skip(super().transfer(target, payload)) 

176 

177 def make_doomed_relation( 

178 self, columns: Set[ColumnTag], messages: Sequence[str], name: str = "0" 

179 ) -> Relation: 

180 # Docstring inherited. 

181 return Select.apply_skip(super().make_doomed_relation(columns, messages, name)) 

182 

183 def make_join_identity_relation(self, name: str = "I") -> Relation: 

184 # Docstring inherited. 

185 return Select.apply_skip(super().make_join_identity_relation(name)) 

186 

187 def get_join_identity_payload(self) -> Payload[_L]: 

188 # Docstring inherited. 

189 select_columns: list[sqlalchemy.sql.ColumnElement] = [] 

190 self.handle_empty_columns(select_columns) 

191 return Payload[_L](sqlalchemy.sql.select(*select_columns).subquery()) 

192 

193 def get_doomed_payload(self, columns: Set[ColumnTag]) -> Payload[_L]: 

194 # Docstring inherited. 

195 select_columns = [sqlalchemy.sql.literal(None).label(tag.qualified_name) for tag in columns] 

196 self.handle_empty_columns(select_columns) 

197 subquery = sqlalchemy.sql.select(*select_columns).subquery() 

198 return Payload( 

199 subquery, 

200 where=[sqlalchemy.sql.literal(False)], 

201 columns_available=self.extract_mapping(columns, subquery.columns), 

202 ) 

203 

204 def append_unary(self, operation: UnaryOperation, target: Relation) -> Select: 

205 # Docstring inherited. 

206 conformed_target = self.conform(target) 

207 return self._append_unary_to_select(operation, conformed_target) 

208 

209 def _append_unary_to_select(self, operation: UnaryOperation, select: Select) -> Select: 

210 """Internal recursive implementation of `append_unary`. 

211 

212 This method should not be called by external code, but it may be 

213 overridden and called recursively by derived engine classes. 

214 

215 Parameters 

216 ---------- 

217 operation : `UnaryOperation` 

218 Operation to add to the tree. 

219 select : `Select` 

220 Existing already-conformed relation tree. 

221 

222 Returns 

223 ------- 

224 appended : `Select` 

225 Conformed relation tree that includes the given operation. 

226 """ 

227 match operation: 

228 case Calculation(tag=tag): 

229 if select.is_compound: 

230 # This Select wraps a Chain operation in order to represent 

231 # a SQL UNION or UNION ALL, and we trust the user's intent 

232 # in putting those upstream of this operation, so we also 

233 # add a nested subquery here. 

234 return Select.apply_skip(operation._finish_apply(select)) 

235 elif select.has_projection: 

236 return select.reapply_skip( 

237 after=operation, 

238 projection=Projection(frozenset(select.columns | {tag})), 

239 ) 

240 else: 

241 return select.reapply_skip(after=operation) 

242 case Deduplication(): 

243 if not select.has_deduplication: 

244 if select.has_slice: 

245 # There was a Slice upstream, which needs to be applied 

246 # before this Deduplication, so we nest the existing 

247 # subquery within a new one that has deduplication. 

248 return Select.apply_skip(select, deduplication=operation) 

249 else: 

250 # Move the Deduplication into the Select's 

251 # operations. 

252 return select.reapply_skip(deduplication=operation) 

253 else: 

254 # There was already another (redundant) 

255 # Deduplication upstream. 

256 return select 

257 case Projection(): 

258 if select.has_deduplication: 

259 # There was a Duplication upstream, so we need to ensure 

260 # that is applied before this Projection via a nested 

261 # subquery. Instead of applying the existing subquery 

262 # operations, though, we apply one without any sort or 

263 # slice that might exist, and save those for the new outer 

264 # query, since putting those in a subquery would destroy 

265 # the ordering. 

266 subquery = select.reapply_skip(sort=None, slice=None) 

267 return Select.apply_skip( 

268 subquery, 

269 projection=operation, 

270 sort=select.sort, 

271 slice=select.slice, 

272 ) 

273 else: 

274 # No Deduplication, so we can just add the Projection to 

275 # the existing Select and reapply it. 

276 match select.skip_to: 

277 case BinaryOperationRelation(operation=Chain() as chain, lhs=lhs, rhs=rhs): 

278 # ... unless the skip_to relation is a Chain; we 

279 # want to move the Projection inside the Chain, to 

280 # make it easier to avoid subqueries in UNION [ALL] 

281 # constructs later. 

282 return select.reapply_skip( 

283 skip_to=chain._finish_apply(operation.apply(lhs), operation.apply(rhs)), 

284 projection=None, 

285 ) 

286 case _: 

287 return select.reapply_skip(projection=operation) 

288 case Selection(): 

289 if select.has_slice: 

290 # There was a Slice upstream, which needs to be applied 

291 # before this Selection via a nested subquery (which means 

292 # we just apply the Selection to target, then add a new 

293 # empty subquery marker). 

294 return Select.apply_skip(operation._finish_apply(select)) 

295 elif select.is_compound: 

296 # This Select wraps a Chain operation in order to represent 

297 # a SQL UNION or UNION ALL, and we trust the user's intent 

298 # in putting those upstream of this operation, so we also 

299 # add a nested subquery here. 

300 return Select.apply_skip(operation._finish_apply(select)) 

301 else: 

302 return select.reapply_skip(after=operation) 

303 case Slice(): 

304 return select.reapply_skip(slice=select.slice.then(operation)) 

305 case Sort(): 

306 if select.has_slice: 

307 # There was a Slice upstream, which needs to be applied 

308 # before this Sort via a nested subquery (which means we 

309 # apply a new Sort-only Select to 'select' itself). 

310 return Select.apply_skip(select, sort=operation) 

311 else: 

312 return select.reapply_skip(sort=select.sort.then(operation)) 

313 case PartialJoin(binary=binary, fixed=fixed, fixed_is_lhs=fixed_is_lhs): 

314 if fixed_is_lhs: 

315 return self.append_binary(binary, fixed, select) 

316 else: 

317 return self.append_binary(binary, select, fixed) 

318 case Identity(): 

319 return select 

320 raise NotImplementedError(f"Unsupported operation type {operation} for engine {self}.") 

321 

322 def append_binary(self, operation: BinaryOperation, lhs: Relation, rhs: Relation) -> Select: 

323 # Docstring inherited. 

324 conformed_lhs = self.conform(lhs) 

325 conformed_rhs = self.conform(rhs) 

326 return self._append_binary_to_select(operation, conformed_lhs, conformed_rhs) 

327 

328 def _append_binary_to_select(self, operation: BinaryOperation, lhs: Select, rhs: Select) -> Select: 

329 """Internal recursive implementation of `append_binary`. 

330 

331 This method should not be called by external code, but it may be 

332 overridden and called recursively by derived engine classes. 

333 

334 Parameters 

335 ---------- 

336 operation : `UnaryOperation` 

337 Operation to add to the tree. 

338 lhs : `Select` 

339 Existing already-conformed relation tree. 

340 rhs : `Select` 

341 The other existing already-conformed relation tree. 

342 

343 Returns 

344 ------- 

345 appended : `Select` 

346 Conformed relation tree that includes the given operation. 

347 """ 

348 if lhs.has_sort and not lhs.has_slice: 

349 raise RelationalAlgebraError( 

350 f"Applying binary operation {operation} to relation {lhs} will not preserve row order." 

351 ) 

352 if rhs.has_sort and not rhs.has_slice: 

353 raise RelationalAlgebraError( 

354 f"Applying binary operation {operation} to relation {rhs} will not preserve row order." 

355 ) 

356 match operation: 

357 case Chain(): 

358 # Chain operands should always be Selects, but they can't have 

359 # Sorts or Slices, at least not directly; if those exist we 

360 # need to move them into subqueries at which point we might as 

361 # well move any Projection or Dedulication there as well). But 

362 # because of the check on Sorts-without-Slices above, there has 

363 # to be a Slice here for there to be a Sort here. 

364 if lhs.has_slice: 

365 lhs = Select.apply_skip(lhs) 

366 if rhs.has_slice: 

367 rhs = Select.apply_skip(rhs) 

368 return Select.apply_skip(operation._finish_apply(lhs, rhs)) 

369 case Join(): 

370 # For Joins, we move Projections after the Join, unless we have 

371 # another reason to make a subquery before the join, and the 

372 # operands are only Selects if they need to be subqueries. 

373 new_lhs, new_lhs_needs_projection = lhs.strip() 

374 new_rhs, new_rhs_needs_projection = rhs.strip() 

375 if new_lhs_needs_projection or new_rhs_needs_projection: 

376 projection = Projection(frozenset(lhs.columns | rhs.columns)) 

377 else: 

378 projection = None 

379 return Select.apply_skip(operation._finish_apply(new_lhs, new_rhs), projection=projection) 

380 case IgnoreOne(ignore_lhs=ignore_lhs): 

381 if ignore_lhs: 

382 return rhs 

383 else: 

384 return lhs 

385 raise AssertionError(f"Match on {operation} should be exhaustive and all branches return..") 

386 

387 def extract_mapping( 

388 self, tags: Iterable[ColumnTag], sql_columns: sqlalchemy.sql.ColumnCollection 

389 ) -> dict[ColumnTag, _L]: 

390 """Extract a mapping with `.ColumnTag` keys and logical column values 

391 from a SQLAlchemy column collection. 

392 

393 Parameters 

394 ---------- 

395 tags : `Iterable` 

396 Set of `.ColumnTag` objects whose logical columns should be 

397 extracted. 

398 sql_columns : `sqlalchemy.sql.ColumnCollection` 

399 SQLAlchemy collection of columns, such as 

400 `sqlalchemy.sql.FromClause.columns`. 

401 

402 Returns 

403 ------- 

404 logical_columns : `dict` 

405 Dictionary mapping `.ColumnTag` to logical column type. 

406 

407 Notes 

408 ----- 

409 This method must be overridden to support a custom logical columns. 

410 """ 

411 return {tag: cast(_L, sql_columns[tag.qualified_name]) for tag in tags} 

412 

413 def select_items( 

414 self, 

415 items: Iterable[tuple[ColumnTag, _L]], 

416 sql_from: sqlalchemy.sql.FromClause, 

417 *extra: sqlalchemy.sql.ColumnElement, 

418 ) -> sqlalchemy.sql.Select: 

419 """Construct a SQLAlchemy representation of a SELECT query. 

420 

421 Parameters 

422 ---------- 

423 items : `Iterable` [ `tuple` ] 

424 Iterable of (`.ColumnTag`, logical column) pairs. This is 

425 typically the ``items()`` of a mapping returned by 

426 `extract_mapping` or obtained from `Payload.columns_available`. 

427 sql_from : `sqlalchemy.sql.FromClause` 

428 SQLAlchemy representation of a FROM clause, such as a single table, 

429 aliased subquery, or join expression. Must provide all columns 

430 referenced by ``items``. 

431 *extra : `sqlalchemy.sql.ColumnElement` 

432 Additional SQL column expressions to include. 

433 

434 Returns 

435 ------- 

436 select : `sqlalchemy.sql.Select` 

437 SELECT query. 

438 

439 Notes 

440 ----- 

441 This method is responsible for handling the case where ``items`` is 

442 empty, typically by delegating to `handle_empty_columns`. 

443 

444 This method must be overridden to support a custom logical columns. 

445 """ 

446 select_columns = [ 

447 cast(sqlalchemy.sql.ColumnElement, logical_column).label(tag.qualified_name) 

448 for tag, logical_column in items 

449 ] 

450 select_columns.extend(extra) 

451 self.handle_empty_columns(select_columns) 

452 return sqlalchemy.sql.select(*select_columns).select_from(sql_from) 

453 

454 def handle_empty_columns(self, columns: list[sqlalchemy.sql.ColumnElement]) -> None: 

455 """Handle the edge case where a SELECT statement has no columns, by 

456 adding a literal column that should be ignored. 

457 

458 Parameters 

459 ---------- 

460 columns : `list` [ `sqlalchemy.sql.ColumnElement` ] 

461 List of SQLAlchemy column objects. This may have no elements when 

462 this method is called, and must always have at least one element 

463 when it returns. 

464 """ 

465 if not columns: 

466 columns.append(sqlalchemy.sql.literal(True).label(self.EMPTY_COLUMNS_NAME)) 

467 

468 def to_executable( 

469 self, relation: Relation, extra_columns: Iterable[sqlalchemy.sql.ColumnElement] = () 

470 ) -> sqlalchemy.sql.expression.SelectBase: 

471 """Convert a relation tree to an executable SQLAlchemy expression. 

472 

473 Parameters 

474 ---------- 

475 relation : `Relation` 

476 The relation tree to convert. 

477 extra_columns : `~collections.abc.Iterable` 

478 Iterable of additional SQLAlchemy column objects to include 

479 directly in the ``SELECT`` clause. 

480 

481 Returns 

482 ------- 

483 select : `sqlalchemy.sql.expression.SelectBase` 

484 A SQLAlchemy ``SELECT`` or compound ``SELECT`` query. 

485 

486 Notes 

487 ----- 

488 This method requires all relations in the tree to have the same engine 

489 (``self``). It also cannot handle `.Materialization` operations 

490 unless they have already been processed once already (and hence have 

491 a payload attached). Use the `.Processor` function to handle both of 

492 these cases. 

493 """ 

494 if relation.engine != self: 

495 raise EngineError( 

496 f"Engine {self!r} cannot operate on relation {relation} with engine {relation.engine!r}. " 

497 "Use lsst.daf.relation.Processor to evaluate transfers first." 

498 ) 

499 relation = self.conform(relation) 

500 return self._select_to_executable(relation, extra_columns) 

501 

502 def _select_to_executable( 

503 self, 

504 select: Select, 

505 extra_columns: Iterable[sqlalchemy.sql.ColumnElement], 

506 ) -> sqlalchemy.sql.expression.SelectBase: 

507 """Internal recursive implementation of `to_executable`. 

508 

509 This method should not be called by external code, but it may be 

510 overridden and called recursively by derived engine classes. 

511 

512 Parameters 

513 ---------- 

514 columns : `~collections.abc.Set` [ `ColumnTag` ] 

515 Columns to include in the ``SELECT`` clause. 

516 select : `Select` 

517 Already-conformed relation tree to convert. 

518 extra_columns : `~collections.abc.Iterable` 

519 Iterable of additional SQLAlchemy column objects to include 

520 directly in the ``SELECT`` clause. 

521 

522 Returns 

523 ------- 

524 executable : `sqlalchemy.sql.expression.SelectBase` 

525 SQLAlchemy ``SELECT`` or compound ``SELECT`` object. 

526 

527 Notes 

528 ----- 

529 This method handles trees terminated with `Select`, the operation 

530 relation types managed by `Select`, and `Chain` operations nested 

531 directly as the `skip_to` target of a `Select`. It delegates to 

532 `to_payload` for all other relation types. 

533 """ 

534 columns_available: Mapping[ColumnTag, _L] | None = None 

535 match select.skip_to: 

536 case BinaryOperationRelation(operation=Chain(), lhs=lhs, rhs=rhs): 

537 lhs_executable = self._select_to_executable(cast(Select, lhs), extra_columns) 

538 rhs_executable = self._select_to_executable(cast(Select, rhs), extra_columns) 

539 if select.has_deduplication: 

540 executable = sqlalchemy.sql.union(lhs_executable, rhs_executable) 

541 else: 

542 executable = sqlalchemy.sql.union_all(lhs_executable, rhs_executable) 

543 case _: 

544 if select.skip_to.payload is not None: 

545 payload: Payload[_L] = select.skip_to.payload 

546 else: 

547 payload = self.to_payload(select.skip_to) 

548 if not select.columns and not extra_columns: 

549 extra_columns = list(extra_columns) 

550 self.handle_empty_columns(extra_columns) 

551 columns_available = payload.columns_available 

552 columns_projected = {tag: columns_available[tag] for tag in select.columns} 

553 executable = self.select_items(columns_projected.items(), payload.from_clause, *extra_columns) 

554 if len(payload.where) == 1: 

555 executable = executable.where(payload.where[0]) 

556 elif payload.where: 

557 executable = executable.where(sqlalchemy.sql.and_(*payload.where)) 

558 if select.has_deduplication: 

559 executable = executable.distinct() 

560 if select.has_sort: 

561 if columns_available is None: 

562 columns_available = self.extract_mapping(select.skip_to.columns, executable.selected_columns) 

563 executable = executable.order_by( 

564 *[self.convert_sort_term(term, columns_available) for term in select.sort.terms] 

565 ) 

566 if select.slice.start: 

567 executable = executable.offset(select.slice.start) 

568 if select.slice.limit is not None: 

569 executable = executable.limit(select.slice.limit) 

570 return executable 

571 

572 def to_payload(self, relation: Relation) -> Payload[_L]: 

573 """Internal recursive implementation of `to_executable`. 

574 

575 This method should not be called by external code, but it may be 

576 overridden and called recursively by derived engine classes. 

577 

578 Parameters 

579 ---------- 

580 relation : `Relation` 

581 Relation to convert. This method handles all operation relation 

582 types other than `Chain` and the operations managed by `Select`. 

583 

584 Returns 

585 ------- 

586 payload : `Payload` 

587 Struct containing a SQLAlchemy represenation of a simple ``SELECT`` 

588 query. 

589 """ 

590 assert relation.engine == self, "Should be guaranteed by callers." 

591 if relation.payload is not None: # Should cover all LeafRelations 

592 return relation.payload 

593 match relation: 

594 case UnaryOperationRelation(operation=operation, target=target): 

595 match operation: 

596 case Calculation(tag=tag, expression=expression): 

597 result = self.to_payload(target).copy() 

598 result.columns_available[tag] = self.convert_column_expression( 

599 expression, result.columns_available 

600 ) 

601 return result 

602 case Selection(predicate=predicate): 

603 result = self.to_payload(target).copy() 

604 result.where.extend( 

605 self.convert_flattened_predicate(predicate, result.columns_available) 

606 ) 

607 return result 

608 case BinaryOperationRelation( 

609 operation=Join(predicate=predicate, common_columns=common_columns), lhs=lhs, rhs=rhs 

610 ): 

611 lhs_payload = self.to_payload(lhs) 

612 rhs_payload = self.to_payload(rhs) 

613 assert common_columns is not None, "Guaranteed by Join.apply and PartialJoin.apply." 

614 on_terms: list[sqlalchemy.sql.ColumnElement] = [] 

615 if common_columns: 

616 on_terms.extend( 

617 lhs_payload.columns_available[tag] == rhs_payload.columns_available[tag] 

618 for tag in common_columns 

619 ) 

620 columns_available = {**lhs_payload.columns_available, **rhs_payload.columns_available} 

621 if predicate.as_trivial() is not True: 

622 on_terms.extend(self.convert_flattened_predicate(predicate, columns_available)) 

623 on_clause: sqlalchemy.sql.ColumnElement 

624 if not on_terms: 

625 on_clause = sqlalchemy.sql.literal(True) 

626 elif len(on_terms) == 1: 

627 on_clause = on_terms[0] 

628 else: 

629 on_clause = sqlalchemy.sql.and_(*on_terms) 

630 return Payload( 

631 from_clause=lhs_payload.from_clause.join(rhs_payload.from_clause, onclause=on_clause), 

632 where=lhs_payload.where + rhs_payload.where, 

633 columns_available=columns_available, 

634 ) 

635 case Select(): 

636 from_clause = self._select_to_executable(relation, ()).subquery() 

637 columns_available = self.extract_mapping(relation.columns, from_clause.columns) 

638 return Payload(from_clause, columns_available=columns_available) 

639 case Materialization(name=name): 

640 raise EngineError( 

641 f"Cannot persist materialization {name!r} during SQL conversion; " 

642 "use `lsst.daf.relation.Processor` first to handle this operation." 

643 ) 

644 case Transfer(target=target): 

645 raise EngineError( 

646 f"Cannot handle transfer from {target.engine} during SQL conversion; " 

647 "use `lsst.daf.relation.Processor` first to handle this operation." 

648 ) 

649 raise NotImplementedError(f"Unsupported relation type {relation} for engine {self}.") 

650 

651 def convert_column_expression( 

652 self, expression: ColumnExpression, columns_available: Mapping[ColumnTag, _L] 

653 ) -> _L: 

654 """Convert a `.ColumnExpression` to a logical column. 

655 

656 Parameters 

657 ---------- 

658 expression : `.ColumnExpression` 

659 Expression to convert. 

660 columns_available : `~collections.abc.Mapping` 

661 Mapping from `.ColumnTag` to logical column, typically produced by 

662 `extract_mapping` or obtained from `Payload.columns_available`. 

663 

664 Returns 

665 ------- 

666 logical_column 

667 SQLAlchemy expression object or other logical column value. 

668 

669 See Also 

670 -------- 

671 :ref:`lsst.daf.relation-sql-logical-columns` 

672 """ 

673 match expression: 

674 case ColumnLiteral(value=value): 

675 return self.convert_column_literal(value) 

676 case ColumnReference(tag=tag): 

677 return columns_available[tag] 

678 case ColumnFunction(name=name, args=args): 

679 sql_args = [self.convert_column_expression(arg, columns_available) for arg in args] 

680 if (function := self.get_function(name)) is not None: 

681 return function(*sql_args) 

682 return getattr(sql_args[0], name)(*sql_args[1:]) 

683 raise AssertionError( 

684 f"matches should be exhaustive and all branches should return; got {expression!r}." 

685 ) 

686 

687 def convert_column_literal(self, value: Any) -> _L: 

688 """Convert a Python literal value to a logical column. 

689 

690 Parameters 

691 ---------- 

692 value 

693 Python value to convert. 

694 

695 Returns 

696 ------- 

697 logical_column 

698 SQLAlchemy expression object or other logical column value. 

699 

700 Notes 

701 ----- 

702 This method must be overridden to support a custom logical columns. 

703 

704 See Also 

705 -------- 

706 :ref:`lsst.daf.relation-sql-logical-columns` 

707 """ 

708 return sqlalchemy.sql.literal(value) 

709 

710 def expect_column_scalar(self, logical_column: _L) -> sqlalchemy.sql.ColumnElement: 

711 """Convert a logical column value to a SQLAlchemy expression. 

712 

713 Parameters 

714 ---------- 

715 logical_column 

716 SQLAlchemy expression object or other logical column value. 

717 

718 Returns 

719 ------- 

720 sql : `sqlalchemy.sql.ColumnElement` 

721 SQLAlchemy expression object. 

722 

723 Notes 

724 ----- 

725 The default implementation assumes the logical column type is just a 

726 SQLAlchemy type and returns the given object unchanged. Subclasses 

727 with a custom logical column type should override to at least assert 

728 that the value is in fact a SQLAlchemy expression. This is only called 

729 in contexts where true SQLAlchemy expressions are required, such as in 

730 ``ORDER BY`` or ``WHERE`` clauses. 

731 

732 See Also 

733 -------- 

734 :ref:`lsst.daf.relation-sql-logical-columns` 

735 """ 

736 return cast(sqlalchemy.sql.ColumnElement, logical_column) 

737 

738 def convert_flattened_predicate( 

739 self, predicate: Predicate, columns_available: Mapping[ColumnTag, _L] 

740 ) -> list[sqlalchemy.sql.ColumnElement]: 

741 """Flatten all logical AND operators in a `.Predicate` and convert each 

742 to a boolean SQLAlchemy expression. 

743 

744 Parameters 

745 ---------- 

746 predicate : `.Predicate` 

747 Predicate to convert. 

748 columns_available : `~collections.abc.Mapping` 

749 Mapping from `.ColumnTag` to logical column, typically produced by 

750 `extract_mapping` or obtained from `Payload.columns_available`. 

751 

752 Returns 

753 ------- 

754 sql : `list` [ `sqlalchemy.sql.ColumnElement` ] 

755 List of boolean SQLAlchemy expressions to be combined with the 

756 ``AND`` operator. 

757 """ 

758 if (flattened := flatten_logical_and(predicate)) is False: 

759 return [sqlalchemy.sql.literal(False)] 

760 else: 

761 return [self.convert_predicate(p, columns_available) for p in flattened] 

762 

763 def convert_predicate( 

764 self, predicate: Predicate, columns_available: Mapping[ColumnTag, _L] 

765 ) -> sqlalchemy.sql.ColumnElement: 

766 """Convert a `.Predicate` to a SQLAlchemy expression. 

767 

768 Parameters 

769 ---------- 

770 predicate : `.Predicate` 

771 Predicate to convert. 

772 columns_available : `~collections.abc.Mapping` 

773 Mapping from `.ColumnTag` to logical column, typically produced by 

774 `extract_mapping` or obtained from `Payload.columns_available`. 

775 

776 Returns 

777 ------- 

778 sql : `sqlalchemy.sql.ColumnElement` 

779 Boolean SQLAlchemy expression. 

780 """ 

781 match predicate: 

782 case PredicateFunction(name=name, args=args): 

783 sql_args = [self.convert_column_expression(arg, columns_available) for arg in args] 

784 if (function := self.get_function(name)) is not None: 

785 return function(*sql_args) 

786 return getattr(sql_args[0], name)(*sql_args[1:]) 

787 case LogicalAnd(operands=operands): 

788 if not operands: 

789 return sqlalchemy.sql.literal(True) 

790 if len(operands) == 1: 

791 return self.convert_predicate(operands[0], columns_available) 

792 else: 

793 return sqlalchemy.sql.and_( 

794 *[self.convert_predicate(operand, columns_available) for operand in operands] 

795 ) 

796 case LogicalOr(operands=operands): 

797 if not operands: 

798 return sqlalchemy.sql.literal(False) 

799 if len(operands) == 1: 

800 return self.convert_predicate(operands[0], columns_available) 

801 else: 

802 return sqlalchemy.sql.or_( 

803 *[self.convert_predicate(operand, columns_available) for operand in operands] 

804 ) 

805 case LogicalNot(operand=operand): 

806 return sqlalchemy.sql.not_(self.convert_predicate(operand, columns_available)) 

807 case PredicateReference(tag=tag): 

808 return self.expect_column_scalar(columns_available[tag]) 

809 case PredicateLiteral(value=value): 

810 return sqlalchemy.sql.literal(value) 

811 case ColumnInContainer(item=item, container=container): 

812 sql_item = self.expect_column_scalar(self.convert_column_expression(item, columns_available)) 

813 match container: 

814 case ColumnRangeLiteral(value=range(start=start, stop=stop_exclusive, step=step)): 

815 # The convert_column_literal calls below should just 

816 # call sqlalchemy.sql.literal(int), which would also 

817 # happen automatically internal to any of the other 

818 # sqlalchemy function calls, but they get the typing 

819 # right, reflecting the fact that the derived engine is 

820 # supposed to have final say over how we convert 

821 # literals. 

822 stop_inclusive = stop_exclusive - 1 

823 if start == stop_inclusive: 

824 return sql_item == self.convert_column_literal(start) 

825 else: 

826 target = sqlalchemy.sql.between( 

827 sql_item, 

828 self.convert_column_literal(start), 

829 self.convert_column_literal(stop_inclusive), 

830 ) 

831 if step != 1: 

832 return sqlalchemy.sql.and_( 

833 *[ 

834 target, 

835 sql_item % self.convert_column_literal(step) 

836 == self.convert_column_literal(start % step), 

837 ] 

838 ) 

839 else: 

840 return target 

841 case ColumnExpressionSequence(items=items): 

842 return sql_item.in_( 

843 [self.convert_column_expression(item, columns_available) for item in items] 

844 ) 

845 raise AssertionError( 

846 f"matches should be exhaustive and all branches should return; got {predicate!r}." 

847 ) 

848 

849 def convert_sort_term( 

850 self, term: SortTerm, columns_available: Mapping[ColumnTag, _L] 

851 ) -> sqlalchemy.sql.ColumnElement: 

852 """Convert a `.SortTerm` to a SQLAlchemy expression. 

853 

854 Parameters 

855 ---------- 

856 term : `.SortTerm` 

857 Sort term to convert. 

858 columns_available : `~collections.abc.Mapping` 

859 Mapping from `.ColumnTag` to logical column, typically produced by 

860 `extract_mapping` or obtained from `Payload.columns_available`. 

861 

862 Returns 

863 ------- 

864 sql : `sqlalchemy.sql.ColumnElement` 

865 Scalar SQLAlchemy expression. 

866 """ 

867 result = self.expect_column_scalar(self.convert_column_expression(term.expression, columns_available)) 

868 if term.ascending: 

869 return result 

870 else: 

871 return result.desc()