Skip to content

Commit

Permalink
Improve querying performance
Browse files Browse the repository at this point in the history
  • Loading branch information
@oliver-ni
oliver-ni committed Dec 2, 2021
1 parent 20164dc commit 250e0db
Show file tree
Hide file tree
Showing 5 changed files with 59 additions and 42 deletions.
15 changes: 14 additions & 1 deletion cogs/auctions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -493,8 +493,21 @@ async def search(self, ctx, **flags):
if flags["page"] < 1:
return await ctx.send("Page must be positive!")

def map_field(field):
if field in {
"_id",
"guild_id",
"user_id",
"current_bid",
"bid_increment",
"bidder_id",
"ends",
}:
return field
return f"pokemon.{field}"

aggregations = await self.bot.get_cog("Pokemon").create_filter(
flags, ctx, order_by=flags["order"]
flags, ctx, order_by=flags["order"], map_field=map_field
)

if aggregations is None:
Expand Down
7 changes: 6 additions & 1 deletion cogs/market.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,8 +67,13 @@ async def market(self, ctx, **flags):
async def search(self, ctx, **flags):
"""Search pokémon from the marketplace."""

def map_field(field):
if field == "_id":
return f"market_data._id"
return field

aggregations = await self.bot.get_cog("Pokemon").create_filter(
flags, ctx, order_by=flags["order"]
flags, ctx, order_by=flags["order"], map_field=map_field
)

if aggregations is None:
Expand Down
9 changes: 2 additions & 7 deletions cogs/mongo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -515,10 +515,7 @@ async def fetch_pokemon_list(self, member: discord.Member, aggregations=[]):
], # TODO Replace
}
},
{"$sort": {"idx": 1}},
{"$project": {"pokemon": "$$ROOT", "idx": "$idx"}},
*aggregations,
{"$replaceRoot": {"newRoot": "$pokemon"}},
],
allowDiskUse=True,
):
Expand All @@ -537,7 +534,6 @@ async def fetch_pokemon_count(self, member: discord.Member, aggregations=[]):
], # TODO Replace
}
},
{"$project": {"pokemon": "$$ROOT"}},
*aggregations,
{"$count": "num_matches"},
],
Expand Down Expand Up @@ -627,16 +623,15 @@ async def fetch_pokemon(self, member: discord.Member, idx: int):
}
},
{"$sort": {"idx": -1}},
{"$project": {"pokemon": "$$ROOT", "idx": "$idx"}},
{"$limit": 1},
],
allowDiskUse=True,
).to_list(None)

