def test__create_partitions(self, spark_session, spark_context):
# arrange
start = datetime.datetime(year=1970, month=1, day=1)
end = datetime.datetime(year=2020, month=12, day=31)
random_dates = [
(
lambda: start + datetime.
timedelta(seconds=random.randint( # noqa: S311
0, int(
(end - start).total_seconds()))))().date().isoformat()
for _ in range(10000)
]
data = [{"timestamp": date} for date in random_dates]
input_df = spark_session.read.json(spark_context.parallelize(data, 1),
schema="timestamp timestamp")
writer = HistoricalFeatureStoreWriter()
# act
result_df = writer._create_partitions(input_df)
# assert
assert result_df.select("year", "month",
"day").distinct().count() == len(
set(random_dates))
def test__repartition_df(self, spark_session, spark_context):
# arrange
start = datetime.datetime(year=1970, month=1, day=1)
end = datetime.datetime(year=2020, month=12, day=31)
random_dates = [
(
lambda: start + datetime.
timedelta(seconds=random.randint( # noqa: S311
0, int(
(end - start).total_seconds()))))().date().isoformat()
for _ in range(10000)
]
data = [{"timestamp": date} for date in random_dates]
input_df = spark_session.read.json(spark_context.parallelize(data, 1),
schema="timestamp timestamp")
writer = HistoricalFeatureStoreWriter()
# act
result_df = writer._create_partitions(input_df)
# assert
# Only one partition id, meaning data is not partitioned
assert input_df.select(spark_partition_id()).distinct().count() == 1
# Desired number of partitions
assert result_df.select(spark_partition_id()).distinct().count() == 200
