pytrilogy 0.0.1.104__tar.gz → 0.0.1.106__tar.gz

{pytrilogy-0.0.1.104/pytrilogy.egg-info → pytrilogy-0.0.1.106}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytrilogy
-Version: 0.0.1.104
+Version: 0.0.1.106
 Summary: Declarative, typed query language that compiles to SQL.
 Home-page: 
 Author: 

{pytrilogy-0.0.1.104 → pytrilogy-0.0.1.106/pytrilogy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytrilogy
-Version: 0.0.1.104
+Version: 0.0.1.106
 Summary: Declarative, typed query language that compiles to SQL.
 Home-page: 
 Author: 

{pytrilogy-0.0.1.104 → pytrilogy-0.0.1.106}/pytrilogy.egg-info/SOURCES.txt

@@ -43,6 +43,7 @@ trilogy/core/functions.py
 trilogy/core/graph_models.py
 trilogy/core/internal.py
 trilogy/core/models.py
+trilogy/core/optimization.py
 trilogy/core/query_processor.py
 trilogy/core/processing/__init__.py
 trilogy/core/processing/concept_strategies_v3.py

{pytrilogy-0.0.1.104 → pytrilogy-0.0.1.106}/tests/test_select.py

@@ -125,4 +125,5 @@ def test_modifiers():
     generator = BigqueryDialect()
 
     text = generator.compile_statement(query)
-    assert "`b` = 2" in text
+    assert "2 = 2" in text
+    assert "as `b`" not in text

{pytrilogy-0.0.1.104 → pytrilogy-0.0.1.106}/trilogy/__init__.py

@@ -2,7 +2,8 @@ from trilogy.core.models import Environment
 from trilogy.dialect.enums import Dialects
 from trilogy.executor import Executor
 from trilogy.parser import parse
+from trilogy.constants import CONFIG
 
-__version__ = "0.0.1.104"
+__version__ = "0.0.1.106"
 
-__all__ = ["parse", "Executor", "Dialects", "Environment"]
+__all__ = ["parse", "Executor", "Dialects", "Environment", "CONFIG"]

{pytrilogy-0.0.1.104 → pytrilogy-0.0.1.106}/trilogy/constants.py

@@ -23,6 +23,7 @@ NULL_VALUE = MagicConstants.NULL
 class Config:
     strict_mode: bool = True
     human_identifiers: bool = True
+    inline_datasources: bool = True
 
 
 CONFIG = Config()

{pytrilogy-0.0.1.104 → pytrilogy-0.0.1.106}/trilogy/core/models.py

@@ -34,7 +34,11 @@ from pydantic import (
 from lark.tree import Meta
 from pathlib import Path
 from trilogy.constants import logger, DEFAULT_NAMESPACE, ENV_CACHE_NAME, MagicConstants
-from trilogy.core.constants import ALL_ROWS_CONCEPT, INTERNAL_NAMESPACE
+from trilogy.core.constants import (
+    ALL_ROWS_CONCEPT,
+    INTERNAL_NAMESPACE,
+    CONSTANT_DATASET,
+)
 from trilogy.core.enums import (
     InfiniteFunctionArgs,
     Purpose,
@@ -113,15 +117,27 @@ NAMESPACED_TYPES = Union[
 
 
 class Namespaced(ABC):
-    pass
 
     def with_namespace(self, namespace: str):
         raise NotImplementedError
 
 
-class SelectGrain(ABC):
-    pass
+class ConceptArgs(ABC):
+
+    @property
+    def concept_arguments(self) -> List["Concept"]:
+        raise NotImplementedError
+
+    @property
+    def existence_arguments(self) -> List["Concept"]:
+        return []
+
+    @property
+    def row_arguments(self) -> List["Concept"]:
+        return self.concept_arguments
+
 
+class SelectGrain(ABC):
     def with_select_grain(self, grain: Grain):
         raise NotImplementedError
 
@@ -1395,16 +1411,11 @@ class MultiSelectStatement(Namespaced, BaseModel):
         return output
 
     def find_source(self, concept: Concept, cte: CTE) -> Concept:
-        all = []
         for x in self.align.items:
             if concept.name == x.alias:
                 for c in x.concepts:
                     if c.address in cte.output_lcl:
-                        all.append(c)
-
-        if len(all) == 1:
-            return all[0]
-
+                        return c
         raise SyntaxError(
             f"Could not find upstream map for multiselect {str(concept)} on cte ({cte})"
         )
@@ -1427,6 +1438,7 @@ class MultiSelectStatement(Namespaced, BaseModel):
 
 class Address(BaseModel):
     location: str
+    is_query: bool = False
 
 
 class Query(BaseModel):
@@ -1522,10 +1534,27 @@ class Datasource(Namespaced, BaseModel):
         default_factory=lambda: DatasourceMetadata(freshness_concept=None)
     )
 
+    @property
+    def condition(self):
+        return None
+
     @cached_property
     def output_lcl(self) -> LooseConceptList:
         return LooseConceptList(concepts=self.output_concepts)
 
+    @property
+    def can_be_inlined(self) -> bool:
+        if isinstance(self.address, Address) and self.address.is_query:
+            return False
+        for x in self.columns:
+            if not isinstance(x.alias, str):
+                return False
+        return True
+
+    @property
+    def non_partial_concept_addresses(self) -> set[str]:
+        return set([c.address for c in self.full_concepts])
+
     @field_validator("namespace", mode="plain")
     @classmethod
     def namespace_validation(cls, v):
@@ -1739,7 +1768,7 @@ class QueryDatasource(BaseModel):
     input_concepts: List[Concept]
     output_concepts: List[Concept]
     source_map: Dict[str, Set[Union[Datasource, "QueryDatasource", "UnnestJoin"]]]
-    datasources: Sequence[Union[Datasource, "QueryDatasource"]]
+    datasources: List[Union[Datasource, "QueryDatasource"]]
     grain: Grain
     joins: List[BaseJoin | UnnestJoin]
     limit: Optional[int] = None
@@ -1790,10 +1819,7 @@ class QueryDatasource(BaseModel):
             c.address for c in values["input_concepts"]
         )
         seen = set()
-        for k, val in v.items():
-            # if val:
-            #     if len(val) != 1:
-            #         raise SyntaxError(f"source map {k} has multiple values {len(val)}")
+        for k, _ in v.items():
             seen.add(k)
         for x in expected:
             if x not in seen:
@@ -1915,18 +1941,18 @@ class QueryDatasource(BaseModel):
         )
 
     def get_alias(
-        self, concept: Concept, use_raw_name: bool = False, force_alias: bool = False
+        self,
+        concept: Concept,
+        use_raw_name: bool = False,
+        force_alias: bool = False,
+        source: str | None = None,
     ):
-        # if we should use the raw datasource name to access
-        use_raw_name = (
-            True
-            if (len(self.datasources) == 1 or use_raw_name) and not force_alias
-            # if ((len(self.datasources) == 1 and isinstance(self.datasources[0], Datasource)) or use_raw_name) and not force_alias
-            else False
-        )
         for x in self.datasources:
             # query datasources should be referenced by their alias, always
             force_alias = isinstance(x, QueryDatasource)
+            use_raw_name = isinstance(x, Datasource) and not force_alias
+            if source and x.identifier != source:
+                continue
             try:
                 return x.get_alias(
                     concept.with_grain(self.grain),
@@ -1960,8 +1986,7 @@ class Comment(BaseModel):
 
 class CTE(BaseModel):
     name: str
-    source: "QueryDatasource"  # TODO: make recursive
-    # output columns are what are selected/grouped by
+    source: "QueryDatasource"
     output_columns: List[Concept]
     source_map: Dict[str, str | list[str]]
     grain: Grain
@@ -1972,6 +1997,10 @@ class CTE(BaseModel):
     condition: Optional[Union["Conditional", "Comparison", "Parenthetical"]] = None
     partial_concepts: List[Concept] = Field(default_factory=list)
     join_derived_concepts: List[Concept] = Field(default_factory=list)
+    order_by: Optional[OrderBy] = None
+    limit: Optional[int] = None
+    requires_nesting: bool = True
+    base_name_override: Optional[str] = None
 
     @computed_field  # type: ignore
     @property
@@ -1982,6 +2011,40 @@ class CTE(BaseModel):
     def validate_output_columns(cls, v):
         return unique(v, "address")
 
+    def inline_parent_datasource(self, parent: CTE) -> bool:
+        qds_being_inlined = parent.source
+        ds_being_inlined = qds_being_inlined.datasources[0]
+        if not isinstance(ds_being_inlined, Datasource):
+            return False
+        self.source.datasources = [
+            ds_being_inlined,
+            *[
+                x
+                for x in self.source.datasources
+                if x.identifier != qds_being_inlined.identifier
+            ],
+        ]
+        # need to identify this before updating joins
+        if self.base_name == parent.name:
+            self.base_name_override = ds_being_inlined.safe_location
+
+        for join in self.joins:
+            if isinstance(join, InstantiatedUnnestJoin):
+                continue
+            if join.left_cte.name == parent.name:
+                join.left_cte = ds_being_inlined
+            if join.right_cte.name == parent.name:
+                join.right_cte = ds_being_inlined
+        for k, v in self.source_map.items():
+            if isinstance(v, list):
+                self.source_map[k] = [
+                    ds_being_inlined.name if x == parent.name else x for x in v
+                ]
+            elif v == parent.name:
+                self.source_map[k] = ds_being_inlined.name
+        self.parent_ctes = [x for x in self.parent_ctes if x.name != parent.name]
+        return True
+
     def __add__(self, other: "CTE"):
         logger.info('Merging two copies of CTE "%s"', self.name)
         if not self.grain == other.grain:
@@ -2024,16 +2087,25 @@ class CTE(BaseModel):
     def relevant_base_ctes(self):
         return self.parent_ctes
 
+    @property
+    def is_root_datasource(self) -> bool:
+        return (
+            len(self.source.datasources) == 1
+            and isinstance(self.source.datasources[0], Datasource)
+            and not self.source.datasources[0].name == CONSTANT_DATASET
+        )
+
     @property
     def base_name(self) -> str:
+        if self.base_name_override:
+            return self.base_name_override
         # if this cte selects from a single datasource, select right from it
         valid_joins: List[Join] = [
             join for join in self.joins if isinstance(join, Join)
         ]
-        if len(self.source.datasources) == 1 and isinstance(
-            self.source.datasources[0], Datasource
-        ):
+        if self.is_root_datasource:
             return self.source.datasources[0].safe_location
+
         # if we have multiple joined CTEs, pick the base
         # as the root
         elif len(self.source.datasources) == 1 and len(self.parent_ctes) == 1:
@@ -2059,11 +2131,10 @@ class CTE(BaseModel):
 
     @property
     def base_alias(self) -> str:
+
+        if self.is_root_datasource:
+            return self.source.datasources[0].identifier
         relevant_joins = [j for j in self.joins if isinstance(j, Join)]
-        if len(self.source.datasources) == 1 and isinstance(
-            self.source.datasources[0], Datasource
-        ):
-            return self.source.datasources[0].full_name.replace(".", "_")
         if relevant_joins:
             return relevant_joins[0].left_cte.name
         elif self.relevant_base_ctes:
@@ -2072,12 +2143,16 @@ class CTE(BaseModel):
             return self.parent_ctes[0].name
         return self.name
 
-    def get_alias(self, concept: Concept) -> str:
+    def get_alias(self, concept: Concept, source: str | None = None) -> str:
         for cte in self.parent_ctes:
             if concept.address in [x.address for x in cte.output_columns]:
+                if source and source != cte.name:
+                    continue
                 return concept.safe_address
         try:
-            source = self.source.get_alias(concept)
+            source = self.source.get_alias(concept, source=source)
+            if not source:
+                raise ValueError("No source found")
             return source
         except ValueError as e:
             return f"INVALID_ALIAS: {str(e)}"
@@ -2090,6 +2165,11 @@ class CTE(BaseModel):
             and not self.group_to_grain
         ):
             return False
+        if (
+            len(self.source.datasources) == 1
+            and self.source.datasources[0].name == CONSTANT_DATASET
+        ):
+            return False
         return True
 
     @property
@@ -2123,19 +2203,43 @@ class JoinKey(BaseModel):
 
 
 class Join(BaseModel):
-    left_cte: CTE
-    right_cte: CTE
+    left_cte: CTE | Datasource
+    right_cte: CTE | Datasource
     jointype: JoinType
     joinkeys: List[JoinKey]
 
+    @property
+    def left_name(self) -> str:
+        if isinstance(self.left_cte, Datasource):
+            return self.left_cte.identifier
+        return self.left_cte.name
+
+    @property
+    def right_name(self) -> str:
+        if isinstance(self.right_cte, Datasource):
+            return self.right_cte.identifier
+        return self.right_cte.name
+
+    @property
+    def left_ref(self) -> str:
+        if isinstance(self.left_cte, Datasource):
+            return f"{self.left_cte.safe_location} as {self.left_cte.identifier}"
+        return self.left_cte.name
+
+    @property
+    def right_ref(self) -> str:
+        if isinstance(self.right_cte, Datasource):
+            return f"{self.right_cte.safe_location} as {self.right_cte.identifier}"
+        return self.right_cte.name
+
     @property
     def unique_id(self) -> str:
-        return self.left_cte.name + self.right_cte.name + self.jointype.value
+        return self.left_name + self.right_name + self.jointype.value
 
     def __str__(self):
         return (
-            f"{self.jointype.value} JOIN {self.left_cte.name} and"
-            f" {self.right_cte.name} on {','.join([str(k) for k in self.joinkeys])}"
+            f"{self.jointype.value} JOIN {self.left_name} and"
+            f" {self.right_name} on {','.join([str(k) for k in self.joinkeys])}"
         )
 
 
@@ -2546,7 +2650,7 @@ class LazyEnvironment(Environment):
         return super().__getattribute__(name)
 
 
-class Comparison(Namespaced, SelectGrain, BaseModel):
+class Comparison(ConceptArgs, Namespaced, SelectGrain, BaseModel):
     left: Union[
         int,
         str,
@@ -2598,7 +2702,7 @@ class Comparison(Namespaced, SelectGrain, BaseModel):
         return f"{str(self.left)} {self.operator.value} {str(self.right)}"
 
     def with_namespace(self, namespace: str):
-        return Comparison(
+        return self.__class__(
             left=(
                 self.left.with_namespace(namespace)
                 if isinstance(self.left, Namespaced)
@@ -2613,7 +2717,7 @@ class Comparison(Namespaced, SelectGrain, BaseModel):
         )
 
     def with_select_grain(self, grain: Grain):
-        return Comparison(
+        return self.__class__(
             left=(
                 self.left.with_select_grain(grain)
                 if isinstance(self.left, SelectGrain)
@@ -2632,7 +2736,9 @@ class Comparison(Namespaced, SelectGrain, BaseModel):
         output: List[Concept] = []
         if isinstance(self.left, (Concept,)):
             output += [self.left]
-        if isinstance(self.left, (Conditional, Parenthetical)):
+        if isinstance(
+            self.left, (Comparison, SubselectComparison, Conditional, Parenthetical)
+        ):
             output += self.left.input
         if isinstance(self.left, FilterItem):
             output += self.left.concept_arguments
@@ -2641,7 +2747,9 @@ class Comparison(Namespaced, SelectGrain, BaseModel):
 
         if isinstance(self.right, (Concept,)):
             output += [self.right]
-        if isinstance(self.right, (Conditional, Parenthetical)):
+        if isinstance(
+            self.right, (Comparison, SubselectComparison, Conditional, Parenthetical)
+        ):
             output += self.right.input
         if isinstance(self.right, FilterItem):
             output += self.right.concept_arguments
@@ -2658,8 +2766,31 @@ class Comparison(Namespaced, SelectGrain, BaseModel):
         return output
 
 
+class SubselectComparison(Comparison):
+
+    @property
+    def row_arguments(self) -> List[Concept]:
+        return get_concept_arguments(self.left)
+
+    @property
+    def existence_arguments(self) -> List[Concept]:
+        return get_concept_arguments(self.right)
+
+    def with_select_grain(self, grain: Grain):
+        # there's no need to pass the select grain through to a subselect comparison
+        return self.__class__(
+            left=(
+                self.left.with_select_grain(grain)
+                if isinstance(self.left, SelectGrain)
+                else self.left
+            ),
+            right=self.right,
+            operator=self.operator,
+        )
+
+
 class CaseWhen(Namespaced, SelectGrain, BaseModel):
-    comparison: Conditional | Comparison
+    comparison: Conditional | SubselectComparison | Comparison
     expr: "Expr"
 
     @property
@@ -2726,7 +2857,7 @@ class CaseElse(Namespaced, SelectGrain, BaseModel):
         )
 
 
-class Conditional(Namespaced, SelectGrain, BaseModel):
+class Conditional(ConceptArgs, Namespaced, SelectGrain, BaseModel):
     left: Union[
         int,
         str,
@@ -2821,6 +2952,32 @@ class Conditional(Namespaced, SelectGrain, BaseModel):
         output += get_concept_arguments(self.right)
         return output
 
+    @property
+    def row_arguments(self) -> List[Concept]:
+        output = []
+        if isinstance(self.left, ConceptArgs):
+            output += self.left.row_arguments
+        else:
+            output += get_concept_arguments(self.left)
+        if isinstance(self.right, ConceptArgs):
+            output += self.right.row_arguments
+        else:
+            output += get_concept_arguments(self.right)
+        return output
+
+    @property
+    def existence_arguments(self) -> List[Concept]:
+        output = []
+        if isinstance(self.left, ConceptArgs):
+            output += self.left.existence_arguments
+        else:
+            output += get_concept_arguments(self.left)
+        if isinstance(self.right, ConceptArgs):
+            output += self.right.existence_arguments
+        else:
+            output += get_concept_arguments(self.right)
+        return output
+
 
 class AggregateWrapper(Namespaced, SelectGrain, BaseModel):
     function: Function
@@ -2864,8 +3021,8 @@ class AggregateWrapper(Namespaced, SelectGrain, BaseModel):
         return AggregateWrapper(function=self.function.with_select_grain(grain), by=by)
 
 
-class WhereClause(Namespaced, SelectGrain, BaseModel):
-    conditional: Union[Comparison, Conditional, "Parenthetical"]
+class WhereClause(ConceptArgs, Namespaced, SelectGrain, BaseModel):
+    conditional: Union[SubselectComparison, Comparison, Conditional, "Parenthetical"]
 
     @property
     def input(self) -> List[Concept]:
@@ -2875,6 +3032,14 @@ class WhereClause(Namespaced, SelectGrain, BaseModel):
     def concept_arguments(self) -> List[Concept]:
         return self.conditional.concept_arguments
 
+    @property
+    def row_arguments(self) -> List[Concept]:
+        return self.conditional.row_arguments
+
+    @property
+    def existence_arguments(self) -> List[Concept]:
+        return self.conditional.existence_arguments
+
     def with_namespace(self, namespace: str) -> WhereClause:
         return WhereClause(conditional=self.conditional.with_namespace(namespace))
 
@@ -3062,7 +3227,7 @@ class RowsetItem(Namespaced, BaseModel):
         return [self.content]
 
 
-class Parenthetical(Namespaced, SelectGrain, BaseModel):
+class Parenthetical(ConceptArgs, Namespaced, SelectGrain, BaseModel):
     content: "Expr"
 
     def __str__(self):
@@ -3106,6 +3271,18 @@ class Parenthetical(Namespaced, SelectGrain, BaseModel):
             base.append(x)
         return base
 
+    @property
+    def row_arguments(self) -> List[Concept]:
+        if isinstance(self.content, ConceptArgs):
+            return self.content.row_arguments
+        return self.concept_arguments
+
+    @property
+    def existence_arguments(self) -> List[Concept]:
+        if isinstance(self.content, ConceptArgs):
+            return self.content.existence_arguments
+        return self.concept_arguments
+
     @property
     def input(self):
         base = []

pytrilogy-0.0.1.106/trilogy/core/optimization.py

@@ -0,0 +1,141 @@
+from trilogy.core.models import (
+    CTE,
+    SelectStatement,
+    PersistStatement,
+    Datasource,
+    MultiSelectStatement,
+)
+from trilogy.core.enums import PurposeLineage
+from trilogy.constants import logger
+from abc import ABC
+
+
+class OptimizationRule(ABC):
+
+    def optimize(self, cte: CTE) -> bool:
+        raise NotImplementedError
+
+    def log(self, message: str):
+        logger.info(f"[Optimization][{self.__class__.__name__}] {message}")
+
+
+class InlineDatasource(OptimizationRule):
+
+    def optimize(self, cte: CTE) -> bool:
+        if not cte.parent_ctes:
+            return False
+
+        optimized = False
+        self.log(
+            f"Checking {cte.name} for consolidating inline tables with {len(cte.parent_ctes)} parents"
+        )
+        to_inline: list[CTE] = []
+        for parent_cte in cte.parent_ctes:
+            if not parent_cte.is_root_datasource:
+                self.log(f"parent {parent_cte.name} is not root")
+                continue
+            if parent_cte.parent_ctes:
+                self.log(f"parent {parent_cte.name} has parents")
+                continue
+            raw_root = parent_cte.source.datasources[0]
+            if not isinstance(raw_root, Datasource):
+                self.log(f"parent {parent_cte.name} is not datasource")
+                continue
+            root: Datasource = raw_root
+            if not root.can_be_inlined:
+                self.log(f"parent {parent_cte.name} datasource is not inlineable")
+                continue
+            root_outputs = {x.address for x in root.output_concepts}
+            cte_outputs = {x.address for x in parent_cte.output_columns}
+            if not cte_outputs.issubset(root_outputs):
+                self.log(f"Not all {parent_cte.name} outputs are found on datasource")
+                continue
+
+            to_inline.append(parent_cte)
+
+        for replaceable in to_inline:
+            self.log(f"Inlining parent {replaceable.name}")
+            cte.inline_parent_datasource(replaceable)
+
+        return optimized
+
+
+REGISTERED_RULES: list[OptimizationRule] = [InlineDatasource()]
+
+
+def filter_irrelevant_ctes(input: list[CTE], root_cte: CTE):
+    relevant_ctes = set()
+
+    def recurse(cte: CTE):
+        relevant_ctes.add(cte.name)
+        for cte in cte.parent_ctes:
+            recurse(cte)
+
+    recurse(root_cte)
+    return [cte for cte in input if cte.name in relevant_ctes]
+
+
+def is_direct_return_eligible(
+    cte: CTE, select: SelectStatement | PersistStatement | MultiSelectStatement
+) -> bool:
+    if isinstance(select, (PersistStatement, MultiSelectStatement)):
+        return False
+    derived_concepts = [
+        c for c in cte.source.output_concepts if c not in cte.source.input_concepts
+    ]
+    eligible = True
+    conditions = (
+        set(x.address for x in select.where_clause.concept_arguments)
+        if select.where_clause
+        else set()
+    )
+    if conditions and select.limit:
+        return False
+    for x in derived_concepts:
+        if x.derivation == PurposeLineage.WINDOW:
+            return False
+        if x.derivation == PurposeLineage.AGGREGATE:
+            if x.address in conditions:
+                return False
+    logger.info(
+        f"Upleveling output select to final CTE with derived_concepts {[x.address for x in derived_concepts]}"
+    )
+    return eligible
+
+
+def sort_select_output(cte: CTE, query: SelectStatement | MultiSelectStatement):
+    hidden_addresses = [c.address for c in query.hidden_components]
+    output_addresses = [
+        c.address for c in query.output_components if c.address not in hidden_addresses
+    ]
+
+    mapping = {x.address: x for x in cte.output_columns}
+
+    new_output = []
+    for x in output_addresses:
+        new_output.append(mapping[x])
+    cte.output_columns = new_output
+
+
+def optimize_ctes(
+    input: list[CTE], root_cte: CTE, select: SelectStatement | MultiSelectStatement
+):
+    complete = False
+
+    while not complete:
+        actions_taken = False
+        for rule in REGISTERED_RULES:
+            for cte in input:
+                actions_taken = rule.optimize(cte)
+        complete = not actions_taken
+
+    if 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
+        )
+        root_cte.requires_nesting = False
+        sort_select_output(cte, select)
+
+    return filter_irrelevant_ctes(input, root_cte)