EMPTY_NODE = _EmptyNode() class Field(EvaluatableLeafNode):
    '''A reference to a column in a table.

    For example, in a table Table with columns A, B, and C,
    a valid Field path would be ('B',) or ('Table', 'A').
    '''
    def __init__(self, path):
        # type: (Tuple[str, ...]) -> None
        '''Set up Field node

        Args:
            path: A tuple of strings describing the path to this field
                (just the column name, or table.column)
        '''
        self.path = path

    def strexpr(self):
        # type: () -> str
        return '.'.join(self.path)

    def mark_grouped_by(self, group_by_paths, context):
        # type: (Sequence[Tuple[str, ...]], EvaluationContext) -> EvaluatableNode
        if context.get_canonical_path(self.path) in group_by_paths:
            return GroupedBy(self)
        return self

    def name(self):
        # type: () -> Optional[str]
        return self.path[-1]

    def _evaluate_leaf_node(self, context):
        # type: (EvaluationContext) -> TypedSeries
        # Retrieves the column from the context
        result = context.lookup(self.path)
        result.series.name = self.name()
        return result

    def __eq__(self, other):
        # type: (Any) -> bool
        if isinstance(other, Field):
            return self.path == other.path
        return False

    def is_constant(self):
        # type: () -> bool
        return False def binary_operator_expression_rule(subexpression_rule):
    # type: (RuleType) -> Callable[[List[str]], AppliedRuleOutputType]
    """Returns a rule for parsing binary expressions given a rule that the operators separate.

    Args:
        subexpression_rule: A grammar rule for some kind of expressions, i.e. a function
            taking tokens and returning a parsed abstract syntax tree node plus leftover tokens.
    Returns:
        A grammar rule that parses expressions where one or more subexpressions (as recognized by
            the rule passed in) are separated by binary operators, parsed according to operator
            precedence.
    """
    return separated_sequence(subexpression_rule,
                              sorted(BINARY_OPERATOR_INFO.keys(), key=len, reverse=True),
                              wrapper=_reparse_binary_expression,
                              keep_separator=True) class QueryExpression(DataframeNode):
    '''Highest level definition of a query.

    https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#sql-syntax
    (see query_expr)
    '''

    def __init__(self,
                 with_clauses,  # type: Union[_EmptyNode, List[Tuple[str, DataframeNode]]]
                 base_query,  # type: DataframeNode
                 order_by,  # type: Union[_EmptyNode, _OrderByType]
                 limit,  # type: Union[_EmptyNode, _LimitType]
                 ):
        # type: (...) -> None
        '''Set up QueryExpression node.

        Args:
            with_clauses: Optional WITH expression
            base_query: Main part of query
            order_by: Expression by which to order results
            limit: Number of rows to return, possibly with an offset
        '''

        self.with_clauses = with_clauses
        self.base_query = base_query
        self.order_by = order_by
        self.limit = limit

    def _order_by(self, order_by, typed_dataframe, table_name, table_context):
        # type: (_OrderByType, TypedDataFrame, Optional[str], TableContext) -> TypedDataFrame
        '''If ORDER BY is specified, sort the data by the given column(s)
        in the given direction(s).

        Args:
            typed_dataframe: The currently resolved query as a TypedDataFrame
            table_name: Resolved name of current typed_dataframe
            table_context: A representation of the state of available tables
        Returns:
            A new TypedDataFrame that is ordered by the given criteria
        '''
        context = EvaluationContext(table_context)
        context.add_table_from_dataframe(typed_dataframe, table_name, EMPTY_NODE)

        # order_by is a list of (field, direction) tuples to sort by
        fields = []
        directions = []  # ascending = True, descending = False
        for field, direction in order_by:
            if isinstance(field, Field):
                path = '.'.join(context.get_canonical_path(field.path))
                fields.append(path)
            elif isinstance(field, Value):
                if not isinstance(field.value, int):
                    raise ValueError('Attempt to order by a literal non-integer constant {}'
                                     .format(field.value))
                index = field.value - 1  # order by 1 means the first field, i.e. index 0
                fields.append(context.table.dataframe.columns[index])
            else:
                raise ValueError('Invalid field specification {}'.format(field))

            if direction == 'DESC':
                directions.append(False)
            else:
                # Default sort order in Standard SQL is ASC
                directions.append(True)
        return TypedDataFrame(
            context.table.dataframe.sort_values(fields, ascending=directions),
            context.table.types)

    def _limit(self, limit, typed_dataframe):
        # type: (_LimitType, TypedDataFrame) -> TypedDataFrame
        '''If limit is specified, only return that many rows.
        If offset is specified, start at that row number, not the first row.

        Args:
            typed_dataframe: The currently resolved query as a TypedDataFrame
        Returns:
            A new TypedDataFrame that conforms to the given limit and offset
        '''
        limit_expression, offset_expression = limit

        # Use empty context because the limit is a constant
        limit_value = limit_expression.evaluate(EMPTY_CONTEXT)
        if not isinstance(limit_value, TypedSeries):
            raise ValueError("invalid limit expression {}".format(limit_expression))
        limit, = limit_value.series
        if offset_expression is not EMPTY_NODE:
            # Use empty context because the offset is also a constant
            offset_value = offset_expression.evaluate(EMPTY_CONTEXT)
            if not isinstance(offset_value, TypedSeries):
                raise ValueError("invalid offset expression {}".format(offset_expression))
            offset, = offset_value.series
        else:
            offset = 0
        return TypedDataFrame(
            typed_dataframe.dataframe[offset:limit + offset],
            typed_dataframe.types)

    def get_dataframe(self, table_context, outer_context=None):
        # type: (TableContext, Optional[EvaluationContext]) -> Tuple[TypedDataFrame, Optional[str]]
        '''See parent, DataframeNode'''
        if not isinstance(self.with_clauses, _EmptyNode):
            name_list = [name for name, _ in self.with_clauses]
            if len(name_list) > len(set(name_list)):
                raise ValueError("Duplicate names in WITH clauses are not allowed: {}"
                                 .format(name_list))
            for name, dataframe_node in self.with_clauses:
                table_context = _WithTableContext(name,
                                                  dataframe_node.get_dataframe(table_context)[0],
                                                  table_context)

        typed_dataframe, table_name = self.base_query.get_dataframe(table_context, outer_context)

        if not isinstance(self.order_by, _EmptyNode):
            typed_dataframe = self._order_by(
                self.order_by, typed_dataframe, table_name, table_context)

        if not isinstance(self.limit, _EmptyNode):
            typed_dataframe = self._limit(self.limit, typed_dataframe)

        return typed_dataframe, DEFAULT_TABLE_NAME class Select(MarkerSyntaxTreeNode, DataframeNode):
    '''SELECT query to retrieve rows from a table(s).

    https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#select-list
    '''

    def __init__(self, modifier,  # type: AbstractSyntaxTreeNode
                 fields,  # type: Sequence[Union[Selector, StarSelector]]
                 from_,  # type: Union[_EmptyNode, DataSource]
                 where,  # type: Union[_EmptyNode, EvaluatableNode]
                 group_by,  # type: Union[_EmptyNode, List[Union[Value, Field]]]
                 having  # type: Union[_EmptyNode, EvaluatableNode]
                 ):
        # type: (...) -> None
        '''Set up SELECT node.

        Args:
            modifier: Optional ALL or DISTINCT modifier
            fields: Columns to return
            from_: Table/expression from which to retrieve rows
            where: WHERE filter condition, if any
            group_by: GROUP BY grouping condition, if any
            having: HAVING filter condition, if any
        '''
        self.modifier = modifier
        self.fields = fields
        for i, field in enumerate(self.fields):
            field.position = i + 1  # position is 1-up, i.e the first selector is position #1.
        self.from_ = from_
        self.where = where
        if isinstance(group_by, _EmptyNode):
            self.group_by = group_by  # type: Union[_EmptyNode, List[Field]]
        else:
            self.group_by = []
            for grouper in group_by:
                if isinstance(grouper, Value):
                    if not isinstance(grouper.value, int):
                        raise ValueError('Attempt to group by a literal non-integer constant {}'
                                         .format(grouper.value))
                    # GROUP BY 3 means group by the third field in the select, the field at index 2,
                    # i.e. we have to subtract one from the user-specified value to get the index.
                    # We construct a one-element field path just as if they'd specified the name
                    # of the corresponding field.
                    nth_field = self.fields[grouper.value - 1]
                    if not isinstance(nth_field, Selector):
                        raise ValueError('cannot GROUP BY {}th selector'.format(grouper.value))
                    grouper_path = (nth_field.name(),)
                    self.group_by.append(Field(grouper_path))
                else:
                    self.group_by.append(grouper)
        self.having = having

    def get_dataframe(self, table_context, outer_context=None):
        # type: (TableContext, Optional[EvaluationContext]) -> Tuple[TypedDataFrame, Optional[str]]
        '''Scope the given datasets by the criteria specified in the
        instance's fields.

        Args:
            table_context: All the tables in the database
            outer_context: The context of the outer query, if this Select is a subquery;
                otherwise None
        Returns:
            Tuple of the resulting table (TypedDataFrame) and a name for
            this table
        '''

        if isinstance(self.from_, _EmptyNode):
            context = EvaluationContext(table_context)
        else:
            context = self.from_.create_context(table_context)

        if outer_context:
            context.add_subcontext(outer_context)

        # Expand out any * fields so that we have a list just of selectors.
        expanded_fields = list(itertools.chain(*[
            [selector] if isinstance(selector, Selector)
            else selector.get_selectors(context)
            for selector in self.fields]))

        context.selector_names = [
                selector.name() for selector in self.fields if isinstance(selector, Selector)]

        if not isinstance(self.where, _EmptyNode):
            # Filter table by WHERE condition
            rows_to_keep = self.where.evaluate(context)
            if not isinstance(rows_to_keep, TypedSeries):
                raise ValueError("Invalid WHERE expression {}".format(rows_to_keep))
            context.table = TypedDataFrame(
                context.table.dataframe.loc[rows_to_keep.series],
                context.table.types)

        if not isinstance(self.group_by, _EmptyNode):
            fields_for_evaluation = context.do_group_by(
                expanded_fields, self.group_by)  # type: Sequence[EvaluatableNode]
        elif any(field.is_aggregated() for field in expanded_fields):
            fields_for_evaluation = context.do_group_by(expanded_fields, [])
        else:
            fields_for_evaluation = expanded_fields
        result = _evaluate_fields_as_dataframe(fields_for_evaluation, context)

        if not isinstance(self.having, _EmptyNode):
            having_context = EvaluationContext(table_context)
            having_context.add_table_from_dataframe(result, None, EMPTY_NODE)
            having_context.add_subcontext(context)
            having_context.group_by_paths = context.group_by_paths
            having = self.having.mark_grouped_by(context.group_by_paths, having_context)
            rows_to_keep = having.evaluate(having_context)
            if not isinstance(rows_to_keep, TypedSeries):
                raise ValueError("Invalid HAVING expression {}".format(rows_to_keep))
            result = TypedDataFrame(result.dataframe.loc[rows_to_keep.series], result.types)

        if self.modifier == 'DISTINCT':
            result = TypedDataFrame(result.dataframe.drop_duplicates(), result.types)

        return result, DEFAULT_TABLE_NAME class SetOperation(DataframeNode):
    '''Represents a set operation between two other query expressions - UNION, INTERSECT, etc.'''

    def __init__(self, left_query, set_operator, right_query):
        # type: (DataframeNode, str, DataframeNode) -> None
        self.left_query = left_query
        self.set_operator = set_operator
        self.right_query = right_query

    def get_dataframe(self, table_context, outer_context=None):
        # type: (TableContext, Optional[EvaluationContext]) -> Tuple[TypedDataFrame, Optional[str]]
        '''See parent, DataframeNode'''
        left_dataframe, unused_left_name = self.left_query.get_dataframe(
            table_context, outer_context)
        right_dataframe, unused_right_name = self.right_query.get_dataframe(
            table_context, outer_context)
        num_left_columns = len(left_dataframe.types)
        num_right_columns = len(right_dataframe.types)
        if num_left_columns != num_right_columns:
            raise ValueError("Queries in {} ALL have mismatched column count: {} vs {}"
                             .format(self.set_operator, num_left_columns, num_right_columns))
        combined_types = [implicitly_coerce(left_type, right_type)
                          for left_type, right_type in zip(left_dataframe.types,
                                                           right_dataframe.types)]
        if self.set_operator == 'UNION_ALL':
            return TypedDataFrame(
                pd.concat([left_dataframe.dataframe,
                           # Rename second table to use first table's column names
                           right_dataframe.dataframe.rename(
                               columns=dict(zip(right_dataframe.dataframe.columns,
                                                left_dataframe.dataframe.columns)))]),
                combined_types), DEFAULT_TABLE_NAME
        else:
            raise NotImplementedError("set operation {} not implemented".format(self.set_operator)) class TableReference(DataframeNode):
    '''A table reference specified as Project.Dataset.Table (or possibly
    Dataset.Table or just Table if there is only one project and/or dataset).
    '''

    def __init__(self, path):
        # type: (Tuple[str, ...]) -> None

        # If the table reference is specified with backticks, it will be parsed
        # as one element, so we need to split into parts here.
        if len(path) == 1:
            split_path = path[0].split('.')  # type: List[str]
            path = tuple(split_path)
        self.path = path  # type: Tuple[str, ...]

    def get_dataframe(self, table_context, outer_context=None):
        # type: (TableContext, Optional[EvaluationContext]) -> Tuple[TypedDataFrame, Optional[str]]
        '''See parent, DataframeNode'''
        del outer_context  # Unused
        return table_context.lookup(self.path) class Unnest(DataframeNode, MarkerSyntaxTreeNode):
    '''An expression unnesting an array into a column of data.'''

    def __init__(self, array_node):
        # type: (Array) -> None

        self.array_node = array_node

    def get_dataframe(self, table_context, outer_context=None):
        # type: (TableContext, Optional[EvaluationContext]) -> Tuple[TypedDataFrame, Optional[str]]
        '''See parent, DataframeNode'''
        del outer_context  # Unused
        context = EvaluationContext(table_context)
        result = self.array_node.evaluate(context)
        if isinstance(result, TypedDataFrame):
            raise ValueError('UNNEST({}) did not result in one column'.format(self.array_node))
        result_type, = result.types

        if not isinstance(result_type, BQArray):
            raise ValueError("UNNESTing a non-array-typed value: {}".format(result_type))

        contained_type = result_type.type_

        if len(result.series) != 1:
            raise ValueError('UNNEST({}) did not result in one row'.format(self.array_node))
        result_array, = result.to_list()
        if not isinstance(result_array, tuple):
            raise ValueError("UNNEST({}) resulted in {!r} rather than an array"
                             .format(self.array_node, result_array))
        if isinstance(contained_type, BQStructType):
            i = 0
            columns = []
            for field in contained_type.fields:
                if field:
                    columns.append(field)
                else:
                    columns.append('f{}_'.format(i))
                    i += 1
            result_dataframe = TypedDataFrame(
                pd.DataFrame([[cell for cell in row] for row in result_array], columns=columns),
                contained_type.types)
        else:
            result_dataframe = TypedDataFrame(
                pd.DataFrame([[cell] for cell in result_array], columns=['f0_']), [contained_type])
        return result_dataframe, None class Array(EvaluatableNodeWithChildren):
    '''An ARRAY expression, grouping several columns into a column of ARRAYs.'''

    def __init__(self, maybe_type, maybe_values):
        # type: (Union[BQArray, _EmptyNode], Union[Sequence[EvaluatableNode], _EmptyNode]) -> None
        '''Creates an array from a literal in the grammar.

        Args:
            maybe_type: If the node is present, this is a string giving the type of the array.
            maybe_values: If the node is present, this is a sequence of expressions giving the
                contents of the array.
        '''
        self.array_type = maybe_type if not isinstance(maybe_type, _EmptyNode) else None
        if isinstance(maybe_values, _EmptyNode):
            # If this is an empty array, we need a placeholder child so that at
            # evaluation time, we know how many rows of empty arrays to produce.
            self.children = [Value(1, BQScalarType.INTEGER)]
            self.empty = True
        else:
            self.children = maybe_values
            self.empty = False

    def copy(self, new_children):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return Array(self.array_type or EMPTY_NODE, EMPTY_NODE if self.empty else new_children)

    def _evaluate_node(self, values):
        # type: (List[TypedSeries]) -> TypedSeries
        value_type = implicitly_coerce(*[value.type_ for value in values])
        if self.array_type:
            join_type = implicitly_coerce(value_type, self.array_type.type_)
            if join_type != self.array_type.type_:
                raise ValueError("Array specifies type {}, incompatible with values of type {}"
                                 .format(self.array_type.type_, value_type))
            array_type = self.array_type
        elif not isinstance(value_type, (BQScalarType, BQStructType)):
            raise ValueError("Cannot create arrays of arrays")
        else:
            array_type = BQArray(value_type)
        if self.empty:
            # If this is an empty array, we have a placeholder child so that
            # we know how many rows of empty arrays to produce.
            arrays = [()] * len(values[0].series)
        else:
            arrays = list(zip(*[value.series for value in values]))
        if not isinstance(array_type.type_, (BQScalarType, BQStructType)):
            raise ValueError("Cannot create arrays of type {}".format(array_type.type_))
        return TypedSeries(pd.Series(arrays, index=values[0].series.index), array_type) class Array_agg(_AggregatingFunction):
    '''An ARRAY_AGG function, aggregating a column of results into an ARRAY-valued cell.'''

    @classmethod
    def create_function_call(
                 cls,
                 distinct,  # type: Union[_EmptyNode, str]
                 expression,  # type: EvaluatableNode
                 nulls,  # type: AbstractSyntaxTreeNode
                 order_by,  # type: AbstractSyntaxTreeNode
                 limit,  # type: AbstractSyntaxTreeNode
                 over_clause,  # type: _OverClauseType
                 ):
        # type: (...) -> EvaluatableNode
        '''Creates an ARRAY_AGG function call based on the inputs from the grammar.

        Args:
            distinct: 'DISTINCT' if the function should return only distinct rows,
                empty otherwise.
            expression: The column to aggregate
            nulls: 'IGNORE' if NULLs should not be returned; 'RESPECT' or empty otherwise.
            order_by: A syntax subtree describing how to order the rows.  Not implemented.
            limit: A syntax subtree giving how many rows to return at most.  Not implemented.
            over_clause: The window to evaluate the function over as an analytic function.
        '''
        function = Array_agg(distinct, nulls, order_by, limit)
        if isinstance(over_clause, _EmptyNode):
            return _AggregatingFunctionCall(function, [expression])
        else:
            return _AnalyticFunctionCall(function, [expression], over_clause)

    def __init__(self,
                 distinct,  # type: Union[_EmptyNode, str]
                 nulls,  # type: AbstractSyntaxTreeNode
                 order_by,  # type: AbstractSyntaxTreeNode
                 limit,  # type: AbstractSyntaxTreeNode
                 ):
        # type: (...) -> None
        if distinct == 'DISTINCT':
            self.distinct = True
        elif distinct is EMPTY_NODE:
            self.distinct = False
        else:
            raise ValueError("Invalid syntax: ARRAY_AGG({}...)".format(distinct))

        if nulls == 'IGNORE':
            self.drop_nulls = True
        elif nulls in ('RESPECT', EMPTY_NODE):
            self.drop_nulls = False
        else:
            raise ValueError("Invalid Syntax: ARRAY_AGG(...{}...)".format(nulls))

        self.order_by = order_by
        self.limit = limit

    def compute_result_type(self, argument_types):
        # type: (Sequence[BQType]) -> BQType
        computed_argument_type = implicitly_coerce(*argument_types)
        if not isinstance(computed_argument_type, BQScalarType):
            raise ValueError("ARRAYs are only supported of scalar types")
            # TODO: support ARRAYs of structs.
        return BQArray(computed_argument_type)

    def aggregating_function(self, values):
        # type: (List[pd.Series]) -> Tuple
        if len(values) != 1:
            raise ValueError("ARRAY_AGG takes exactly one argument, not {}"
                             .format(len(values)))
        result = values[0]
        if self.distinct:
            result = result.drop_duplicates()
        if self.drop_nulls:
            result = result.dropna()
        if self.order_by is not EMPTY_NODE:
            raise NotImplementedError("ARRAY_AGG(ORDER BY) is not implemented")
        if self.limit is not EMPTY_NODE:
            raise NotImplementedError("ARRAY_AGG(LIMIT) is not implemented")

        return tuple(result) class Case(MarkerSyntaxTreeNode, EvaluatableNodeWithChildren):
    '''A CASE expression, similar to an IF.  For example:
    CASE WHEN a > 0 THEN "positive" WHEN a < 0 THEN "negative" ELSE "zero" END
    CASE a WHEN 1 THEN "one" WHEN 2 THEN "two" END
    '''
    def __init__(self, comparand,  # type: AbstractSyntaxTreeNode
                 whens,  # type: List[Tuple[AbstractSyntaxTreeNode, EvaluatableNode]]
                 else_  # type: Union[EvaluatableNode, _EmptyNode]
                 ):
        # type: (...) -> None
        '''Set up Case node

        Args:
            comparand: An expression that should be compared to each of the
                provided WHENs, or an empty node if the whens are expressions
                that evaluate to booleans
            whens: A list of tuples describing each WHEN. The first value is
                either a condition that evaluates to a boolean, or a value to
                compare against the comparand. The second value is the
                expression that follows the THEN.
            else_: An optional ELSE expression
        '''
        self.children = []  # type: List[EvaluatableNode]
        if not whens:
            raise ValueError("Must provide at least one WHEN for a CASE")
        for when, then in whens:
            if not isinstance(when, EvaluatableNode):
                raise ValueError("Invalid CASE expression; WHEN {} must be an expression."
                                 .format(comparand))
            if not isinstance(comparand, _EmptyNode):
                # If there's a comparand, generate conditions by comparing the
                # comparand with each WHEN
                # This allows us to consistently treat each when as a binary expression,
                # and not have to separately store the comparand.
                if not isinstance(comparand, EvaluatableNode):
                    raise ValueError("Invalid CASE expression; comparand {} must be an expression."
                                     .format(comparand))
                when = BinaryExpression(comparand, '=', when)
            self.children.append(when)
            self.children.append(then)
        # The fallback value, if no conditions are true, will be the ELSE if provided else NULL
        self.children.append(else_ if not isinstance(else_, _EmptyNode) else Value(None, None))

    def copy(self, new_children):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        new_children = list(new_children)
        new_else_ = new_children.pop()
        return Case(
            comparand=EMPTY_NODE,
            whens=zip(new_children[::2], new_children[1::-2]),
            else_=new_else_)

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries

        # Evaluated_children has the following structure:
        # when0, then0, when1, then1, ... else
        # that is, a condition which, if true, is followed by what to return for that condition,
        # followed by another condition, etc.  The final value is what is returned when no
        # conditions are true.  We evaluate this as if it were a nested series of IF clauses,
        # i.e. IF(when0, then0, IF(when1, then1, ... IF(whenn, thenn, else_) ... ))
        #
        # First pop off the else column off the end.

        result = evaluated_children.pop()

        # Next, evaluate each nested IF/ELSE by walking backward through the whens
        for condition, then in zip(evaluated_children[-2::-2], evaluated_children[-1::-2]):
            if condition.type_ != BQScalarType.BOOLEAN:
                raise ValueError("CASE condition isn't boolean! Found: {!r}".format(condition))
            result = TypedSeries(then.series.where(condition.series, result.series),
                                 implicitly_coerce(then.type_, result.type_))
        return result class Cast(MarkerSyntaxTreeNode, EvaluatableNodeWithChildren):
    '''An expression that converts an expression of one type to another type.  For example:
    CAST(1 AS STRING)
    '''
    def __init__(self, expression, type_):
        # type: (EvaluatableNode, str) -> None
        '''Set up Cast node

        Args:
            expression: The expression whose type to change
            type_: The type to cast to (string must resolve to BQScalarType)
        '''
        self.children = [expression]  # type: List[EvaluatableNode]
        self.type_ = BQScalarType.from_string(type_)

    def copy(self, new_arguments):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return Cast(new_arguments[0], self.type_.value)

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries
        evaluated_expression = evaluated_children[0]
        return TypedSeries(
            evaluated_expression.series.astype(self.type_.to_dtype()), self.type_) class Count(_AggregatingFunction):
    '''A COUNT function, for example:
    SELECT COUNT(*) FROM Table
    '''

    _result_type = BQScalarType.INTEGER

    @classmethod
    def create_count_function_call(cls,
                                   countee,  # type: _CounteeType
                                   over_clause  # type: _OverClauseType
                                   ):
        # type: (...) -> EvaluatableNode
        '''Creates a Count expression.

        COUNT has a factory creation method, unlike other _Function subtypes, because it has
        unique syntax (COUNT(*), COUNT(DISTINCT expr)), and because it has state in addition
        to its child nodes (whether or not to count only distinct rows).  This method is
        called by the grammar to create a FunctionCall syntax tree node of the appropriate type
        (aggregating or analytic) with a Count expression having the appropriate state.

        Args:
            countee: Either a single string '*' or a tuple: an optional
                string 'DISTINCT' and a required expression to count
            over_clause: An optional OVER clause
        Returns:
            An AggregatingFunctionCall expression, if the over clause isn't present,
            or an _AnalyticFunctionCall expression, if it is, either way using this function.
        '''
        # Treat count(*) as if it were count(1), which is equivalent.
        if isinstance(countee, str):
            if countee != '*':
                raise ValueError("Invalid expression COUNT({})".format(countee))
            countee = (EMPTY_NODE, Value(1, BQScalarType.INTEGER))
        maybe_distinct, argument = countee
        if maybe_distinct == 'DISTINCT':
            distinct = True
        elif maybe_distinct == EMPTY_NODE:
            distinct = False
        else:
            raise NotImplementedError("Non-DISTINCT modifiers for COUNT are not implemented:"
                                      " {}".format(maybe_distinct))

        if isinstance(over_clause, _EmptyNode):
            return _AggregatingFunctionCall(Count(distinct), [argument])
        else:
            return _AnalyticFunctionCall(Count(distinct), [argument], over_clause)

    def __init__(self, distinct):
        # type: (bool) -> None
        self.distinct = distinct

    def aggregating_function(self, values):
        # type: (List[pd.Series]) -> int
        value, = values
        value = value.dropna()  # COUNT counts non-NULL rows
        if self.distinct:
            value = set(value)
        return len(value) class Exists(MarkerSyntaxTreeNode, EvaluatableLeafNode):
    '''An expression that returns TRUE if the subquery produces one or more rows.  For example:
    EXISTS(SELECT a FROM table WHERE a=1)
    '''
    def __init__(self, subquery):
        # type: (DataframeNode) -> None
        '''An EXISTS node

        Args:
            subquery: A subquery
        '''
        self.subquery = subquery

    def mark_grouped_by(self, group_by_paths, context):
        # type: (Sequence[Tuple[str, ...]], EvaluationContext) -> EvaluatableNode
        return self

    def _evaluate_leaf_node(self, context):
        # type: (EvaluationContext) -> TypedSeries
        # We need to calculate the Exists for each row in the current
        # context's table.  We also need to provide the current context to the
        # Select query, because it can reference fields from the current
        # context.  For example:
        # SELECT a, EXISTS(SELECT * FROM table_b WHERE table_a.a = table_b.b)
        # FROM table_a
        # This returns a row with: `a` and whether there exists a `b` that
        # equals it.  The inner Select query needs to know about table_a in
        # order to compare `a` to `b`.
        results = []  # type: List[bool]
        for index, row in context.table.dataframe.iterrows():
            # Create a new context just for this one row
            single_row_df = TypedDataFrame(
                pd.DataFrame([row], index=context.table.dataframe.index), context.table.types)
            row_context = EvaluationContext.clone_context_new_table(single_row_df, context)
            typed_df, df_name = self.subquery.get_dataframe(context.table_context, row_context)
            results.append(len(typed_df.dataframe) > 0)
        # Construct a Series that contains each of the individual result rows
        return TypedSeries(pd.Series(results, index=_get_index(context.table.dataframe)),
                           BQScalarType.BOOLEAN) class Extract(MarkerSyntaxTreeNode, EvaluatableNodeWithChildren):
    '''An expression that returns the value corresponding to the specified date part.
    For example:
    EXTRACT(DAY FROM timestamp)

    Reference:
    https://cloud.google.com/bigquery/docs/reference/standard-sql/date_functions#extract
    '''
    def __init__(self, part, date_expression):
        # type: (str, EvaluatableNode) -> None
        '''An EXTRACT node

        Args:
            part: The part of the date object to retrieve (day, month, year, etc)
            date_expression: The timestamp from which to extract the part
                (must evaluate to a pd.Timestamp)
        '''
        self.part = part.upper()
        self.children = [date_expression]

        weekdays = ['SUNDAY', 'MONDAY', 'TUESDAY', 'WEDNESDAY', 'THURSDAY', 'FRIDAY', 'SATURDAY']
        valid_parts = (['DAYOFWEEK', 'DAY', 'DAYOFYEAR', 'WEEK', 'ISOWEEK',
                        'MONTH', 'QUARTER', 'YEAR', 'ISOYEAR'] +
                       ['WEEK({})'.format(day) for day in weekdays])
        if self.part not in valid_parts:
            raise ValueError('Not a valid date part to retrieve: {}'.format(self.part))

    def copy(self, new_arguments):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return Extract(self.part, new_arguments[0])

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries
        date = evaluated_children[0].series
        # date_expression will evaluate to a pd.Timestamp, not datetime.datetime
        isoparts = ['ISOYEAR', 'ISOWEEK']
        if self.part in isoparts:
            # isocalendar() returns a tuple where first two elements are the
            # isoyear and the isoweek
            return TypedSeries(date.apply(lambda dt: dt.isocalendar()[isoparts.index(self.part)]),
                               BQScalarType.INTEGER)
        elif self.part.startswith('WEEK('):
            # WEEK(<weekday>) means the number of this week, assuming weeks
            # start with the given weekday.  There isn't an easy pandas analog
            # to this, and I don't think it's widely used
            raise NotImplementedError('{} not implemented'.format(self.part))
        else:
            return TypedSeries(date.apply(lambda dt: getattr(dt, self.part.lower())),
                               BQScalarType.INTEGER) class FunctionCall(object):
    '''Abstract base class for function call expressions. Subclasses are aggregating or not.'''

    _FUNCTION_MAP = {function_info.name(): function_info()
                     for function_info in (Min, Max, Sum, Mod, Concat, Timestamp, Current_Timestamp,
                                           Row_Number)}

    @classmethod
    def create(cls,
               function_name,  # type: str
               expression,  # type: Union[_EmptyNode, List[EvaluatableNode]]
               over_clause  # type: Union[_EmptyNode, _OverClauseType]
               ):
        # type: (...) -> EvaluatableNode
        function_name = function_name.upper()
        if function_name not in cls._FUNCTION_MAP:
            raise NotImplementedError('Function {} not implemented'.format(function_name))
        function_info = cls._FUNCTION_MAP[function_name]

        # The implementation of aggregating functions currently requires exactly one argument.
        # This works for sum, min, max; it excludes a few functions like corr.  The reason
        # for the restriction is that in a group by context, the aggregating function is called
        # by applying it to a single column.  Lifting the restriction will require combining
        # evaluated SeriesGroupBys into a single DataFrameGroupBy.
        if (isinstance(function_info, _AggregatingFunction) and
                (isinstance(expression, _EmptyNode) or len(expression) != 1)):
            raise NotImplementedError("Aggregating functions are only supported with 1 argument")

        if not isinstance(over_clause, _EmptyNode):
            return _AnalyticFunctionCall(function_info, expression, over_clause)
        elif isinstance(function_info, _AggregatingFunction):
            return _AggregatingFunctionCall(function_info, expression)
        elif isinstance(function_info, _NonAggregatingFunction):
            return _NonAggregatingFunctionCall(function_info, expression)
        else:
            raise ValueError("Invalid function info {}".format(function_info))

    def _evaluate(self, arguments, function, compute_result_type):
        # type: (List[TypedSeries], _FunctionType, Callable[[List[BQType]], BQType]) -> TypedSeries
        '''Evaluates arguments using function.

        Args:
            arguments: A list of TypedSeries (columns of data)
            function: A function to apply to the arguments.
            compute_result_type: A function to call to compute the result type based on the types
                of the arguments.
        Returns:
            The result of applying the function to the arguments.
        '''
        argument_values = [argument.series for argument in arguments]
        argument_types = [argument.type_ for argument in arguments]
        return TypedSeries(function(argument_values), compute_result_type(argument_types)) class If(MarkerSyntaxTreeNode, EvaluatableNodeWithChildren):
    '''A conditional expression, for example:

    IF(a > 0, a, 0)
        condition: a > 0
        then: a
        else_: 0
    Will return a if a is greater than zero, otherwise it will return zero.
    '''
    def __init__(self, condition, then, else_):
        # type: (EvaluatableNode, EvaluatableNode, EvaluatableNode) -> None
        '''Set up an If node

        Args:
            condition: A conditional expression, evaluates to a boolean
            then: Value to return if the condition is true
            else_: Value to return if the condition is false
        '''
        self.children = [condition, then, else_]

    def copy(self, new_arguments):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return If(*new_arguments)

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries
        condition, then, else_ = evaluated_children
        if condition.type_ != BQScalarType.BOOLEAN:
            raise ValueError("IF condition isn't boolean! Found: {!r}".format(condition))
        result_type = implicitly_coerce(then.type_, else_.type_)
        return TypedSeries(then.series.where(condition.series, else_.series), result_type) class InCheck(EvaluatableNodeWithChildren):
    '''Expression that checks whether element is in or not in a particular selection'''
    def __init__(self, expression, direction, elements):
        # type: (EvaluatableNode, str, Tuple[EvaluatableNode, ...]) -> None
        '''Set up InCheck node

        Args:
            expression: Expression to check
            direction: 'IN' or 'NOT_IN'
            elements: Set of values to be in (or not in)
        '''
        self.children = [expression] + list(elements)

        if direction == 'IN':
            self.direction = True
        elif direction == 'NOT_IN':
            self.direction = False
        else:
            raise ValueError("Invalid direction for InCheck, not IN or NOT_IN: {}"
                             .format(direction))

    def copy(self, new_children):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return InCheck(new_children[0],
                       'IN' if self.direction else 'NOT_IN',
                       tuple(new_children[1:]))

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries
        expression_value = evaluated_children[0]
        element_values = evaluated_children[1:]
        contained = pd.Series([False] * len(expression_value.series),
                              index=expression_value.series.index)
        for element_value in element_values:
            contained = contained | (expression_value.series == element_value.series)
        if self.direction:
            return TypedSeries(contained, BQScalarType.BOOLEAN)
        else:
            # ~ is the boolean NOT operator
            return TypedSeries(~contained, BQScalarType.BOOLEAN) class Not(MarkerSyntaxTreeNode, EvaluatableNodeWithChildren):
    '''Expression that negates the boolean series, such as turning [True] into [False]'''
    def __init__(self, expression):
        # type: (EvaluatableNode) -> None
        '''Set up NOT node

        Args:
            expression: The expression to negate (expression must evaluate to a series)
        '''
        self.children = [expression]

    def copy(self, new_arguments):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return Not(new_arguments[0])

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries
        evaluated, = evaluated_children
        if not isinstance(evaluated, TypedSeries):
            raise ValueError("NOT expression must evaluate to a series")
        if evaluated.type_ != BQScalarType.BOOLEAN:
            raise ValueError("NOT accepts only booleans but was given type: {}".format(
                evaluated.type_))
        return TypedSeries(~evaluated.series, BQScalarType.BOOLEAN) class NullCheck(EvaluatableNodeWithChildren):
    '''Expression that checks whether element is null or not'''

    def __init__(self, expression, direction):
        # type: (EvaluatableNode, str) -> None
        '''Set up NullCheck node

        Args:
            expression: Expression to check
            direction: 'IS_NULL' or 'IS_NOT_NULL'
        '''
        self.children = [expression]

        if direction == 'IS_NULL':
            self.direction = True
        elif direction == 'IS_NOT_NULL':
            self.direction = False
        else:
            raise ValueError("Invalid direction for NullCheck, not IS_NULL or IS_NOT_NULL: {}"
                             .format(direction))

    def copy(self, new_arguments):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return NullCheck(new_arguments[0], 'IS_NULL' if self.direction else 'IS_NOT_NULL')

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries
        value = evaluated_children[0].series
        if self.direction:
            return TypedSeries(value.isnull(), BQScalarType.BOOLEAN)
        else:
            return TypedSeries(value.notnull(), BQScalarType.BOOLEAN) class Selector(EvaluatableNodeWithChildren):
    '''Corresponds to one of SELECT's fields, i.e. Selector is a column that
    should be returned from a SELECT query.
    Also includes the alias for this field, if specified.
    '''

    def __init__(self,
                 selector,  # type: EvaluatableNode
                 alias  # type: Union[_EmptyNode, str]
                 ):
        # type: (...) -> None
        '''Set up Selector node

        Args:
            selector: Column/expression to select from table.  Must be a tuple of
                an evaluatable expression.
            alias: Name to assign to this column/expression, if any.
        '''
        self.children = [selector]
        self.alias = alias

        # The position in the select list; 1-up (i.e. the first selector will have position=1).
        # This field will be populated by the Select object this Selector is passed to.
        self.position = None  # type: Optional[int]

    def copy(self, new_children):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return Selector(new_children[0], self.name())

    def name(self):
        # type: () -> str
        if not isinstance(self.alias, _EmptyNode):
            return self.alias
        if self.position is None:
            raise RuntimeError("Accessing name of Selector {} before position is populated."
                               .format(self))
        return self.children[0].name() or '_f{}'.format(self.position)

    def mark_grouped_by(self, group_by_paths, context):
        # type: (Sequence[Tuple[str, ...]], EvaluationContext) -> EvaluatableNode
        if context.get_canonical_path((self.name(),)) in group_by_paths:
            return GroupedBy(self)
        return self.copy([self.children[0].mark_grouped_by(group_by_paths, context)])

    def _evaluate_node(self, evaluated_arguments):
        # type: (List[TypedSeries]) -> TypedSeries
        result, = evaluated_arguments
        result.series.name = self.name()
        return result class StarSelector(AbstractSyntaxTreeNode):
    '''Corresponds a wildcard field of SELECT, like SELECT * or SELECT table.*'''

    StarSelectorType = Tuple[AbstractSyntaxTreeNode,  # expression
                             AbstractSyntaxTreeNode,  # exception
                             AbstractSyntaxTreeNode   # replacement
                             ]

    def __init__(self,
                 expression,  # type: Union[_EmptyNode, EvaluatableNode]
                 exception,  # type: Union[_EmptyNode, Tuple[str, ...]]
                 replacement,  # type: Union[_EmptyNode, Tuple[Tuple[EvaluatableNode, Any, str], ...]]  # noqa: E501
                 ):
        # type: (...) -> None
        '''Sets up StarSelector node

        Args:
            selector: Columns to select from table.  Must be a tuple of (expression, exception,
            replacement).
        '''
        self.expression = expression
        self.exception = exception
        self.replacement = replacement

        # The position in the select list; 1-up (i.e. the first selector will have position=1).
        # This field will be populated by the Select object this Selector is passed to.
        self.position = None  # type: Optional[int]

    def get_selectors(self, context):
        # type: (EvaluationContext) -> List[Selector]
        '''See parent, EvaluatableNode'''
        # First: get the list of selectors, ignoring exception and replacement.

        # Case 1: SELECT *
        if isinstance(self.expression, _EmptyNode):
            selectors = [Selector(Field(path), path[-1])
                         for col in context.table.dataframe.columns
                         for path in (col.split('.'),)]

        # Case 2: SELECT table.*
        elif isinstance(self.expression, Field) and len(self.expression.path) == 1:
            table_identifier, = self.expression.path
            columns = context.table_to_column_ids[table_identifier]
            selectors = [Selector(Field((table_identifier, column)), column) for column in columns]

        # CASE 3: SELECT struct-expression.* (currently unimplemented)
        # This 'else' would also also match nonsense like SELECT 123.*
        else:
            # TODO: Implement select (struct expression).*
            raise ValueError("Cannot select {}.*".format(self.expression))

        if not isinstance(self.exception, _EmptyNode):
            selectors = [selector for selector in selectors if selector.alias not in self.exception]
        if not isinstance(self.replacement, _EmptyNode):
            replacement_map = {name: expression for expression, _, name in self.replacement}
            selectors = [selector
                         if (not isinstance(selector.alias, str)
                             or selector.alias not in replacement_map)
                         else Selector(replacement_map[selector.alias], selector.alias)
                         for selector in selectors]
        return selectors class Struct(EvaluatableNodeWithChildren):
    '''A STRUCT expression.'''

    def __init__(self, type_, expressions):
        # type: (BQStructType, Sequence[EvaluatableNode]) -> None
        self.type_ = type_
        self.children = expressions

    def copy(self, new_children):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return Struct(self.type_, new_children)

    def _evaluate_node(self, values):
        # type: (List[TypedSeries]) -> TypedSeries
        value_types = [value.type_ for value in values]
        type_ = implicitly_coerce(self.type_, BQStructType([None] * len(values), value_types))
        if self.type_:
            if not isinstance(type_, BQStructType):
                raise RuntimeError("STRUCT types coerced to non-STRUCT type {}".format(type_))
            for i, (declared_type, coerced_type) in enumerate(zip(self.type_.types, type_.types)):
                if declared_type and declared_type != coerced_type:
                    raise ValueError('Struct field {} has type {} which does not coerce to {}'
                                     .format(i + 1,  # enumerate is 0-up, we want to report 1-up
                                             coerced_type, declared_type))

        structs = list(zip(*[value.series for value in values]))
        return TypedSeries(pd.Series(structs, index=values[0].series.index), type_)

    @classmethod
    def create_from_typeless(cls, maybe_named_fields):
        # type: (Sequence[Tuple[EvaluatableNode, Union[_EmptyNode, str]]]) -> Struct
        '''Creates a Struct from the typeless grammar syntax STRUCT(value [AS name], ...)

        Args:
            maybe_named_fields: A list of pairs.  The first element is a field's value, an
                expression (evaluatable node), the second is an optional name for the field.
        Returns:
            A Struct abstract syntax tree node.
        '''
        children, maybe_names = list(zip(*maybe_named_fields))
        maybe_names = tuple(maybe_name if not isinstance(maybe_name, _EmptyNode) else None
                            for maybe_name in maybe_names)
        return Struct(BQStructType(maybe_names, [None] * len(maybe_named_fields)), children)

    @classmethod
    def create_from_typed(cls, type_, expressions):
        # type: (BQStructType, Sequence[EvaluatableNode]) -> Struct
        '''Creates a Struct from the typed grammar syntax STRUCT<[name] type, ...>(value, ...).

        Args:
            type_: A declared STRUCT type for this structure.
            expressions: The values of the fields of the structure.

        Returns:
            A Struct abstract syntax tree node.
        '''
        return Struct(type_, expressions)

    @classmethod
    def create_from_tuple(cls, first_expression, other_expressions):
        # type: (EvaluatableNode, Sequence[EvaluatableNode]) -> Struct
        '''Creates a Struct from the tuple grammar syntax (value, value, ...)

        Args:
            first_expression: The evaluatable node of the first field in the structure.
            other_expressions: The evaluatable nodes of the rest of the fields in the structure
                (if any).
        Returns:
            A Struct abstract syntax tree node.
        '''
        return Struct(None, [first_expression] + list(other_expressions)) class UnaryNegation(EvaluatableNodeWithChildren):
    '''Expression that negates a series, such as turning [1, 2] into [-1, -2]'''
    def __init__(self, expression):
        # type: (EvaluatableNode) -> None
        ''' Set up unary negation

        Args:
            expression: The expression to negate; must evaluate to a series of
                either integers or floats.
        '''
        self.children = [expression]

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries
        typed_value, = evaluated_children
        if not isinstance(typed_value, TypedSeries):
            raise ValueError("UnaryNegation expression must evaluate to a series")
        if typed_value.type_ not in [BQScalarType.INTEGER, BQScalarType.FLOAT]:
            raise TypeError("UnaryNegation expression supports only integers and floats, got: {}"
                            .format(typed_value.type_))
        return TypedSeries(-typed_value.series, typed_value.type_)

    def copy(self, new_arguments):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return UnaryNegation(new_arguments[0]) _CROSS_KEY = '__cross_key__'
    BIGQUERY_TO_PANDAS_JOIN_TYPE = {
        EMPTY_NODE: 'inner',
        'CROSS': 'inner',  # Cross join behaves like inner if join conditions are provided.
        'INNER': 'inner',
        'LEFT': 'left',
        'LEFT_OUTER': 'left',
        'FULL': 'outer',
        'FULL_OUTER': 'outer',
        'RIGHT': 'right',
        'RIGHT_OUTER': 'right',
    }
