class BinaryExpression(EvaluatableNodeWithChildren):

    def __init__(self, left, operator_str, right):
        # type: (EvaluatableNode, str, EvaluatableNode) -> None
        self.children = [left, right]
        self.operator_info = BINARY_OPERATOR_INFO[operator_str]

    def copy(self, new_children):
        # type: (Sequence[EvaluatableNode]) -> BinaryExpression
        new_left, new_right = new_children
        return BinaryExpression(new_left, self.operator_info.operator, new_right)

    def strexpr(self):
        # type: () -> str
        """Returns a prefix-expression serialization for testing purposes."""
        left, right = self.children
        return '({} {} {})'.format(self.operator_info.operator,
                                   left.strexpr(),
                                   right.strexpr())

    def _evaluate_node(self, evaluated_children):
        # type: (List[TypedSeries]) -> TypedSeries
        left_value, right_value = evaluated_children
        # We need to know the type of the result of the operation.  Some operators specify their
        # result type (e.g. comparators have a boolean output regardless of input types).  Some
        # operators keep the type of their inputs (e.g. multiplication keeps the input numerical
        # types); this is notated in the BINARY_OPERATOR_INFO table by specifying the result_type
        # as None.  The right thing to do if the types of the inputs aren't the same is to apply
        # specific logic depending on the input types: see here for the logic for arithmetic
        # operators:
        # https://cloud.google.com/bigquery/docs/reference/standard-sql/operators#arithmetic_operators
        result_type = self.operator_info.result_type
        if not result_type:
            result_type = implicitly_coerce(left_value.type_, right_value.type_)
        return TypedSeries(self.operator_info.function(left_value.series, right_value.series),
                           result_type) _SELECTOR_TABLE = '__selector__'
