# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from collections import defaultdict
from copy import deepcopy
from itertools import ifilter
from logging import getLogger
from random import shuffle, choice, randint, randrange

from tests.comparison.common import TableExprList, ValExpr, ValExprList, Table, Column
from tests.comparison.query_profile import DefaultProfile
from tests.comparison.funcs import (
    AGG_FUNCS,
    AggFunc,
    ANALYTIC_FUNCS,
    AnalyticFunc,
    And,
    Coalesce,
    Equals,
    Func,
    FUNCS,
    PartitionByClause,
    WindowBoundary,
    WindowClause)
from tests.comparison.db_types import (
    Boolean,
    Int,
    Float,
    TYPES)
from tests.comparison.query import (
    FromClause,
    GroupByClause,
    HavingClause,
    InlineView,
    JoinClause,
    LimitClause,
    OrderByClause,
    Query,
    SelectClause,
    SelectItem,
    Subquery,
    UnionClause,
    WhereClause,
    WithClause,
    WithClauseInlineView)

UNBOUNDED_PRECEDING = WindowBoundary.UNBOUNDED_PRECEDING
PRECEDING = WindowBoundary.PRECEDING
CURRENT_ROW = WindowBoundary.CURRENT_ROW
FOLLOWING = WindowBoundary.FOLLOWING
UNBOUNDED_FOLLOWING = WindowBoundary.UNBOUNDED_FOLLOWING

LOG = getLogger(__name__)