def _extract_simple_comparison(node):
def _cross_join(left_table, right_table):
def _get_common_columns(left_table, right_table):
def _get_join_using(join_condition, right_table_id, context):
def _get_join_on_equality_comparisons(
        join_comparisons,  # type: List[Tuple[Field, Field]]
        right_table_id,  # type: str
        context  # type: EvaluationContext
):
def _join_on_arbitrary_condition(left_table,  # type: TypedDataFrame
                                 right_table,  # type: TypedDataFrame
                                 join_condition,  # type: EvaluatableNode
                                 context,  # type: EvaluationContext
                                 pandas_join_type  # type: str
                                 ):
    def __init__(self,
                 first_from,  # type: FromItemType
                 joins  # type: List[Join]
                 ):
    def _get_joined_table(self,
                          context,  # type: EvaluationContext
                          table,  # type: TypedDataFrame
                          join_type,  # type: Union[str, _EmptyNode]
                          join_with_alias,  # type: FromItemType
                          join_condition  # type: ConditionsType
                          ):
    def create_context(self, table_context):
class DataSource(AbstractSyntaxTreeNode): def separated_sequence(rule,  # RuleType
                       separator_rule,  # RuleType
                       wrapper=tuple,  # Callable[[List[AbstractSyntaxTreeNode]], AbstractSyntaxTreeNode] # noqa: E501
                       keep_separator=False  # bool
                       ):
    def check_sequence(tokens):