EMPTY_NODE = _EmptyNode()
EMPTY_CONTEXT = EvaluationContext(TableContext())
class AbstractSyntaxTreeNode(object):
class MarkerSyntaxTreeNode(AbstractSyntaxTreeNode):
class EvaluatableNode(AbstractSyntaxTreeNode):
class EvaluatableLeafNode(EvaluatableNode):
class EvaluatableNodeWithChildren(EvaluatableNode):
class EvaluatableNodeThatAggregatesOrGroups(EvaluatableNodeWithChildren):
class Result(object):
class DataframeNode(AbstractSyntaxTreeNode):
class GroupedBy(EvaluatableNodeThatAggregatesOrGroups):
class Field(EvaluatableLeafNode):
class _EmptyNode(AbstractSyntaxTreeNode):
class EvaluationContext:
    def __repr__(self):
    def strexpr(self):
    def literal(cls):
    def literal(cls):
    def evaluate(self, context):
    def pre_group_by_partially_evaluate(self, context):
    def name(self):
    def mark_grouped_by(self, group_by_paths, context):
    def is_constant(self):
    def is_aggregated(self):
    def evaluate(self, context):
    def pre_group_by_partially_evaluate(self, context):
    def _evaluate_leaf_node(self, context):
    def __init__(self, children):
    def pre_group_by_partially_evaluate(self, context):
    def _ensure_fully_evaluated(self, evaluated_children):
    def evaluate(self, context):
    def _evaluate_node(self, evaluated_children):
    def copy(self, new_children):
    def mark_grouped_by(self, group_by_paths, context):
    def is_constant(self):
    def is_aggregated(self):
    def pre_group_by_partially_evaluate(self, context):
    def mark_grouped_by(self, group_by_paths, context):
    def is_constant(self):
    def is_aggregated(self):
    def __init__(self, statement_type, path=None, table=None):
    def get_dataframe(self, table_context, outer_context=None):
    def execute(self, table_context):
    def __init__(self, expression):
    def copy(self, new_children):
    def _evaluate_node(self, evaluated_children):
    def __init__(self, path):
    def strexpr(self):
    def mark_grouped_by(self, group_by_paths, context):
    def name(self):
    def _evaluate_leaf_node(self, context):
    def __eq__(self, other):
    def is_constant(self):
    def strexpr(self):
    def __init__(self, table_context):
    def add_subcontext(self, subcontext):
    def maybe_add_column(self, column):
    def _partially_evaluate(self, selector, group_by_paths):
    def do_group_by(self, selectors, group_by):
    def add_table_from_node(self, from_item, alias):
    def add_table_from_dataframe(self, table,  # type: TypedDataFrame
                                 table_id,  # type: Optional[str]
                                 alias  # type: Union[_EmptyNode, str]
                                 ):
    def clone_context_new_table(cls, table, old_context):
    def get_all_canonical_paths(self, path):
    def get_canonical_path(self, path):
    def lookup(self, path): class BQScalarType(BQType, enum.Enum):
    """Representation of basic (scalar) BigQuery types as an enum."""

    __metaclass__ = AbstractEnumMeta

    BOOLEAN = 'BOOLEAN'
    DATE = 'DATE'
    DATETIME = 'DATETIME'
    INTEGER = 'INTEGER'
    FLOAT = 'FLOAT'
    STRING = 'STRING'
    TIMESTAMP = 'TIMESTAMP'

    def to_dtype(self):
        # type: () -> np.dtype
        """Converts this BigQuery type to a NumPy dtype.

        Returns:
           A NumPy dtype corresponding to this BigQuery type (e.g. np.dtype('int64') for INTEGER)
        """
        return np.dtype(_BQ_SCALAR_TYPE_TO_NUMPY_TYPE[self])

    @classmethod
    def from_string(cls, typename):
        # type: (str) -> BQScalarType
        """Reads this type from a string representation.

        A Factory method constructing an instance of BQScalarType corresponding to typename.
        The reason for not just using the inherited Enum constructor is to allow Standard BigQuery
        typenames to be aliases for Legacy typenames.

        Args:
            typename: A BigQuery type name, either Legacy (INTEGER, FLOAT, ...) or Standard
            (INT64, FLOAT64, ...).

        Returns:
            The corresponding BQScalarType enum.
        """
        if isinstance(typename, str):
            typename = typename.upper()
        if typename in _LEGACY_BQ_SCALAR_TYPE_FROM_BQ_SCALAR_TYPE:
            return _LEGACY_BQ_SCALAR_TYPE_FROM_BQ_SCALAR_TYPE[typename]
        return cls(typename)

    def to_schema_field(self, name):
        # type: (str) -> SchemaField
        """Converts this type to a BigQuery SchemaField.

        Args:
            name: The name of the column.  This class represents a type; SchemaField represents
            a column, so it includes the type and also the name of the column.

        Returns:
            A SchemaField object corresponding to a column containing this class' type.
        """
        return SchemaField(name=name, field_type=self.value)

    def __repr__(self):
        return 'BQScalarType.{}'.format(self.value)

    def convert(self, element):
        # type: (PandasType) -> ScalarPythonType
        """Converts a pandas Series element to a Python type corresponding to this BigQuery type.

        Args:
            element: One cell of a Pandas DataFrame or Series.  Will have numpy types like np.int64
            rather than a Python type like int.

        Returns:
            The element's value, cast to a corresponding Python type.
        """
        if pd.isnull(element):
            return None
        return _BQ_SCALAR_TYPE_TO_PYTHON_TYPE[self](element) class TypedDataFrame(object):
    def __init__(self, dataframe, types):
        # type: (Union[pd.DataFrame, pd.DataFrameGroupBy], Sequence[BQType]) -> None
        if isinstance(dataframe, pd.DataFrame) and len(dataframe.columns) != len(types):
            raise ValueError(
                "Trying to create TypedDataFrame with mismatching type set; columns {!r} types {!r}"
                .format(list(dataframe.columns), types))
        self._dataframe = dataframe
        self._types = list(types)

    @property
    def dataframe(self):
        # type: () -> Union[pd.DataFrame, pd.DataFrameGroupBy]
        """Returns the underlying DataFrame.

        This is a property so that it's immutable.

        Returns:
           The actual tabular data.
        """
        return self._dataframe

    @property
    def types(self):
        # type: () -> List[BQType]
        """Returns the data types, in the same order as the table's columns."""
        return self._types

    def __repr__(self):
        return 'TypedDataFrame({!r}, {!r})'.format(self.dataframe, self.types)

    def to_bq_schema(self):
        # type: () -> List[SchemaField]
        """Returns a BigQuery schema (list of schema fields) matching this object's types."""
        return [type_.to_schema_field(name)
                for name, type_ in zip(self.dataframe.columns, self.types)]

    def to_list_of_lists(self):
        # type: () -> List[List[PythonType]]
        """Returns the data as a list of rows, each row a list of Python-typed objects."""
        rows = []
        for unused_index, row in self.dataframe.iterrows():
            rows.append([type_.convert(element) for element, type_ in zip(list(row), self.types)])
        return rows 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 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 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 Value(EvaluatableLeafNode):
    '''A node representing a literal value (number, string, boolean, null).'''

    def __init__(self, value, type_):
        # type: (Optional[LiteralType], Optional[BQScalarType]) -> None
        if (value is None) != (type_ is None):
            raise ValueError(
                "Value(None, None) means NULL; Value({}, {}) is not allowed"
                .format(value, type_))
        self.value = value
        self.type_ = type_

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

    def strexpr(self):
        # type: () -> str
        return repr(self.value)

    def __eq__(self, other):
        # type: (Any) -> bool
        if isinstance(other, Value):
            return (self.type_ == other.type_) and (self.value == other.value)
        return False

    def _evaluate_leaf_node(self, context):
        # type: (EvaluationContext) -> TypedSeries
        '''See parent, EvaluatableNode'''
        return TypedSeries(pd.Series([self.value] * len(context.table.dataframe),
                                     index=_get_index(context.table.dataframe)),
                           self.type_) def query_expression(tokens):
    # type: (List[str]) -> AppliedRuleOutputType
    '''This is the highest-level grammar method.  It is called by query.execute_query().

    The "raw" rule syntax is supposedly this:
    https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax

    query_expression = ([with_expression, None],
                        [select,
                         ('(', query_expression, ')'),
                         (query_expression, set_operator, query_expression)],
                        maybe_order_by,
                        maybe_limit)

    However, (a) this is left-recursive, so we would need to juggle things to avoid
    infinite recursion, and (b) this permits things like
    WITH query1 as (SELECT 1) WITH query2 as (SELECT 2) SELECT 3 UNION ALL SELECT 4
    that (1) seem kind of strange (two successive with clauses)
    and (2) give an error on prod BigQuery.

    So we use the following simpler syntax.

    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.
    '''
    core_query_expression = (
        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) 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) class DataSource(AbstractSyntaxTreeNode):
    '''Node representing JOIN operations.

    Not a child of EvaluatableNode because this evaluate() has a different input and
    output type than the rest of the evaluate()s.
    '''
    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 __init__(self,
                 first_from,  # type: FromItemType
                 joins  # type: List[Join]
                 ):
        # type: (...) -> None
        '''Set up JOIN node.

        Args:
            first_from: The first table that was specified after FROM
            joins: The rest of the tables, if any, to join to the first
                This is a list, where each element is a tuple of 3 elements:
                join_type: Key into self.BIGQUERY_TO_PANDAS_JOIN_TYPE
                join_with_alias: Tuple of a TableReference and an optional alias
                join_condition: Tuple of fields to join on, where each is
                    either a string (field name) or tuple of 2 Field nodes
        '''
        self.first_from = first_from
        self.joins = joins

    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
                          ):
        # type: (...) -> TypedDataFrame
        """Gets the TypedDataFrame after performing a join.

        Args:
            context: EvaluationContext in which to evaluate the join.
            table: The table to join to.
            join_type: If provided, the type of join (e.g. INNER, OUTER, etc.)
            join_with_alias: The table to be joined in, with an optional alias
            join_condition: The specified conditions on the join.

        Returns:
            The TypedDataFrame after the join is done.
        """
        join_table, join_table_id = context.add_table_from_node(*join_with_alias)

        join_type = join_type.upper() if isinstance(join_type, str) else join_type

        pandas_join_type = self.BIGQUERY_TO_PANDAS_JOIN_TYPE.get(join_type)
        if pandas_join_type is None:
            raise NotImplementedError("Join type {} is not supported".format(join_type))

        # Now we execute the JOIN operation by determining the type of join condition -- the
        # user-specified condition of which rows from `table' are joined with which rows from
        # `join_table` -- and taking the appropriate action.  As much as possible, we want to use
        # the pandas DataFrame.merge method, which means we the action will be to convert the
        # user-specified join condition into a list of columns on the left and columns on the right
        # (left_ons and right_ons) and call pandas.merge.
        #
        # Two kinds of joins -- unconditional cross joins, and joins on an arbitrary boolean
        # expression -- are not supported by the merge method, and so we have separate logic that
        # directly calculates and returns the merged table.
        if join_type == 'CROSS' and join_condition is EMPTY_NODE:
            return _cross_join(table, join_table)

        # If no specific join condition is given, we join on columns that are common between
        # the two tables.
        elif isinstance(join_condition, _EmptyNode):
            left_ons, right_ons = _get_common_columns(table, join_table)

        # If join USING(list of fields) is specified.
        elif isinstance(join_condition, tuple):
            left_ons, right_ons = _get_join_using(join_condition, join_table_id, context)

        # If join ON is specified
        else:
            join_comparisons = _extract_simple_comparison(join_condition)

            # If join ON (a = b AND c = d AND ...) is specified.
            if join_comparisons is not None:
                left_ons, right_ons = _get_join_on_equality_comparisons(
                        join_comparisons, join_table_id, context)

            # The user can also join ON some arbitrary boolean condition, e.g. JOIN ON (a+b < c).
            else:
                return _join_on_arbitrary_condition(table, join_table, join_condition, context,
                                                    pandas_join_type)

        return TypedDataFrame(table.dataframe.merge(
                join_table.dataframe, how=pandas_join_type, left_on=left_ons, right_on=right_ons),
                              table.types + join_table.types)

    def create_context(self, table_context):
        # type: (TableContext) -> EvaluationContext
        '''Given a representation of the entire database, add the specified
        table(s) to the query's context.
        '''
        context = EvaluationContext(table_context)
        table, _ = context.add_table_from_node(*self.first_from)

        for join_type, join_with_alias, join_condition in self.joins:
            table = self._get_joined_table(
                    context, table, join_type, join_with_alias, join_condition)

        context.table = table
        return context class DatasetTableContext(TableContext):
    '''A TableContext containing a set of datasets.'''

    def __init__(self, datasets):
        # type: (DatasetType) -> None
        '''Construct the TableContext.

        Args:
            datasets: A series of nested dictionaries mapping to a TypedDataFrame.
            For example, {'my_project': {'my_dataset': {'table1': t1, 'table2': t2}}},
            where t1 and t2 are two-dimensional TypeDataFrames representing a table.
        '''
        self.datasets = datasets

    def lookup(self, path):
        # type: (Sequence[str]) -> Tuple[TypedDataFrame, Optional[str]]
        '''See TableContext.lookup for docstring.'''
        if not self.datasets:
            raise ValueError("Attempt to look up path {} with no projects/datasets/tables given."
                             .format(path))
        if len(path) < 3:
            # Table not fully qualified - attempt to resolve
            if len(self.datasets) != 1:
                raise ValueError("Non-fully-qualified table {} with multiple possible projects {}"
                                 .format(path, sorted(self.datasets.keys())))
            project, = self.datasets.keys()

            if len(path) == 1:
                # No dataset specified, only table
                if len(self.datasets[project]) != 1:
                    raise ValueError(
                            "Non-fully-qualified table {} with multiple possible datasets {}"
                            .format(path, sorted(self.datasets[project].keys())))
                dataset, = self.datasets[project].keys()
                path = (project, dataset) + tuple(path)
            else:
                # Dataset and table both specified
                path = (project,) + tuple(path)

        if len(path) > 3:
            raise ValueError("Invalid path has more than three parts: {}".format(path))
        project_id, dataset_id, table_id = path

        return self.datasets[project_id][dataset_id][table_id], table_id

    def set(self, path, table):
        # type: (Tuple[str, ...], TypedDataFrame) -> None
        '''Sets a path to refer to a table.

        Args:
            path: A tuple of strings representing a period-separated path to a table, like
                projectname.datasetname.tablename
            table: The table to add
         '''
        project_id, dataset_id, table_id = path
        if project_id not in self.datasets:
            raise ValueError("Attempting to create {!r} but project {!r} not created"
                             .format(path, project_id))
        self.datasets[project_id].setdefault(dataset_id, {})[table_id] = table def tokenize(query):
    # type: (str) -> List[str]
    return _COMBINED_PATTERN.findall(remove_comments(query))
