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) EMPTY_NODE = _EmptyNode() class EvaluatableNode(AbstractSyntaxTreeNode):
    '''Abstract base class for syntax tree nodes that can be evaluated to return a column of data'''

    __metaclass__ = ABCMeta

    @abstractmethod
    def evaluate(self, context):
        # type: (EvaluationContext) -> Union[TypedDataFrame, TypedSeries]
        '''Generates a new table or column based on applying the instance's
        fields to the given context.

        This method should never be overridden by concrete subclasses.

        Args:
            context: The tables that are involved in this query

        Returns:
            A new table (TypedDataFrame) or column (TypedSeries)
        '''

    @abstractmethod
    def pre_group_by_partially_evaluate(self, context):
        # type: (EvaluationContext) -> Union[TypedSeries, EvaluatableNode]
        '''Partially evaluates an expression to prepare for an upcoming GROUP BY.

        This function evaluates expressions that are contained within an aggregation or a grouped by
        field, and stores the results of these in the context to be referred to later.  It is used
        as part of a two-pass evaluation to implement GROUP BY -- see EvaluationContext.do_group_by
        for more detail.

        This function should never be overriden by concrete subclasses.

        Args:
            context: The tables that are involved in this query
        Returns:
            A fully evaluated column or an unevaluated expression (an EvaluatableNode)
        '''

    def name(self):
        # type: () -> Optional[str]
        '''Returns a name for this expression, or None if no name can be inferred.

        Returns:
            The DataFrame returned from a select statement needs a name for every column.  This
            function returns a name that will be used for this expression if the user doesn't
            provide an explicit alias with AS.  If the expression doesn't have an obvious name (most
            expressions), it should return None, and the outermost expression (Selector) will
            generate one.
        '''
        return None

    @abstractmethod
    def mark_grouped_by(self, group_by_paths, context):
        # type: (Sequence[Tuple[str, ...]], EvaluationContext) -> EvaluatableNode
        '''Returns a new syntax tree rooted at the current one, marking fields that are grouped by.

        Args:
            group_by_paths: Canonicalized paths of the columns that are grouped by
            context: Context to evaluate in (for canonicalizing)
        '''

    def is_constant(self):
        # type: () -> bool
        '''Returns true if this expression is the same when evaluated in any context.'''
        return True

    def is_aggregated(self):
        # type: () -> bool
        '''Returns true if this expression contains any aggregation.'''
        return False 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 class _EmptyNode(AbstractSyntaxTreeNode):
    '''An Empty syntax tree node; represents an optional element not present.'''

    def strexpr(self):
        # type: () -> str
        return 'null' 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) DEFAULT_TABLE_NAME = None
class _WithTableContext(TableContext):
class QueryExpression(DataframeNode):
class SetOperation(DataframeNode):
class Select(MarkerSyntaxTreeNode, DataframeNode):
class TableReference(DataframeNode):
class Unnest(DataframeNode, MarkerSyntaxTreeNode):
    def __init__(self, name, table, parent_context):
    def lookup(self, path):
    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]
                 ):
    def _order_by(self, order_by, typed_dataframe, table_name, table_context):
    def _limit(self, limit, typed_dataframe):
    def get_dataframe(self, table_context, outer_context=None):
    def __init__(self, left_query, set_operator, right_query):
    def get_dataframe(self, table_context, outer_context=None):
