pytrilogy 0.0.2.8__py3-none-any.whl → 0.0.2.10__py3-none-any.whl

trilogy/core/processing/nodes/__init__.py

@@ -6,7 +6,7 @@ from .window_node import WindowNode
 from .base_node import StrategyNode, NodeJoin
 from .unnest_node import UnnestNode
 from pydantic import BaseModel, Field, ConfigDict
-from trilogy.core.models import Concept, Environment
+from trilogy.core.models import Concept, Environment, WhereClause
 
 
 class History(BaseModel):
@@ -15,23 +15,42 @@ class History(BaseModel):
     started: set[str] = Field(default_factory=set)
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def _concepts_to_lookup(self, search: list[Concept], accept_partial: bool) -> str:
+    def _concepts_to_lookup(
+        self,
+        search: list[Concept],
+        accept_partial: bool,
+        conditions: WhereClause | None = None,
+    ) -> str:
+        if conditions:
+            return (
+                "-".join([c.address for c in search])
+                + str(accept_partial)
+                + str(conditions)
+            )
         return "-".join([c.address for c in search]) + str(accept_partial)
 
     def search_to_history(
-        self, search: list[Concept], accept_partial: bool, output: StrategyNode | None
+        self,
+        search: list[Concept],
+        accept_partial: bool,
+        output: StrategyNode | None,
+        conditions: WhereClause | None = None,
     ):
-        self.history[self._concepts_to_lookup(search, accept_partial)] = output
+        self.history[
+            self._concepts_to_lookup(search, accept_partial, conditions=conditions)
+        ] = output
 
     def get_history(
         self,
         search: list[Concept],
+        conditions: WhereClause | None = None,
         accept_partial: bool = False,
         parent_key: str = "",
     ) -> StrategyNode | None | bool:
         key = self._concepts_to_lookup(
             search,
             accept_partial,
+            conditions,
         )
         if parent_key and parent_key == key:
             raise ValueError(
@@ -48,11 +67,13 @@ class History(BaseModel):
         self,
         search: list[Concept],
         accept_partial: bool = False,
+        conditions: WhereClause | None = None,
     ):
         self.started.add(
             self._concepts_to_lookup(
                 search,
-                accept_partial,
+                accept_partial=accept_partial,
+                conditions=conditions,
             )
         )
 
@@ -60,11 +81,13 @@ class History(BaseModel):
         self,
         search: list[Concept],
         accept_partial: bool = False,
+        conditions: WhereClause | None = None,
     ):
         return (
             self._concepts_to_lookup(
                 search,
                 accept_partial,
+                conditions=conditions,
             )
             in self.started
         )
@@ -76,6 +99,7 @@ class History(BaseModel):
         accept_partial: bool,
         fail_if_not_found: bool,
         accept_partial_optional: bool,
+        conditions: WhereClause | None = None,
     ) -> str:
         return (
             str(main.address)
@@ -84,6 +108,7 @@ class History(BaseModel):
             + str(accept_partial)
             + str(fail_if_not_found)
             + str(accept_partial_optional)
+            + str(conditions)
         )
 
     def gen_select_node(
@@ -97,6 +122,7 @@ class History(BaseModel):
         fail_if_not_found: bool = False,
         accept_partial: bool = False,
         accept_partial_optional: bool = False,
+        conditions: WhereClause | None = None,
     ) -> StrategyNode | None:
         from trilogy.core.processing.node_generators.select_node import gen_select_node
 
@@ -105,7 +131,8 @@ class History(BaseModel):
             local_optional,
             accept_partial,
             fail_if_not_found,
-            accept_partial_optional,
+            accept_partial_optional=accept_partial_optional,
+            conditions=conditions,
         )
         if fingerprint in self.select_history:
             return self.select_history[fingerprint]
@@ -119,6 +146,7 @@ class History(BaseModel):
             accept_partial=accept_partial,
             accept_partial_optional=accept_partial_optional,
             source_concepts=source_concepts,
+            conditions=conditions,
         )
         self.select_history[fingerprint] = gen
         return gen

trilogy/core/processing/nodes/base_node.py

