def test_should_select_cheap_plan(self):
# mocking the cost model
def side_effect_func(value):
if value is grp_expr1:
return 1
elif value is grp_expr2:
return 2
cm = CostModel()
cm.calculate_cost = MagicMock(side_effect=side_effect_func)
opt_cxt = OptimizerContext(cm)
grp_expr1 = GroupExpression(MagicMock())
grp_expr1.opr.is_logical = lambda: False
grp_expr2 = GroupExpression(MagicMock())
grp_expr2.opr.is_logical = lambda: False
opt_cxt.memo.add_group_expr(grp_expr1)
opt_cxt.memo.add_group_expr(grp_expr2, grp_expr1.group_id)
grp = opt_cxt.memo.get_group_by_id(grp_expr1.group_id)
opt_cxt.task_stack.push(OptimizeGroup(grp, opt_cxt))
self.execute_task_stack(opt_cxt.task_stack)
plan = PlanGenerator().build_optimal_physical_plan(
grp_expr1.group_id, opt_cxt)
self.assertEqual(plan, grp_expr1.opr)
self.assertEqual(grp.get_best_expr_cost(PropertyType.DEFAULT), 1)
def _xform_opr_to_group_expr(self, opr: Operator) -> GroupExpression:
"""
Note: Internal function Generate a group expressions from a
logical operator tree. Caller is responsible for assigning
the group to the returned GroupExpression.
"""
# Go through the children first.
child_ids = []
for child_opr in opr.children:
if isinstance(child_opr, Dummy):
child_ids.append(child_opr.group_id)
else:
child_expr = self._xform_opr_to_group_expr(opr=child_opr)
# add the expr to memo
# handles duplicates and assigns group id
memo_expr = self.memo.add_group_expr(child_expr)
child_ids.append(memo_expr.group_id)
# Group Expression only needs the operator content. Remove
# the opr children as parent-child relationship is captured
# by the group expressions.
# Hack: Shallow copy all the content except children and
# manually clearing the children as we don't need the
# dependency. Better fix is to rewrite the operator class to
# support exposing only the content
opr_copy = copy.copy(opr)
opr_copy.clear_children()
expr = GroupExpression(opr=opr_copy, children=child_ids)
return expr
def add_group_expr(self,
expr: GroupExpression,
group_id: int = UNDEFINED_GROUP_ID) -> GroupExpression:
"""
Add an expression into the memo.
If expr exists, we return it.
If group_id is not specified, creates a new group
Otherwise, inserts the expr into specified group.
"""
# check duplicate expression
duplicate_expr = self.find_duplicate_expr(expr)
if duplicate_expr is not None:
return duplicate_expr
# did not find a dulpicate expression
expr.group_id = group_id
if expr.group_id == UNDEFINED_GROUP_ID:
self._create_new_group(expr)
else:
self._insert_expr(expr, group_id)
assert expr.group_id is not UNDEFINED_GROUP_ID, '''Expr
should have a
valid group
id'''
return expr
def test_add_group_expr_cost(self):
grp = Group(0)
prpty = Property(PropertyType(1))
grp_expr1 = GroupExpression(MagicMock(), 1)
grp_expr1.opr.is_logical = lambda: True
grp_expr2 = GroupExpression(MagicMock())
grp_expr2.opr.is_logical = lambda: False
grp.add_expr(grp_expr1)
grp.add_expr_cost(grp_expr1, prpty, 1)
grp.add_expr(grp_expr2)
grp.add_expr_cost(grp_expr2, prpty, 0)
self.assertEqual(grp.get_best_expr(prpty), grp_expr2)
self.assertEqual(grp.get_best_expr_cost(prpty), 0)
def test_add_group_expr_with_unmatched_group_id(self):
grp = Group(0)
grp_expr1 = GroupExpression(MagicMock(), 1)
grp_expr1.opr.is_logical = lambda: True
grp.add_expr(grp_expr1)
self.assertEquals(len(grp.logical_exprs), 0)
self.assertEquals(len(grp.physical_exprs), 0)
def add_expr(self, expr: GroupExpression):
if expr.group_id == UNDEFINED_GROUP_ID:
expr.group_id = self.group_id
if expr.group_id != self.group_id:
logger.error('Expected group id {}, found {}'.format(
self.group_id, expr.group_id))
return
if expr.opr.is_logical():
self._add_logical_expr(expr)
else:
self._add_physical_expr(expr)
def test_simple_add_group_expr(self):
grp = Group(0)
grp_expr1 = GroupExpression(MagicMock())
grp_expr1.opr.is_logical = lambda: True
grp_expr2 = GroupExpression(MagicMock())
grp_expr2.opr.is_logical = lambda: False
grp_expr3 = GroupExpression(MagicMock(), 0)
grp_expr3.opr.is_logical = lambda: True
grp.add_expr(grp_expr1)
self.assertEquals(len(grp.logical_exprs), 1)
grp.add_expr(grp_expr2)
self.assertEquals(len(grp.logical_exprs), 1)
self.assertEquals(len(grp.physical_exprs), 1)
grp.add_expr(grp_expr3)
self.assertEquals(len(grp.logical_exprs), 2)
self.assertEquals(len(grp.physical_exprs), 1)
def test_should_select_cheap_plan_with_tree(self):
# mocking the cost model
def side_effect_func(value):
cost = dict({
grp_expr00: 1,
grp_expr01: 2,
grp_expr10: 4,
grp_expr11: 3,
grp_expr20: 5
})
return cost[value]
cm = CostModel()
cm.calculate_cost = MagicMock(side_effect=side_effect_func)
opt_cxt = OptimizerContext(cm)
# group 0
grp_expr00 = GroupExpression(Operator(MagicMock()))
grp_expr00.opr.is_logical = lambda: False
grp_expr01 = GroupExpression(Operator(MagicMock()))
grp_expr01.opr.is_logical = lambda: False
opt_cxt.memo.add_group_expr(grp_expr00)
opt_cxt.memo.add_group_expr(grp_expr01, grp_expr00.group_id)
# group 1
grp_expr10 = GroupExpression(Operator(MagicMock()))
grp_expr10.opr.is_logical = lambda: False
opt_cxt.memo.add_group_expr(grp_expr10)
grp_expr11 = GroupExpression(Operator(MagicMock()))
grp_expr11.opr.is_logical = lambda: False
opt_cxt.memo.add_group_expr(grp_expr11, grp_expr10.group_id)
# group 2
grp_expr20 = GroupExpression(Operator(MagicMock()))
grp_expr20.opr.is_logical = lambda: False
opt_cxt.memo.add_group_expr(grp_expr20)
grp = opt_cxt.memo.get_group_by_id(grp_expr20.group_id)
# tree: 2->1->0
grp_expr10.children = [grp_expr01.group_id]
grp_expr11.children = [grp_expr01.group_id]
grp_expr20.children = [grp_expr10.group_id]
opt_cxt.task_stack.push(OptimizeGroup(grp, opt_cxt))
self.execute_task_stack(opt_cxt.task_stack)
plan = PlanGenerator().build_optimal_physical_plan(
grp_expr20.group_id, opt_cxt)
subplan = copy(grp_expr11.opr)
subplan.children = [copy(grp_expr01.opr)]
expected_plan = copy(grp_expr20.opr)
expected_plan.children = [subplan]
self.assertEqual(plan, expected_plan)
self.assertEqual(grp.get_best_expr_cost(PropertyType.DEFAULT), 9)