def _evaluate_fields_as_dataframe(fields, context):
    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]
                 ):
    def get_dataframe(self, table_context, outer_context=None):
    def __init__(self, path):
    def get_dataframe(self, table_context, outer_context=None):
    def __init__(self, array_node):
    def get_dataframe(self, table_context, outer_context=None): 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 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 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]) 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_) class _AggregatingFunctionCall(FunctionCall, EvaluatableNodeThatAggregatesOrGroups):
    '''A function call that aggregates rows, e.g. sum.'''

    def __init__(self, function_info, expression):
        # type: (_AggregatingFunction, Union[_EmptyNode, Sequence[EvaluatableNode]]) -> None
        self.function_info = function_info
        if not isinstance(expression, _EmptyNode):
            self.children = expression
        else:
            self.children = []

    def copy(self, new_arguments):
        # type: (Sequence[EvaluatableNode]) -> EvaluatableNode
        return _AggregatingFunctionCall(self.function_info, new_arguments)

    def _evaluate_node(self, arguments):
        # type: (List[TypedSeries]) -> TypedSeries
        return self._evaluate(
            arguments,
            lambda values: values[0].apply(lambda x: self.function_info.aggregating_function([x])),
            self.function_info.compute_result_type) def core_expression(tokens):
    # type: (List[str]) -> AppliedRuleOutputType
    """Grammar rule for a core set of expressions that can be nested inside other expressions.

    The current set of handled expressions are:
    - Array
    - Array_agg
    - Count
    - Function call
    - A field (column)
    - Case
    - A literal (number, string, etc)
    - If
    - Cast
    - Exists
    - Extract
    - Another expression nested in parentheses
    - Not
    - Unary negation
    """
    return apply_rule(
        [
            # COUNT(*), COUNT(DISTINCT expression), COUNT(expression)
            wrap(Count.create_count_function_call,
                 ('COUNT', '(', ['*', (['DISTINCT', None], expression)], ')', [over_clause, None])),

            wrap(Array_agg.create_function_call,
                 ('ARRAY_AGG', '(', ['DISTINCT', None], expression,
                  [(['IGNORE', 'RESPECT'], 'NULLS'), None],
                  [(grammar_literal('ORDER', 'BY'),
                    separated_sequence(identifier, ['ASC', 'DESC', None], ',')),
                   None],
                  [('LIMIT', literal), None],
                  ')',
                  [over_clause, None])),

            wrap(FunctionCall.create,
                 (identifier, '(', [separated_sequence(expression, ','), None], ')', [over_clause,
                                                                                      None])),
            field,

            (Case, [expression, None], 'WHEN',
                separated_sequence((expression, 'THEN', expression), 'WHEN'),
                [('ELSE', expression), None], 'END'),

            literal,

            (If, '(', expression, ',', expression, ',', expression, ')'),

            (Cast, '(', expression, 'AS', scalar_type, ')'),

            (Exists, '(', query_expression, ')'),

            array_expression,

            (Extract, '(', identifier, 'FROM', expression, ')'),

            # Typeless struct syntax
            wrap(Struct.create_from_typeless,
                 ('STRUCT', '(', separated_sequence((expression, [('AS', identifier), None]), ','),
                  ')',)),

            # Typed struct syntax
            wrap(Struct.create_from_typed,
                 (struct_type, '(', separated_sequence(expression, ','), ')')),

            # Tuple syntax for structs (anonymous column names, implicit types)
            wrap(Struct.create_from_tuple,
                 ('(', expression, ',', separated_sequence(expression, ','), ')')),

            ('(', expression, ')'),

            (Not, expression),

            (UnaryNegation, '-', expression),
        ],
        tokens) def data_source(orig_tokens):
    # type: (List[str]) -> AppliedRuleOutputType
    '''Includes the initial FROM expression as well as any following JOINs.

    This describes everything that comes after a FROM, essentially in the form:
    from_item JOIN from_item ON on_expression JOIN from_item ON on_expression JOIN ...

    The first from_item is called first_from and is required.

    After that is any number of repetititions of ('JOIN', from_item, (ON, on_expression)),
    and is parsed into an array (possibly empty) called joins.

    Args:
        orig_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.
    '''
    first_from, tokens = apply_rule(from_item, orig_tokens)
    if not first_from:
        return None, tokens
    joins = []
    while True:
        next_join, tokens = apply_rule([(',', from_item),  # shorthand for cross-join
                                        (join_type,
                                         'JOIN',
                                         from_item,
                                         [('ON', expression),
                                          ('USING', '(', separated_sequence(identifier, ','), ')'),
                                          None])],
                                       tokens)
        if next_join:
            # This case is triggered by the shorthand cross-join above where the table to be joined
            # is just specified separated by a comma, with no join type or condition specified.
            if not isinstance(next_join, tuple):
                raise RuntimeError("Internal error; join rule above must result in tuple not {}"
                                   .format(next_join))
            if len(next_join) == 2:
                joins.append(Join('CROSS', cast(FromItemType, next_join), EMPTY_NODE))
            else:
                joins.append(Join(*next_join))
        else:
            break
    return DataSource(cast(FromItemType, first_from), joins), tokens def post_expression(tokens):
    """Grammar rule for expressions that occur only after a core expression.

    For example:
    <core_expression> IS NULL
    <core_expression> IN (a, b, c)

    If the query has none of these, it can still match a plain `core_expression`,
    the last item in the list.  [Currently this is the only thing implemented.]
    """
    return apply_rule(
        [
            (NullCheck, core_expression, [grammar_literal('IS', 'NULL'),
                                          grammar_literal('IS', 'NOT', 'NULL')]),
            (InCheck, core_expression, ['IN', grammar_literal('NOT', 'IN')],
             '(', separated_sequence(expression, ','), ')'),
            core_expression,
        ],
        tokens) 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)
import unittest
from typing import List, Tuple, Union  # 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
                                                 _EmptyNode)
from purplequery.bq_types import BQScalarType
from purplequery.dataframe_node import DataSource, QueryExpression, Select, TableReference
from purplequery.evaluatable_node import (Case, Cast, Count, Exists, Extract, If, InCheck, Not,
                                          NullCheck, Selector, StarSelector, UnaryNegation, Value,
                                          _AggregatingFunctionCall)
