Spaces:
Running
Running
| # 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 codecs | |
| import decimal | |
| from functools import partial | |
| import itertools | |
| import sys | |
| import unicodedata | |
| import numpy as np | |
| import pyarrow as pa | |
| KILOBYTE = 1 << 10 | |
| MEGABYTE = KILOBYTE * KILOBYTE | |
| DEFAULT_NONE_PROB = 0.3 | |
| def _multiplicate_sequence(base, target_size): | |
| q, r = divmod(target_size, len(base)) | |
| return [base] * q + [base[:r]] | |
| def get_random_bytes(n, seed=42): | |
| """ | |
| Generate a random bytes object of size *n*. | |
| Note the result might be compressible. | |
| """ | |
| rnd = np.random.RandomState(seed) | |
| # Computing a huge random bytestring can be costly, so we get at most | |
| # 100KB and duplicate the result as needed | |
| base_size = 100003 | |
| q, r = divmod(n, base_size) | |
| if q == 0: | |
| result = rnd.bytes(r) | |
| else: | |
| base = rnd.bytes(base_size) | |
| result = b''.join(_multiplicate_sequence(base, n)) | |
| assert len(result) == n | |
| return result | |
| def get_random_ascii(n, seed=42): | |
| """ | |
| Get a random ASCII-only unicode string of size *n*. | |
| """ | |
| arr = np.frombuffer(get_random_bytes(n, seed=seed), dtype=np.int8) & 0x7f | |
| result, _ = codecs.ascii_decode(arr) | |
| assert isinstance(result, str) | |
| assert len(result) == n | |
| return result | |
| def _random_unicode_letters(n, seed=42): | |
| """ | |
| Generate a string of random unicode letters (slow). | |
| """ | |
| def _get_more_candidates(): | |
| return rnd.randint(0, sys.maxunicode, size=n).tolist() | |
| rnd = np.random.RandomState(seed) | |
| out = [] | |
| candidates = [] | |
| while len(out) < n: | |
| if not candidates: | |
| candidates = _get_more_candidates() | |
| ch = chr(candidates.pop()) | |
| # XXX Do we actually care that the code points are valid? | |
| if unicodedata.category(ch)[0] == 'L': | |
| out.append(ch) | |
| return out | |
| _1024_random_unicode_letters = _random_unicode_letters(1024) | |
| def get_random_unicode(n, seed=42): | |
| """ | |
| Get a random non-ASCII unicode string of size *n*. | |
| """ | |
| indices = np.frombuffer(get_random_bytes(n * 2, seed=seed), | |
| dtype=np.int16) & 1023 | |
| unicode_arr = np.array(_1024_random_unicode_letters)[indices] | |
| result = ''.join(unicode_arr.tolist()) | |
| assert len(result) == n, (len(result), len(unicode_arr)) | |
| return result | |
| class BuiltinsGenerator(object): | |
| def __init__(self, seed=42): | |
| self.rnd = np.random.RandomState(seed) | |
| def sprinkle(self, lst, prob, value): | |
| """ | |
| Sprinkle *value* entries in list *lst* with likelihood *prob*. | |
| """ | |
| for i, p in enumerate(self.rnd.random_sample(size=len(lst))): | |
| if p < prob: | |
| lst[i] = value | |
| def sprinkle_nones(self, lst, prob): | |
| """ | |
| Sprinkle None entries in list *lst* with likelihood *prob*. | |
| """ | |
| self.sprinkle(lst, prob, None) | |
| def generate_int_list(self, n, none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of Python ints with *none_prob* probability of | |
| an entry being None. | |
| """ | |
| data = list(range(n)) | |
| self.sprinkle_nones(data, none_prob) | |
| return data | |
| def generate_float_list(self, n, none_prob=DEFAULT_NONE_PROB, | |
| use_nan=False): | |
| """ | |
| Generate a list of Python floats with *none_prob* probability of | |
| an entry being None (or NaN if *use_nan* is true). | |
| """ | |
| # Make sure we get Python floats, not np.float64 | |
| data = list(map(float, self.rnd.uniform(0.0, 1.0, n))) | |
| assert len(data) == n | |
| self.sprinkle(data, none_prob, value=float('nan') if use_nan else None) | |
| return data | |
| def generate_bool_list(self, n, none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of Python bools with *none_prob* probability of | |
| an entry being None. | |
| """ | |
| # Make sure we get Python bools, not np.bool_ | |
| data = [bool(x >= 0.5) for x in self.rnd.uniform(0.0, 1.0, n)] | |
| assert len(data) == n | |
| self.sprinkle_nones(data, none_prob) | |
| return data | |
| def generate_decimal_list(self, n, none_prob=DEFAULT_NONE_PROB, | |
| use_nan=False): | |
| """ | |
| Generate a list of Python Decimals with *none_prob* probability of | |
| an entry being None (or NaN if *use_nan* is true). | |
| """ | |
| data = [decimal.Decimal('%.9f' % f) | |
| for f in self.rnd.uniform(0.0, 1.0, n)] | |
| assert len(data) == n | |
| self.sprinkle(data, none_prob, | |
| value=decimal.Decimal('nan') if use_nan else None) | |
| return data | |
| def generate_object_list(self, n, none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of generic Python objects with *none_prob* | |
| probability of an entry being None. | |
| """ | |
| data = [object() for i in range(n)] | |
| self.sprinkle_nones(data, none_prob) | |
| return data | |
| def _generate_varying_sequences(self, random_factory, n, min_size, | |
| max_size, none_prob): | |
| """ | |
| Generate a list of *n* sequences of varying size between *min_size* | |
| and *max_size*, with *none_prob* probability of an entry being None. | |
| The base material for each sequence is obtained by calling | |
| `random_factory(<some size>)` | |
| """ | |
| base_size = 10000 | |
| base = random_factory(base_size + max_size) | |
| data = [] | |
| for i in range(n): | |
| off = self.rnd.randint(base_size) | |
| if min_size == max_size: | |
| size = min_size | |
| else: | |
| size = self.rnd.randint(min_size, max_size + 1) | |
| data.append(base[off:off + size]) | |
| self.sprinkle_nones(data, none_prob) | |
| assert len(data) == n | |
| return data | |
| def generate_fixed_binary_list(self, n, size, none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of bytestrings with a fixed *size*. | |
| """ | |
| return self._generate_varying_sequences(get_random_bytes, n, | |
| size, size, none_prob) | |
| def generate_varying_binary_list(self, n, min_size, max_size, | |
| none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of bytestrings with a random size between | |
| *min_size* and *max_size*. | |
| """ | |
| return self._generate_varying_sequences(get_random_bytes, n, | |
| min_size, max_size, none_prob) | |
| def generate_ascii_string_list(self, n, min_size, max_size, | |
| none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of ASCII strings with a random size between | |
| *min_size* and *max_size*. | |
| """ | |
| return self._generate_varying_sequences(get_random_ascii, n, | |
| min_size, max_size, none_prob) | |
| def generate_unicode_string_list(self, n, min_size, max_size, | |
| none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of unicode strings with a random size between | |
| *min_size* and *max_size*. | |
| """ | |
| return self._generate_varying_sequences(get_random_unicode, n, | |
| min_size, max_size, none_prob) | |
| def generate_int_list_list(self, n, min_size, max_size, | |
| none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of lists of Python ints with a random size between | |
| *min_size* and *max_size*. | |
| """ | |
| return self._generate_varying_sequences( | |
| partial(self.generate_int_list, none_prob=none_prob), | |
| n, min_size, max_size, none_prob) | |
| def generate_tuple_list(self, n, none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of tuples with random values. | |
| Each tuple has the form `(int value, float value, bool value)` | |
| """ | |
| dicts = self.generate_dict_list(n, none_prob=none_prob) | |
| tuples = [(d.get('u'), d.get('v'), d.get('w')) | |
| if d is not None else None | |
| for d in dicts] | |
| assert len(tuples) == n | |
| return tuples | |
| def generate_dict_list(self, n, none_prob=DEFAULT_NONE_PROB): | |
| """ | |
| Generate a list of dicts with random values. | |
| Each dict has the form | |
| `{'u': int value, 'v': float value, 'w': bool value}` | |
| """ | |
| ints = self.generate_int_list(n, none_prob=none_prob) | |
| floats = self.generate_float_list(n, none_prob=none_prob) | |
| bools = self.generate_bool_list(n, none_prob=none_prob) | |
| dicts = [] | |
| # Keep half the Nones, omit the other half | |
| keep_nones = itertools.cycle([True, False]) | |
| for u, v, w in zip(ints, floats, bools): | |
| d = {} | |
| if u is not None or next(keep_nones): | |
| d['u'] = u | |
| if v is not None or next(keep_nones): | |
| d['v'] = v | |
| if w is not None or next(keep_nones): | |
| d['w'] = w | |
| dicts.append(d) | |
| self.sprinkle_nones(dicts, none_prob) | |
| assert len(dicts) == n | |
| return dicts | |
| def get_type_and_builtins(self, n, type_name): | |
| """ | |
| Return a `(arrow type, list)` tuple where the arrow type | |
| corresponds to the given logical *type_name*, and the list | |
| is a list of *n* random-generated Python objects compatible | |
| with the arrow type. | |
| """ | |
| size = None | |
| if type_name in ('bool', 'decimal', 'ascii', 'unicode', 'int64 list'): | |
| kind = type_name | |
| elif type_name.startswith(('int', 'uint')): | |
| kind = 'int' | |
| elif type_name.startswith('float'): | |
| kind = 'float' | |
| elif type_name.startswith('struct'): | |
| kind = 'struct' | |
| elif type_name == 'binary': | |
| kind = 'varying binary' | |
| elif type_name.startswith('binary'): | |
| kind = 'fixed binary' | |
| size = int(type_name[6:]) | |
| assert size > 0 | |
| else: | |
| raise ValueError("unrecognized type %r" % (type_name,)) | |
| if kind in ('int', 'float'): | |
| ty = getattr(pa, type_name)() | |
| elif kind == 'bool': | |
| ty = pa.bool_() | |
| elif kind == 'decimal': | |
| ty = pa.decimal128(9, 9) | |
| elif kind == 'fixed binary': | |
| ty = pa.binary(size) | |
| elif kind == 'varying binary': | |
| ty = pa.binary() | |
| elif kind in ('ascii', 'unicode'): | |
| ty = pa.string() | |
| elif kind == 'int64 list': | |
| ty = pa.list_(pa.int64()) | |
| elif kind == 'struct': | |
| ty = pa.struct([pa.field('u', pa.int64()), | |
| pa.field('v', pa.float64()), | |
| pa.field('w', pa.bool_())]) | |
| factories = { | |
| 'int': self.generate_int_list, | |
| 'float': self.generate_float_list, | |
| 'bool': self.generate_bool_list, | |
| 'decimal': self.generate_decimal_list, | |
| 'fixed binary': partial(self.generate_fixed_binary_list, | |
| size=size), | |
| 'varying binary': partial(self.generate_varying_binary_list, | |
| min_size=3, max_size=40), | |
| 'ascii': partial(self.generate_ascii_string_list, | |
| min_size=3, max_size=40), | |
| 'unicode': partial(self.generate_unicode_string_list, | |
| min_size=3, max_size=40), | |
| 'int64 list': partial(self.generate_int_list_list, | |
| min_size=0, max_size=20), | |
| 'struct': self.generate_dict_list, | |
| 'struct from tuples': self.generate_tuple_list, | |
| } | |
| data = factories[kind](n) | |
| return ty, data | |