@@ -17,6 +17,7 @@ from trilogy.core.models import (
 from trilogy.core.enums import Purpose, JoinType, PurposeLineage, Granularity
 from trilogy.utility import unique
 from dataclasses import dataclass
+from trilogy.core.enums import BooleanOperator
 
 
 def concept_list_to_grain(
@@ -71,7 +72,6 @@ def resolve_concept_map(
                 concept_map[concept.address].add(input)
             elif concept.address not in concept_map:
                 concept_map[concept.address].add(input)
-
     # second loop, include partials
     for input in inputs:
         for concept in input.output_concepts:
@@ -92,14 +92,31 @@ def resolve_concept_map(
     return concept_map
 
 
-def get_all_parent_partial(all_concepts: List[Concept], parents: List["StrategyNode"]):
-    return [
-        c
-        for c in all_concepts
-        if len([c.address in [x.address for x in p.partial_concepts] for p in parents])
-        >= 1
-        and all([c.address in [x.address for x in p.partial_concepts] for p in parents])
-    ]
+def get_all_parent_partial(
+    all_concepts: List[Concept], parents: List["StrategyNode"]
+) -> List[Concept]:
+    return unique(
+        [
+            c
+            for c in all_concepts
+            if len(
+                [
+                    p
+                    for p in parents
+                    if c.address in [x.address for x in p.partial_concepts]
+                ]
+            )
+            >= 1
+            and all(
+                [
+                    c.address in p.partial_lcl
+                    for p in parents
+                    if c.address in p.output_lcl
+                ]
+            )
+        ],
+        "address",
+    )
 
 
 class StrategyNode:
@@ -137,7 +154,7 @@ class StrategyNode:
         self.partial_concepts = partial_concepts or get_all_parent_partial(
             self.output_concepts, self.parents
         )
-        self.partial_lcl = LooseConceptList(concepts=self.partial_concepts)
+
         self.depth = depth
         self.conditions = conditions
         self.grain = grain
@@ -146,13 +163,50 @@ class StrategyNode:
         self.hidden_concepts = hidden_concepts or []
         self.existence_concepts = existence_concepts or []
         self.virtual_output_concepts = virtual_output_concepts or []
+        self.validate_parents()
+
+    def add_parents(self, parents: list["StrategyNode"]):
+        self.parents += parents
+        self.validate_parents()
+
+    def add_condition(self, condition: Conditional | Comparison | Parenthetical):
+        if self.conditions:
+            self.conditions = Conditional(
+                left=self.conditions, right=condition, operator=BooleanOperator.AND
+            )
+        else:
+            self.conditions = condition
+
+    def validate_parents(self):
+        # validate parents exist
+        # assign partial values where needed
         for parent in self.parents:
             if not parent:
                 raise SyntaxError("Unresolvable parent")
 
+        # TODO: make this accurate
+        if self.parents:
+            self.partial_concepts = get_all_parent_partial(
+                self.output_concepts, self.parents
+            )
+
+        self.partial_lcl = LooseConceptList(concepts=self.partial_concepts)
+
     def add_output_concepts(self, concepts: List[Concept]):
         for concept in concepts:
-            self.output_concepts.append(concept)
+            if concept.address not in self.output_lcl.addresses:
+                self.output_concepts.append(concept)
+        self.output_lcl = LooseConceptList(concepts=self.output_concepts)
+        self.rebuild_cache()
+
+    def add_existence_concepts(self, concepts: List[Concept]):
+        for concept in concepts:
+            if concept.address not in [x.address for x in self.output_concepts]:
+                self.existence_concepts.append(concept)
+        self.rebuild_cache()
+
+    def set_output_concepts(self, concepts: List[Concept]):
+        self.output_concepts = concepts
         self.output_lcl = LooseConceptList(concepts=self.output_concepts)
         self.rebuild_cache()
 
@@ -195,8 +249,8 @@ class StrategyNode:
         grain = self.grain if self.grain else Grain(components=self.output_concepts)
         source_map = resolve_concept_map(
             parent_sources,
-            self.output_concepts,
-            self.input_concepts + self.existence_concepts,
+            targets=self.output_concepts,
+            inherited_inputs=self.input_concepts + self.existence_concepts,
         )
         return QueryDatasource(
             input_concepts=self.input_concepts,

trilogy/core/processing/nodes/filter_node.py

@@ -36,6 +36,7 @@ class FilterNode(StrategyNode):
         partial_concepts: List[Concept] | None = None,
         force_group: bool | None = False,
         grain: Grain | None = None,
+        existence_concepts: List[Concept] | None = None,
     ):
         super().__init__(
             output_concepts=output_concepts,
@@ -49,6 +50,7 @@ class FilterNode(StrategyNode):
             partial_concepts=partial_concepts,
             force_group=force_group,
             grain=grain,
+            existence_concepts=existence_concepts,
         )
 
     def copy(self) -> "FilterNode":
@@ -64,4 +66,5 @@ class FilterNode(StrategyNode):
             partial_concepts=list(self.partial_concepts),
             force_group=self.force_group,
             grain=self.grain,
+            existence_concepts=list(self.existence_concepts),
         )

trilogy/core/processing/nodes/group_node.py

@@ -39,6 +39,7 @@ class GroupNode(StrategyNode):
         partial_concepts: Optional[List[Concept]] = None,
         force_group: bool | None = None,
         conditions: Conditional | Comparison | Parenthetical | None = None,
+        existence_concepts: List[Concept] | None = None,
     ):
         super().__init__(
             input_concepts=input_concepts,
@@ -51,6 +52,7 @@ class GroupNode(StrategyNode):
             partial_concepts=partial_concepts,
             force_group=force_group,
             conditions=conditions,
+            existence_concepts=existence_concepts,
         )
 
     def _resolve(self) -> QueryDatasource:
@@ -128,7 +130,7 @@ class GroupNode(StrategyNode):
                     if self.conditions
                     else self.output_concepts
                 ),
-                inherited_inputs=self.input_concepts,
+                inherited_inputs=self.input_concepts + self.existence_concepts,
             ),
             joins=[],
             grain=grain,
