chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,276 @@
|
||||
# ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
|
||||
# ┃ ██████ ██████ ██████ █ █ █ █ █ █▄ ▀███ █ ┃
|
||||
# ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█ ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄ ▀█ █ ▀▀▀▀▀ ┃
|
||||
# ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄ █ ▄▄▄▄▄ ┃
|
||||
# ┃ █ ██████ █ ▀█▄ █ ██████ █ ███▌▐███ ███████▄ █ ┃
|
||||
# ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫
|
||||
# ┃ Copyright (c) 2017, the Perspective Authors. ┃
|
||||
# ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃
|
||||
# ┃ This file is part of the Perspective library, distributed under the terms ┃
|
||||
# ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃
|
||||
# ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
|
||||
|
||||
from datetime import datetime, date, timezone
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
from pytest import fixture
|
||||
from random import random, randint, choice
|
||||
from faker import Faker
|
||||
|
||||
|
||||
# Our tests construct naive datetimes everywhere
|
||||
# so setting it here is an easy way to fix it globally.
|
||||
import os
|
||||
|
||||
# Perspective used to support datetime.date and datetime.datetime
|
||||
# as Table() constructor arguments, but now we forward the parameters
|
||||
# directly to JSON.loads. So to make sure the tests dont need to be
|
||||
# so utterly transmogrified, we have this little hack :)
|
||||
import json
|
||||
|
||||
|
||||
os.environ["TZ"] = "UTC"
|
||||
|
||||
|
||||
def new_encoder(self, obj):
|
||||
if isinstance(obj, datetime):
|
||||
return str(obj)
|
||||
elif isinstance(obj, date):
|
||||
return str(obj)
|
||||
else:
|
||||
return old(self, obj)
|
||||
|
||||
|
||||
old = json.JSONEncoder.default
|
||||
json.JSONEncoder.default = new_encoder
|
||||
|
||||
fake = Faker()
|
||||
|
||||
|
||||
def _make_date_time_index(size, time_unit):
|
||||
return pd.date_range("2000-01-01", periods=size, freq=time_unit)
|
||||
|
||||
|
||||
def _make_period_index(size, time_unit):
|
||||
return pd.period_range(start="2000", periods=size, freq=time_unit)
|
||||
|
||||
|
||||
def _make_dataframe(index, size=10):
|
||||
"""Create a new random dataframe of `size` and with a DateTimeIndex of
|
||||
frequency `time_unit`.
|
||||
"""
|
||||
return pd.DataFrame(
|
||||
index=index,
|
||||
data={
|
||||
"a": np.random.rand(size),
|
||||
"b": np.random.rand(size),
|
||||
"c": np.random.rand(size),
|
||||
"d": np.random.rand(size),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class Util:
|
||||
@staticmethod
|
||||
def make_arrow(names, data, types=None, legacy=False):
|
||||
"""Create an arrow binary that can be loaded and manipulated from memory.
|
||||
|
||||
Args:
|
||||
names (list): a list of str column names
|
||||
data (list): a list of lists containing data for each column
|
||||
types (list): an optional list of `pyarrow.type` function references.
|
||||
Types will be inferred if not provided.
|
||||
legacy (bool): if True, use legacy IPC format (pre-pyarrow 0.15). Defaults to False.
|
||||
|
||||
Returns:
|
||||
bytes : a bytes object containing the arrow-serialized output.
|
||||
"""
|
||||
stream = pa.BufferOutputStream()
|
||||
arrays = []
|
||||
|
||||
for idx, column in enumerate(data):
|
||||
# only apply types if array is present
|
||||
kwargs = {}
|
||||
if types:
|
||||
kwargs["type"] = types[idx]
|
||||
arrays.append(pa.array(column, **kwargs))
|
||||
|
||||
batch = pa.RecordBatch.from_arrays(arrays, names)
|
||||
table = pa.Table.from_batches([batch])
|
||||
writer = pa.RecordBatchStreamWriter(
|
||||
stream, table.schema
|
||||
)
|
||||
|
||||
writer.write_table(table)
|
||||
writer.close()
|
||||
return stream.getvalue().to_pybytes()
|
||||
|
||||
@staticmethod
|
||||
def make_arrow_from_pandas(df, schema=None, legacy=False):
|
||||
"""Create a pyarrow Table from a Pandas dataframe.
|
||||
|
||||
Args:
|
||||
df (:obj:`pandas.DataFrame`)
|
||||
schema (:obj:`pyarrow.Schema`)
|
||||
legacy (bool): unused; retained for backwards compatibility.
|
||||
|
||||
Returns:
|
||||
pyarrow.Table
|
||||
"""
|
||||
return pa.Table.from_pandas(df, schema=schema)
|
||||
|
||||
@staticmethod
|
||||
def make_dictionary_arrow(names, data, types=None, legacy=False):
|
||||
"""Create an arrow binary that can be loaded and manipulated from memory, with
|
||||
each column being a dictionary array of `str` values and `int` indices.
|
||||
|
||||
Args:
|
||||
names (list): a list of str column names
|
||||
data (list:tuple): a list of tuples, the first value being a list of indices,
|
||||
and the second value being a list of values.
|
||||
types (list:list:pyarrow.func): a list of lists, containing the indices type and
|
||||
dictionary value type for each array.
|
||||
legacy (bool): if True, use legacy IPC format (pre-pyarrow 0.15). Defaults to False.
|
||||
|
||||
Returns:
|
||||
bytes : a bytes object containing the arrow-serialized output.
|
||||
"""
|
||||
stream = pa.BufferOutputStream()
|
||||
|
||||
arrays = []
|
||||
for idx, column in enumerate(data):
|
||||
indice_type = pa.int64()
|
||||
value_type = pa.string()
|
||||
|
||||
if types is not None:
|
||||
indice_type = types[idx][0]
|
||||
value_type = types[idx][1]
|
||||
|
||||
indices = pa.array(column[0], type=indice_type)
|
||||
values = pa.array(column[1], type=value_type)
|
||||
parray = pa.DictionaryArray.from_arrays(indices, values)
|
||||
arrays.append(parray)
|
||||
|
||||
batch = pa.RecordBatch.from_arrays(arrays, names)
|
||||
table = pa.Table.from_batches([batch])
|
||||
writer = pa.RecordBatchStreamWriter(
|
||||
stream, table.schema
|
||||
)
|
||||
|
||||
writer.write_table(table)
|
||||
writer.close()
|
||||
return stream.getvalue().to_pybytes()
|
||||
|
||||
@staticmethod
|
||||
def to_timestamp(obj):
|
||||
"""Return an integer timestamp based on a date/datetime object.
|
||||
|
||||
A `date` is a timezone-agnostic calendar day and serializes to epoch
|
||||
ms at *UTC* midnight regardless of the host process timezone, so it
|
||||
is converted here in UTC. A naive `datetime` is a wall-clock reading
|
||||
in the process-local timezone (callers express expected instants
|
||||
this way, e.g. `aware.astimezone(TZ).replace(tzinfo=None)`), so it is
|
||||
converted via local `timestamp()`.
|
||||
"""
|
||||
classname = obj.__class__.__name__
|
||||
if classname == "date":
|
||||
return int(
|
||||
datetime(
|
||||
obj.year, obj.month, obj.day, tzinfo=timezone.utc
|
||||
).timestamp()
|
||||
* 1000
|
||||
)
|
||||
elif classname == "datetime":
|
||||
return int(obj.timestamp() * 1000)
|
||||
else:
|
||||
return -1
|
||||
|
||||
@staticmethod
|
||||
def make_dataframe(size=10, freq="D"):
|
||||
index = _make_date_time_index(size, freq)
|
||||
return _make_dataframe(index, size)
|
||||
|
||||
@staticmethod
|
||||
def make_period_dataframe(size=10):
|
||||
index = _make_period_index(size, "M")
|
||||
return _make_dataframe(index, size)
|
||||
|
||||
@staticmethod
|
||||
def make_series(size=10, freq="D"):
|
||||
index = _make_date_time_index(size, freq)
|
||||
return pd.Series(data=np.random.rand(size), index=index)
|
||||
|
||||
|
||||
class Sentinel(object):
|
||||
"""Generic sentinel class for testing side-effectful code in Python 2 and
|
||||
3.
|
||||
"""
|
||||
|
||||
def __init__(self, value):
|
||||
self.value = value
|
||||
|
||||
def get(self):
|
||||
return self.value
|
||||
|
||||
def set(self, new_value):
|
||||
self.value = new_value
|
||||
|
||||
|
||||
@fixture()
|
||||
def sentinel():
|
||||
"""Pass `sentinel` into a test and call it with `value` to create a new
|
||||
instance of the Sentinel class.
|
||||
|
||||
Example:
|
||||
>>> def test_with_sentinel(self, sentinel):
|
||||
>>> s = sentinel(True)
|
||||
>>> s.set(False)
|
||||
>>> s.get() # returns False
|
||||
"""
|
||||
|
||||
def _sentinel(value):
|
||||
return Sentinel(value)
|
||||
|
||||
return _sentinel
|
||||
|
||||
|
||||
@fixture
|
||||
def util():
|
||||
"""Pass the `Util` class in to a test."""
|
||||
return Util
|
||||
|
||||
|
||||
@fixture
|
||||
def superstore(count=100):
|
||||
data = []
|
||||
for id in range(count):
|
||||
dat = {}
|
||||
dat["Row ID"] = id
|
||||
dat["Order ID"] = "{}-{}".format(fake.ein(), fake.zipcode())
|
||||
dat["Order Date"] = fake.date_this_year()
|
||||
dat["Ship Date"] = fake.date_between_dates(dat["Order Date"]).strftime(
|
||||
"%Y-%m-%d"
|
||||
)
|
||||
dat["Order Date"] = dat["Order Date"].strftime("%Y-%m-%d")
|
||||
dat["Ship Mode"] = choice(["First Class", "Standard Class", "Second Class"])
|
||||
dat["Ship Mode"] = choice(["First Class", "Standard Class", "Second Class"])
|
||||
dat["Customer ID"] = fake.zipcode()
|
||||
dat["Segment"] = choice(["A", "B", "C", "D"])
|
||||
dat["Country"] = "US"
|
||||
dat["City"] = fake.city()
|
||||
dat["State"] = fake.state()
|
||||
dat["Postal Code"] = fake.zipcode()
|
||||
dat["Region"] = choice(["Region %d" % i for i in range(5)])
|
||||
dat["Product ID"] = fake.bban()
|
||||
sector = choice(["Industrials", "Technology", "Financials"])
|
||||
industry = choice(["A", "B", "C"])
|
||||
dat["Category"] = sector
|
||||
dat["Sub-Category"] = industry
|
||||
dat["Sales"] = randint(1, 100) * 100
|
||||
dat["Quantity"] = randint(1, 100) * 10
|
||||
dat["Discount"] = round(random() * 100, 2)
|
||||
dat["Profit"] = round(random() * 1000, 2)
|
||||
data.append(dat)
|
||||
return pd.DataFrame(data)
|
||||
Reference in New Issue
Block a user