if len(result) == 0 or "pokemon" not in result[0]:
if len(result) == 0:
result = None
else:
result = result[0]["pokemon"]
result = result[0]
else:
result = await self.db.pokemon.find_one(
{
Expand Down
34 changes: 19 additions & 15 deletions cogs/pokemon.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -583,7 +583,7 @@ def parse_numerical_flag(self, text):

return ops

async def create_filter(self, flags, ctx, order_by=None):
async def create_filter(self, flags, ctx, order_by=None, map_field=lambda x: x):
aggregations = []

if "mine" in flags and flags["mine"]:
Expand All @@ -597,40 +597,44 @@ async def create_filter(self, flags, ctx, order_by=None):
if x in flags and flags[x]:
rarity += getattr(self.bot.data, f"list_{x}")
if rarity:
aggregations.append({"$match": {"pokemon.species_id": {"$in": rarity}}})
aggregations.append({"$match": {map_field("species_id"): {"$in": rarity}}})

for x in ("alolan", "galarian", "mega", "event"):
if x in flags and flags[x]:
aggregations.append(
{"$match": {"pokemon.species_id": {"$in": getattr(self.bot.data, f"list_{x}")}}}
{
"$match": {
map_field("species_id"): {"$in": getattr(self.bot.data, f"list_{x}")}
}
}
)

if "type" in flags and flags["type"]:
all_species = [i for x in flags["type"] for i in self.bot.data.list_type(x)]
aggregations.append({"$match": {"pokemon.species_id": {"$in": all_species}}})
aggregations.append({"$match": {map_field("species_id"): {"$in": all_species}}})

if "region" in flags and flags["region"]:
all_species = [i for x in flags["region"] for i in self.bot.data.list_region(x)]
aggregations.append({"$match": {"pokemon.species_id": {"$in": all_species}}})
aggregations.append({"$match": {map_field("species_id"): {"$in": all_species}}})

if "favorite" in flags and flags["favorite"]:
aggregations.append({"$match": {"pokemon.favorite": True}})
aggregations.append({"$match": {map_field("favorite"): True}})

if "shiny" in flags and flags["shiny"]:
aggregations.append({"$match": {"pokemon.shiny": True}})
aggregations.append({"$match": {map_field("shiny"): True}})

if "name" in flags and flags["name"] is not None:
all_species = [
i for x in flags["name"] for i in self.bot.data.find_all_matches(" ".join(x))
]

aggregations.append({"$match": {"pokemon.species_id": {"$in": all_species}}})
aggregations.append({"$match": {map_field("species_id"): {"$in": all_species}}})

if "nickname" in flags and flags["nickname"] is not None:
aggregations.append(
{
"$match": {
"pokemon.nickname": {
map_field("nickname"): {
"$regex": "("
+ ")|(".join(" ".join(x) for x in flags["nickname"])
+ ")",
Expand All @@ -641,7 +645,7 @@ async def create_filter(self, flags, ctx, order_by=None):
)

if "embedcolor" in flags and flags["embedcolor"]:
aggregations.append({"$match": {"pokemon.has_color": True}})
aggregations.append({"$match": {map_field("has_color"): True}})

if "ends" in flags and flags["ends"] is not None:
aggregations.append({"$match": {"ends": {"$lt": datetime.utcnow() + flags["ends"]}}})
Expand All @@ -662,15 +666,15 @@ async def create_filter(self, flags, ctx, order_by=None):

if ops[0] == "<":
aggregations.append(
{"$match": {expr: {"$lt": math.ceil(ops[1])}}},
{"$match": {map_field(expr): {"$lt": math.ceil(ops[1])}}},
)
elif ops[0] == "=":
aggregations.append(
{"$match": {expr: {"$eq": round(ops[1])}}},
{"$match": {map_field(expr): {"$eq": round(ops[1])}}},
)
elif ops[0] == ">":
aggregations.append(
{"$match": {expr: {"$gt": math.floor(ops[1])}}},
{"$match": {map_field(expr): {"$gt": math.floor(ops[1])}}},
)

for flag, amt in constants.FILTER_BY_DUPLICATES.items():
Expand All @@ -679,7 +683,7 @@ async def create_filter(self, flags, ctx, order_by=None):

# Processing combinations
combinations = [
{field: iv for field in combo}
{map_field(field): iv for field in combo}
for combo in itertools.combinations(constants.IV_FIELDS, amt)
]
aggregations.append({"$match": {"$or": combinations}})
Expand All @@ -691,7 +695,7 @@ async def create_filter(self, flags, ctx, order_by=None):
else:
asc = -1 if order_by in constants.DEFAULT_DESCENDING else 1

aggregations.append({"$sort": {constants.SORTING_FUNCTIONS[order_by]: asc}})
aggregations.append({"$sort": {map_field(constants.SORTING_FUNCTIONS[order_by]): asc}})

if "skip" in flags and flags["skip"] is not None:
aggregations.append({"$skip": flags["skip"]})
Expand Down
36 changes: 18 additions & 18 deletions helpers/constants.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,10 +129,10 @@

SORTING_FUNCTIONS = {
"number": "idx",
"iv": "pokemon.iv_total",
"level": "pokemon.level",
"pokedex": "pokemon.species_id",
"price": "price",
"iv": "iv_total",
"level": "level",
"pokedex": "species_id",
"price": "market_data.price",
"bid": "current_bid",
"ends": "ends",
"id": "_id",
Expand All @@ -141,14 +141,14 @@
DEFAULT_DESCENDING = {"iv", "level"}

FILTER_BY_NUMERICAL = {
"iv": "pokemon.iv_total",
"level": "pokemon.level",
"hpiv": "pokemon.iv_hp",
"atkiv": "pokemon.iv_atk",
"defiv": "pokemon.iv_defn",
"spatkiv": "pokemon.iv_satk",
"spdefiv": "pokemon.iv_sdef",
"spdiv": "pokemon.iv_spd",
"iv": "iv_total",
"level": "level",
"hpiv": "iv_hp",
"atkiv": "iv_atk",
"defiv": "iv_defn",
"spatkiv": "iv_satk",
"spdefiv": "iv_sdef",
"spdiv": "iv_spd",
}

FILTER_BY_DUPLICATES = {
Expand All @@ -159,12 +159,12 @@
}

IV_FIELDS = [
"pokemon.iv_hp",
"pokemon.iv_atk",
"pokemon.iv_defn",
"pokemon.iv_satk",
"pokemon.iv_sdef",
"pokemon.iv_spd",
"iv_hp",
"iv_atk",
"iv_defn",
"iv_satk",
"iv_sdef",
"iv_spd",
]

NATURE_MULTIPLIERS = {
Expand Down

0 comments on commit 250e0db

Please sign in to comment.