text: # 
  #  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 pandas.api.types import is_list_like
  import itertools
  import heapq
  from functools import total_ordering
  
  @total_ordering
  class Query:
      def __init__(self, qid, idx, query, parameters, param_idx=0):
          self.data_idx = idx
          self.qid = qid
          self.query = query
          self.parameters = parameters
          self.param_idx = param_idx # used to order queries
  
      def __iter__(self):
          yield self.data_idx
          yield self.qid
          yield self.query
          yield self.parameters
      
      def info(self):
          return {
              'query_idx': self.data_idx,
              'qid': self.qid,
              'query': self.query,
              'query_parameters': self.parameters,
          }
      
      def __str__(self):
          return f"{self.query} {self.parameters} {self.data_idx}"
      
      def __lt__(self, b):        
          return self.data_idx < b.data_idx
      
      def __eq__(self,b):
          return self.data_idx == b.data_idx
          
  class QueryGenerator:
      name = "generic_QG"
      END = 1000000000 
      
      def __init__(self):
          self.prepared = False
          
      def connectDataGenerator(self, data_generator):
          self.data_generator = data_generator
  
      def getName(self):
          return self.name
              
      def getID(self):
          return self.name
      
  class DataGeneratorSeq:
      def __init__(self, data_generator=None, length=None, by=None):
          self.data_generator = data_generator
          self.length = length
          self.by = by
      
      def genSeq(self):
          if self.length is not None:
              n = len(self.data_generator)
              assert(n > 1)
              by = (n-1) / self.length
          else:
              n = QueryGenerator.END
              by = self.by
          
          i = 1
          while i*by < len(self.data_generator):
              yield int(i*by)
              i += 1
          if i*by != n-1:
              yield QueryGenerator.END
              
              
          
  class ConfigQueryGenerator(QueryGenerator):
      name = 'config_QG'
  
      def __init__(self, queries, indices=[QueryGenerator.END], parameters=None):
          self.queries = queries
          self.indices = indices
          self.query_parameters = parameters
      
      def genQueries(self):
          qid = 0
          if isinstance(self.queries, list):
              queries = self.queries
          else:
              queries = [self.queries]
              
          if is_list_like(self.query_parameters) and not isinstance(self.query_parameters, dict):
              query_parameters = self.query_parameters
          else:
              query_parameters = [self.query_parameters]
              
          if isinstance(self.indices, DataGeneratorSeq):
              indices = self.indices.genSeq()
          else:
              indices = self.indices
          
          for q, idx, params in itertools.product(queries, indices, query_parameters):
              yield Query(qid, idx, q, params)
              qid += 1
      
      def connectDataGenerator(self, data_generator):
          self.data_generator = data_generator
          if isinstance(self.indices, DataGeneratorSeq):
              self.indices.data_generator = data_generator
          
  class ChainQueryGenerators(QueryGenerator):
      name = 'chained_QG'
  
      def __init__(self, generators=[]):
          super().__init__()
          self.query_generators = generators
          self.heap = []
          
      def genQueries(self):
          # reassign the qid's and ensure queries are ordered by idx
          iters = [qg.genQueries() for qg in self.query_generators]
              
          for i, it in enumerate(iters):
              q = next(it, None)
              if q is not None:
                  self.heap.append((q, i))
              
          heapq.heapify(self.heap)
          qid = 0
          while self.heap:
              q, i = heapq.heappop(self.heap)
              q.qid = qid
              qid += 1
              yield q
              
              nextq = next(iters[i], None)
              if nextq is not None:
                  heapq.heappush(self.heap, (nextq, i))
              
      def connectDataGenerator(self, data_generator):
          for qg in self.query_generators:
              qg.connectDataGenerator(data_generator)
          
  
  class TopKQueryGenerator(ConfigQueryGenerator):
      def __init__(self, k_values, num_queries, data_generator=None, **kwargs):
          super().__init__(queries="topk", 
                           indices=DataGeneratorSeq(data_generator, length=num_queries),
                           parameters=k_values)
          