@@ -137,10 +139,6 @@ class GroupNode(StrategyNode):
         )
         # if there is a condition on a group node and it's not scalar
         # inject an additional CTE
-        if self.conditions:
-            logger.info("CONDITIONS")
-            logger.info(str(self.conditions))
-            logger.info(is_scalar_condition(self.conditions))
         if self.conditions and not is_scalar_condition(self.conditions):
             base.condition = None
             base.output_concepts = self.output_concepts + self.conditions.row_arguments
@@ -173,4 +171,5 @@ class GroupNode(StrategyNode):
             partial_concepts=list(self.partial_concepts),
             force_group=self.force_group,
             conditions=self.conditions,
+            existence_concepts=list(self.existence_concepts),
         )

trilogy/core/processing/nodes/merge_node.py

@@ -138,16 +138,6 @@ class MergeNode(StrategyNode):
                     continue
                 final_joins.append(join)
             self.node_joins = final_joins
-        partial_lookup: list[Concept] = []
-        non_partial: List[Concept] = []
-        for node in parents or []:
-            partial_lookup += node.partial_concepts
-            non_partial += [
-                x for x in node.output_concepts if x not in node.partial_concepts
-            ]
-
-        final_partial = [x for x in partial_lookup if x not in non_partial]
-        self.partial_concepts = final_partial
 
     def translate_node_joins(self, node_joins: List[NodeJoin]) -> List[BaseJoin]:
         joins = []
@@ -265,7 +255,6 @@ class MergeNode(StrategyNode):
             for x in final_datasets
             if all([y in self.existence_concepts for y in x.output_concepts])
         ]
-
         if len(merged.keys()) == 1:
             final: QueryDatasource | Datasource = list(merged.values())[0]
             if (
@@ -378,4 +367,5 @@ class MergeNode(StrategyNode):
             node_joins=self.node_joins,
             join_concepts=list(self.join_concepts) if self.join_concepts else None,
             force_join_type=self.force_join_type,
+            existence_concepts=list(self.existence_concepts),
         )

trilogy/core/processing/nodes/select_node_v2.py

@@ -166,6 +166,7 @@ class SelectNode(StrategyNode):
             resolution = self.resolve_from_provided_datasource()
             if resolution:
                 return resolution
