Skip to content

Commit

Permalink
Account for storage costs
Browse files Browse the repository at this point in the history
  • Loading branch information
@geoffxy
geoffxy committed Nov 13, 2023
1 parent ef82fad commit a5852f2
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 25 deletions.
35 changes: 22 additions & 13 deletions src/brad/planner/beam/query_based_candidate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -254,8 +254,20 @@ def add_query(
planner_config=ctx.planner_config,
)

# Table movement costs that this query imposes.
# Storage costs and table movement that this query imposes.
for name, next_placement in table_diffs:
# If we added a table to Athena or Aurora, we need to take into
# account its storage costs.
if (next_placement & EngineBitmapValues[Engine.Athena]) != 0:
# We added the table to Athena.
self.storage_cost += compute_single_athena_table_cost(name, ctx)

if (next_placement & EngineBitmapValues[Engine.Aurora]) != 0:
# Added table to Aurora.
# You only pay for 1 copy of the table on Aurora, regardless of
# how many read replicas you have.
self.storage_cost += compute_single_aurora_table_cost(name, ctx)

curr = ctx.current_blueprint.table_locations_bitmap()[name]
if ((~curr) & next_placement) == 0:
# This table was already present on the engine.
Expand All @@ -267,18 +279,6 @@ def add_query(
self.table_movement_trans_cost += result.movement_cost
self.table_movement_trans_time_s += result.movement_time_s

# If we added a table to Athena or Aurora, we need to take into
# account its storage costs.
if (((~curr) & next_placement) & (EngineBitmapValues[Engine.Athena])) != 0:
# We added the table to Athena.
self.storage_cost += compute_single_athena_table_cost(name, ctx)

if (((~curr) & next_placement) & (EngineBitmapValues[Engine.Aurora])) != 0:
# Added table to Aurora.
# You only pay for 1 copy of the table on Aurora, regardless of
# how many read replicas you have.
self.storage_cost += compute_single_aurora_table_cost(name, ctx)

# Adding a new query can affect the feasibility of the provisioning.
self.feasibility = BlueprintFeasibility.Unchecked
self.explored_provisionings = False
Expand Down Expand Up @@ -365,6 +365,7 @@ def add_query_last_step(

def add_transactional_tables(self, ctx: ScoringContext) -> None:
referenced_tables = set()
newly_added = set()

# Make sure that tables referenced in transactions are present on
# Aurora.
Expand All @@ -373,9 +374,17 @@ def add_transactional_tables(self, ctx: ScoringContext) -> None:
if tbl not in self.table_placements:
# This is a CTE.
continue
orig = self.table_placements[tbl]
self.table_placements[tbl] |= EngineBitmapValues[Engine.Aurora]
referenced_tables.add(tbl)

if ((~orig) & self.table_placements[tbl]) != 0:
newly_added.add(tbl)

# Account for storage costs (Aurora only charges for 1 copy).
for tbl in newly_added:
self.storage_cost += compute_single_aurora_table_cost(tbl, ctx)

# Update the table movement score if needed.
for tbl in referenced_tables:
cur = ctx.current_blueprint.table_locations_bitmap()[tbl]
Expand Down
33 changes: 21 additions & 12 deletions src/brad/planner/beam/table_based_candidate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,18 @@ def add_placement(
changed_tables.append(table_name)
changed = True

# If we added the table to Athena or Aurora, we need to take into
# account its storage costs.
if (((~cur) & nxt) & (EngineBitmapValues[Engine.Athena])) != 0:
# We added the table to Athena.
self.storage_cost += compute_single_athena_table_cost(table_name, ctx)

if (((~cur) & nxt) & (EngineBitmapValues[Engine.Aurora])) != 0:
# Added table to Aurora.
# You only pay for 1 copy of the table on Aurora, regardless of
# how many read replicas you have.
self.storage_cost += compute_single_aurora_table_cost(table_name, ctx)

# Update movement scoring.
self._update_movement_score(changed_tables, ctx)

Expand Down Expand Up @@ -295,6 +307,7 @@ async def _route_queries_compute_scan_stats(

def add_transactional_tables(self, ctx: ScoringContext) -> None:
referenced_tables = set()
newly_added = set()

# Make sure that tables referenced in transactions are present on
# Aurora.
Expand All @@ -303,9 +316,17 @@ def add_transactional_tables(self, ctx: ScoringContext) -> None:
if tbl not in self.table_placements:
# This is a CTE.
continue
orig = self.table_placements[tbl]
self.table_placements[tbl] |= EngineBitmapValues[Engine.Aurora]
referenced_tables.add(tbl)

if ((~orig) & self.table_placements[tbl]) != 0:
newly_added.add(tbl)

for tbl in newly_added:
# Aurora only charges for 1 copy of the data.
self.storage_cost += compute_single_aurora_table_cost(tbl, ctx)

# If we made a change to the table placement, see if it corresponds to a
# table score change.
self._update_movement_score(referenced_tables, ctx)
Expand All @@ -323,18 +344,6 @@ def _update_movement_score(
self.table_movement_trans_cost += result.movement_cost
self.table_movement_trans_time_s += result.movement_time_s

# If we added the table to Athena or Aurora, we need to take into
# account its storage costs.
if (((~cur) & nxt) & (EngineBitmapValues[Engine.Athena])) != 0:
# We added the table to Athena.
self.storage_cost += compute_single_athena_table_cost(tbl, ctx)

if (((~cur) & nxt) & (EngineBitmapValues[Engine.Aurora])) != 0:
# Added table to Aurora.
# You only pay for 1 copy of the table on Aurora, regardless of
# how many read replicas you have.
self.storage_cost += compute_single_aurora_table_cost(tbl, ctx)

def try_to_make_feasible_if_needed(self, ctx: ScoringContext) -> None:
"""
Checks if this blueprint is already feasible, and if so, does nothing
Expand Down

0 comments on commit a5852f2

Please sign in to comment.