from purplequery.grammar import (core_expression, data_source, post_expression, query_expression,
                                 select)
                         ):
        # type: (...) -> None
        ast, leftover = data_source(tokens)

        assert isinstance(ast, DataSource)
        self.assertEqual(leftover, [])
        assert isinstance(ast.first_from[0], TableReference)
        self.assertEqual(ast.first_from[0].path, first_from_path)
        self.assertEqual(ast.first_from[1], first_from_alias)
        self.assertEqual(len(ast.joins), num_joins)
        if num_joins > 0:
            join = ast.joins[0]
            self.assertEqual(join[0], join_type)
            assert isinstance(join[1][0], TableReference)
            self.assertEqual(join[1][0].path, join_table_path)
            self.assertEqual(join[1][1], join_table_alias)
            self.assertEqual(repr(join[2]), repr(join_conditions))

    def test_core_expression(self):
        # type: () -> None
        ast, leftover = core_expression(['a'])

        self.assertEqual(leftover, [])
        assert isinstance(ast, Field)
        self.assertEqual(ast.path, ('a',))

    @data(
        (
            ['SELECT', 'a', 'FROM', '`my_project.my_dataset.my_table`'],
            ('a',)
        ),
        (
            ['SELECT', 'my_table', '.', 'a', 'FROM', '`my_project.my_dataset.my_table`'],
            ('my_table', 'a',)
        )
    )
    @unpack
    def test_select_field(self, tokens, path):
        # type: (List[str], Tuple[str, ...]) -> None
        ast, leftover = select(tokens)

        self.assertEqual(leftover, [])
        assert isinstance(ast, Select)
        assert isinstance(ast.fields[0], Selector)
        field = ast.fields[0].children[0]
        assert isinstance(field, Field)
        self.assertEqual(field.path, path)

    def test_order_by(self):
        # type: () -> None
        ast, leftover = query_expression(['SELECT', '*', 'FROM', 'my_table',
                                          'ORDER', 'BY', 'a', 'ASC'])

        assert isinstance(ast, QueryExpression)
        self.assertEqual(leftover, [])
        assert not isinstance(ast.order_by, _EmptyNode)
        self.assertEqual(len(ast.order_by), 1)
        (field, direction) = ast.order_by[0]
        self.assertEqual(field.path, ('a',))
        self.assertEqual(direction, 'ASC')

    def test_limit_offset(self):
        # type: () -> None
        ast, leftover = query_expression(['SELECT', '*', 'FROM', 'my_table',
                                          'LIMIT', '5', 'OFFSET', '10'])
        assert isinstance(ast, QueryExpression)
        self.assertEqual(leftover, [])
        assert not isinstance(ast.limit, _EmptyNode)
        limit_expression, offset_expression = ast.limit
        assert isinstance(limit_expression, Value)
        self.assertEqual(limit_expression.value, 5)
        assert isinstance(offset_expression, Value)
        self.assertEqual(offset_expression.value, 10)

    def test_group_by(self):
        # type: () -> None
        ast, leftover = select(['SELECT', '*', 'FROM', 'my_table', 'GROUP', 'BY', 'a'])

        self.assertEqual(leftover, [])
        assert isinstance(ast, Select)
        assert not isinstance(ast.group_by, _EmptyNode)
        self.assertEqual(len(ast.group_by), 1)
        field = ast.group_by[0]
        assert isinstance(field, Field)
        self.assertEqual(field.path, ('a',))

    def test_selector_alias(self):
        # type: () -> None
        ast, leftover = select(['SELECT', 'a', 'AS', 'field_name', 'FROM', 'my_table'])

        self.assertEqual(leftover, [])
        assert isinstance(ast, Select)
        self.assertEqual(len(ast.fields), 1)
        selector = ast.fields[0]
        assert isinstance(selector, Selector)
        alias = selector.alias
        self.assertEqual(alias, 'field_name')

    @data(
        dict(tokens=['COUNT', '(', '*', ')'],
             countee=Value(1, BQScalarType.INTEGER),
             distinct=False),
        dict(tokens=['COUNT', '(', 'DISTINCT', 'a', ')'],
             countee=Field(('a',)),
             distinct=True),
        dict(tokens=['COUNT', '(', '1', ')'],
             countee=Value(1, BQScalarType.INTEGER),
             distinct=False),
    )
    @unpack
    def test_count(self, tokens,  # type: List[str]
                   countee,  # type: EvaluatableNode
                   distinct,
                   ):
        # type: (...) -> None
        ast, leftover = core_expression(tokens)

        self.assertEqual(leftover, [])
        assert isinstance(ast, _AggregatingFunctionCall)
        self.assertEqual(ast.children, [countee])

        assert isinstance(ast.function_info, Count)