pytrilogy 0.0.1.109__py3-none-any.whl → 0.0.1.111__py3-none-any.whl

trilogy/core/optimization.py

@@ -4,24 +4,32 @@ from trilogy.core.models import (
     PersistStatement,
     Datasource,
     MultiSelectStatement,
+    Conditional,
+    BooleanOperator,
 )
 from trilogy.core.enums import PurposeLineage
-from trilogy.constants import logger
+from trilogy.constants import logger, CONFIG
 from abc import ABC
 
 
+REGISTERED_RULES: list["OptimizationRule"] = []
+
+
 class OptimizationRule(ABC):
 
-    def optimize(self, cte: CTE) -> bool:
+    def optimize(self, cte: CTE, inverse_map: dict[str, list[CTE]]) -> bool:
         raise NotImplementedError
 
     def log(self, message: str):
         logger.info(f"[Optimization][{self.__class__.__name__}] {message}")
 
+    def debug(self, message: str):
+        logger.debug(f"[Optimization][{self.__class__.__name__}] {message}")
+
 
 class InlineDatasource(OptimizationRule):
 
-    def optimize(self, cte: CTE) -> bool:
+    def optimize(self, cte: CTE, inverse_map: dict[str, list[CTE]]) -> bool:
         if not cte.parent_ctes:
             return False
 
@@ -30,6 +38,7 @@ class InlineDatasource(OptimizationRule):
             f"Checking {cte.name} for consolidating inline tables with {len(cte.parent_ctes)} parents"
         )
         to_inline: list[CTE] = []
+        force_group = False
         for parent_cte in cte.parent_ctes:
             if not parent_cte.is_root_datasource:
                 self.log(f"parent {parent_cte.name} is not root")
@@ -47,23 +56,114 @@ class InlineDatasource(OptimizationRule):
                 continue
             root_outputs = {x.address for x in root.output_concepts}
             cte_outputs = {x.address for x in parent_cte.output_columns}
+            grain_components = {x.address for x in root.grain.components}
             if not cte_outputs.issubset(root_outputs):
                 self.log(f"Not all {parent_cte.name} outputs are found on datasource")
                 continue
-
+            if not grain_components.issubset(cte_outputs):
+                self.log("Not all datasource components in cte outputs, forcing group")
+                force_group = True
             to_inline.append(parent_cte)
 
         for replaceable in to_inline:
             self.log(f"Inlining parent {replaceable.name}")
-            cte.inline_parent_datasource(replaceable)
+            cte.inline_parent_datasource(replaceable, force_group=force_group)
 
         return optimized
 
 
-REGISTERED_RULES: list[OptimizationRule] = [InlineDatasource()]
+# This will be used in the future for more complex condition decomposition
+def decompose_condition(conditional: Conditional):
+    chunks = []
+    if conditional.operator == BooleanOperator.AND:
+        for val in [conditional.left, conditional.right]:
+            if isinstance(val, Conditional):
+                chunks.extend(decompose_condition(val))
+            else:
+                chunks.append(val)
+    else:
+        chunks.append(conditional)
+    return chunks
+
+
+def is_child_of(a, comparison):
+    if isinstance(comparison, Conditional):
+        return (
+            is_child_of(a, comparison.left) or is_child_of(a, comparison.right)
+        ) and comparison.operator == BooleanOperator.AND
+    return comparison == a
+
+
+class PredicatePushdown(OptimizationRule):
 
+    def optimize(self, cte: CTE, inverse_map: dict[str, list[CTE]]) -> bool:
 
