| # |
| # 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. |
| # |
| |
| import random |
| import scipy |
| import scipy.stats |
| import numpy as np |
| import itertools |
| import pandas as pd |
| import os |
| |
| |
| class DataGenerator: |
| rng = None |
| data = None # iteratable providing the data stream |
| queries = None # iterable providing (index, query) tuples where answer is the desired answer after processing index data items |
| name = "GenericData" |
| |
| MATERIALIZE = False |
| |
| def __init__(self, *args, **kwargs): |
| self.seed = kwargs.get('seed', 0) |
| self.current_seed = None |
| self.rng = random.Random(self.seed) |
| self.prepared = False |
| |
| def prepareData(self, size=None, *args, **kwargs): |
| pass |
| |
| def genData(self): |
| raise Exception |
| |
| def getName(self): |
| return self.name |
| |
| def getID(self): |
| return f"{self.getName()}_{str(self.seed)}" |
| |
| def __len__(self): |
| """ |
| Return the length of the data stream |
| """ |
| raise Exception |
| |
| # These should be filled out in the future to serialize workloads |
| # This would allow tests to be run/written in another language but share the same workload/answers/evaluation |
| # |
| # def writeToCache(self): |
| # pass |
| |
| # def prepareFromCache(self): |
| # pass |
| |
| # def getCached(self): |
| # return self |
| |
| def prepareForPickle(self): |
| """ |
| This should remove any large objects |
| """ |
| pass |
| |
| def reset(self, seed): |
| """ |
| Update the seed on the data generator |
| Typically this generates a new sequence for synthetic data and |
| permutes the input order for data from a file. |
| """ |
| self.seed = seed |
| self.rng = random.Random(self.seed) |
| |
| |
| class Workload: |
| name = "BaseWorkload" |
| data_generator = None |
| query_generator = None |
| |
| def __init__(self, data_generator = None, query_generator=None, **kwargs): |
| self.prepared = False |
| self.data_generator = data_generator |
| self.query_generator = query_generator |
| |
| def getName(self, data=True, query=False): |
| data_name = self.data_generator.getName() |
| query_name = self.query_generator.getName() |
| if data and query: |
| return f"Data_{data_name}_Queries_{query_name}" |
| elif data: |
| return data_name |
| elif query: |
| return query_name |
| else: |
| return self.name |
| |
| def getID(self): |
| data_name = self.data_generator.getID() |
| query_name = self.query_generator.getID() |
| return f"Data_{data_name}_Queries_{query_name}" |
| |
| def genData(self): |
| return self.data_generator.genData() |
| |
| def __len__(self): |
| return len(self.data_generator) |
| |
| def genQueries(self): |
| return self.query_generator.genQueries() |
| |
| def prepareData(self): |
| if self.data_generator is None: |
| raise Exception |
| |
| self.data_generator.prepareData() |
| |
| def prepare(self): |
| if self.prepared: |
| return |
| |
| self.prepareData() |
| self.query_generator.connectDataGenerator(self.data_generator) |
| |
| def reset(self, seed=None): |
| self.data_generator.reset(seed=seed) |
| self.prepared = False |
| |
| def info(self): |
| info = {'workload': self.getName(), |
| 'data_seed': self.data_generator.seed, |
| } |
| return info |
| |
| def prepareForPickle(self): |
| self.data_generator.prepareForPickle() |
| |
| |
| ########################################################################################## |
| |
| from random import randint |
| |
| class DistributionDataGenerator(DataGenerator): |
| """ |
| takes a scipy.stats distribution and generates data using it |
| In particular, the data generator is assigned its own rng with a specified seed |
| """ |
| def __init__(self, length, distribution, name, seed=0, params={}, dim=1, *args, **kwargs): |
| super().__init__(**kwargs) # need to cooperate with other classes for multiple inheritance |
| self.size = length |
| self.distribution = distribution |
| self.seed = seed |
| self.params = params |
| self.dim = dim |
| self.name = name |
| self.chunksize = 1000 |
| |
| def __len__(self): |
| return self.size |
| |
| def prepareData(self): |
| pass |
| |
| def genDistributionSequence(self, dim=1): |
| d = self.distribution # (**self.params) |
| d.random_state = np.random.default_rng(seed=self.seed) |
| chunksize = self.chunksize |
| while True: |
| self.buffer = d.rvs(chunksize*dim) |
| if dim==1: |
| for x in self.buffer: |
| yield x |
| else: |
| for i in range(chunksize): |
| yield self.buffer[i:(i+dim)] |
| |
| self.buffer = None |
| |
| def genData(self): |
| return itertools.islice(self.genDistributionSequence(dim=self.dim), self.size) |
| |
| def prepareForPickle(self): |
| self.buffer = None |
| |
| class FileDataGenerator(DataGenerator): |
| """ |
| Takes a file of tokens and plays it back as a stream. |
| """ |
| def __init__(self, filename=None, **kwargs): |
| super().__init__(**kwargs) |
| self.filename = filename |
| self.data = None |
| self.name = os.path.basename(filename) |
| |
| def prepareData(self, **kwargs): |
| if self.prepared: |
| return |
| |
| self.data = [] |
| with open(self.filename, 'r') as f: |
| for line in f: |
| tokens = line.rstrip().split() |
| self.data.extend(tokens) |
| |
| self.data = pd.Series(self.data) |
| self.prepared = True |
| print(f"finished reading {self.filename}") |
| |
| def __len__(self): |
| if self.data is None: |
| self.prepareData() |
| |
| return len(self.data) |
| |
| def genData(self): |
| # don't support permutations yet. |
| return self.data.sample(frac=1, random_state=self.seed) |
| |
| def prepareForPickle(self): |
| self.data = None |
| self.prepared = False |
| |
| class CSVFileDataGenerator(FileDataGenerator): |
| """ |
| Take a csv file and turns the specified column into a permuted stream |
| """ |
| |
| MATERIALIZE = True |
| |
| def __init__(self, filename=None, column=None, filetype='csv', *args, **kwargs): |
| super().__init__(**kwargs) |
| self.filename = filename |
| self.filetype = filetype |
| self.column = column |
| self.name = filename |
| self.data = None |
| self.cache_file = None |
| |
| def prepareData(self, **kwargs): |
| |
| if self.prepared: |
| return |
| |
| try: |
| df = pd.read_csv(self.filename) #, self.filetype) |
| self.data = df[self.column] |
| except Exception as e: |
| print(e) |
| |
| self.prepared = True |
| |
| |
| # Given a distribution d, returns |
| # a random binary vector with X ~ d non-zero entries |
| class BinaryVecDataGenerator(DataGenerator): |
| """ |
| takes a scipy.stats distribution |
| and assigns it its own rng with a specified seed |
| """ |
| def __init__(self, length, distribution, name, seed=0, params={}, dim=1, *args, **kwargs): |
| super().__init__(**kwargs) # need to cooperate with other classes for multiple inheritance |
| self.size = length |
| self.distribution = distribution |
| self.seed = seed |
| self.params = params |
| self.dim = dim |
| self.name = name |
| |
| def __len__(self): |
| return self.size |
| |
| def prepareData(self): |
| pass |
| |
| def genData(self): |
| d = self.distribution # (**self.params) |
| d.random_state = np.random.default_rng(seed=self.seed) |
| np_rng = np.random.RandomState(seed=self.seed) |
| for i in range(self.size): |
| x = d.rvs(1)[0] |
| while x > self.dim: |
| x = d.rvs(1)[0] |
| pi = np_rng.permutation(self.dim) |
| idx = pi[:x] |
| z = np.zeros(self.dim) + 0.1 |
| z[idx] = 1.0 |
| yield z |
| |