def _apply_rule_str(rule, tokens):
def _apply_rule_tuple(rule, tokens):
def _apply_rule_list(rule, tokens):
def _apply_rule_none(rule, tokens):
def _apply_rule_method(rule, tokens):
def apply_rule(rule, tokens):
def wrap(wrapper, rule):
    def wrapped_rule(tokens): def apply_rule(rule, tokens):
    # type: (RuleType, List[str]) -> AppliedRuleOutputType
    """Apply the given rule to tokens if possible.

    Args:
        rule: The object (string, tuple, list, method, None) that represents the rule.
        tokens: A list of tokens.
    Returns:
        A tuple of the Abstract Syntax Tree nodes representing the result of applying the rule
        to the tokens, and the remaining unmatched tokens.
    """

    if isinstance(rule, str):
        return _apply_rule_str(rule, tokens)
    elif isinstance(rule, tuple):
        return _apply_rule_tuple(rule, tokens)
    elif isinstance(rule, list):
        return _apply_rule_list(rule, tokens)
    elif rule is None:
        return _apply_rule_none(rule, tokens)
    else:
        return _apply_rule_method(rule, tokens) def separated_sequence(rule,  # RuleType
                       separator_rule,  # RuleType
                       wrapper=tuple,  # Callable[[List[AbstractSyntaxTreeNode]], AbstractSyntaxTreeNode] # noqa: E501
                       keep_separator=False  # bool
                       ):
    # type: (...) -> Callable[[List[str]], AppliedRuleOutputType]
    """Return a grammar rule that is a sequence of `rule`s with a separator.

    The returned rule `result` is equivalent to

    result :: = rule separator_rule result | rule

    Args:
        rule: A grammar rule.
        separator_rule: A grammar rule, used as a separator.
        wrapper: An optional class; if provided, it will be called with the matching
           syntax tree nodes as arguments to create the return value.
        keep_separator: Nodes representing the separators are discarded unless this is True.

    Returns:
        A grammar rule, i.e. a function taking a list of tokens, and returning a result
        and all non-matching tokens.  That result is:

            None (if the rule doesn't match)
            An object of type class, constructed with all matching nodes (if cls is provided)
            node (if one token matches `rule` and the separator isn't a literal string)
            a tuple of nodes matching (rule, rule, ... rule) (if separator is a literal string)
            a tuple of nodes matching (rule, separator, rule, separator, ... rule) (otherwise)
    """
    def check_sequence(tokens):
        # type: (List[str]) -> AppliedRuleOutputType
        nodes = []
        # Initial value must be non-None but also have the correct type
        maybe_separator = AbstractSyntaxTreeNode()  # type: AppliedRuleNode
        tokens_after_separator = tokens

        while maybe_separator:
            # Try to match the given rule
            node, tokens = apply_rule(rule, tokens_after_separator)
            if not node:
                break
            nodes.append(node)

            # Try to match the given separator
            maybe_separator, tokens_after_separator = apply_rule(separator_rule, tokens)
            if maybe_separator and keep_separator:
                nodes.append(maybe_separator)
        if not nodes:
            return None, tokens
        return wrapper(nodes), tokens
    return check_sequence def wrap(wrapper, rule):
    '''Returns a grammar rule that, if `rule' matches, wraps the result by calling `wrapper'

    Args:
        wrapper: A function to wrap the rule's result.
        rule: Any grammar rule
    Returns:
        If the rule matches, the result of calling wrapper on the rule's result, plus leftover
        unmatched tokens.  Otherwise, returns None and all the tokens.
    '''
    def wrapped_rule(tokens):
        # type: (List[str]) -> AppliedRuleOutputType
        result, new_tokens = apply_rule(rule, tokens)
        if result is None:
            return None, tokens
        if (isinstance(rule, tuple)
                and len([subrule for subrule in rule if not isinstance(subrule, str)]) > 1):
            return wrapper(*result), new_tokens
        else:
            return wrapper(result), new_tokens
    return wrapped_rule def grammar_literal(*words):
    # type: (*str) -> Callable[[List[str]], Tuple[Optional[str], List[str]]]
    """Checks if the first token(s) is a grammar literal. This includes all the reserved words
    above, as well as binary operators, parentheses, brackets, etc.

    Args:
        words: Grammatic literal(s) to match.
    Returns:
        Function that checks whether given tokens match given grammatic literals.
    """

    def match_strings_to_tokens(tokens):
        # type: (List[str]) -> Tuple[Optional[str], List[str]]
        """Checks if the next token(s) match the previously given grammatic literal(s).

        Args:
            tokens: List of tokens, to decide whether the first token matches a grammatic literal.
        Returns:
            Tuple of grammatic literal(s) (if there was a match, joined by an underscore if there
            is more than one), and the rest of the unidentified tokens.
        """

        # Convert all requested words to upper case
        words_upper = tuple(word.upper() for word in words)

        # If there are enough tokens for the given reserved word(s) AND
        # the first n tokens match the given reserved word(s)
        if (len(tokens) >= len(words) and
                words_upper == tuple(token.upper() for token in tokens[:len(words)])):

            # Put underscores between the found reserved words, and return the
            # unused tokens
            return '_'.join(words_upper), tokens[len(words):]
        return None, tokens

    return match_strings_to_tokens def identifier(tokens):
    # type: (List[str]) -> Tuple[Optional[str], List[str]]
    """Checks if the first token is an identifier.

    Args:
        tokens: List of tokens, to decide whether tokens[0] is an identifier.
    Returns:
        Tuple of identifier (if there was a match) and rest of unidentified tokens.
    """
    if tokens:
        maybe_id = tokens[0]
        if (re.search(r'^' + IDENTIFIER_PATTERN + '$', maybe_id)
                and maybe_id.upper() not in RESERVED_WORDS):
            return maybe_id, tokens[1:]
        # If enclosed in backticks, identifiers can contain any character,
        # including spaces, and can even be reserved words.
        if re.search(r'^' + BACKTICK_PATTERN + '$', maybe_id):
            # Remove backticks before returning
            return maybe_id[1:-1], tokens[1:]
    return None, tokens def literal(tokens):
    # type: (List[str]) -> Tuple[Optional[Value], List[str]]
    """Checks if the first token is a literal (number, string, boolean, null).

    Args:
        tokens: List of tokens, to decide whether tokens[0] is a literal.
    Returns:
        Tuple of parsed literal (if there was a match) and rest of unidentified tokens.
    """
    if tokens:
        token = tokens[0]
        rest_of_tokens = tokens[1:]
        float_pattern = '|'.join([r'^' + pattern + '$' for pattern in FLOAT_LITERAL_PATTERNS])
        if re.search(r'^' + float_pattern + '$', token):
            return Value(float(token), BQScalarType.FLOAT), rest_of_tokens
        if re.search(r'^' + INT_LITERAL_PATTERN + '$', token):
            return Value(int(tokens[0]), BQScalarType.INTEGER), rest_of_tokens
        str_pattern = '|'.join([r'^' + pattern + '$' for pattern in STR_LITERAL_PATTERNS])
        if re.search(r'^' + str_pattern + '$', token):
            # Remove quotes before returning
            return Value(token[1:-1], BQScalarType.STRING), rest_of_tokens
        token = token.upper()
        if token in _CONSTANTS:
            return Value(*_CONSTANTS[token]), rest_of_tokens
    return None, tokens def bigquery_type(tokens):
