AQuery/engine/projection.py

from engine.ast import ColRef, TableInfo, ast_node, Context, include
from engine.groupby import groupby
from engine.join import join
from engine.expr import expr
from engine.orderby import orderby
from engine.scan import filter
from engine.utils import base62uuid, enlist, base62alp
from engine.ddl import outfile
import copy

class projection(ast_node):
    name='select'
    def __init__(self, parent:ast_node, node, context:Context = None, outname = None, disp = True):
        self.disp = disp
        self.outname = outname
        self.group_node = None
        ast_node.__init__(self, parent, node, context)
    def init(self, _):
        if self.outname is None:
            self.outname = self.context.gen_tmptable()

    def produce(self, node):
        p = node['select']
        self.projections = p if type(p) is list else [p]
        print(node)

    def spawn(self, node):
        self.datasource = None
        if 'from' in node:
            from_clause = node['from']
            if type(from_clause) is list:
                # from joins
                join(self, from_clause)
            elif type(from_clause) is dict:
                if 'value' in from_clause:
                    value = from_clause['value']
                    if type(value) is dict:
                        if 'select' in value:
                            # from subquery
                            projection(self, from_clause, disp = False)
                        else:
                            # TODO: from func over table
                            print(f'from func over table{node}')
                    elif type(value) is str:
                        self.datasource = self.context.tables_byname[value]
                if 'assumptions' in from_clause:
                    for assumption in enlist(from_clause['assumptions']):
                        orderby(self, assumption)

            elif type(from_clause) is str:
                self.datasource = self.context.tables_byname[from_clause]
            
            if self.datasource is None:
                raise ValueError('spawn error: from clause')
           
        if self.datasource is not None:
            self.datasource_changed = True
            self.prev_datasource = self.context.datasource
            self.context.datasource = self.datasource            
        if 'where' in node:
            self.datasource = filter(self, node['where'], True).output
            self.context.datasource = self.datasource            

        if 'groupby' in node:
            self.group_node = groupby(self, node['groupby'])
            self.datasource = copy.copy(self.datasource) # shallow copy
            self.datasource.groupinfo = self.group_node
        else:
            self.group_node = None
            
    def consume(self, node):
        self.inv = True
        disp_varname = 'd'+base62uuid(7)
        has_groupby = False
        if self.group_node is not None:
            # There is group by;
            has_groupby = True
        cexpr = f'(' 
        flatten = False
        cols = []
        self.out_table = TableInfo('out_'+base62uuid(4), [], self.context)
        if 'outfile' in node:
            flatten = True
            
        for i, proj in enumerate(self.projections):
            cname = ''
            compound = False
            self.datasource.rec = []
            if type(proj) is dict:
                if 'value' in proj:
                    e = proj['value']
                    if type(e) is str:
                        cname = e # TODO: deal w/ alias
                        cexpr += (f"{self.datasource.parse_tablenames(proj['value'])}")
                    elif type(e) is dict:
                        p_expr = expr(self, e)
                        cname = p_expr.cexpr
                        compound = True
                        cexpr += f"{cname}"
                    cname = ''.join([a if a in base62alp else '' for a in cname])
                    cexpr += ';'if i < len(self.projections)-1 else ''

            compound = compound and has_groupby and self.datasource.rec not in self.group_node.referenced
            
            cols.append(ColRef(f'{disp_varname}[{i}]', 'generic', self.out_table, 0, None, cname, i, compound=compound))
        self.out_table.add_cols(cols, False)
        
        cexpr += ')'
        if has_groupby:
            self.group_node.finalize(cexpr, disp_varname)
        else:
            self.emit(f'auto {disp_varname} = {cexpr};')
        self.datasource.group_node = None

        has_orderby = 'orderby' in node

        if has_orderby:
            self.datasource = self.out_table
            self.context.datasource = self.out_table # discard current ds
            orderby_node = orderby(self, node['orderby'])
            self.context.datasource.materialize_orderbys()
            self.emit_no_ln(f"{f'{disp_varname}:+' if flatten else ''}(")
            
        if self.disp or has_orderby:
            if len(self.projections) > 1:
                self.emit_no_ln(f"{'+' if self.inv else ''}{disp_varname}")
            else:
                self.emit_no_ln(f'print({disp_varname});')
            if flatten:
                self.emit_no_ln(f'{disp_varname}')
        if has_orderby:
            self.emit(f')[{orderby_node.view}]')
        else:
            self.context.emit_flush()
        if flatten:
            if len(self.projections) > 1 and not self.inv:
                self.emit(f"{disp_varname}:+{disp_varname}")
            outfile(self, node['outfile'])

        if self.datasource_changed:
            self.context.datasource = self.prev_datasource


import sys
include(sys.modules[__name__])