pytrilogy 0.0.3.92__py3-none-any.whl → 0.0.3.94__py3-none-any.whl

trilogy/core/models/build_environment.py

@@ -90,9 +90,7 @@ def validate_concepts(v) -> BuildEnvironmentConceptDict:
     if isinstance(v, BuildEnvironmentConceptDict):
         return v
     elif isinstance(v, dict):
-        return BuildEnvironmentConceptDict(
-            **{x: BuildConcept.model_validate(y) for x, y in v.items()}
-        )
+        return BuildEnvironmentConceptDict(**{x: y for x, y in v.items()})
     raise ValueError
 
 
@@ -100,9 +98,7 @@ def validate_datasources(v) -> BuildEnvironmentDatasourceDict:
     if isinstance(v, BuildEnvironmentDatasourceDict):
         return v
     elif isinstance(v, dict):
-        return BuildEnvironmentDatasourceDict(
-            **{x: BuildDatasource.model_validate(y) for x, y in v.items()}
-        )
+        return BuildEnvironmentDatasourceDict(**{x: y for x, y in v.items()})
     raise ValueError
 
 

trilogy/core/models/execute.py

@@ -859,7 +859,7 @@ class RecursiveCTE(CTE):
                     ),
                     BuildCaseElse(expr=False),
                 ],
-                output_datatype=DataType.BOOL,
+                output_data_type=DataType.BOOL,
                 output_purpose=Purpose.KEY,
             ),
         )
@@ -884,7 +884,7 @@ class RecursiveCTE(CTE):
                     ),
                     BuildCaseElse(expr=False),
                 ],
-                output_datatype=DataType.BOOL,
+                output_data_type=DataType.BOOL,
                 output_purpose=Purpose.KEY,
             ),
         )
@@ -909,7 +909,7 @@ class RecursiveCTE(CTE):
                     ),
                     BuildCaseElse(expr=right_recurse_concept),
                 ],
-                output_datatype=recursive_derived.datatype,
+                output_data_type=recursive_derived.datatype,
                 output_purpose=recursive_derived.purpose,
             ),
         )

trilogy/core/processing/node_generators/group_node.py