import unittest
import pandas as pd
from typing import Any, List, Tuple  # noqa: F401
from ddt import data, ddt, unpack
from purplequery.binary_expression import BinaryExpression
from purplequery.bq_abstract_syntax_tree import (EMPTY_NODE, EvaluatableNode, Field,  # noqa: F401
                                                 TableContext)
from purplequery.bq_types import BQScalarType, TypedDataFrame
from purplequery.dataframe_node import QueryExpression, Select, TableReference
from purplequery.evaluatable_node import Selector, StarSelector, Value
from purplequery.grammar import query_expression as query_expression_rule
from purplequery.grammar import select as select_rule
from purplequery.join import DataSource
from purplequery.storage import DatasetTableContext
from purplequery.tokenizer import tokenize
    @data(
        dict(
            query_expression=(
                'WITH q1 as (SELECT a+1 as b FROM my_table),'
                'q1 as (SELECT a+2 as b FROM my_table)'
                'SELECT * from q1'),
            error='Duplicate names in WITH clauses are not allowed'),
    )
    @unpack
    def test_with_clause_error(self, query_expression, error):
        # type: (str, str) -> None
        query_expression_node, leftover = query_expression_rule(tokenize(query_expression))
        self.assertFalse(leftover)
        assert isinstance(query_expression_node, QueryExpression)
        with self.assertRaisesRegexp(ValueError, error):
            query_expression_node.get_dataframe(self.table_context)

    # constants for use in data-driven syntax tests below
    with_clause = 'with foo as (select 1), bar as (select 2)'
    select_clause = 'select 3'
    order_by_clause = 'order by a, b'
    limit_clause = 'limit 5'

    @data(
        # core query expression grammar (no recursion)
        dict(query=(with_clause, select_clause, order_by_clause, limit_clause)),
        dict(query=(with_clause, select_clause, limit_clause)),
        dict(query=(with_clause, select_clause, order_by_clause)),
        dict(query=(with_clause, select_clause)),
        dict(query=(select_clause, order_by_clause, limit_clause)),
        dict(query=(select_clause, limit_clause)),
        dict(query=(select_clause, order_by_clause)),
        dict(query=(select_clause,)),

        # parenthesized recursion
        dict(query=(with_clause,
                    '(', with_clause, select_clause, order_by_clause, limit_clause, ')',
                    order_by_clause, limit_clause)),
        dict(query=(with_clause,
                    '(', with_clause, select_clause, order_by_clause, limit_clause, ')',
                    limit_clause)),
        dict(query=(with_clause,
                    '(', with_clause, select_clause, order_by_clause, limit_clause, ')',
                    order_by_clause)),
        dict(query=(with_clause,
                    '(', with_clause, select_clause, order_by_clause, limit_clause, ')')),

        # set operations
        dict(query=(select_clause, 'UNION ALL', select_clause)),  # simplest case
        dict(query=(
                # first query expression
                with_clause, select_clause, order_by_clause, limit_clause,
                'UNION ALL',
                # second query expression
                with_clause, select_clause, order_by_clause, limit_clause)),
    )
    @unpack
    def test_query_expression_syntax(self, query):
        query_expression = ' '.join(query)
        query_expression_node, leftover = query_expression_rule(tokenize(query_expression))
        self.assertFalse(leftover)

    @data(
        dict(query_expression='select * from my_table order by a DESC',
             expected_result=[[3], [2], [1]]),
        dict(query_expression='select a from my_table order by 1 DESC',
             expected_result=[[3], [2], [1]]),
    )
    @unpack
    def test_query_expression_order_by(self, query_expression, expected_result):
        # type: (str, List[List[int]]) -> None
        query_expression_node, leftover = query_expression_rule(tokenize(query_expression))
        assert isinstance(query_expression_node, QueryExpression)
        dataframe, unused_table_name = query_expression_node.get_dataframe(self.table_context)
        self.assertFalse(leftover)
        self.assertEqual(dataframe.to_list_of_lists(), expected_result)

    def test_select(self):
        # type: () -> None
        from_ = DataSource((TableReference(('my_project', 'my_dataset', 'my_table')),
                            EMPTY_NODE), [])
        selector = StarSelector(EMPTY_NODE, EMPTY_NODE, EMPTY_NODE)
        select = Select(EMPTY_NODE, [selector], from_, EMPTY_NODE, EMPTY_NODE, EMPTY_NODE)
        dataframe, table_name = select.get_dataframe(self.table_context)

        self.assertEqual(table_name, None)
        self.assertEqual(dataframe.to_list_of_lists(), [[1], [2], [3]])
        self.assertEqual(list(dataframe.dataframe), ['a'])
        self.assertEqual(dataframe.types, [BQScalarType.INTEGER])

    @data(
        dict(select='select c from my_table group by c', expected_result=[[3]]),
        dict(select='select my_table.c from my_table group by c', expected_result=[[3]]),
        dict(select='select c from my_table group by my_table.c', expected_result=[[3]]),
        dict(select='select my_table.c from my_table group by my_table.c', expected_result=[[3]]),
        dict(select='select c from my_table group by 1', expected_result=[[3]]),
        dict(select='select max(b), c from my_table group by 2', expected_result=[[3, 3]]),
        dict(select='select c+1 from my_table group by c', expected_result=[[4]]),
        dict(select='select c+1 from my_table group by 1', expected_result=[[4]]),
        dict(select='select c+1 from my_table group by c, 1', expected_result=[[4]]),
        dict(select='select c+1 as d from my_table group by c,d', expected_result=[[4]]),
        dict(select='select c+1 as d from my_table group by c,1', expected_result=[[4]]),
        dict(select='select c+1 from my_table group by c,1', expected_result=[[4]]),
        # Naively evaluating the below after the group by has occurred fails with
        # Column(s) my_table.c already selected
        dict(select='select c+1 as d, count(*) from my_table group by d', expected_result=[[4, 2]]),
        # Naively evaluating the below after the group by has occurred tries and fails to add
        # two SeriesGroupbys
        dict(select='select a+c as d from my_table group by d', expected_result=[[4]]),
    )
    @unpack
    def test_select_group_by(self, select, expected_result):
        # type: (str, List[List[int]]) -> None
        group_table_context = DatasetTableContext(
            {'my_project': {'my_dataset': {'my_table': TypedDataFrame(
                pd.DataFrame(
                    [[1, 2, 3], [1, 3, 3]],
                    columns=['a', 'b', 'c']),
                types=[BQScalarType.INTEGER, BQScalarType.INTEGER, BQScalarType.INTEGER]
            )}}})

        select_node, leftover = select_rule(tokenize(select))