def visit_individual_node(
self, node: IndividualNode[Entity]) -> EquivalenceGraph:
ret: EquivalenceGraph = {}
entity = get_entity(entity_from_node(node))
for relationship in entity.get_all_join_relationships().values():
if relationship.rhs_entity in self.__entities:
for equivalence in relationship.equivalences:
self.__add_relationship(
ret,
entity_from_node(node),
equivalence.left_col,
relationship.rhs_entity,
equivalence.right_col,
)
return ret
def visit_join_clause(self, node: JoinClause[Entity]) -> EquivalenceGraph:
ret: EquivalenceGraph = {}
mapping = node.get_alias_node_map()
for condition in node.keys:
self.__add_relationship(
ret,
entity_from_node(mapping[condition.left.table_alias]),
condition.left.column,
entity_from_node(mapping[condition.right.table_alias]),
condition.right.column,
)
def merge_into_graph(node: JoinNode[Entity]) -> None:
for col, equivalences in node.accept(self).items():
ret[col] = ret.get(col, set()) | equivalences
merge_into_graph(node.left_node)
merge_into_graph(node.right_node)
return ret
def get_equivalent_columns(
join: JoinClause[Entity], ) -> Mapping[QualifiedCol, Set[QualifiedCol]]:
"""
Given a Join, it returns the set of all the semantically equivalent
columns across the entities involved in the join.
This is obtained by generating, through EquivalenceExtractor, the
graph of all equivalences.
We then have the sets of semantically equivalent columns by
generating the list of connected components in the equivalence graph
Each node in a connected component of the equivalence graph is by
definition, semantically equivalent to all the nodes of the same
connected component (directly if there is an edge between two columns
or transitively).
The connected components are returned as a Mapping of nodes to the
set of the nodes in the same component, which means the nodes in the
same connected component
"""
def traverse_graph(node: QualifiedCol,
visited_nodes: Set[QualifiedCol]) -> Set[QualifiedCol]:
"""
Traverse the whole connected component in with a depth first
algorithm starting from the node provided.
"""
if node in visited_nodes:
return visited_nodes
visited_nodes.add(node)
for n in adjacency_sets.get(node, set()):
visited_nodes = traverse_graph(n, visited_nodes)
return visited_nodes
entities_in_join = {
entity_from_node(node)
for node in join.get_alias_node_map().values()
}
adjacency_sets = join.accept(EquivalenceExtractor(entities_in_join))
connected_components: MutableMapping[QualifiedCol, Set[QualifiedCol]] = {}
for node in adjacency_sets:
if node not in connected_components:
component = traverse_graph(node, set())
for node in component:
equivalent_nodes = copy(component)
equivalent_nodes.remove(node)
connected_components[node] = equivalent_nodes
return connected_components
def add_equivalent_conditions(query: CompositeQuery[Entity]) -> None:
"""
Finds conditions in a join query on columns that have a semantic
equivalent in another entity in the join and add the same condition
on the equivalent column.
Example: In a join between events and groupedmessage, if there is
a condition on events.project_id, it would replicate the same
condition on groupedmessage.project_id as this is a semantically
equivalent column.
The goal is to reduce the amount of data that is loaded by clickhouse
for each subquery by adding all the conditions we can to all
subqueries.
Cases we skip:
- top level conditions that include columns in multiple tables.
These cannot be pushed down to subqueries.
- top level conditions containing multiple columns as some may
not have a semantic equivalent. TODO: This can be extended by
supporting conditions that contain multiple column which all
have an equivalent in the same entity
"""
from_clause = query.get_from_clause()
if isinstance(from_clause, CompositeQuery):
add_equivalent_conditions(from_clause)
return
elif isinstance(from_clause, ProcessableQuery):
return
# Now this has to be a join, so we can work with it.
alias_to_entity = {
alias: entity_from_node(node)
for alias, node in from_clause.get_alias_node_map().items()
}
entity_to_alias: MutableMapping[EntityKey, Set[str]] = {}
for alias, entity in alias_to_entity.items():
entity_to_alias.setdefault(entity, set()).add(alias)
column_equivalence = get_equivalent_columns(from_clause)
condition = query.get_condition()
if condition is None:
return
and_components = get_first_level_and_conditions(condition)
conditions_to_add = []
for sub_condition in and_components:
# We duplicate only the top level conditions that reference one
# and only one column that has a semantic equivalent.
# This excludes top level conditions that contains columns from
# multiple entities, and cannot be pushed down to subqueries.
#
# TODO: Address top level conditions that contain multiple
# columns each of which has an equivalent in the same entity.
sole_column = _classify_single_column_condition(
sub_condition, alias_to_entity)
if sole_column is not None:
column_in_condition, table_alias_in_condition = sole_column
for equivalent_table_alias in entity_to_alias[
column_in_condition.entity]:
if equivalent_table_alias != table_alias_in_condition:
# There are multiple occurrences of the entity found.
# Apply the same condition everywhere.
replacer = partial(
_replace_col,
table_alias_in_condition,
column_in_condition.column,
equivalent_table_alias,
column_in_condition.column,
)
conditions_to_add.append(sub_condition.transform(replacer))
for equivalent in column_equivalence.get(column_in_condition, []):
# There are equivalent column on different entities
# in the query. Transform the condition and add it
# to all entities.
equivalent_aliases = entity_to_alias.get(
equivalent.entity, set())
for table_alias in equivalent_aliases:
replacer = partial(
_replace_col,
table_alias_in_condition,
column_in_condition.column,
table_alias,
equivalent.column,
)
conditions_to_add.append(sub_condition.transform(replacer))
query.set_ast_condition(
combine_and_conditions([*and_components, *conditions_to_add]))