def _create_struct_type(typed_fields):
from typing import List, cast  # noqa: F401
from .bq_abstract_syntax_tree import EMPTY_NODE, Field
from .bq_operator import binary_operator_expression_rule
from .dataframe_node import QueryExpression, Select, SetOperation, TableReference, Unnest
from .evaluatable_node import (Array, Array_agg, Case, Cast, Count, Exists, Extract, FunctionCall,
                               If, InCheck, Not, NullCheck, Selector, StarSelector, Struct,
                               UnaryNegation)
from .join import DataSource, FromItemType, Join
from .query_helper import AppliedRuleOutputType  # noqa: F401
from .query_helper import apply_rule, separated_sequence, wrap
from .terminals import grammar_literal, identifier, literal
from .type_grammar import array_type, scalar_type, struct_type
        QueryExpression,
        [with_expression, None],
        [select,
         ('(', query_expression, ')')],

        # order_by
        maybe_order_by,

        # limit
        maybe_limit)

    return apply_rule(
        [(SetOperation, core_query_expression, set_operator, query_expression),
         core_query_expression],
        tokens)


with_expression = ('WITH', separated_sequence((identifier, 'AS', '(', query_expression, ')'), ','))


def alias(tokens):
    # type: (List[str]) -> AppliedRuleOutputType
    '''An optional alias to rename a field or table.

    Args:
        tokens: Parts of the user's query (split by spaces into tokens) that
        are not yet parsed
    Returns:
        A tuple of the matched alias identifier (if any), and the remaining tokens.
    '''
    alias_node, new_tokens = apply_rule([
        (['AS', None], identifier),
        None
    ], tokens)

    if alias_node == EMPTY_NODE:
        return EMPTY_NODE, tokens

    if not (isinstance(alias_node, tuple) and len(alias_node) == 2):
        raise RuntimeError("Internal parse error: alias rule returned {!r}".format(alias_node))

    # The alias node will be a tuple of ('AS', new_name), so we get rid of the
    # 'AS' and just return the new name (alias identifier)
    as_token, alias_identifier = alias_node
    return alias_identifier, new_tokens