class QueryGenerator(object):

  def __init__(self, query_profile):
    self.profile = query_profile
    self.queries_under_construction = list()
    self.max_nested_query_count = None
    self.cur_id = 0

  def get_next_id(self):
    self.cur_id += 1
    return 'a' + str(self.cur_id)

  @property
  def current_query(self):
    if self.queries_under_construction:
      return self.queries_under_construction[-1]

  @property
  def root_query(self):
    if self.queries_under_construction:
      return self.queries_under_construction[0]

  def clear_state(self):
    # This function clears the state from the previous query.
    self.cur_id = 0
    self.profile.query = None
    self.max_nested_query_count = None

  def generate_statement(self,
      table_exprs,
      allow_with_clause=True,
      allow_union_clause=True,
      select_item_data_types=None,
      required_select_item_type=None,
      required_table_expr_col_type=None,
      require_aggregate=None,
      dml_table=None,
      table_alias_prefix='t'):
    '''Create a random query using various language features.

       The initial call to this method should only use tables in the table_exprs
       parameter, and not  inline views or WITH definitions. The other types of
       table exprs may be added as part of the query generation.

       Due to implementation limitations, nested queries are not distributed evenly by
       default even if the query profile assigns an equal likelihood of use to each
       possible nested query decision. This is because the implementation chooses nested
       queries in a fixed order (WITH -> FROM -> WHERE). For example if only one nested
       query were allowed and each clause had a 50% chance of using a subquery, the
       resulting set of generated queries would contain a subquery in the WITH clause
       50% of the time, in the FROM 25% of the time, and WHERE 12.5% of the time. If the
       max nested queries were two, it's most likely that the generated query would
       contain a WITH clause which has the second nested query within it....

       If select_item_data_types is specified it must be a sequence or iterable of
       DataType. The generated query.select_clause.select will have data types suitable
       for use in a UNION.

       required_select_item_type may be set to a DataType to force at least one of the
       SELECT items to be of the given type. This can be used to ensure that inline
       views will have at least one joinable column.

       required_table_expr_col_type may be set to ensure that at least one of the
       TableExprs used in the FROM clause will have a column of the given type. This
       can be used to unsure that a correlation condition can be made for a Subquery.

       require_aggregate can be set to True or False to force or disable the creation
       of an aggregate query. This is used during Subquery creation where the context
       may require an aggregate or non-aggregate.

       dml_table exists for kwargs compatibility with other statement generators but is
       ignored
    '''
    if not table_exprs:
      raise Exception("At least one TableExpr is needed")

    query = Query()
    query.parent = self.current_query
    self.queries_under_construction.append(query)
    self.profile.query = query

    # Make a copy so tables can be added if a WITH clause is used
    table_exprs = TableExprList(table_exprs)

    if self.max_nested_query_count is None:
      self.max_nested_query_count = self.profile.get_max_nested_query_count()
    elif self.max_nested_query_count == 0:
      raise Exception('Maximum nested query count exceeded')
    else:
      self.max_nested_query_count -= 1

    with_clause = None
    if allow_with_clause \
        and self.allow_more_nested_queries \
        and self.profile.use_with_clause():
      with_clause = self._create_with_clause(table_exprs)
      table_exprs.extend(with_clause.table_exprs)
      query.with_clause = with_clause

    from_clause = self._create_from_clause(
        table_exprs, table_alias_prefix, required_table_expr_col_type)
    query.from_clause = from_clause

    select_clause = self._create_select_clause(
        from_clause.visible_table_exprs,
        select_item_data_types,
        required_select_item_type,
        require_aggregate)
    query.select_clause = select_clause

    if self.profile.use_where_clause():
      query.where_clause = self._create_where_clause(
          from_clause.visible_table_exprs, table_alias_prefix)

    # If agg and non-agg SELECT items are present then a GROUP BY is required otherwise
    # it's optional and is effectively a "SELECT DISTINCT".
    if (select_clause.agg_items and select_clause.basic_items) \
        or (require_aggregate is None \
            and select_clause.basic_items \
            and not select_clause.analytic_items \
            and not select_clause.contains_approximate_types \
            and self.profile.use_group_by_clause()):
      group_by_items = [deepcopy(item) for item in select_clause.basic_items
                        if not item.val_expr.is_constant]
      # TODO: What if there are select_clause.analytic_items?
      if group_by_items:
        query.group_by_clause = GroupByClause(group_by_items)

    # Impala doesn't support DISTINCT with analytics or "SELECT DISTINCT" when
    # GROUP BY is used.
    if not select_clause.analytic_items \
        and not query.group_by_clause \
        and not select_clause.contains_approximate_types \
        and self.profile.use_distinct():
      if select_clause.agg_items:
        self._enable_distinct_on_random_agg_items(select_clause.agg_items)
      else:
        select_clause.distinct = True

    if self.profile.use_having_clause() \
        and (query.group_by_clause
             or (self.profile.use_having_without_groupby()
             and select_clause.agg_items)):
      basic_select_item_exprs = \
          ValExprList(item.val_expr for item in select_clause.basic_items)
      query.having_clause = self._create_having_clause(
          from_clause.visible_table_exprs, basic_select_item_exprs)

    if allow_union_clause \
        and self.allow_more_nested_queries \
        and self.profile.use_union_clause():
      data_type_candidates_by_base_type = defaultdict(list)
      for data_type in TYPES:
        data_type_candidates_by_base_type[data_type.get_base_type()].append(data_type)
      select_item_data_types = list()
      for select_item in select_clause.items:
        select_item_data_types.append(
            choice(data_type_candidates_by_base_type[select_item.val_expr.base_type]))
      query.union_clause = UnionClause(self.generate_statement(
          table_exprs,
          allow_with_clause=False,
          select_item_data_types=select_item_data_types))
      if select_clause.contains_approximate_types or any(
          q.select_clause.contains_approximate_types for q in query.union_clause.queries):
        query.union_clause.all = True
      else:
        query.union_clause.all = self.profile.use_union_all()

    self.queries_under_construction.pop()
    if self.queries_under_construction:
      self.profile.query = self.queries_under_construction[-1]
    else:
      self.clear_state()

    return query

  @property
  def allow_more_nested_queries(self):
    return self.max_nested_query_count > 0

  def _create_with_clause(self, table_exprs):
    # Make a copy so newly created tables can be added and made available for use in
    # future table definitions.
    table_exprs = TableExprList(table_exprs)
    with_clause_inline_views = TableExprList()
    for with_clause_inline_view_idx \
        in xrange(self.profile.get_with_clause_table_ref_count()):
      query = self.generate_statement(table_exprs,
                                      allow_with_clause=self.profile.use_nested_with())
      with_clause_alias_count = getattr(self.root_query, 'with_clause_alias_count', 0) + 1
      self.root_query.with_clause_alias_count = with_clause_alias_count
      with_clause_inline_view = \
          WithClauseInlineView(query, 'with_%s' % with_clause_alias_count)
      table_exprs.append(with_clause_inline_view)
      with_clause_inline_views.append(with_clause_inline_view)
      if not self.allow_more_nested_queries:
        break
    return WithClause(with_clause_inline_views)

  def _create_select_clause(self,
      table_exprs,
      select_item_data_types,
      required_select_item_type,
      require_aggregate):
    if select_item_data_types:
      select_item_data_types = tuple(select_item_data_types)
    if select_item_data_types \
        and required_select_item_type \
        and not issubclass(required_select_item_type, select_item_data_types):
      raise Exception('Required select item type is not in allowed types')

    # Generate column types for the select clause if not specified.
    if not select_item_data_types:
      select_item_data_types = []
      desired_item_count = self.profile.get_select_item_count()
      if required_select_item_type:
        select_item_data_types.append(required_select_item_type)
      while len(select_item_data_types) < desired_item_count:
        select_item_data_types.append(self.profile.choose_type(table_exprs.col_types))
      shuffle(select_item_data_types)
      select_item_data_types = tuple(select_item_data_types)

    # We want to assign BASIC, AGG or ANALYTIC to each column. We want to prevent GROUP BY
    # float_col, because the groupings will differ across databases because floats are
    # approximate values.
    column_categories = ['-' for _ in select_item_data_types]
    if require_aggregate:
      column_categories[randint(0, len(column_categories) - 1)] = 'AGG'
    # Assign AGG and ANALYTIC some columns
    for i in range(len(column_categories)):
      item_category = self.profile._choose_from_weights(
          self.profile.weights('SELECT_ITEM_CATEGORY'))
      if item_category in ('AGG', 'ANALYTIC'):
        column_categories[i] = item_category
    agg_present = 'AGG' in column_categories
    # Assign BASIC to the remaining columns
    for i in range(len(column_categories)):
      if column_categories[i] == '-':
        if select_item_data_types[i].is_approximate() and agg_present:
          column_categories[i] = 'AGG'
        else:
          # If AGG column is present, BASIC can only be assigned to a column if it's not a
          # Float.
          column_categories[i] = 'BASIC'

    select_items = []

    for i, column_category in enumerate(column_categories):
      if column_category == 'BASIC':
        select_items.append(self._create_basic_select_item(
          table_exprs, select_item_data_types[i]))
      else:
        select_items.append((column_category, select_item_data_types[i]))

    analytic_count = sum(1 for c in column_category if c == 'ANALYTIC')

    select_item_exprs = \
        ValExprList(item.val_expr for item in select_items if type(item) == SelectItem)
    for idx, item in enumerate(select_items):
      if type(item) == tuple and item[0] == 'AGG':
        select_items[idx] = \
            self._create_agg_select_item(table_exprs, select_item_exprs, item[1])

    for item in select_items:
      if type(item) == tuple:
        continue
      if item.is_agg:
        select_item_exprs.append(item.val_expr)
    for idx, item in enumerate(select_items):
      if type(item) == tuple:
        select_items[idx] = self._create_analytic_select_item(
            table_exprs,
            select_item_exprs,
            len(select_item_exprs) == 0 and analytic_count == 1,
            item[1])

    # So far all the SELECT items are defined and set but none of them have aliases. If
    # an item is a simple column reference, then it will only get an alias if there is a
    # conflict with another simple column ref. All other item types, such as functions
    # or constants, will always have an alias.
    #
    # item_names is the set of final names or aliases for each item in the SELECT list.
    item_names = set()
    for item in select_items:
      if not item.alias and not item.val_expr.is_col:
        continue
      if item.name in item_names:
        item.alias = '*CONFLICT*'
      else:
        item_names.add(item.name)

    # base_item_name_counts stores the number of conflicts that occurred for a name,
    # and hence the number of name to skip forward to create a non-conflicting name.
    base_item_name_counts = defaultdict(int)
    for item in select_items:
      if item.alias == '*CONFLICT*' or (not item.val_expr.is_col and not item.alias):
        # Use names close to the Impala functional test database so that bugs in
        # resolution will be more likely to surface.
        alias = base_alias = '%s_col' % item.type.__name__.lower()
        while alias in item_names:
          base_item_name_counts[base_alias] += 1
          alias = base_alias + '_' + str(base_item_name_counts[base_alias])
        item.alias = alias
        item_names.add(alias)
    return SelectClause(select_items)

  def _create_basic_select_item(self, table_exprs, return_type):
    max_children = self.profile.choose_nested_expr_count()
    if max_children:
      value = self.create_func_tree(return_type)
      value = self.populate_func_with_vals(value, table_exprs)
    elif return_type in table_exprs.col_types:
      value = self.profile.choose_val_expr(table_exprs.cols_by_type[return_type])
    else:
      value = self.profile.choose_constant(return_type)
    return SelectItem(value)

  def create_func_tree(self, return_type, allow_subquery=False):
    '''Returns an instance of a basic function that has all of it's arguments either set
       to None or another instance of a function that has it's arguments set likewise. The
       caller should replace the None values with column references or constants as
       desired. The depth of the tree is determined by the query profile (self.profile).
    '''
    signatures = self._funcs_to_allowed_signatures(FUNCS)
    root_signatures = self._find_matching_signatures(
        signatures, returns=return_type, allow_subquery=allow_subquery)
    root_signature = self.profile.choose_func_signature(root_signatures)
    func = root_signature.func(root_signature)   # An instance of a function
    max_children = self.profile.choose_nested_expr_count()
    if max_children:
      # Impala does not allow functions that contain subqueries to have arguments
      # that contain subqueries. Ex: ... WHERE (int_col IN (SELECT 1)) IN (SELECT TRUE)
      subquery_allowed_null_args = list()
      subquery_not_allowed_null_args = list()
      if func.contains_subquery:
        null_args = subquery_not_allowed_null_args
      else:
        null_args = subquery_allowed_null_args
      null_args.extend((func, idx) for idx, arg in enumerate(func.args)
                       if type(arg) != list and arg.val is None)
      while max_children \
          and (subquery_allowed_null_args or subquery_not_allowed_null_args):
        idx = randrange(
            len(subquery_allowed_null_args) + len(subquery_not_allowed_null_args))
        if idx < len(subquery_allowed_null_args):
          null_args = subquery_allowed_null_args
        else:
          null_args = subquery_not_allowed_null_args
        shuffle(null_args)
        parent_func, parent_arg_idx = null_args.pop()
        child_signatures = self._find_matching_signatures(
            signatures,
            returns=parent_func.args[parent_arg_idx].type,
            allow_subquery=(allow_subquery and null_args == subquery_allowed_null_args))
        child_signature = self.profile.choose_func_signature(child_signatures)
        child_func = child_signature.func(child_signature)
        parent_func.args[parent_arg_idx] = child_func
        if child_func.contains_subquery:
          null_args = subquery_not_allowed_null_args
        else:
          null_args = subquery_allowed_null_args
        null_args.extend((child_func, idx) for idx, arg in enumerate(child_func.args)
                         if type(arg) != list and arg.val is None)
        max_children -= 1
    return func

  def _funcs_to_allowed_signatures(self, funcs):
    '''Return a list of the signatures contained in "funcs" that are eligible for use
       based on the query profile.
    '''
    return [signature for func in funcs for signature in func.signatures()
            if self.profile.allow_func_signature(signature)]

  def _find_matching_signatures(self,
      signatures,
      returns=None,
      accepts=None,
      accepts_only=None,
      allow_subquery=False):
    '''Returns the subset of signatures matching the given criteria.
         returns: The signature must return this type.
         accepts: The signature must have at least one argument of this type.
         accepts_only: The signature must have arguments of only this type.
         allow_subquery: If False, the signature cannot contain a subquery.
    '''
    matching_signatures = list()
    for signature in signatures:
      if returns and not issubclass(signature.return_type, returns):
        continue
      if accepts and not any(not arg.is_subquery and issubclass(arg.type, accepts)
                             for arg in signature.args):
        continue
      if accepts_only and (not signature.args or any(
          not arg.is_subquery and not issubclass(arg.type, accepts_only)
          for arg in signature.args)):
        continue
      if not allow_subquery and any(arg.is_subquery for arg in signature.args):
        continue
      matching_signatures.append(signature)
    return matching_signatures

  def populate_func_with_vals(self,
      func,
      table_exprs=TableExprList(),
      val_exprs=ValExprList(),
      table_alias_prefix='',
      allow_subquery=False,
      _allow_table_exprs=None):
    if not _allow_table_exprs and func.is_agg:
      _allow_table_exprs = True
    elif _allow_table_exprs is None and func.contains_agg:
      _allow_table_exprs = False
    elif _allow_table_exprs is None:
      _allow_table_exprs = True
    # If a function's return type depends on some of its args then at least one of those
    # args must not be the NULL literal. Example: IF(false, NULL, NULL) is considered
    # invalid because the return type cannot be determined.
    has_non_null_literal_arg = False
    for idx, arg in enumerate(func.args):
      signature_arg = func.signature.args[idx]
      if signature_arg.is_subquery \
          or (allow_subquery \
              and self.allow_more_nested_queries \
              and self.profile.use_scalar_subquery()):
        usage = self.profile.choose_subquery_predicate_category(
            func.name(),
            self.current_query.from_clause.table_exprs.joinable_cols_by_type)
        if usage is not None \
            and self.allow_more_nested_queries \
            and (usage[1] == 'UNCORRELATED'
                or self.current_query.from_clause.table_exprs.joinable_cols_by_type):
          use_scalar_subquery = (usage[0] == 'Scalar')
          use_agg_subquery = (usage[1] == 'AGG')
          use_correlated_subquery = (usage[2] == 'CORRELATED')
          if use_correlated_subquery:
            # TODO: Sometimes this causes an exception because the list is empty
            join_expr_type = self.profile.choose_type(list(
                self.current_query.from_clause.table_exprs.joinable_cols_by_type))
          else:
            join_expr_type = None
          select_item_data_types = \
              [signature_arg.type] if use_scalar_subquery else signature_arg.type
          query = self.generate_statement(
              table_exprs,
              select_item_data_types=select_item_data_types,
              required_table_expr_col_type=join_expr_type,
              require_aggregate=use_agg_subquery,
              # Don't use UNION + LIMIT; IMPALA-1379
              allow_union_clause=(not signature_arg.is_subquery),
              table_alias_prefix=(table_alias_prefix +
                  ('t' if use_correlated_subquery else '')),
              allow_with_clause=self.profile.use_nested_with())
          if use_scalar_subquery and not use_agg_subquery:
            # Impala will assume the query will return more than one row unless a LIMIT 1
            # is added. An ORDER BY will also be added under the assumption that we want
            # deterministic results.
            query.order_by_clause = OrderByClause([Int(1)])
            query.limit_clause = LimitClause(Int(1))
          if use_correlated_subquery:
            outer_table_expr = choice(
                self.current_query.from_clause.table_exprs.by_col_type[join_expr_type])
            correlation_condition = self._create_relational_join_condition(
                outer_table_expr,
                query.from_clause.table_exprs.by_col_type[join_expr_type])
            if query.where_clause:
              query.where_clause.boolean_expr = And.create_from_args(
                  query.where_clause.boolean_expr, correlation_condition)
            else:
              query.where_clause = WhereClause(correlation_condition)
          func.args[idx] = Subquery(query)
        else:
          replacement_func = self.create_func_tree(func.type)
          return self.populate_func_with_vals(
              replacement_func,
              table_exprs=table_exprs,
              val_exprs=val_exprs,
              table_alias_prefix=table_alias_prefix,
              allow_subquery=allow_subquery,
              _allow_table_exprs=_allow_table_exprs)
      else:
        if arg.is_constant and arg.val is None:
          candidate_val_exprs = ValExprList()
          if val_exprs:
            candidate_val_exprs.extend(val_exprs.by_type[arg.type])
          if _allow_table_exprs:
            candidate_val_exprs.extend(table_exprs.cols_by_type[arg.type])
          if candidate_val_exprs:
            val = self.profile.choose_val_expr(candidate_val_exprs)
          else:
            val = self.profile.choose_constant(
                return_type=arg.type,
                allow_null=(signature_arg.can_be_null \
                    and signature_arg.can_be_null_literal \
                    and (has_non_null_literal_arg \
                        or not signature_arg.determines_signature)))
          func.args[idx] = val
          arg = val
        elif arg.is_func:
          func.args[idx] = self.populate_func_with_vals(
              arg,
              table_exprs=table_exprs,
              val_exprs=val_exprs,
              _allow_table_exprs=_allow_table_exprs)
        if not signature_arg.can_be_null and not arg.is_constant:
          val = self.profile.choose_constant(return_type=arg.type, allow_null=False)
          func.args[idx] = Coalesce.create_from_args(arg, val)
        if not arg.is_constant or arg.val is not None:
          has_non_null_literal_arg = True
    return func

  def _create_agg_select_item(self, table_exprs, basic_select_item_exprs, return_type):
    value = self._create_agg_func_tree(return_type)
    value = self.populate_func_with_vals(value, table_exprs, basic_select_item_exprs)
    return SelectItem(value)

  def _create_agg_func_tree(self, return_type):
    return self._create_agg_or_analytic_tree(return_type, agg_funcs=AGG_FUNCS)

  def _create_agg_or_analytic_tree(self, return_type, agg_funcs=[], analytic_funcs=[],
                                   basic_funcs=FUNCS):
    '''Returns an instance of a function that is guaranteed to either be or contain an
       aggregate or analytic function. The arguments of the returned function will either
       be None or an instance of a function as in create_func_tree.

       The chosen aggregate or analytic functions will be restricted to the list of
       functions in agg_funcs and analytic_funcs.

       If analytic_funcs is non-empty the returned function will be guaranteed to
       be an analytic or contain at least on analytic function.

       return_type must be set and refers to the data type of the function output.
       agg_funcs and analytic_funcs should be used to determine the class of the
       returned function. The caller is responsible for restricting the return_type
       to types that can be generated by permutations of available functions. If the max
       nested expr count in the query profile is at least one, then any return_type
       should be possible to generate but this is not guaranteed.

       basic_funcs is a list of allowed basic functions that will be used in the function
       tree. By default, all basic funcs defined in funcs.py will be allowed. A basic
       function is any non-aggregate, non-analytic function.
    '''
    # The creation of aggregate and analytic functions is so similar that they are
    # combined here. "basic" function creation is much simpler so that is kept separate.
    # What's going to happen is there will be essentially two important data structures:
    #
    #   1) A tree of functions, the root of which will be returned. The leaves of the
    #      tree are "place holders" which are actually instances of a concrete DataType,
    #      such as Int, with a value of None. In other words, the arguments to all
    #      functions are either other functions or SQL NULL.
    #
    #   2) A mapping to place holders from the type of function that they may be replaced
    #      with. The actual data structure is
    #      dict<func class> -> list<tuple<tree node, index of place holder in tree node>>
    #      where "func class" is one of "AggFunc", "AnalyticFunc" or "Func".
    #
    # This means once a child function is generated, a spot where it can be placed into
    # the tree can easily be identified. Although the work is actually done in reverse
    # order, a place holder is chosen, then a replacement is generated.
    if not agg_funcs and not analytic_funcs:
      raise Exception('At least one analytic or aggregate function is required')

    basic_signatures_by_return_type = self._group_signatures_by_return_type(
        self._find_matching_signatures(
            self._funcs_to_allowed_signatures(basic_funcs),
            allow_subquery=False))
    agg_signatures_by_return_type = self._group_signatures_by_return_type(
        self._funcs_to_allowed_signatures(agg_funcs))
    analytic_signatures_by_return_type = self._group_signatures_by_return_type(
        self._funcs_to_allowed_signatures(analytic_funcs))
    if analytic_funcs:
      return_class = AnalyticFunc
      return_signatures_by_return_type = analytic_signatures_by_return_type
    else:
      return_class = AggFunc
      return_signatures_by_return_type = agg_signatures_by_return_type

    min_children, max_children = self.profile.bounds('MAX_NESTED_EXPR_COUNT')
    if max_children == 0 and return_type not in return_signatures_by_return_type:
      raise Exception(('At least one child expr is required to create a %s expr'
          + ' using an %s function')
          % (return_type.__name__,
             'aggregate' if return_class == AggFunc else 'analytic'))
    if min_children == 0 and return_type not in return_signatures_by_return_type:
      min_children = 1
    max_children = self.profile._choose_from_bounds(min_children, max_children)

    if not max_children:
      signature = self.profile.choose_func_signature(
          return_signatures_by_return_type[return_type])
      return signature.func(signature)

    root_func = None

    # Every time a function is created its arguments will initially be NULL. Those NULL
    # arguments may be replaced by other functions up to "max_children". The types of
    # valid child functions depends on the type of parent function.
    #
    #  * Analytics may not contain other analytics
    #  * Analytics may contain aggregates or basic functions
    #  * Aggregates may not contain other aggregates or analytics
    #  * Aggregates may contain basic functions
    #  * Basic functions may contain any function so long as the above conditions are
    #    not broken (indirectly).
    null_args_by_func_allowed = {AggFunc: list(), AnalyticFunc: list(), Func: list()}
    while max_children:
      null_arg_pool = None
      parent_func, parent_arg_idx = None, None
      chosen_signature=None

      # Since aggregate (and let's assume analytic functions) return a limited set of
      # types, some prep work may be needed if the return type isn't in the set of
      # directly producible types. For example if a Boolean is desired and there are
      # only two children remaining, the next child must be something that accepts an
      # aggregate and returns a Boolean, such as GreaterThan.
      if (not root_func and max_children == 1) \
          or (root_func and max_children == 2 \
              and ((return_class == AggFunc and not root_func.contains_agg)
                  or (return_class == AnalyticFunc and not root_func.contains_analytic))):
        if root_func:
          idx = randrange(len(null_args_by_func_allowed[return_class])
              + len(null_args_by_func_allowed[Func]))
          use_return_class = idx < len(null_args_by_func_allowed[return_class])
          if use_return_class:
            null_arg_pool = null_args_by_func_allowed[return_class]
            signature_pools_by_return_type = return_signatures_by_return_type
          else:
            null_arg_pool = null_args_by_func_allowed[Func]
            signature_pools_by_return_type = basic_signatures_by_return_type
          shuffle(null_arg_pool)
          parent_func, parent_arg_idx = null_arg_pool.pop()
          parent_arg_type = parent_func.args[parent_arg_idx].type
          if use_return_class:
            signature_pool = signature_pools_by_return_type[parent_arg_type]
          else:
            # This is a basic functions so it needs to accept one of the types returned
            # by the desired return_class.
            signature_pool = self._find_matching_signatures(
                signature_pools_by_return_type[parent_arg_type],
                accepts=tuple(return_signatures_by_return_type))
        else:
          signature_pool = list()
          if return_type in return_signatures_by_return_type:
            signature_pool.extend(return_signatures_by_return_type[return_type])
          if return_type in basic_signatures_by_return_type:
            signature_pool.extend(self._find_matching_signatures(
                basic_signatures_by_return_type[return_type],
                accepts=tuple(return_signatures_by_return_type)))
        chosen_signature = self.profile.choose_func_signature(signature_pool)
      elif max_children == 1 \
          and ((return_class == AggFunc and not root_func.contains_agg)
              or (return_class == AnalyticFunc and not root_func.contains_analytic)):
        # This is the last iteration and the generated function still doesn't contain
        # an instance of the desired function class. From the setup above, at least
        # one of the available place holder arguments can be replaced by an aggregate or
        # analytic.
        null_arg_pool = null_args_by_func_allowed[return_class]
        if not null_arg_pool:
          raise Exception(
              'No leaves in the expr tree may be replaced by an %s function'
              % ('aggregate' if return_class == AggFunc else 'analytic'))
        shuffle(null_arg_pool)
        return_types = tuple(return_signatures_by_return_type)
        while null_arg_pool:
          parent_func, parent_arg_idx = null_arg_pool.pop()
          parent_arg_type = parent_func.args[parent_arg_idx].type
          if issubclass(parent_arg_type, return_types):
            break
          if not null_arg_pool:
            raise Exception('No functions could accept an %s function'
                % ('aggregate' if return_class == AggFunc else 'analytic'))
        chosen_signature = self.profile.choose_func_signature(
            return_signatures_by_return_type[parent_arg_type])
      elif not root_func:
        # No restrictions, just choose a root_func that returns the needed type.
        signature_pool = list()
        if return_type in return_signatures_by_return_type:
          signature_pool.extend(return_signatures_by_return_type[return_type])
        if return_type in basic_signatures_by_return_type:
          signature_pool.extend(basic_signatures_by_return_type[return_type])
        chosen_signature= self.profile.choose_func_signature(signature_pool)
      else:
        # A root_func was chosen and it's children are in one or more of the
        # null_args_by_func_allowed pools. A pool will be chosen, then a child function.
        null_arg_counts_by_pool = dict((pool_category, len(pool)) for pool_category, pool
                                       in null_args_by_func_allowed.iteritems())
        # There is a special case that would lead to a dead end. If there is only one
        # distinct place holder across all the pools and an analytic is still needed,
        # then that place holder cannot be replaced by an aggregate since aggregates
        # are not allowed to contain analytics. Example: an analytic is impossible once
        # the tree is FLOOR(AVG(NULL)).
        if return_class == AnalyticFunc \
            and not root_func.contains_analytic \
            and len(null_args_by_func_allowed[AnalyticFunc]) == 1 \
            and null_args_by_func_allowed[AnalyticFunc][0] \
                in null_args_by_func_allowed[AggFunc]:
          null_arg_counts_by_pool[AggFunc] = 0

        null_arg_pool_category = \
            self.profile._choose_from_weights(null_arg_counts_by_pool)
        null_arg_pool = null_args_by_func_allowed[null_arg_pool_category]
        shuffle(null_arg_pool)
        parent_func, parent_arg_idx = null_arg_pool.pop()
        parent_arg_type = parent_func.args[parent_arg_idx].type

        if null_arg_pool_category == AggFunc:
          signature_pool = agg_signatures_by_return_type[parent_arg_type]
        elif null_arg_pool_category == AnalyticFunc:
          signature_pool = analytic_signatures_by_return_type[parent_arg_type]
        else:
          signature_pool = basic_signatures_by_return_type[parent_arg_type]
        chosen_signature = self.profile.choose_func_signature(signature_pool)

      chosen_func = chosen_signature.func(chosen_signature)
      if root_func:
        max_children -= 1
      else:
        root_func = chosen_func

      if parent_func:
        # Remove the place holder from all of the other pools.
        for pool_category, pool in null_args_by_func_allowed.iteritems():
          for null_arg_idx, (func, arg_idx) in enumerate(pool):
            if func is parent_func and arg_idx == parent_arg_idx:
              del pool[null_arg_idx]

        # Replace the place holder with the child function
        parent_func.args[parent_arg_idx] = chosen_func
        parent_func.validate()
        chosen_func.parent = parent_func
      else:
        chosen_func.parent = None

      # Check if the func_tree contains any analytic functions returned by
      # profile.get_analytic_funcs_that_cannot_contain_aggs; if it does then aggregate
      # functions are not allowed in child operators
      allow_agg_func = not self._func_tree_contains_funcs(
          chosen_func, self.profile.get_analytic_funcs_that_cannot_contain_aggs())

      # Place the args of the chosen function into the appropriate pools. Aggregate and
      # analytic functions have different rules about which functions they accept as
      # arguments. If the chosen function is neither then the tree must be inspected to
      # find the first aggregate or analytic ancestor.
      child_null_arg_pools = set()
      node = chosen_func
      while True:
        if node.is_analytic:
          if allow_agg_func:
            child_null_arg_pools.add(AggFunc)
          child_null_arg_pools.add(Func)
          break
        elif node.is_agg:
          child_null_arg_pools.add(Func)
          break
        node = node.parent
        if not node:
          # This means the root_func is a non-aggregate and non-analytic, it can accept
          # anything.
          child_null_arg_pools.add(AggFunc)
          child_null_arg_pools.add(Func)
          if return_class == AnalyticFunc:
            child_null_arg_pools.add(AnalyticFunc)
          break
      for func_allowed in child_null_arg_pools:
        null_args = null_args_by_func_allowed[func_allowed]
        for idx, arg in enumerate(chosen_func.args):
          if arg.val is not None:
            # Some functions come with unusual arguments pre-populated.
            # Ex: The analytic function LEAD(foo, constant) will have "constant" non-NULL.
            continue
          if func_allowed == AggFunc:
            returnable_types = agg_signatures_by_return_type
          elif func_allowed == AnalyticFunc:
            returnable_types = analytic_signatures_by_return_type
          else:
            returnable_types = basic_signatures_by_return_type
          # Skip place holders that are impossible to replace. The function selection
          # logic depends on this.
          if arg.type not in returnable_types:
            continue
          null_args.append((chosen_func, idx))

      if not any(null_args_by_func_allowed.itervalues()):
        # Some analytic functions take no arguments. Ex: ROW_NUM()
        break

    if (return_class == AggFunc and not root_func.contains_agg) \
        or (return_class == AnalyticFunc and not root_func.contains_analytic):
      raise Exception('Unable to create an %s function'
          % ('aggregate' if return_class == AggFunc else 'analytic'))

    return root_func

  def _func_tree_contains_funcs(self, func_tree, funcs):
    """
    Returns True if any of the funcs are in the given func_tree. A func_tree must have
    a parent property defined that links to the parent function. This method compares
    functions based on their name.
    """
    if not funcs:
      return False
    node = func_tree
    func_names = set(func.name() for func in funcs)
    while node:
      if node.name() in func_names:
        return True
      node = node.parent
    return False

  def _create_analytic_select_item(self,
      table_exprs,
      select_item_exprs,
      is_only_item,
      return_type):
    '''Create and return a list of SelectItems where each item includes an anayltic
       function in its val_expr.

       table_exprs must be a non-empty collection of tables exprs to choose columns from.

       select_items may be an empty collection. If not empty, anayltic function
       expressions will be restricted so that PARTITION BY and ORDER BY values are
       chosen from within this set. If empty, any random expr may be chosen.

       is_only_item should be True if this analytic item will be the only SelectItem in
       the SELECT clause. When True, an additional analytic variation that would
       otherwise produce a non-deterministic result set is enabled.
    '''
    # Results may not be deterministic when analytics have an ORDER BY. The code below
    # will assume we want deterministic results though we know this may not be true
    # for certain use cases such as just searching for crashes where results will be
    # thrown away.
    #
    # These cases will produce a deterministic result when ORDER BY is used
    #
    #   1) There is only one SELECT item and the analytic function has no argument
    #      or the ORDER BY includes the argument.
    #
    #   2) The analytic has set of ORDER BY expressions that has a deterministic
    #      result. Ex: ORDER BY primary_key. This is only reliable if there are no
    #      JOINs.
    #
    #   3) A window is explicitly specified as UNBOUNDED PRECEDING to UNBOUNDED
    #      FOLLOWING and either the analytic function is an aggregate or the
    #      analytic function is FIRST_VALUE or LAST_VALUE.
    #
    #   4) RANK is a special case, it is intrinsically deterministic
    excluded_designs = list()
    if len(self.queries_under_construction) > 1:
      excluded_designs.append('TOP_LEVEL_QUERY_WITHOUT_LIMIT')
    if not is_only_item:
      excluded_designs.append('ONLY_ITEM')
    if len(table_exprs) == 1:
      unique_col_combos = []
      if table_exprs[0].unique_cols:
        if select_item_exprs:
          select_cols = set(val_expr for val_expr in select_item_exprs if val_expr.is_col)
          for unique_cols in table_exprs[0].unique_cols:
            if unique_cols <= select_cols:
              unique_col_combos.append(unique_cols)
        else:
          unique_col_combos = table_exprs[0].unique_cols
      if not unique_col_combos:
        excluded_designs.append('DETERMINISTIC_ORDER')
    else:
      excluded_designs.append('DETERMINISTIC_ORDER')

    allow_agg = any(ifilter(lambda expr: expr.contains_agg, select_item_exprs))
    value = self._create_analytic_func_tree(return_type, excluded_designs, allow_agg)
    value = self.populate_func_with_vals(
        value,
        table_exprs=TableExprList(),
        val_exprs=ValExprList())
    self._populate_analytic_clauses(value, table_exprs, select_item_exprs)
    return SelectItem(value)

  def _create_analytic_func_tree(self, return_type, excluded_designs, allow_agg):
    # The results of queries with analytic functions may not be deterministic. Let's
    # assume deterministic results are required. To generate deterministic results, the
    # structure of queries will be limited to one of several "designs". The goal of each
    # design is to ensure that results are deterministic. See the notes in
    # _create_analytic_select_item for more info on designs.
    designs = set(self.profile.allowed_analytic_designs())
    for design in excluded_designs:
      if design in designs:
        designs.remove(design)
    window_weigts = self.profile.weights('ANALYTIC_WINDOW')
    if 'UNBOUNDED_WINDOW' in designs \
        and window_weigts[('ROWS', UNBOUNDED_PRECEDING, UNBOUNDED_FOLLOWING)] == 0 \
        and window_weigts[('RANGE', UNBOUNDED_PRECEDING, UNBOUNDED_FOLLOWING)] == 0:
      designs.remove('UNBOUNDED_WINDOW')
    if not designs:
      raise Exception('No options for creating a deterministic analytic function are'
          ' available')

    available_funcs = list()
    for func in ANALYTIC_FUNCS:
      if 'TOP_LEVEL_QUERY_WITHOUT_LIMIT' in designs \
          or 'DETERMINISTIC_ORDER_BY' in designs \
          or 'ONLY_SELECT_ITEM' in designs \
          or ('RANK_FUNC' in designs and func.name() == 'Rank') \
          or ('NO_ORDER_BY' in designs and not func.REQUIRES_ORDER_BY) \
          or ('UNBOUNDED_WINDOW' in designs and func.SUPPORTS_WINDOWING):
        available_funcs.append(func)
    if not available_funcs:
      raise Exception('No analytic functions available')

    # In general this may be a tree-like val expr that is guaranteed to contain at least
    # one analytic function. The arguments to the function are already outlined but the
    # analytic clause is still undetermined.
    root_func = self._create_agg_or_analytic_tree(
        return_type,
        analytic_funcs=available_funcs,
        agg_funcs=(AGG_FUNCS if allow_agg else []))

    # The following designs are preferred because they place no restrictions on the
    # analytic. From now on, if "designs" is non-empty this means that a design must
    # be chosen.
    if 'TOP_LEVEL_QUERY_WITHOUT_LIMIT' in designs \
        or 'DETERMINISTIC_ORDER_BY' in designs \
        or 'ONLY_SELECT_ITEM' in designs:
      designs.clear()
    else:
      designs = list(designs)

    # Set the analytic clauses, such as PARTITION BY, etc, for each analytic function.
    analytic_funcs = list()
    if root_func.is_analytic:
      analytic_funcs.append(root_func)
    for func in root_func.iter_exprs(lambda expr: expr.is_analytic):
      analytic_funcs.append(func)
    for func in analytic_funcs:
      if designs:
        func.design = choice([design for design in designs
                              if design != 'RANK' or func.name() != 'Rank'])
      else:
        func.design = 'NO_RESTRICTIONS'

      use_window = (func.design == 'UNBOUNDED_WINDOW') \
          or (func.design != 'NO_ORDER_BY'
              and (func.SUPPORTS_WINDOWING and self.profile.use_window_in_analytic()))
      if use_window:
        window_range_or_rows, window_start_boundary_type, window_end_boundary_type \
            = self.profile.choose_window_type()
        if window_start_boundary_type in (PRECEDING, FOLLOWING):
          start_val = Int(self.profile._choose_from_bounds('ANALYTIC_WINDOW_OFFSET'))
        else:
          start_val = None
        if window_end_boundary_type in (PRECEDING, FOLLOWING):
          end_val = Int(self.profile.get_window_offset())
          if window_start_boundary_type == PRECEDING \
              and window_end_boundary_type == PRECEDING:
            if start_val.val < end_val.val:
              start_val, end_val = end_val, start_val
            elif start_val.val == end_val.val:
              end_val.val -= 1
          elif window_start_boundary_type == FOLLOWING \
              and window_end_boundary_type == FOLLOWING:
            if start_val.val > end_val.val:
              start_val, end_val = end_val, start_val
            elif start_val.val == end_val.val:
              start_val.val -= 1
        else:
          end_val = None
        func.window_clause = WindowClause(
            window_range_or_rows,
            WindowBoundary(window_start_boundary_type, start_val),
            WindowBoundary(window_end_boundary_type, end_val) \
                if window_end_boundary_type else None)

      if self.profile.use_partition_by_clause_in_analytic():
        func.partition_by_clause = PartitionByClause([])

      if (func.design == 'DETERMINISTIC_ORDER') \
          or use_window \
          or func.REQUIRES_ORDER_BY \
          or (func.design != 'NO_ORDER_BY' \
              and self.profile.use_order_by_clause_in_analytic()):
        func.order_by_clause = OrderByClause([])

      if func.name() in ('Lead', 'Lag') and len(func.args) == 2:
        func.args[1].val = self.profile.get_offset_for_analytic_lead_or_lag()

    return root_func

  def _group_signatures_by_return_type(self, signatures):
    groups = defaultdict(list)
    for signature in signatures:
      groups[signature.return_type].append(signature)
    return dict(groups)

  def _populate_analytic_clauses(self, func, table_exprs, val_exprs):
    for arg in func.args:
      if arg.is_func:
        self._populate_analytic_clauses(arg, table_exprs, val_exprs)
    if func.is_analytic:
      if func.design == 'DETERMINISTIC_ORDER':
        unique_cols = shuffle(list(choice(table_exprs[0].unique_cols)))
      else:
        unique_cols = list()

      if func.partition_by_clause:
        if len(unique_cols) > 1:
          func.partition_by_clause.val_exprs.append(unique_cols.pop())
        elif len(val_exprs) > 1:
          shuffle(val_exprs)
          func.partition_by_clause.val_exprs.append(val_exprs[0])
        else:
          func.partition_by_clause = None

      if func.order_by_clause:
        if unique_cols:
          order_by_exprs = unique_cols
        elif val_exprs:
          order_by_exprs = val_exprs[:2]
        else:
          cols = list(table_exprs.cols)
          shuffle(cols)
          order_by_exprs = cols[:2]
        for order_by_expr in order_by_exprs:
          func.order_by_clause.exprs_to_order.append(
              (order_by_expr, choice([None, 'ASC', 'DESC', 'DESC'])))

  def _create_from_clause(self,
      table_exprs,
      table_alias_prefix,
      required_table_expr_col_type):
    from_clause = None
    table_count = self.profile.get_table_count()

    candidate_table_exprs = TableExprList(table_exprs)
    candidate_table_exprs.extend(candidate_table_exprs.collections)
    # If filter out table expressions with no scalar columns.
    candidate_table_exprs = TableExprList(
        table_expr for table_expr in candidate_table_exprs if table_expr.cols)

    if required_table_expr_col_type:
      candidate_table_exprs = \
          candidate_table_exprs.by_col_type[required_table_expr_col_type]
      if not candidate_table_exprs:
        raise Exception('No tables have the required column type')

    if table_count > 1:
      # If there are multiple multiple tables, the first chosen table must have joinable
      # columns.
      candidate_table_exprs = TableExprList(
          table_expr for table_expr in candidate_table_exprs
          if table_expr.joinable_cols_by_type)
      if not candidate_table_exprs:
        raise Exception('No tables have any joinable types')
      required_type = self.profile.choose_type(
          candidate_table_exprs.joinable_cols_by_type)
    else:
      required_type = None
    table_expr = self._create_table_expr(candidate_table_exprs,
        required_type=required_type)
    table_expr.alias = self.get_next_id()
    from_clause = FromClause(table_expr)

    for idx in xrange(1, table_count):
      join_clause = self._create_join_clause(from_clause, table_exprs)
      join_clause.table_expr.alias = self.get_next_id()
      from_clause.join_clauses.append(join_clause)

    return from_clause

  def _create_table_expr(self, table_exprs, required_type=None):
    if not table_exprs:
      raise Exception('At least one table_expr is required')
    if self.allow_more_nested_queries and self.profile.use_inline_view():
      return self._create_inline_view(table_exprs, required_type=required_type)
    if required_type:
      candidiate_table_exprs = TableExprList()
      for table_expr in table_exprs:
        for col in table_expr.cols:
          if issubclass(col.type, required_type):
            candidiate_table_exprs.append(table_expr)
            break
      if not candidiate_table_exprs:
        raise Exception('No tables have a column with the required data type')
      table_exprs = candidiate_table_exprs
    return deepcopy(self.profile.choose_table(table_exprs))

  def _create_inline_view(self, table_exprs, required_type=None):
    return InlineView(self.generate_statement(
        table_exprs, required_select_item_type=required_type,
        allow_with_clause=self.profile.use_nested_with()))

  def _create_join_clause(self, from_clause, table_exprs):
    join_type = self.profile.choose_join_type(JoinClause.JOINS_TYPES)
    use_lateral = self.profile.use_lateral_join()
    correlated_collections = from_clause.collections
    if use_lateral and correlated_collections:
      # TODO: Add correlated collections from ancestor queries
      if self.allow_more_nested_queries and self.profile.use_inline_view():
        collection = self._create_inline_view(correlated_collections)
      else:
        collection = deepcopy(choice(correlated_collections))
      join_clause = JoinClause(join_type, table_expr=collection)
      join_clause.is_lateral_join = True
      for _ in range(self.profile.get_num_boolean_exprs_for_lateral_join()):
        predicate = self._create_single_join_condition(collection, from_clause.table_exprs)
        if not predicate:
          break
        join_clause.boolean_expr = And.create_from_args(join_clause.boolean_expr,
            predicate) if join_clause.boolean_expr else predicate
      return join_clause
    else:
      candidate_table_exprs = TableExprList(table_exprs)
      candidate_table_exprs.extend(candidate_table_exprs.collections)
      if join_type == 'CROSS':
        table_expr = self._create_table_expr(candidate_table_exprs)
      else:
        available_join_expr_types = set(from_clause.table_exprs.joinable_cols_by_type) \
            & set(candidate_table_exprs.col_types)
        if not available_join_expr_types:
          raise Exception('No tables have any columns eligible for joining')
        join_expr_type = self.profile.choose_type(tuple(available_join_expr_types))
        table_expr = self._create_table_expr(
            candidate_table_exprs, required_type=join_expr_type)

      if join_type in ('LEFT ANTI', 'LEFT SEMI'):
        table_expr.is_visible = False

      join_clause = JoinClause(join_type, table_expr)

      if join_type == 'CROSS':
        return join_clause

      join_table_expr_candidates = from_clause.table_exprs.by_col_type[join_expr_type]
      if not join_table_expr_candidates:
        raise Exception('table_expr has no common joinable columns')
      join_clause.boolean_expr = self._create_relational_join_condition(
          table_expr, join_table_expr_candidates)
      return join_clause

  def _get_joinable_cols_in_common(self, table_expr, other_table_expr):
    common_col_types = set(table_expr.joinable_cols_by_type) \
        & set(other_table_expr.joinable_cols_by_type)
    return common_col_types

  _has_joinable_cols_in_common = _get_joinable_cols_in_common

  def _create_single_join_condition(self, cur_table_expr, available_table_exprs):
    candidates = []
    for possible_table_expr in available_table_exprs:
      if self._has_joinable_cols_in_common(cur_table_expr, possible_table_expr):
        candidates.append(possible_table_expr)

    if not candidates:
      return None

    target_table_expr = choice(candidates)
    common_col_types = set(cur_table_expr.joinable_cols_by_type) \
        & set(target_table_expr.joinable_cols_by_type)

    col_type = choice(list(common_col_types))
    left_col = choice(cur_table_expr.joinable_cols_by_type[col_type])
    right_col = choice(target_table_expr.joinable_cols_by_type[col_type])
    return Equals.create_from_args(left_col, right_col)

  def _create_relational_join_condition(self,
      right_table_expr, possible_left_table_exprs):
    assert possible_left_table_exprs

    table_exprs_by_col_types = defaultdict(list)
    candidiate_table_exprs = list()

    for left_table_expr in possible_left_table_exprs:
      joinable_col_types = self._get_joinable_cols_in_common(
          right_table_expr, left_table_expr)
      if joinable_col_types:
        for col_type in joinable_col_types:
          table_exprs_by_col_types[col_type].append(left_table_expr)
          candidiate_table_exprs.append(left_table_expr)
    if not table_exprs_by_col_types:
      raise Exception('Tables have no joinable columns in common')

    join_signatures = self._funcs_to_allowed_signatures(FUNCS)
    if self.profile.only_use_equality_join_predicates():
      join_signatures = [sig for sig in join_signatures if not
          self.profile.is_non_equality_join_predicate(sig.func)]
    else:
      join_signatures = self.profile.get_allowed_join_signatures(join_signatures)

    root_predicate, relational_predicates = self._create_boolean_func_tree(
      require_relational_func=True,
      relational_col_types=table_exprs_by_col_types.keys(),
      allowed_signatures=join_signatures)

    for predicate in relational_predicates:
      left_table_expr = choice(candidiate_table_exprs)
      joinable_col_types = self._get_joinable_cols_in_common(
          right_table_expr, left_table_expr)

      left_null_args = self._populate_null_args_list(
          predicate, tuple(joinable_col_types), arg_stop_idx=1)
      arg_idx, func = choice(left_null_args)
      arg_type = func.args[arg_idx].type
      func.args[arg_idx] = choice(left_table_expr.cols_by_type[arg_type])

      right_null_args = self._populate_null_args_list(
          predicate, tuple(joinable_col_types), arg_start_idx=1)
      arg_idx, func = choice(right_null_args)
      arg_type = func.args[arg_idx].type
      func.args[arg_idx] = choice(right_table_expr.cols_by_type[arg_type])

    table_exprs = TableExprList(possible_left_table_exprs)
    table_exprs.append(right_table_expr)
    self.populate_func_with_vals(root_predicate, table_exprs=table_exprs)
    return root_predicate

  def _populate_null_args_list(self,
      func,
      allowed_types,  # Only allow args of this type in the result list.
      arg_start_idx=0,  # Start at this arg index, only applies to "func".
      arg_stop_idx=None,  # Stop at this arg index, only applies to "func".
      null_args=None):  # Used internally for recursion.
    '''Walks through func and creates a list of function arguments that are null. The
       returned list contains tuples with values (arg index, parent func).
    '''
    if arg_stop_idx is None:
      arg_stop_idx = len(func.args)
    if arg_stop_idx < arg_start_idx:
      raise Exception("'arg_stop_idx' (%s) cannot be less than 'arg_start_idx' (%s)"
          % (arg_stop_idx, arg_start_idx))
    if null_args is None:
      null_args = list()
    for idx in xrange(arg_start_idx, arg_stop_idx):
      arg = func.args[idx]
      if arg.is_constant and issubclass(arg.type, allowed_types):
        assert arg.val is None
        null_args.append((idx, func))
      elif arg.is_func:
        self._populate_null_args_list(arg, allowed_types, null_args=null_args)
    return null_args

  def _create_where_clause(self,
      table_exprs,
      table_alias_prefix):
    predicate, _ = self._create_boolean_func_tree(allow_subquery=True)
    predicate = self.populate_func_with_vals(
        predicate,
        table_exprs=table_exprs,
        table_alias_prefix=table_alias_prefix)
    if predicate.contains_subquery and not table_exprs[0].alias:
      # TODO: Figure out if an alias is really needed.
      table_exprs[0].alias = self.get_next_id()
    return WhereClause(predicate)

  def _create_having_clause(self, table_exprs, basic_select_item_exprs):
    if self.profile.only_use_aggregates_in_having_clause():
      predicate = self._create_agg_func_tree(Boolean)
    else:
      predicate, _ = self._create_boolean_func_tree()
    predicate = self.populate_func_with_vals(
        predicate, val_exprs=basic_select_item_exprs)
    # IMPALA-1423
    # Make sure any cols used have a table identifier. As of this writing the only
    # single table FROM clauses don't use table aliases. Setting a table alias
    # automatically propagates as a column table identifier ("t1.col" instead of "col").
    for arg in predicate.iter_exprs():
      if isinstance(arg, ValExpr) and arg.is_col and not arg.owner.alias:
        # TODO: Figure out if an alias is really needed.
        arg.owner.alias = self.get_next_id()
    return HavingClause(predicate)

  def _enable_distinct_on_random_agg_items(self, agg_items):
    '''For each agg function, set it to be DISTINCT with some probability'''
    def _enable_distinct_helper(val_expr):
      if val_expr.is_agg:
        if self.profile.use_distinct_in_func():
          val_expr.distinct = True
      elif val_expr.is_func:
        for arg in val_expr.args:
          _enable_distinct_helper(arg)

    for item in agg_items:
      _enable_distinct_helper(item.val_expr)

  def _create_boolean_func_tree(self, require_relational_func=False,
                                relational_col_types=None, allow_subquery=False,
                                allowed_signatures=None):
    '''Creates a function that returns a Boolean. The returned value is a
       Tuple<Function, List<Function>> where the first function is the root of the tree
       and the list contains relational functions (though other relational functions
       may also be in the tree). The root function's arguments may be other functions
       depending on the query profile in use. The leaf arguments of the tree will all be
       set to a NULL literal.

       Args:
         require_relational_func: If True, at least one function in the tree will be a
         relational function such as Equals, LessThanOrEquals, etc.

         relational_col_types: Can be used to restrict the signature of the chosen
         relational function to the given types.

         allow_subquery: If True, a subquery may be added to the function tree.

         allowed_signatures: A list of allowable signatures to be used inside the function
         tree. This list is not definitive, functions outside this list may be added to
         the function tree if they are explicitly added by the QueryProfile.
    '''
    # To create a realistic expression, the nodes nearest the root of the tree will be
    # ANDs and ORs, their children will be relational functions, then other number of
    # functions requested by the profile will be added as children of the relational
    # functions. There is no requirement that all children of ANDs and ORs will be
    # relational functions so a wide variety of queries should be generated.

    # Determine the number of non-leaf nodes in the tree.
    if require_relational_func:
      # Just assume this is for a JOIN.
      func_count = self.profile.choose_join_condition_count()
      if func_count <= 0:
        raise Exception("Relational condition count must be greater than zero")
    else:
      func_count = self.profile.choose_nested_expr_count() + 1

    and_or_fill_ratio = self.profile.choose_conjunct_disjunct_fill_ratio()
    and_or_count = int(and_or_fill_ratio * func_count)
    if and_or_count == func_count and require_relational_func:
      # Leave a space for a relational function.
      and_or_count -= 1

    relational_fill_ratio = self.profile.choose_relational_func_fill_ratio()
    relational_count = int(relational_fill_ratio * (func_count - and_or_count))
    if relational_count > and_or_count + 1:
      # Reduce the AND/OR count to the number of boolean spaces that are guaranteed to
      # exist.
      relational_count = and_or_count + 1
    elif relational_count == 0 and require_relational_func:
      relational_count = 1

    root_func = None   # The root of the function tree.
    relational_funcs = list()  # These will be in the tree somewhere.

    # After a root function is chosen, each additional function will be chosen based on
    # the types required by the leaves (existing unused function args). For example if
    # the current tree is
    #      And
    #     /   \
    #   And  Equals(Int, Int)
    # Then the next function chosen must either return a Boolean or an Int. The dict
    # below tracks which arg types are available.
    #
    # In the example above, assuming Equals was chosen as a relational function, the
    # children of Equals must preserve the Int arg if the caller requested Ints using
    # 'relational_col_types'.
    null_args_by_type = defaultdict(list)

    if not allowed_signatures:
      allowed_signatures = self._funcs_to_allowed_signatures(FUNCS)

    if not relational_col_types:
      relational_col_types = tuple()

    for _ in xrange(func_count):
      is_relational = False

      if and_or_count > 0:
        and_or_count -= 1
        func_class = self.profile.choose_conjunct_disjunct()
        signature = self.profile.choose_func_signature(func_class.signatures())
      elif relational_count > 0:
        relational_count -= 1
        is_relational = True
        relational_signatures = self._find_matching_signatures(
          allowed_signatures, accepts_only=tuple(relational_col_types), returns=Boolean)
        signature = self.profile.choose_relational_func_signature(relational_signatures)
      else:
        if not root_func or Boolean in null_args_by_type:
          # Prefer to replace Boolean leaves to get a more realistic expression.
          return_type = Boolean
        else:
          return_type = choice(null_args_by_type.keys())
        # Rather than track if this is a child of a relational function, in which case
        # the arg type needs to be preserved, just always assume that this is a child of
        # a relational function.
        accepts = return_type if return_type in relational_col_types else None
        matching_signatures = self._find_matching_signatures(
            allowed_signatures, returns=return_type, accepts=accepts,
            allow_subquery=allow_subquery)
        if matching_signatures:
          signature = self.profile.choose_func_signature(matching_signatures)
        else:
          continue

      func = signature.func(signature)
      if root_func:
        null_args = null_args_by_type[signature.return_type]
        shuffle(null_args)
        parent_func, parent_arg_idx = null_args.pop()
        parent_func.args[parent_arg_idx] = func
        if not null_args:
          del null_args_by_type[signature.return_type]
      else:
        root_func = func

      if is_relational:
        relational_funcs.append(func)

      for idx, arg in enumerate(func.args):
        if signature.args[idx].is_subquery:
          continue
        null_args = null_args_by_type[arg.type]
        null_args.append((func, idx))
      if not null_args_by_type:
        # Subqueries might not provide leaves, so stopping early may be necessary. Ex
        # WHERE EXISTS(...).
        assert func.contains_subquery
        break

    return root_func, relational_funcs


