pytrilogy 0.0.3.93__py3-none-any.whl → 0.0.3.95__py3-none-any.whl

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,
@@ -131,16 +123,15 @@ def create_pruned_concept_graph(
         common: set[BuildConcept] = set.intersection(
             *[set(x.output_concepts) for x in ds_list]
         )
-        g.add_node(node_address, datasource=ds_list)
+        g.add_datasource_node(node_address, 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
@@ -227,7 +214,7 @@ def create_pruned_concept_graph(
     relevant = set(relevant_concepts + relevent_datasets)
     for edge in orig_g.edges():
         if edge[0] in relevant and edge[1] in relevant:
-            g.add_edge(edge[0], edge[1])
+            g.add_edge(edge[0], edge[1], fast=True)
     # if we have no ds nodes at all, for non constant, we can't find it
     if not any([n.startswith("ds~") for n in g.nodes]):
         logger.info(
@@ -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/core/query_processor.py

@@ -40,8 +40,8 @@ from trilogy.core.statements.author import (
     PersistStatement,
     SelectStatement,
 )
-from trilogy.core.statements.common import MaterializedDataset
 from trilogy.core.statements.execute import (
+    MaterializedDataset,
     ProcessedCopyStatement,
     ProcessedQuery,
     ProcessedQueryPersist,
@@ -567,4 +567,5 @@ def process_query(
         base=root_cte,
         hidden_columns=set([x for x in statement.hidden_components]),
         local_concepts=statement.local_concepts,
+        locally_derived=statement.locally_derived,
     )

trilogy/core/statements/author.py

@@ -12,6 +12,7 @@ from trilogy.core.enums import (
     IOType,
     Modifier,
     ShowCategory,
+    ValidationScope,
 )
 from trilogy.core.models.author import (
     AggregateWrapper,
@@ -147,11 +148,13 @@ class SelectStatement(HasUUID, SelectTypeMixin, BaseModel):
                     continue
                 if CONFIG.parsing.select_as_definition and not environment.frozen:
                     if x.concept.address not in environment.concepts:
-                        environment.add_concept(x.content.output)
+                        environment.add_concept(x.content.output, add_derived=False)
                     elif x.concept.address in environment.concepts:
                         version = environment.concepts[x.concept.address]
                         if version.metadata.concept_source == ConceptSource.SELECT:
-                            environment.add_concept(x.content.output, force=True)
+                            environment.add_concept(
+                                x.content.output, force=True, add_derived=False
+                            )
                 x.content.output = x.content.output.set_select_grain(
                     output.grain, environment
                 )
@@ -378,6 +381,13 @@ class MultiSelectStatement(HasUUID, SelectTypeMixin, BaseModel):
             output = output.union(select.hidden_components)
         return output
 
+    @property
+    def locally_derived(self) -> set[str]:
+        locally_derived: set[str] = set([x.address for x in self.derived_concepts])
+        for select in self.selects:
+            locally_derived = locally_derived.union(select.locally_derived)
+        return locally_derived
+
 
 class RowsetDerivationStatement(HasUUID, BaseModel):
     name: str
@@ -428,8 +438,13 @@ class PersistStatement(HasUUID, BaseModel):
         return self.datasource.address
 
 
+class ValidateStatement(BaseModel):
+    scope: ValidationScope
+    targets: Optional[List[str]] = None  # list of identifiers
+
+
 class ShowStatement(BaseModel):
-    content: SelectStatement | PersistStatement | ShowCategory
+    content: SelectStatement | PersistStatement | ValidateStatement | ShowCategory
 
 
 class Limit(BaseModel):

trilogy/core/statements/common.py

@@ -4,7 +4,6 @@ from pydantic import BaseModel, Field
 
 from trilogy.core.enums import IOType
 from trilogy.core.models.author import ConceptRef, HavingClause, WhereClause
-from trilogy.core.models.datasource import Address, Datasource
 
 
 class CopyQueryMixin(BaseModel):
@@ -12,15 +11,6 @@ class CopyQueryMixin(BaseModel):
     target_type: IOType
 
 
-class MaterializedDataset(BaseModel):
-    address: Address
-
-
-class PersistQueryMixin(BaseModel):
-    output_to: MaterializedDataset
-    datasource: Datasource
-
-
 class SelectTypeMixin(BaseModel):
     where_clause: Union["WhereClause", None] = Field(default=None)
     having_clause: Union["HavingClause", None] = Field(default=None)

trilogy/core/statements/execute.py

@@ -1,49 +1,104 @@
-from typing import Annotated, List, Optional, Union
+from dataclasses import dataclass, field
+from typing import List, Optional, Union
 
-from pydantic import BaseModel, Field
-from pydantic.functional_validators import PlainValidator
-
-from trilogy.core.models.author import ConceptRef
+from trilogy.core.enums import IOType, ValidationScope
+from trilogy.core.models.author import ConceptRef, HavingClause, WhereClause
 from trilogy.core.models.build import (
     BuildConcept,
     BuildDatasource,
     BuildOrderBy,
 )
-from trilogy.core.models.environment import EnvironmentConceptDict, validate_concepts
+from trilogy.core.models.datasource import Address, Datasource
+from trilogy.core.models.environment import EnvironmentConceptDict
 from trilogy.core.models.execute import CTE, UnionCTE
-from trilogy.core.statements.common import CopyQueryMixin, PersistQueryMixin
 
 
-class ProcessedQuery(BaseModel):
+@dataclass
+class CopyQueryMixin:
+    target: str
+    target_type: IOType
+
+
+@dataclass
+class MaterializedDataset:
+    address: Address
+
+
+@dataclass
+class PersistQueryMixin:
+    output_to: MaterializedDataset
+    datasource: Datasource
+
+
+@dataclass
+class SelectTypeMixin:
+    where_clause: Union["WhereClause", None] = field(default=None)
+    having_clause: Union["HavingClause", None] = field(default=None)
+
+    @property
+    def output_components(self) -> List[ConceptRef]:
+        raise NotImplementedError
+
+
+@dataclass
+class ProcessedQuery:
     output_columns: List[ConceptRef]
     ctes: List[CTE | UnionCTE]
     base: CTE | UnionCTE
-    hidden_columns: set[str] = Field(default_factory=set)
+    hidden_columns: set[str] = field(default_factory=set)
     limit: Optional[int] = None
     order_by: Optional[BuildOrderBy] = None
-    local_concepts: Annotated[
-        EnvironmentConceptDict, PlainValidator(validate_concepts)
-    ] = Field(default_factory=EnvironmentConceptDict)
+    local_concepts: EnvironmentConceptDict = field(
+        default_factory=EnvironmentConceptDict
+    )
+    locally_derived: set[str] = field(default_factory=set)
 
 
+@dataclass
 class ProcessedQueryPersist(ProcessedQuery, PersistQueryMixin):
     pass
 
 
+@dataclass
 class ProcessedCopyStatement(ProcessedQuery, CopyQueryMixin):
     pass
 
 
-class ProcessedRawSQLStatement(BaseModel):
+@dataclass
+class ProcessedRawSQLStatement:
     text: str
 
 
-class ProcessedStaticValueOutput(BaseModel):
+@dataclass
+class ProcessedValidateStatement:
+    scope: ValidationScope
+    targets: Optional[List[str]]
+
+
+@dataclass
+class ProcessedStaticValueOutput:
     values: List[dict]
 
 
-class ProcessedShowStatement(BaseModel):
+@dataclass
+class ProcessedShowStatement:
     output_columns: List[ConceptRef]
     output_values: List[
-        Union[BuildConcept, BuildDatasource, ProcessedQuery, ProcessedStaticValueOutput]
+        Union[
+            BuildConcept,
+            BuildDatasource,
+            ProcessedQuery,
+            ProcessedValidateStatement,
+            ProcessedStaticValueOutput,
+        ]
     ]
+
+
+PROCESSED_STATEMENT_TYPES = (
+    ProcessedCopyStatement
+    | ProcessedQuery
+    | ProcessedRawSQLStatement
+    | ProcessedQueryPersist
+    | ProcessedShowStatement
+    | ProcessedValidateStatement
+)