def test_process_data_source_df(spark):
start_date = '20190101'
exp_8d = Experiment('experiment-with-8-day-cohort', start_date, 8)
data_source_df = _get_data_source_df(spark)
end_date = '20190114'
# Are the fixtures sufficiently complicated that we're actually testing
# things?
assert _simple_return_agg_date(F.min, data_source_df) < start_date
assert _simple_return_agg_date(F.max, data_source_df) > end_date
tl_03 = TimeLimits.for_single_analysis_window(
first_enrollment_date=exp_8d.start_date,
last_date_full_data=end_date,
analysis_start_days=0,
analysis_length_dates=3,
num_dates_enrollment=exp_8d.num_dates_enrollment)
assert tl_03.first_date_data_required == start_date
assert tl_03.last_date_data_required == '20190110'
proc_ds = exp_8d._process_data_source_df(data_source_df, tl_03)
assert _simple_return_agg_date(F.min,
proc_ds) == tl_03.first_date_data_required
assert _simple_return_agg_date(F.max,
proc_ds) == tl_03.last_date_data_required
tl_23 = TimeLimits.for_single_analysis_window(
first_enrollment_date=exp_8d.start_date,
last_date_full_data=end_date,
analysis_start_days=2,
analysis_length_dates=3,
num_dates_enrollment=exp_8d.num_dates_enrollment)
assert tl_23.first_date_data_required == add_days(start_date, 2)
assert tl_23.last_date_data_required == '20190112'
p_ds_2 = exp_8d._process_data_source_df(data_source_df, tl_23)
assert _simple_return_agg_date(F.min,
p_ds_2) == tl_23.first_date_data_required
assert _simple_return_agg_date(F.max,
p_ds_2) == tl_23.last_date_data_required
assert proc_ds.select(F.col('data_source.client_id'))
with pytest.raises(AnalysisException):
assert data_source_df.select(F.col('data_source.client_id'))
def test_get_per_client_data_join(spark):
exp = Experiment('a-stub', '20190101')
enrollments = spark.createDataFrame(
[
['aaaa', 'control', '20190101'],
['bbbb', 'test', '20190101'],
['cccc', 'control', '20190108'],
['dddd', 'test', '20190109'],
['annie-nodata', 'control', '20190101'],
['bob-badtiming', 'test', '20190102'],
['carol-gooddata', 'test', '20190101'],
['derek-lateisok', 'control', '20190110'],
],
[
"client_id",
"branch",
"enrollment_date",
],
)
ex_d = {'a-stub': 'fake-branch-lifes-too-short'}
data_source_df = spark.createDataFrame(
[
# bob-badtiming only has data before/after analysis window
# but missed by `process_data_source`
['bob-badtiming', '20190102', ex_d, 1],
['bob-badtiming', '20190106', ex_d, 2],
# carol-gooddata has data on two days (including a dupe day)
['carol-gooddata', '20190102', ex_d, 3],
['carol-gooddata', '20190102', ex_d, 2],
['carol-gooddata', '20190104', ex_d, 6],
# derek-lateisok has data before and during the analysis window
['derek-lateisok', '20190110', ex_d, 1000],
['derek-lateisok', '20190111', ex_d, 1],
# TODO: exercise the last condition on the join
],
[
"client_id",
"submission_date_s3",
"experiments",
"some_value",
],
)
ds = DataSource.from_dataframe('ds', data_source_df)
metric = Metric.from_col('some_value', agg_sum(data_source_df.some_value),
ds)
res = exp.get_per_client_data(enrollments, [metric],
'20190114',
1,
3,
keep_client_id=True)
# Check that the dataframe has the correct number of rows
assert res.count() == enrollments.count()
# Check that dataless enrollments are handled correctly
annie_nodata = res.filter(res.client_id == 'annie-nodata')
assert annie_nodata.count() == 1
assert annie_nodata.first()['some_value'] == 0
# Check that early and late data were ignored
# i.e. check the join, not just _process_data_source_df
bob_badtiming = res.filter(res.client_id == 'bob-badtiming')
assert bob_badtiming.count() == 1
assert bob_badtiming.first()['some_value'] == 0
# Check that _process_data_source_df didn't do the
# heavy lifting above
time_limits = TimeLimits.for_single_analysis_window(
exp.start_date, '20190114', 1, 3, exp.num_dates_enrollment)
pds = exp._process_data_source_df(data_source_df, time_limits)
assert pds.filter(pds.client_id == 'bob-badtiming').select(
F.sum(pds.some_value).alias('agg_val')).first()['agg_val'] == 3
# Check that relevant data was included appropriately
carol_gooddata = res.filter(res.client_id == 'carol-gooddata')
assert carol_gooddata.count() == 1
assert carol_gooddata.first()['some_value'] == 11
derek_lateisok = res.filter(res.client_id == 'derek-lateisok')
assert derek_lateisok.count() == 1
assert derek_lateisok.first()['some_value'] == 1
# Check that it still works for `data_source`s without an experiments map
ds_df_noexp = data_source_df.drop('experiments')
ds_noexp = DataSource.from_dataframe('ds_noexp', ds_df_noexp)
metric_noexp = Metric.from_col('some_value',
agg_sum(ds_df_noexp.some_value), ds_noexp)
res2 = exp.get_per_client_data(enrollments, [metric_noexp],
'20190114',
1,
3,
keep_client_id=True)
assert res2.count() == enrollments.count()