-def filter_irrelevant_ctes(input: list[CTE], root_cte: CTE):
+        if not cte.parent_ctes:
+            self.debug(f"No parent CTEs for {cte.name}")
+
+            return False
+
+        optimized = False
+        if not cte.condition:
+            self.debug(f"No CTE condition for {cte.name}")
+            return False
+        self.log(
+            f"Checking {cte.name} for predicate pushdown with {len(cte.parent_ctes)} parents"
+        )
+        if isinstance(cte.condition, Conditional):
+            candidates = cte.condition.decompose()
+        else:
+            candidates = [cte.condition]
+        logger.info(f"Have {len(candidates)} candidates to try to push down")
+        for candidate in candidates:
+            conditions = {x.address for x in candidate.concept_arguments}
+            for parent_cte in cte.parent_ctes:
+                materialized = {k for k, v in parent_cte.source_map.items() if v != []}
+                if conditions.issubset(materialized):
+                    if all(
+                        [
+                            is_child_of(candidate, child.condition)
+                            for child in inverse_map[parent_cte.name]
+                        ]
+                    ):
+                        self.log(
+                            f"All concepts are found on {parent_cte.name} and all it's children include same filter; pushing up filter"
+                        )
+                        if parent_cte.condition:
+                            parent_cte.condition = Conditional(
+                                left=parent_cte.condition,
+                                operator=BooleanOperator.AND,
+                                right=candidate,
+                            )
+                        else:
+                            parent_cte.condition = candidate
+                        optimized = True
+                else:
+                    logger.info("conditions not subset of parent materialized")
+
+        if all(
+            [
+                is_child_of(cte.condition, parent_cte.condition)
+                for parent_cte in cte.parent_ctes
+            ]
+        ):
+            self.log("All parents have same filter, removing filter")
+            cte.condition = None
+            optimized = True
+
+        return optimized
+
+
+if CONFIG.optimizations.datasource_inlining:
+    REGISTERED_RULES.append(InlineDatasource())
+if CONFIG.optimizations.predicate_pushdown:
+    REGISTERED_RULES.append(PredicatePushdown())
+
+
+def filter_irrelevant_ctes(
+    input: list[CTE],
+    root_cte: CTE,
+):
     relevant_ctes = set()
 
     def recurse(cte: CTE):
@@ -75,6 +175,16 @@ def filter_irrelevant_ctes(input: list[CTE], root_cte: CTE):
     return [cte for cte in input if cte.name in relevant_ctes]
 
 
+def gen_inverse_map(input: list[CTE]) -> dict[str, list[CTE]]:
+    inverse_map: dict[str, list[CTE]] = {}
+    for cte in input:
+        for parent in cte.parent_ctes:
+            if parent.name not in inverse_map:
+                inverse_map[parent.name] = []
+            inverse_map[parent.name].append(cte)
+    return inverse_map
+
+
 def is_direct_return_eligible(
     cte: CTE, select: SelectStatement | PersistStatement | MultiSelectStatement
 ) -> bool:
