def visit_query_engine_node()

in smallpond/logical/optimizer.py [0:0]

• 15 lines of code
• 7 McCabe index (conditional complexity)

    def visit_query_engine_node(self, node: SqlEngineNode, depth: int) -> Node:
        # fuse consecutive SqlEngineNodes
        if len(node.input_deps) == 1 and isinstance(child := self.visit(node.input_deps[0], depth + 1), SqlEngineNode):
            fused = copy.copy(node)
            fused.input_deps = child.input_deps
            fused.udfs = node.udfs + child.udfs
            fused.cpu_limit = max(node.cpu_limit, child.cpu_limit)
            fused.gpu_limit = max(node.gpu_limit, child.gpu_limit)
            fused.memory_limit = (
                max(node.memory_limit, child.memory_limit)
                if node.memory_limit is not None and child.memory_limit is not None
                else node.memory_limit or child.memory_limit
            )
            # merge the sql queries
            # example:
            # ```
            # child.sql_queries = ["select * from {0}"]
            #  node.sql_queries = ["select a, b from {0}"]
            # fused.sql_queries = ["select a, b from (select * from {0})"]
            # ```
            fused.sql_queries = child.sql_queries[:-1] + [query.format(f"({child.sql_queries[-1]})") for query in node.sql_queries]
            return fused
        return self.generic_visit(node, depth)