+
         required = [c.address for c in self.all_concepts]
         raise NoDatasourceException(
             f"Could not find any way to resolve datasources for required concepts {required} with derivation {[x.derivation for x in self.all_concepts]}"

trilogy/core/processing/utility.py

@@ -20,7 +20,7 @@ from trilogy.core.models import (
     DataType,
 )
 
-from trilogy.core.enums import Purpose, Granularity
+from trilogy.core.enums import Purpose, Granularity, BooleanOperator
 from trilogy.core.constants import CONSTANT_DATASET
 from enum import Enum
 from trilogy.utility import unique
@@ -62,7 +62,7 @@ def create_log_lambda(prefix: str, depth: int, logger: Logger):
     pad = padding(depth)
 
     def log_lambda(msg: str):
-        logger.info(f"{pad} {prefix} {msg}")
+        logger.info(f"{pad}{prefix} {msg}")
 
     return log_lambda
 
@@ -328,14 +328,6 @@ def get_disconnected_components(
     for datasource, concepts in concept_map.items():
         graph.add_node(datasource, type=NodeType.NODE)
         for concept in concepts:
-            # TODO: determine if this is the right way to handle things
-            # if concept.derivation in (PurposeLineage.FILTER, PurposeLineage.WINDOW):
-            #     if isinstance(concept.lineage, FilterItem):
-            #         graph.add_node(concept.lineage.content.address, type=NodeType.CONCEPT)
-            #         graph.add_edge(datasource, concept.lineage.content.address)
-            #     if isinstance(concept.lineage, WindowItem):
-            #         graph.add_node(concept.lineage.content.address, type=NodeType.CONCEPT)
-            #         graph.add_edge(datasource, concept.lineage.content.address)
             graph.add_node(concept.address, type=NodeType.CONCEPT)
             graph.add_edge(datasource, concept.address)
             all_concepts.add(concept)
@@ -363,18 +355,58 @@ def is_scalar_condition(
         | MagicConstants
         | DataType
     ),
+    materialized: set[str] | None = None,
 ) -> bool:
     if isinstance(element, Parenthetical):
-        return is_scalar_condition(element.content)
+        return is_scalar_condition(element.content, materialized)
     elif isinstance(element, SubselectComparison):
         return True
     elif isinstance(element, Comparison):
-        return is_scalar_condition(element.left) and is_scalar_condition(element.right)
+        return is_scalar_condition(element.left, materialized) and is_scalar_condition(
+            element.right, materialized
+        )
     elif isinstance(element, Function):
         if element.operator in FunctionClass.AGGREGATE_FUNCTIONS.value:
             return False
+    elif isinstance(element, Concept):
+        if materialized and element.address in materialized:
+            return True
+        if element.lineage and isinstance(element.lineage, AggregateWrapper):
+            return is_scalar_condition(element.lineage, materialized)
+        return True
     elif isinstance(element, AggregateWrapper):
-        return is_scalar_condition(element.function)
+        return is_scalar_condition(element.function, materialized)
     elif isinstance(element, Conditional):
-        return is_scalar_condition(element.left) and is_scalar_condition(element.right)
+        return is_scalar_condition(element.left, materialized) and is_scalar_condition(
+            element.right, materialized
+        )
     return True
+
+
+def decompose_condition(
+    conditional: Conditional | Comparison | Parenthetical,
+) -> list[SubselectComparison | Comparison | Conditional | Parenthetical]:
+    chunks: list[SubselectComparison | Comparison | Conditional | Parenthetical] = []
+    if not isinstance(conditional, Conditional):
+        return [conditional]
+    if conditional.operator == BooleanOperator.AND:
+        if not (
+            isinstance(
+                conditional.left,
+                (SubselectComparison, Comparison, Conditional, Parenthetical),
+            )
+            and isinstance(
+                conditional.right,
+                (SubselectComparison, Comparison, Conditional, Parenthetical),
+            )
+        ):
+            chunks.append(conditional)
+        else:
+            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

trilogy/core/query_processor.py

@@ -4,9 +4,9 @@ from trilogy.core.env_processor import generate_graph
 from trilogy.core.graph_models import ReferenceGraph
 from trilogy.core.constants import CONSTANT_DATASET
 from trilogy.core.processing.concept_strategies_v3 import source_query_concepts
-from trilogy.core.enums import SelectFiltering
+from trilogy.core.enums import SelectFiltering, BooleanOperator
 from trilogy.constants import CONFIG, DEFAULT_NAMESPACE
-from trilogy.core.processing.nodes import GroupNode, SelectNode, StrategyNode
+from trilogy.core.processing.nodes import GroupNode, SelectNode, StrategyNode, History
 from trilogy.core.models import (
     Concept,
     Environment,
@@ -24,6 +24,7 @@ from trilogy.core.models import (
     Datasource,
     BaseJoin,
     InstantiatedUnnestJoin,
+    Conditional,
 )
 
 from trilogy.utility import unique
@@ -260,7 +261,7 @@ def datasource_to_ctes(
 
     human_id = generate_cte_name(query_datasource.full_name, name_map)
     logger.info(
-        f"Finished building source map for {human_id} with {len(parents)} parents, have {source_map}, query_datasource had non-empty keys {[k for k, v in query_datasource.source_map.items() if v]} "
+        f"Finished building source map for {human_id} with {len(parents)} parents, have {source_map}, query_datasource had non-empty keys {[k for k, v in query_datasource.source_map.items() if v]} and existence had non-empty keys {[k for k, v in query_datasource.existence_source_map.items() if v]} "
     )
     final_joins = [
         x
@@ -307,7 +308,10 @@ def datasource_to_ctes(
 
 
 def append_existence_check(
-    node: StrategyNode, environment: Environment, graph: ReferenceGraph
+    node: StrategyNode,
+    environment: Environment,
+    graph: ReferenceGraph,
+    history: History | None = None,
 ):
     # we if we have a where clause doing an existence check
     # treat that as separate subquery
@@ -318,25 +322,22 @@ def append_existence_check(
             logger.info(
                 f"{LOGGER_PREFIX} fetching existance clause inputs {[str(c) for c in subselect]}"
             )
-            eds = source_query_concepts([*subselect], environment=environment, g=graph)
-
-            final_eds = eds.resolve()
-            first_parent = node.resolve()
-            first_parent.datasources.append(final_eds)
-            for x in final_eds.output_concepts:
-                if x.address not in first_parent.existence_source_map:
-                    first_parent.existence_source_map[x.address] = {final_eds}
+            eds = source_query_concepts(
+                [*subselect], environment=environment, g=graph, history=history
+            )
+            node.add_parents([eds])
+            node.add_existence_concepts([*subselect])
 
 
-def get_query_datasources(
+def get_query_node(
     environment: Environment,
     statement: SelectStatement | MultiSelectStatement,
     graph: Optional[ReferenceGraph] = None,
-    hooks: Optional[List[BaseHook]] = None,
-) -> QueryDatasource:
+    history: History | None = None,
+) -> StrategyNode:
     graph = graph or generate_graph(environment)
     logger.info(
-        f"{LOGGER_PREFIX} getting source datasource for query with output {[str(c) for c in statement.output_components]}"
+        f"{LOGGER_PREFIX} getting source datasource for query with filtering {statement.where_clause_category} and output {[str(c) for c in statement.output_components]}"
     )
     if not statement.output_components:
         raise ValueError(f"Statement has no output components {statement}")
@@ -353,22 +354,28 @@ def get_query_datasources(
         )
         nest_where = True
 
-    ods = source_query_concepts(
+    ods: StrategyNode = source_query_concepts(
         search_concepts,
         environment=environment,
         g=graph,
+        conditions=(statement.where_clause if statement.where_clause else None),
+        history=history,
     )
-    ds: GroupNode | SelectNode
+    if not ods:
+        raise ValueError(
+            f"Could not find source query concepts for {[x.address for x in search_concepts]}"
+        )
+    ds: StrategyNode
     if nest_where and statement.where_clause:
         if not all_aggregate:
             ods.conditions = statement.where_clause.conditional
-        ods.output_concepts = search_concepts
+        ods.output_concepts = statement.output_components
         # ods.hidden_concepts = where_delta
         ods.rebuild_cache()
-        append_existence_check(ods, environment, graph)
+        append_existence_check(ods, environment, graph, history)
         ds = GroupNode(
             output_concepts=statement.output_components,
-            input_concepts=search_concepts,
+            input_concepts=statement.output_components,
             parents=[ods],
             environment=ods.environment,
             g=ods.g,
@@ -390,7 +397,26 @@ def get_query_datasources(
 
     else:
         ds = ods
+    if statement.having_clause:
+        if ds.conditions:
+            ds.conditions = Conditional(
+                left=ds.conditions,
+                right=statement.having_clause.conditional,
+                operator=BooleanOperator.AND,
+            )
+        else:
+            ds.conditions = statement.having_clause.conditional
+    return ds
+
+
+def get_query_datasources(
+    environment: Environment,
+    statement: SelectStatement | MultiSelectStatement,
+    graph: Optional[ReferenceGraph] = None,
+    hooks: Optional[List[BaseHook]] = None,
+) -> QueryDatasource:
 
+    ds = get_query_node(environment, statement, graph)
     final_qds = ds.resolve()
     if hooks:
         for hook in hooks:
@@ -475,6 +501,7 @@ def process_query(
         grain=statement.grain,
         limit=statement.limit,
         where_clause=statement.where_clause,
+        having_clause=statement.having_clause,
         output_columns=statement.output_components,
         ctes=final_ctes,
         base=root_cte,

trilogy/dialect/base.py

@@ -2,7 +2,7 @@ from typing import List, Union, Optional, Dict, Any, Sequence, Callable
 
 from jinja2 import Template
 
-from trilogy.core.processing.utility import is_scalar_condition
+from trilogy.core.processing.utility import is_scalar_condition, decompose_condition
 from trilogy.constants import CONFIG, logger, MagicConstants
 from trilogy.core.internal import DEFAULT_CONCEPTS
 from trilogy.core.enums import (
@@ -386,14 +386,19 @@ class BaseDialect:
                         e.right.address,
                         [
                             INVALID_REFERENCE_STRING(
-                                f"Missing source reference to {e.right.name}"
+                                f"Missing source reference to {e.right.address}"
                             )
                         ],
                     )
                 else:
                     lookup = lookup_cte.existence_source_map[e.right.address]
-
-                return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} (select {lookup[0]}.{self.QUOTE_CHARACTER}{e.right.safe_address}{self.QUOTE_CHARACTER} from {lookup[0]} where {lookup[0]}.{self.QUOTE_CHARACTER}{e.right.safe_address}{self.QUOTE_CHARACTER} is not null)"
+                if len(lookup) > 0:
+                    target = lookup[0]
+                else:
+                    target = INVALID_REFERENCE_STRING(
+                        f"Missing source CTE for {e.right.address}"
+                    )
+                return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} (select {target}.{self.QUOTE_CHARACTER}{e.right.safe_address}{self.QUOTE_CHARACTER} from {target} where {target}.{self.QUOTE_CHARACTER}{e.right.safe_address}{self.QUOTE_CHARACTER} is not null)"
             elif isinstance(e.right, (ListWrapper, Parenthetical, list)):
                 return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} {self.render_expr(e.right, cte=cte, cte_map=cte_map)}"
 
@@ -514,7 +519,9 @@ class BaseDialect:
                 ):
                     source = f"{render_unnest(self.UNNEST_MODE, self.QUOTE_CHARACTER, cte.join_derived_concepts[0], self.render_concept_sql, cte)}"
                 # direct - eg DUCK DB - can be directly selected inline
-                elif cte.join_derived_concepts and self.UNNEST_MODE == UnnestMode.DIRECT:
+                elif (
+                    cte.join_derived_concepts and self.UNNEST_MODE == UnnestMode.DIRECT
+                ):
                     source = None
                 else:
                     raise SyntaxError("CTE has joins but no from clause")
@@ -531,6 +538,20 @@ class BaseDialect:
             final_joins = []
         else:
             final_joins = cte.joins or []
+        where: Conditional | Parenthetical | Comparison | None = None
+        having: Conditional | Parenthetical | Comparison | None = None
+        materialized = {x for x, v in cte.source_map.items() if v}
+        if cte.condition:
+            if is_scalar_condition(cte.condition, materialized=materialized):
+                where = cte.condition
+            else:
+                components = decompose_condition(cte.condition)
+                for x in components:
+                    if is_scalar_condition(x, materialized=materialized):
+                        where = where + x if where else x
+                    else:
+                        having = having + x if having else x
+
         return CompiledCTE(
             name=cte.name,
             statement=self.SQL_TEMPLATE.render(
@@ -554,16 +575,8 @@ class BaseDialect:
                     ]
                     if j
                 ],
-                where=(
-                    self.render_expr(cte.condition, cte)
-                    if cte.condition and is_scalar_condition(cte.condition)
-                    else None
-                ),
-                having=(
-                    self.render_expr(cte.condition, cte)
-                    if cte.condition and not is_scalar_condition(cte.condition)
-                    else None
-                ),
+                where=(self.render_expr(where, cte) if where else None),
+                having=(self.render_expr(having, cte) if having else None),
                 order_by=(
                     [self.render_order_item(i, cte) for i in cte.order_by.items]
                     if cte.order_by

trilogy/dialect/duckdb.py

@@ -36,7 +36,7 @@ FUNCTION_MAP = {
 # we may return a static value
 FUNCTION_GRAIN_MATCH_MAP = {
     **FUNCTION_MAP,
-    FunctionType.COUNT: lambda args: "1",
+    FunctionType.COUNT: lambda args: f"{args[0]}",
     FunctionType.SUM: lambda args: f"{args[0]}",
     FunctionType.AVG: lambda args: f"{args[0]}",
 }