@@ -61,13 +61,13 @@ def gen_group_node(
 
     # if the aggregation has a grain, we need to ensure these are the ONLY optional in the output of the select
     output_concepts = [concept]
-
+    grain_components = [environment.concepts[c] for c in concept.grain.components]
     if (
         concept.grain
         and len(concept.grain.components) > 0
         and not concept.grain.abstract
     ):
-        grain_components = [environment.concepts[c] for c in concept.grain.components]
+
         parent_concepts += grain_components
         build_grain_parents = get_aggregate_grain(concept, environment)
         output_concepts += grain_components
@@ -131,10 +131,6 @@ def gen_group_node(
     else:
         parents = []
 
-    # the keys we group by
-    # are what we can use for enrichment
-    group_key_parents = [environment.concepts[c] for c in concept.grain.components]
-
     group_node = GroupNode(
         output_concepts=output_concepts,
         input_concepts=parent_concepts,
@@ -164,7 +160,7 @@ def gen_group_node(
     )
     return gen_enrichment_node(
         group_node,
-        join_keys=group_key_parents,
+        join_keys=grain_components,
         local_optional=local_optional,
         environment=environment,
         g=g,

trilogy/core/processing/node_generators/node_merge_node.py

@@ -6,7 +6,11 @@ from networkx.algorithms import approximation as ax
 from trilogy.constants import logger
 from trilogy.core.enums import Derivation
 from trilogy.core.exceptions import AmbiguousRelationshipResolutionException
-from trilogy.core.graph_models import concept_to_node, prune_sources_for_conditions
+from trilogy.core.graph_models import (
+    ReferenceGraph,
+    concept_to_node,
+    prune_sources_for_conditions,
+)
 from trilogy.core.models.build import BuildConcept, BuildConditional, BuildWhereClause
 from trilogy.core.models.build_environment import BuildEnvironment
 from trilogy.core.processing.nodes import History, MergeNode, StrategyNode
@@ -17,11 +21,12 @@ LOGGER_PREFIX = "[GEN_MERGE_NODE]"
 AMBIGUITY_CHECK_LIMIT = 20
 
 
-def filter_pseudonyms_for_source(ds_graph: nx.DiGraph, node: str):
+def filter_pseudonyms_for_source(
+    ds_graph: nx.DiGraph, node: str, pseudonyms: set[tuple[str, str]]
+):
     to_remove = set()
-
     for edge in ds_graph.edges:
-        if ds_graph.edges[edge].get("pseudonym", False):
+        if edge in pseudonyms:
             lengths = {}
             for n in edge:
                 lengths[n] = nx.shortest_path_length(ds_graph, node, n)
@@ -52,12 +57,14 @@ def filter_unique_graphs(graphs: list[list[str]]) -> list[list[str]]:
     return [list(x) for x in unique_graphs]
 
 
-def extract_ds_components(g: nx.DiGraph, nodelist: list[str]) -> list[list[str]]:
+def extract_ds_components(
+    g: nx.DiGraph, nodelist: list[str], pseudonyms: set[tuple[str, str]]
+) -> list[list[str]]:
     graphs = []
     for node in g.nodes:
         if node.startswith("ds~"):
             local = g.copy()
-            filter_pseudonyms_for_source(local, node)
+            filter_pseudonyms_for_source(local, node, pseudonyms)
             ds_graph: nx.DiGraph = nx.ego_graph(local, node, radius=10).copy()
             graphs.append(
                 [
@@ -78,7 +85,7 @@ def extract_ds_components(g: nx.DiGraph, nodelist: list[str]) -> list[list[str]]
 
 
 def determine_induced_minimal_nodes(
-    G: nx.DiGraph,
+    G: ReferenceGraph,
     nodelist: list[str],
     environment: BuildEnvironment,
     filter_downstream: bool,
@@ -86,23 +93,19 @@ def determine_induced_minimal_nodes(
 ) -> nx.DiGraph | None:
     H: nx.Graph = nx.to_undirected(G).copy()
     nodes_to_remove = []
-    concepts = nx.get_node_attributes(G, "concept")
-
-    for node in G.nodes:
-        if concepts.get(node):
-            lookup: BuildConcept = concepts[node]
-            # inclusion of aggregates can create ambiguous node relation chains
-            # there may be a better way to handle this
-            # can be revisited if we need to connect a derived synonym based on an aggregate
-            if lookup.derivation in (
-                Derivation.CONSTANT,
-                Derivation.AGGREGATE,
-                Derivation.FILTER,
-            ):
-                nodes_to_remove.append(node)
-            # purge a node if we're already looking for all it's parents
-            if filter_downstream and lookup.derivation not in (Derivation.ROOT,):
-                nodes_to_remove.append(node)
+    for node, lookup in G.concepts.items():
+        # inclusion of aggregates can create ambiguous node relation chains
+        # there may be a better way to handle this
+        # can be revisited if we need to connect a derived synonym based on an aggregate
+        if lookup.derivation in (
+            Derivation.CONSTANT,
+            Derivation.AGGREGATE,
+            Derivation.FILTER,
+        ):
+            nodes_to_remove.append(node)
+        # purge a node if we're already looking for all it's parents
+        if filter_downstream and lookup.derivation not in (Derivation.ROOT,):
+            nodes_to_remove.append(node)
     if nodes_to_remove:
         # logger.debug(f"Removing nodes {nodes_to_remove} from graph")
         H.remove_nodes_from(nodes_to_remove)
@@ -259,7 +262,7 @@ def filter_duplicate_subgraphs(
 def resolve_weak_components(
     all_concepts: List[BuildConcept],
     environment: BuildEnvironment,
-    environment_graph: nx.DiGraph,
+    environment_graph: ReferenceGraph,
     filter_downstream: bool = True,
     accept_partial: bool = False,
     search_conditions: BuildWhereClause | None = None,
@@ -316,8 +319,6 @@ def resolve_weak_components(
             ]
             new = [x for x in all_graph_concepts if x.address not in all_concepts]
 
-            new_addresses = set([x.address for x in new if x.address not in synonyms])
-
             if not new:
                 break_flag = True
             # remove our new nodes for the next search path
@@ -329,6 +330,7 @@ def resolve_weak_components(
             # from trilogy.hooks.graph_hook import GraphHook
             # GraphHook().query_graph_built(g, highlight_nodes=[concept_to_node(c.with_default_grain()) for c in all_concepts if "__preql_internal" not in c.address])
             found.append(g)
+            new_addresses = set([x.address for x in new if x.address not in synonyms])
             reduced_concept_sets.append(new_addresses)
 
         except nx.exception.NetworkXNoPath:
@@ -346,7 +348,7 @@ def resolve_weak_components(
     subgraphs: list[list[BuildConcept]] = []
     # components = nx.strongly_connected_components(g)
     node_list = [x for x in g.nodes if x.startswith("c~")]
-    components = extract_ds_components(g, node_list)
+    components = extract_ds_components(g, node_list, environment_graph.pseudonyms)
     logger.debug(f"Extracted components {components} from {node_list}")
     for component in components:
         # we need to take unique again as different addresses may map to the same concept

trilogy/core/processing/node_generators/select_helpers/datasource_injection.py

@@ -128,17 +128,32 @@ def simplify_conditions(
     for condition in conditions:
         if not isinstance(condition, BuildComparison):
             return False
-        if not isinstance(condition.left, REDUCABLE_TYPES) and not isinstance(
-            condition.right, REDUCABLE_TYPES
-        ):
-            return False
-        if not isinstance(condition.left, BuildConcept) and not isinstance(
-            condition.right, BuildConcept
+        left_is_concept = False
+        left_is_reducable = False
+        right_is_concept = False
+        right_is_reducable = False
+        if isinstance(condition.left, BuildConcept):
+            left_is_concept = True
+        elif isinstance(condition.left, REDUCABLE_TYPES):
+            left_is_reducable = True
+
+        if isinstance(condition.right, BuildConcept):
+            right_is_concept = True
+        elif isinstance(condition.right, REDUCABLE_TYPES):
+            right_is_reducable = True
+
+        if not (
+            (left_is_concept and right_is_reducable)
+            or (right_is_concept and left_is_reducable)
         ):
             return False
-        vars = [condition.left, condition.right]
-        concept = [x for x in vars if isinstance(x, BuildConcept)][0]
-        raw_comparison = [x for x in vars if not isinstance(x, BuildConcept)][0]
+        if left_is_concept:
+            concept = condition.left
+            raw_comparison = condition.right
+        else:
+            concept = condition.right
+            raw_comparison = condition.left
+
         if isinstance(raw_comparison, BuildFunction):
             if not raw_comparison.operator == FunctionType.CONSTANT:
                 return False
@@ -154,7 +169,7 @@ def simplify_conditions(
         if not isinstance(comparison, REDUCABLE_TYPES):
             return False
 
-        var = concept
+        var: BuildConcept = concept  # type: ignore
         op = condition.operator
         grouped[var].append((op, comparison))
 
@@ -240,7 +255,6 @@ def get_union_sources(
         assocs[merge_key.address].append(x)
     final: list[list[BuildDatasource]] = []
     for _, dses in assocs.items():
-
         conditions = [c.non_partial_for.conditional for c in dses if c.non_partial_for]
         if simplify_conditions(conditions):
             final.append(dses)

trilogy/core/processing/node_generators/select_merge_node.py

@@ -6,6 +6,7 @@ import networkx as nx
 from trilogy.constants import logger
 from trilogy.core.enums import Derivation
 from trilogy.core.graph_models import (
+    ReferenceGraph,
     concept_to_node,
     get_graph_exact_match,
     prune_sources_for_conditions,
@@ -41,77 +42,68 @@ def extract_address(node: str):
 
 
 def get_graph_partial_nodes(
-    g: nx.DiGraph, conditions: BuildWhereClause | None
+    g: ReferenceGraph, conditions: BuildWhereClause | None
 ) -> dict[str, list[str]]:
-    datasources: dict[str, BuildDatasource | list[BuildDatasource]] = (
-        nx.get_node_attributes(g, "datasource")
-    )
     partial: dict[str, list[str]] = {}
-    for node in g.nodes:
-        if node in datasources:
-            ds = datasources[node]
-            if not isinstance(ds, list):
-
-                if ds.non_partial_for and conditions == ds.non_partial_for:
-                    partial[node] = []
-                    continue
-                partial[node] = [concept_to_node(c) for c in ds.partial_concepts]
-                ds = [ds]
-            # assume union sources have no partial
-            else:
-                partial[node] = []
+    for node, ds in g.datasources.items():
 
+        if not isinstance(ds, list):
+
+            if ds.non_partial_for and conditions == ds.non_partial_for:
+                partial[node] = []
+                continue
+            partial[node] = [concept_to_node(c) for c in ds.partial_concepts]
+        # assume union sources have no partial
+        else:
+            partial[node] = []
     return partial
 
 
-def get_graph_grains(g: nx.DiGraph) -> dict[str, list[str]]:
-    datasources: dict[str, BuildDatasource | list[BuildDatasource]] = (
-        nx.get_node_attributes(g, "datasource")
-    )
+def get_graph_grains(g: ReferenceGraph) -> dict[str, list[str]]:
     grain_length: dict[str, list[str]] = {}
-    for node in g.nodes:
-        if node in datasources:
-            base: set[str] = set()
-            lookup = datasources[node]
-            if not isinstance(lookup, list):
-                lookup = [lookup]
-            assert isinstance(lookup, list)
-            grain_length[node] = reduce(
-                lambda x, y: x.union(y.grain.components), lookup, base  # type: ignore
-            )
+    for node, lookup in g.datasources.items():
+
+        base: set[str] = set()
+        if not isinstance(lookup, list):
+            flookup = [lookup]
+        else:
+            flookup = lookup
+        assert isinstance(flookup, list)
+        grain_length[node] = reduce(
+            lambda x, y: x.union(y.grain.components), flookup, base  # type: ignore
+        )
     return grain_length
 
 
 def subgraph_is_complete(
     nodes: list[str], targets: set[str], mapping: dict[str, str], g: nx.DiGraph
 ) -> bool:
-    mapped = set([mapping.get(n, n) for n in nodes])
-    passed = all([t in mapped for t in targets])
-    if not passed:
+    # Check if all targets are present in mapped nodes
+    mapped = {mapping.get(n, n) for n in nodes}
+    if not targets.issubset(mapped):
         # logger.info(
         #     f"Subgraph {nodes} is not complete, missing targets {targets} - mapped {mapped}"
         # )
         return False
-    # check if all concepts have a datasource edge
-    has_ds_edge = {
-        mapping.get(n, n): any(x.startswith("ds~") for x in nx.neighbors(g, n))
-        for n in nodes
-        if n.startswith("c~")
-    }
-    has_ds_edge = {k: False for k in targets}
-    # check at least one instance of concept has a datasource edge
-    for n in nodes:
-        if n.startswith("c~"):
-            neighbors = nx.neighbors(g, n)
-            for neighbor in neighbors:
-                if neighbor.startswith("ds~"):
-                    has_ds_edge[mapping.get(n, n)] = True
-                    break
-    return all(has_ds_edge.values()) and passed
+
+    # Check if at least one concept node has a datasource edge
+    has_ds_edge = {target: False for target in targets}
+
+    for node in nodes:
+        if node.startswith("c~"):
+            mapped_node = mapping.get(node, node)
+            if mapped_node in targets and not has_ds_edge[mapped_node]:
+                # Only check neighbors if we haven't found a ds edge for this mapped node yet
+                if any(
+                    neighbor.startswith("ds~") for neighbor in nx.neighbors(g, node)
+                ):
+                    has_ds_edge[mapped_node] = True
+
+    return all(has_ds_edge.values())
 
 
 def create_pruned_concept_graph(
-    g: nx.DiGraph,
+    g: ReferenceGraph,
     all_concepts: List[BuildConcept],
     datasources: list[BuildDatasource],
     accept_partial: bool = False,
@@ -133,14 +125,13 @@ def create_pruned_concept_graph(
         )
         g.add_node(node_address, datasource=ds_list)
         for c in common:
-            g.add_edge(node_address, concept_to_node(c))
-            g.add_edge(concept_to_node(c), node_address)
+            cnode = concept_to_node(c)
+            g.add_edge(node_address, cnode)
+            g.add_edge(cnode, node_address)
     prune_sources_for_conditions(g, accept_partial, conditions)
     target_addresses = set([c.address for c in all_concepts])
-    concepts: dict[str, BuildConcept] = nx.get_node_attributes(orig_g, "concept")
-    datasource_map: dict[str, BuildDatasource | list[BuildDatasource]] = (
-        nx.get_node_attributes(orig_g, "datasource")
-    )
+    concepts: dict[str, BuildConcept] = orig_g.concepts
+    datasource_map: dict[str, BuildDatasource] = orig_g.datasources
     relevant_concepts_pre = {
         n: x.address
         for n in g.nodes()
@@ -170,31 +161,27 @@ def create_pruned_concept_graph(
         for edge in to_remove:
             g.remove_edge(*edge)
 
-    for n in g.nodes():
-        if not n.startswith("ds~"):
-            continue
-        actual_neighbors = [
-            x for x in relevant_concepts if x in (nx.all_neighbors(g, n))
-        ]
-        if actual_neighbors:
+    for n in g.datasources:
+        if any([[n, x] in g.edges for x in relevant_concepts]):
             relevent_datasets.append(n)
+            continue
 
     # for injecting extra join concepts that are shared between datasets
     # use the original graph, pre-partial pruning
-    for n in orig_g.nodes:
+    for n in orig_g.concepts:
         # readd ignoring grain
         # we want to join inclusive of all concepts
-        roots: dict[str, set[str]] = {}
-        if n.startswith("c~") and n not in relevant_concepts:
-            root = n.split("@")[0]
-            neighbors = roots.get(root, set())
-            for neighbor in nx.all_neighbors(orig_g, n):
+        if n not in relevant_concepts:
+            n_neighbors = nx.all_neighbors(orig_g, n)
+            # check if the irrelevant concept is a join between
+            # two relevant datasets
+            neighbors = set()
+            for neighbor in n_neighbors:
                 if neighbor in relevent_datasets:
                     neighbors.add(neighbor)
-            if len(neighbors) > 1:
-                relevant_concepts.append(n)
-            roots[root] = set()
-
+                    if len(neighbors) > 1:
+                        relevant_concepts.append(n)
+                        continue
     g.remove_nodes_from(
         [
             n
@@ -238,7 +225,7 @@ def create_pruned_concept_graph(
 
 
 def resolve_subgraphs(
-    g: nx.DiGraph,
+    g: ReferenceGraph,
     relevant: list[BuildConcept],
     accept_partial: bool,
     conditions: BuildWhereClause | None,
@@ -261,7 +248,7 @@ def resolve_subgraphs(
     partial_map = get_graph_partial_nodes(g, conditions)
     exact_map = get_graph_exact_match(g, accept_partial, conditions)
     grain_length = get_graph_grains(g)
-    concepts: dict[str, BuildConcept] = nx.get_node_attributes(g, "concept")
+    concepts: dict[str, BuildConcept] = g.concepts
     non_partial_map = {
         ds: [concepts[c].address for c in subgraphs[ds] if c not in partial_map[ds]]
         for ds in datasources
@@ -460,7 +447,7 @@ def create_select_node(
     ds_name: str,
     subgraph: list[str],
     accept_partial: bool,
-    g,
+    g: ReferenceGraph,
     environment: BuildEnvironment,
     depth: int,
     conditions: BuildWhereClause | None = None,
@@ -486,9 +473,7 @@ def create_select_node(
             preexisting_conditions=conditions.conditional if conditions else None,
         )
 
-    datasource: dict[str, BuildDatasource | list[BuildDatasource]] = (
-        nx.get_node_attributes(g, "datasource")[ds_name]
-    )
+    datasource: BuildDatasource = g.datasources[ds_name]
     if isinstance(datasource, BuildDatasource):
         bcandidate, force_group = create_datasource_node(
             datasource,
@@ -581,6 +566,7 @@ def gen_select_merge_node(
     logger.info(
         f"{padding(depth)}{LOGGER_PREFIX} searching for root source graph for concepts {[c.address for c in all_concepts]} and conditions {conditions}"
     )
+    pruned_concept_graph = None
     for attempt in attempts:
         pruned_concept_graph = create_pruned_concept_graph(
             g,

trilogy/parsing/parse_engine.py

@@ -244,7 +244,7 @@ def unwrap_transformation(
     elif isinstance(input, Parenthetical):
         return unwrap_transformation(input.content, environment)
     else:
-        return Function(
+        return Function.model_construct(
             operator=FunctionType.CONSTANT,
             output_datatype=arg_to_datatype(input),
             output_purpose=Purpose.CONSTANT,