def select(tokens):
    # type: (List[str]) -> AppliedRuleOutputType
    '''Grammar rule matching a select clause.

    This rule is adapted from here:
    https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#select-list

    Args:
        tokens: Parts of the user's query (split by spaces into tokens) that
        are not yet parsed
    Returns:
        A tuple of the Abstract Syntax Tree nodes representing the result of
        applying the rule to the tokens, and the remaining unmatched tokens.
    '''
    return apply_rule(
        (Select,

         # modifier
         ['ALL', 'DISTINCT', None],

         # fields
         separated_sequence(

              # The selector either contains a * or is just an expression
              [
                  # If selector's expression includes a '*', it can also include
                  # exception and replacement
                  (
                      StarSelector,

                      # expression - can be just * or field.* or table.field.*, etc
                      [(expression, '.'), None],

                      '*',

                      # exception
                      [('EXCEPT', '(', separated_sequence(identifier, ','), ')'), None],

                      # replacement
                      [('REPLACE',
                        '(',
                        separated_sequence((expression, ['AS', None], identifier), ','),
                        ')'),
                       None],
                  ),

                  # The selector does not include a * and therefore cannot have
                  # an exception or replacement
                  (Selector, expression, alias),
              ],

              # separator for separated_sequence()
              ','
         ),

         # from
         [('FROM', data_source), None],

         # where
         [('WHERE', expression), None],

         # group by
         [('GROUP', 'BY', separated_sequence([field, literal], ',')), None],

         # having
         [('HAVING', expression), None],
         ),
        tokens)


# Expression following "FROM": a table or another query, followed optionally by an AS alias
# Examples: "SomeTable", "(SELECT a from SomeTable)", "SomeTable AS t"
from_item = ([

    (TableReference, separated_sequence(identifier, '.')),