def test_no_analysis_exception_when_shared_parent_dataframe(spark): # Check that we don't fall victim to # https://issues.apache.org/jira/browse/SPARK-10925 df = spark.createDataFrame( [ # Just need the schema, really ['someone', '20190102', 'fake', 1], ], [ "client_id", "submission_date_s3", "branch", "some_value", ]) enrollments = df.groupby('client_id', 'branch').agg( F.min('submission_date_s3').alias('enrollment_date')) exp = Experiment('a-stub', '20180101') time_limits = TimeLimits.for_single_analysis_window( exp.start_date, last_date_full_data='20190522', analysis_start_days=28, analysis_length_dates=7) enrollments = exp._add_analysis_windows_to_enrollments( enrollments, time_limits) exp._get_results_for_one_data_source( enrollments, df, [F.max(F.col('some_value'))], )
def test_add_analysis_windows_to_enrollments(spark): exp = Experiment('a-stub', '20190101', num_dates_enrollment=8) enrollments = exp.get_enrollments(spark, _get_enrollment_view(slug="a-stub")) assert enrollments.count() == 3 tl = TimeLimits.for_ts( first_enrollment_date=exp.start_date, last_date_full_data='20190114', time_series_period='daily', num_dates_enrollment=exp.num_dates_enrollment, ) assert len(tl.analysis_windows) == 7 new_enrollments = exp._add_analysis_windows_to_enrollments(enrollments, tl) nep = new_enrollments.toPandas() assert len(nep) == enrollments.count() * len(tl.analysis_windows) a = nep[nep['client_id'] == 'aaaa'] assert len(a) == len(tl.analysis_windows) assert (a.mozanalysis_analysis_window_start.sort_values() == np.arange( len(tl.analysis_windows))).all() assert (a.mozanalysis_analysis_window_end.sort_values() == np.arange( len(tl.analysis_windows))).all()
def test_get_results_for_one_data_source(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 = 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", ], ) time_limits = TimeLimits.for_single_analysis_window( exp.start_date, '20190114', 1, 3, ) enrollments = exp._add_analysis_windows_to_enrollments( enrollments, time_limits) res = exp._get_results_for_one_data_source( enrollments, data_source, [ F.coalesce(F.sum(data_source.some_value), F.lit(0)).alias('some_value'), ], ) # 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 bob_badtiming = res.filter(res.client_id == 'bob-badtiming') assert bob_badtiming.count() == 1 assert bob_badtiming.first()['some_value'] == 0 # 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 res2 = exp._get_results_for_one_data_source( enrollments, data_source.drop('experiments'), [ F.coalesce(F.sum(data_source.some_value), F.lit(0)).alias('some_value'), ], ) assert res2.count() == enrollments.count()