def generate_queries_for_manual_inspection():
  '''Generate some queries for manual inspection. The query won't run anywhere because the
     tables used are fake. To make real queries, we'd need to connect to a database and
     read the table metadata and such.
  '''
  from model_translator import SqlWriter

  NUM_TABLES = 5
  NUM_COLS_EACH_TYPE = 10
  NUM_QUERIES = 3000

  tables = list()
  data_types = list(TYPES)
  data_types.remove(Float)
  for table_idx in xrange(NUM_TABLES):
    table = Table('table_%s' % table_idx)
    tables.append(table)
    cols = table.cols
    col_idx = 0
    for _ in xrange(NUM_COLS_EACH_TYPE):
      for col_type in data_types:
        col = Column(table, '%s_col_%s' % (col_type.__name__.lower(), col_idx), col_type)
        cols.append(col)
        col_idx += 1
    table.cols = cols

  query_profile = DefaultProfile()
  query_generator = QueryGenerator(query_profile)
  sql_writer = SqlWriter.create(dialect='IMPALA')
  ref_writer = SqlWriter.create(dialect='POSTGRESQL',
      nulls_order_asc=query_profile.nulls_order_asc())
  for _ in xrange(NUM_QUERIES):
    query = query_generator.generate_statement(tables)
    print("Test db")
    print(sql_writer.write_query(query) + '\n')
    print("Ref db")
    print(ref_writer.write_query(query) + '\n')


if __name__ == '__main__':
  generate_queries_for_manual_inspection()