@@ -94,6 +204,8 @@ def is_direct_return_eligible(
     for x in derived_concepts:
         if x.derivation == PurposeLineage.WINDOW:
             return False
+        if x.derivation == PurposeLineage.UNNEST:
+            return False
         if x.derivation == PurposeLineage.AGGREGATE:
             if x.address in conditions:
                 return False
@@ -126,15 +238,25 @@ def optimize_ctes(
         actions_taken = False
         for rule in REGISTERED_RULES:
             for cte in input:
-                actions_taken = rule.optimize(cte)
+                inverse_map = gen_inverse_map(input)
+                actions_taken = rule.optimize(cte, inverse_map)
         complete = not actions_taken
 
-    if is_direct_return_eligible(root_cte, select):
+    if CONFIG.optimizations.direct_return and is_direct_return_eligible(
+        root_cte, select
+    ):
         root_cte.order_by = select.order_by
         root_cte.limit = select.limit
-        root_cte.condition = (
-            select.where_clause.conditional if select.where_clause else None
-        )
+        if select.where_clause:
+
+            if root_cte.condition:
+                root_cte.condition = Conditional(
+                    left=root_cte.condition,
+                    operator=BooleanOperator.AND,
+                    right=select.where_clause.conditional,
+                )
+            else:
+                root_cte.condition = select.where_clause.conditional
         root_cte.requires_nesting = False
         sort_select_output(root_cte, select)
 

trilogy/core/processing/concept_strategies_v3.py

@@ -1,7 +1,6 @@
 from collections import defaultdict
 from typing import List, Optional, Callable
 
-
 from trilogy.constants import logger
 from trilogy.core.enums import PurposeLineage, Granularity, FunctionType
 from trilogy.core.env_processor import generate_graph
@@ -278,9 +277,10 @@ def generate_node(
         logger.info(
             f"{depth_to_prefix(depth)}{LOGGER_PREFIX} for {concept.address}, generating multiselect node with optional {[x.address for x in local_optional]}"
         )
-        return gen_concept_merge_node(
+        node = gen_concept_merge_node(
             concept, local_optional, environment, g, depth + 1, source_concepts, history
         )
+        return node
     elif concept.derivation == PurposeLineage.CONSTANT:
         logger.info(
             f"{depth_to_prefix(depth)}{LOGGER_PREFIX} for {concept.address}, generating constant node"
@@ -340,13 +340,15 @@ def validate_stack(
     stack: List[StrategyNode],
     concepts: List[Concept],
     accept_partial: bool = False,
-) -> tuple[ValidationResult, set[str], set[str], set[str]]:
+) -> tuple[ValidationResult, set[str], set[str], set[str], set[str]]:
     found_map = defaultdict(set)
     found_addresses: set[str] = set()
     non_partial_addresses: set[str] = set()
     partial_addresses: set[str] = set()
+    virtual_addresses: set[str] = set()
     for node in stack:
-        for concept in node.resolve().output_concepts:
+        resolved = node.resolve()
+        for concept in resolved.output_concepts:
             found_map[str(node)].add(concept)
             if concept not in node.partial_concepts:
                 found_addresses.add(concept.address)
@@ -354,11 +356,20 @@ def validate_stack(
                 # remove it from our partial tracking
                 if concept.address in partial_addresses:
                     partial_addresses.remove(concept.address)
+                if concept.address in virtual_addresses:
+                    virtual_addresses.remove(concept.address)
             if concept in node.partial_concepts:
+                if concept.address in non_partial_addresses:
+                    continue
                 partial_addresses.add(concept.address)
                 if accept_partial:
                     found_addresses.add(concept.address)
                     found_map[str(node)].add(concept)
+        for concept in node.virtual_output_concepts:
+            if concept.address in non_partial_addresses:
+                continue
+            found_addresses.add(concept.address)
+            virtual_addresses.add(concept.address)
     # zip in those we know we found
     if not all([c.address in found_addresses for c in concepts]):
         return (
@@ -366,12 +377,25 @@ def validate_stack(
             found_addresses,
             {c.address for c in concepts if c.address not in found_addresses},
             partial_addresses,
+            virtual_addresses,
         )
     graph_count, graphs = get_disconnected_components(found_map)
     if graph_count in (0, 1):
-        return ValidationResult.COMPLETE, found_addresses, set(), partial_addresses
+        return (
+            ValidationResult.COMPLETE,
+            found_addresses,
+            set(),
+            partial_addresses,
+            virtual_addresses,
+        )
     # if we have too many subgraphs, we need to keep searching
-    return ValidationResult.DISCONNECTED, found_addresses, set(), partial_addresses
+    return (
+        ValidationResult.DISCONNECTED,
+        found_addresses,
+        set(),
+        partial_addresses,
+        virtual_addresses,
+    )
 
 
 def depth_to_prefix(depth: int) -> str:
@@ -404,7 +428,10 @@ def search_concepts(
         accept_partial=accept_partial,
         history=history,
     )
-    history.search_to_history(mandatory_list, accept_partial, result)
+    # a node may be mutated after be cached; always store a copy
+    history.search_to_history(
+        mandatory_list, accept_partial, result.copy() if result else None
+    )
     return result
 
 
@@ -472,13 +499,13 @@ def _search_concepts(
                     skip.add(priority_concept.address)
                 break
         attempted.add(priority_concept.address)
-        complete, found, missing, partial = validate_stack(
+        complete, found, missing, partial, virtual = validate_stack(
             stack, mandatory_list, accept_partial
         )
 
         logger.info(
             f"{depth_to_prefix(depth)}{LOGGER_PREFIX} finished concept loop for {priority_concept} flag for accepting partial addresses is "
-            f" {accept_partial} (complete: {complete}), have {found} from {[n for n in stack]} (missing {missing} partial {partial}), attempted {attempted}"
+            f" {accept_partial} (complete: {complete}), have {found} from {[n for n in stack]} (missing {missing} partial {partial} virtual {virtual}), attempted {attempted}"
         )
         # early exit if we have a complete stack with one node
         # we can only early exit if we have a complete stack
@@ -489,7 +516,7 @@ def _search_concepts(
             break
 
     logger.info(
-        f"{depth_to_prefix(depth)}{LOGGER_PREFIX} finished sourcing loop (complete: {complete}), have {found} from {[n for n in stack]} (missing {all_mandatory - found}), attempted {attempted}"
+        f"{depth_to_prefix(depth)}{LOGGER_PREFIX} finished sourcing loop (complete: {complete}), have {found} from {[n for n in stack]} (missing {all_mandatory - found}), attempted {attempted}, virtual {virtual}"
     )
     if complete == ValidationResult.COMPLETE:
         all_partial = [
@@ -503,24 +530,22 @@ def _search_concepts(
                 ]
             )
         ]
+        non_virtual = [c for c in mandatory_list if c.address not in virtual]
         if len(stack) == 1:
+            output = stack[0]
             logger.info(
-                f"{depth_to_prefix(depth)}{LOGGER_PREFIX} Source stack has single node, returning just that node"
+                f"{depth_to_prefix(depth)}{LOGGER_PREFIX} Source stack has single node, returning that {type(output)}"
             )
-            return stack[0]
+            return output
 
         output = MergeNode(
-            input_concepts=mandatory_list,
-            output_concepts=mandatory_list,
+            input_concepts=non_virtual,
+            output_concepts=non_virtual,
             environment=environment,
             g=g,
             parents=stack,
             depth=depth,
             partial_concepts=all_partial,
-            # always hide merge concepts
-            hidden_concepts=[
-                x for x in mandatory_list if x.derivation == PurposeLineage.MERGE
-            ],
         )
 
         # ensure we can resolve our final merge
@@ -573,7 +598,7 @@ def _search_concepts(
             )
             return partial_search
     logger.error(
-        f"{depth_to_prefix(depth)}{LOGGER_PREFIX} Could not resolve concepts {[c.address for c in mandatory_list]}, network outcome was {complete}, missing {all_mandatory - found}"
+        f"{depth_to_prefix(depth)}{LOGGER_PREFIX} Could not resolve concepts {[c.address for c in mandatory_list]}, network outcome was {complete}, missing {all_mandatory - found},"
     )
     return None
 

trilogy/core/processing/node_generators/basic_node.py

@@ -8,6 +8,7 @@ from trilogy.core.processing.nodes import StrategyNode, History, MergeNode
 from trilogy.core.processing.node_generators.common import (
     resolve_function_parent_concepts,
 )
+from trilogy.utility import unique
 from trilogy.constants import logger
 
 LOGGER_PREFIX = "[GEN_BASIC_NODE]"
@@ -37,6 +38,7 @@ def gen_basic_node(
         attempts.append((parent_concepts + local_optional, local_optional + [concept]))
 
     for attempt, output in reversed(attempts):
+        attempt = unique(attempt, "address")
         parent_node = source_concepts(
             mandatory_list=attempt,
             environment=environment,

trilogy/core/processing/node_generators/common.py

@@ -56,7 +56,9 @@ def resolve_filter_parent_concepts(
     base_existence = []
     base_rows = [direct_parent]
     base_rows += concept.lineage.where.row_arguments
-    base_existence += concept.lineage.where.existence_arguments
+    # TODO: pass tuple groups through
+    for ctuple in concept.lineage.where.existence_arguments:
+        base_existence += list(ctuple)
     if direct_parent.grain:
         base_rows += direct_parent.grain.components_copy
     if (

trilogy/core/processing/node_generators/concept_merge_node.py

@@ -56,6 +56,7 @@ def gen_concept_merge_node(
 
     # get additional concepts that should be merged across the environments
     additional_merge: List[Concept] = [*lineage.concepts]
+    target_namespaces = set(x.namespace for x in [concept] + local_optional)
     for x in local_optional:
         if x.address in environment.merged_concepts:
             ms = environment.merged_concepts[x.address].lineage
@@ -64,6 +65,8 @@ def gen_concept_merge_node(
 
     for select in lineage.concepts:
         # if it's a merge concept, filter it out of the optional
+        if select.namespace not in target_namespaces:
+            continue
         sub_optional = [
             x
             for x in local_optional
@@ -76,6 +79,9 @@ def gen_concept_merge_node(
         ]
         sub_optional += sub_additional_merge
         final: List[Concept] = unique([select] + sub_optional, "address")
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} generating concept merge parent node with {[x.address for x in final]}"
+        )
         snode: StrategyNode = source_concepts(
             mandatory_list=final,
             environment=environment,
@@ -111,17 +117,18 @@ def gen_concept_merge_node(
 
     additional_relevant = [x for x in outputs if x.address in enrichment]
     final_outputs = outputs + additional_relevant + [concept]
+    virtual_outputs = [x for x in final_outputs if x.derivation == PurposeLineage.MERGE]
     node = MergeNode(
         input_concepts=[x for y in base_parents for x in y.output_concepts],
-        output_concepts=[x for x in final_outputs],
-        hidden_concepts=[
-            x for x in final_outputs if x.derivation == PurposeLineage.MERGE
+        output_concepts=[
+            x for x in final_outputs if x.derivation != PurposeLineage.MERGE
         ],
         environment=environment,
         g=g,
         depth=depth,
         parents=base_parents,
         node_joins=node_joins,
+        virtual_output_concepts=virtual_outputs,
     )
 
     qds = node.rebuild_cache()
@@ -149,9 +156,17 @@ def gen_concept_merge_node(
             f"{padding(depth)}{LOGGER_PREFIX} all enriched concepts returned from base merge concept node; exiting early"
         )
         return node
+    missing = [
+        x
+        for x in local_optional
+        if x.address not in [y.address for y in node.output_concepts]
+    ]
+    logger.info(
+        f"{padding(depth)}{LOGGER_PREFIX} generating merge concept enrichment node for missing {[x.address for x in missing]}"
+    )
     enrich_node: MergeNode = source_concepts(  # this fetches the parent + join keys
         # to then connect to the rest of the query
-        mandatory_list=additional_relevant + local_optional,
+        mandatory_list=additional_relevant + missing,
         environment=environment,
         g=g,
         depth=depth + 1,
@@ -159,7 +174,7 @@ def gen_concept_merge_node(
     )
     if not enrich_node:
         logger.info(
-            f"{padding(depth)}{LOGGER_PREFIX} Cannot generate merge concept enrichment node for {concept} with optional {local_optional}, returning just merge concept"
+            f"{padding(depth)}{LOGGER_PREFIX} Cannot generate merge concept enrichment node for {concept.address} with optional {[x.address for x in local_optional]}, returning just merge concept"
         )
         return node
 
@@ -170,12 +185,12 @@ def gen_concept_merge_node(
     return MergeNode(
         input_concepts=enrich_node.output_concepts + node.output_concepts,
         # also filter out the
-        output_concepts=node.output_concepts + local_optional,
-        hidden_concepts=[
+        output_concepts=[
             x
             for x in node.output_concepts + local_optional
-            if x.derivation == PurposeLineage.MERGE
+            if x.derivation != PurposeLineage.MERGE
         ],
+        hidden_concepts=[],
         environment=environment,
         g=g,
         depth=depth,
@@ -195,4 +210,5 @@ def gen_concept_merge_node(
             )
         ],
         partial_concepts=node.partial_concepts,
+        virtual_output_concepts=virtual_outputs,
     )

trilogy/core/processing/node_generators/filter_node.py

@@ -2,11 +2,14 @@ from typing import List
 
 
 from trilogy.core.enums import JoinType
-from trilogy.core.models import (
-    Concept,
-    Environment,
+from trilogy.core.models import Concept, Environment, FilterItem
+from trilogy.core.processing.nodes import (
+    FilterNode,
+    MergeNode,
+    NodeJoin,
+    History,
+    StrategyNode,
 )
-from trilogy.core.processing.nodes import FilterNode, MergeNode, NodeJoin, History
 from trilogy.core.processing.node_generators.common import (
     resolve_filter_parent_concepts,
 )
@@ -25,16 +28,19 @@ def gen_filter_node(
     depth: int,
     source_concepts,
     history: History | None = None,
-) -> MergeNode | FilterNode | None:
+) -> StrategyNode | None:
     immediate_parent, parent_row_concepts, parent_existence_concepts = (
         resolve_filter_parent_concepts(concept)
     )
+    if not isinstance(concept.lineage, FilterItem):
+        raise SyntaxError('Filter node must have a lineage of type "FilterItem"')
+    where = concept.lineage.where
 
     logger.info(
         f"{padding(depth)}{LOGGER_PREFIX} fetching filter node row parents {[x.address for x in parent_row_concepts]}"
     )
     core_parents = []
-    parent = source_concepts(
+    parent: StrategyNode = source_concepts(
         mandatory_list=parent_row_concepts,
         environment=environment,
         g=g,
@@ -43,7 +49,28 @@ def gen_filter_node(
     )
 
     if not parent:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} filter node row parents {[x.address for x in parent_row_concepts]} could not be found"
+        )
         return None
+
+    if not local_optional and not parent_existence_concepts:
+        optimized_pushdown = True
+    else:
+        optimized_pushdown = False
+
+    if optimized_pushdown:
+        if parent.conditions:
+            parent.conditions = parent.conditions + where.conditional
+        else:
+            parent.conditions = where.conditional
+        parent.output_concepts = [concept]
+        parent.rebuild_cache()
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} returning optimized filter node with pushdown to parent with condition {where.conditional}"
+        )
+        return parent
+
     core_parents.append(parent)
     if parent_existence_concepts:
         logger.info(
@@ -57,6 +84,9 @@ def gen_filter_node(
             history=history,
         )
         if not parent_existence:
+            logger.info(
+                f"{padding(depth)}{LOGGER_PREFIX} filter existence node parents could not be found"
+            )
             return None
         core_parents.append(parent_existence)