diff --git a/tests/conftest.py b/tests/conftest.py index d2c5d93559..2d2aa1db13 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,6 +31,9 @@ def pytest_collection_modifyitems(items, config): our_marks = {*marks, *dialect_groups.keys()} for item in items: + # Any test without backend-specific marker gets 'core' marker (by definition) + # Also a dialect marker for each specified dialect + # Does not get {dialect}_only, as these are specifically for non-core tests if not any(marker.name in our_marks for marker in item.iter_markers()): item.add_marker("core") for mark in our_marks: @@ -92,7 +95,10 @@ def df_spark(spark): def fake_1000(): import pyarrow.csv as pv - return pv.read_csv("./tests/datasets/fake_1000_from_splink_demos.csv") + return pv.read_csv( + "./tests/datasets/fake_1000_from_splink_demos.csv", + convert_options=pv.ConvertOptions(strings_can_be_null=True), + ) @pytest.fixture(scope="function") diff --git a/tests/helpers.py b/tests/helpers.py index 8dd843c64c..367b4facbc 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -22,10 +22,6 @@ def DatabaseAPI(self): def db_api_args(self): return {} - # def extra_linker_args(self): - # # create fresh api each time - # return {"db_api": self.DatabaseAPI(**self.db_api_args())} - @property def date_format(self): return "yyyy-mm-dd" @@ -33,19 +29,18 @@ def date_format(self): def db_api(self): return self.DatabaseAPI(**self.db_api_args()) - @abstractmethod - def convert_frame(self, df): - pass - def load_frame_from_csv(self, path): - import pandas as pd + import pyarrow.csv as pv - return pd.read_csv(path) + return pv.read_csv( + path, + convert_options=pv.ConvertOptions(strings_can_be_null=True), + ) def load_frame_from_parquet(self, path): - import pandas as pd + import pyarrow.parquet as pq - return pd.read_parquet(path) + return pq.read_table(path) @property def arrays_from(self) -> int: @@ -82,9 +77,6 @@ def DatabaseAPI(self): def db_api_args(self): return {"connection": self.con} - def convert_frame(self, df): - return df - def load_frame_from_csv(self, path): return self.con.read_csv(path) @@ -113,11 +105,6 @@ def db_api_args(self): "break_lineage_method": "parquet", } - def convert_frame(self, df): - spark_frame = self.spark.createDataFrame(df) - spark_frame.persist() - return spark_frame - def load_frame_from_csv(self, path): df = self.spark.read.csv(path, header=True) df.persist() @@ -153,17 +140,6 @@ def _get_input_name(cls): cls._frame_counter += 1 return name - def convert_frame(self, df): - name = self._get_input_name() - df.to_sql(name, self.con, if_exists="replace") - return name - - def load_frame_from_csv(self, path): - return self.convert_frame(super().load_frame_from_csv(path)) - - def load_frame_from_parquet(self, path): - return self.convert_frame(super().load_frame_from_parquet(path)) - class PostgresTestHelper(TestHelper): _frame_counter = 0 @@ -188,32 +164,6 @@ def _get_input_name(cls): cls._frame_counter += 1 return name - def convert_frame(self, df): - from sqlalchemy.dialects import postgresql - from sqlalchemy.types import INTEGER, TEXT - - name = self._get_input_name() - # workaround to handle array column conversion - # manually mark any list columns so type is handled correctly - dtypes = {} - for colname, values in df.items(): - # TODO: will fail if first value is null - if isinstance(values[0], list): - # TODO: will fail if first array is empty - initial_array_val = values[0][0] - if isinstance(initial_array_val, int): - dtypes[colname] = postgresql.ARRAY(INTEGER) - elif isinstance(initial_array_val, str): - dtypes[colname] = postgresql.ARRAY(TEXT) - df.to_sql(name, con=self.engine, if_exists="replace", dtype=dtypes) - return name - - def load_frame_from_csv(self, path): - return self.convert_frame(super().load_frame_from_csv(path)) - - def load_frame_from_parquet(self, path): - return self.convert_frame(super().load_frame_from_parquet(path)) - class SplinkTestException(Exception): pass diff --git a/tests/test_chunking.py b/tests/test_chunking.py index 2abbda13df..0820002df7 100644 --- a/tests/test_chunking.py +++ b/tests/test_chunking.py @@ -434,10 +434,7 @@ def test_chunked_predict_link_only(test_helpers, dialect): df1_pd = df_pd[df_pd.index % 2 == 0].copy().reset_index(drop=True) df2_pd = df_pd[df_pd.index % 2 == 1].copy().reset_index(drop=True) - df1 = helper.convert_frame(df1_pd) - df2 = helper.convert_frame(df2_pd) - - linker = helper.linker_with_registration([df1, df2], settings) + linker = helper.linker_with_registration([df1_pd, df2_pd], settings) # Get baseline predictions predictions_baseline = linker.inference.predict(threshold_match_weight=-10) @@ -488,11 +485,7 @@ def test_chunked_predict_link_only_three_datasets(test_helpers, dialect): df2_pd = df_pd[df_pd.index % 3 == 1].copy().reset_index(drop=True) df3_pd = df_pd[df_pd.index % 3 == 2].copy().reset_index(drop=True) - df1 = helper.convert_frame(df1_pd) - df2 = helper.convert_frame(df2_pd) - df3 = helper.convert_frame(df3_pd) - - linker = helper.linker_with_registration([df1, df2, df3], settings) + linker = helper.linker_with_registration([df1_pd, df2_pd, df3_pd], settings) # Get baseline predictions predictions_baseline = linker.inference.predict(threshold_match_weight=-10) @@ -539,10 +532,7 @@ def test_chunked_predict_link_and_dedupe(test_helpers, dialect): df1_pd = df_pd[df_pd.index % 2 == 0].copy().reset_index(drop=True) df2_pd = df_pd[df_pd.index % 2 == 1].copy().reset_index(drop=True) - df1 = helper.convert_frame(df1_pd) - df2 = helper.convert_frame(df2_pd) - - linker = helper.linker_with_registration([df1, df2], settings) + linker = helper.linker_with_registration([df1_pd, df2_pd], settings) # Get baseline predictions predictions_baseline = linker.inference.predict(threshold_match_weight=-10) @@ -593,11 +583,7 @@ def test_chunked_predict_link_and_dedupe_three_datasets(test_helpers, dialect): df2_pd = df_pd[df_pd.index % 3 == 1].copy().reset_index(drop=True) df3_pd = df_pd[df_pd.index % 3 == 2].copy().reset_index(drop=True) - df1 = helper.convert_frame(df1_pd) - df2 = helper.convert_frame(df2_pd) - df3 = helper.convert_frame(df3_pd) - - linker = helper.linker_with_registration([df1, df2, df3], settings) + linker = helper.linker_with_registration([df1_pd, df2_pd, df3_pd], settings) # Get baseline predictions predictions_baseline = linker.inference.predict(threshold_match_weight=-10) diff --git a/tests/test_clustering.py b/tests/test_clustering.py index 6d974bcd36..bd7f613cf6 100644 --- a/tests/test_clustering.py +++ b/tests/test_clustering.py @@ -60,8 +60,7 @@ def test_clustering(test_helpers, dialect, link_type, input_pd_tables): block_on("dob"), ], ) - linker_input = list(map(helper.convert_frame, input_pd_tables)) - linker = helper.linker_with_registration(linker_input, settings) + linker = helper.linker_with_registration(input_pd_tables, settings) df_predict = linker.inference.predict() linker.clustering.cluster_pairwise_predictions_at_threshold(df_predict, 0.95) @@ -121,8 +120,7 @@ def test_clustering_no_edges(test_helpers, dialect): ], unique_id_column_name="id", ) - linker_input = helper.convert_frame(df) - linker = helper.linker_with_registration(linker_input, settings) + linker = helper.linker_with_registration(df, settings) # due to blocking rules, df_predict will be empty df_predict = linker.inference.predict() diff --git a/tests/test_column_expression.py b/tests/test_column_expression.py index ea24df53af..6cb70c15b6 100644 --- a/tests/test_column_expression.py +++ b/tests/test_column_expression.py @@ -20,9 +20,8 @@ def test_access_extreme_array_element(test_helpers, dialect): helper = test_helpers[dialect] # register table with backend table_name = "arr_tab" - table = helper.convert_frame(df_arr) db_api = helper.DatabaseAPI(**helper.db_api_args()) - arr_tab = db_api.register(table, table_name) + arr_tab = db_api.register(df_arr, table_name) # construct a SQL query from ColumnExpressions and run it against backend splink_dialect = SplinkDialect.from_string(dialect) @@ -64,9 +63,8 @@ def test_nullif(test_helpers, dialect): helper = test_helpers[dialect] # register table with backend table_name = "nully_name_table" - table = helper.convert_frame(df_arr) db_api = helper.DatabaseAPI(**helper.db_api_args()) - nully_table = db_api.register(table, table_name) + nully_table = db_api.register(df_arr, table_name) # construct a SQL query from ColumnExpressions and run it against backend splink_dialect = SplinkDialect.from_string(dialect) diff --git a/tests/test_completeness.py b/tests/test_completeness.py index f41d857afa..3bafaea122 100644 --- a/tests/test_completeness.py +++ b/tests/test_completeness.py @@ -20,19 +20,12 @@ def test_completeness_chart(dialect, test_helpers): @mark_with_dialects_excluding("sqlite") -def test_completeness_chart_mismatched_columns(dialect, test_helpers): +def test_completeness_chart_mismatched_columns(dialect, test_helpers, fake_1000): helper = test_helpers[dialect] db_api = helper.db_api() - df_l = helper.load_frame_from_csv( - "./tests/datasets/fake_1000_from_splink_demos.csv" - ) - df_r = pd.read_csv("./tests/datasets/fake_1000_from_splink_demos.csv") - df_r.rename(columns={"surname": "surname_2"}, inplace=True) - df_r = helper.convert_frame(df_r) - - df_l_sdf = db_api.register(df_l) - df_r_sdf = db_api.register(df_r) + df_l_sdf = db_api.register(fake_1000) + df_r_sdf = db_api.register(fake_1000.rename_columns(names={"surname": "surname_2"})) with raises(SplinkException): completeness_chart([df_l_sdf, df_r_sdf]) @@ -56,7 +49,6 @@ def test_completeness_chart_complex_columns(dialect, test_helpers): ], } ) - df = helper.convert_frame(df) df_sdf = db_api.register(df) first = helper.arrays_from # check completeness when we have more complicated column constructs, such as @@ -70,6 +62,5 @@ def test_completeness_chart_source_dataset(dialect, test_helpers): db_api = helper.db_api() df_pd = pd.read_csv("./tests/datasets/fake_1000_from_splink_demos.csv") df_pd["source_dataset"] = "fake_1000" - df = helper.convert_frame(df_pd) - df_sdf = db_api.register(df) + df_sdf = db_api.register(df_pd) completeness_chart(df_sdf) diff --git a/tests/test_estimate_prob_two_rr_match.py b/tests/test_estimate_prob_two_rr_match.py index c9bdbc3695..d764cd1e9f 100644 --- a/tests/test_estimate_prob_two_rr_match.py +++ b/tests/test_estimate_prob_two_rr_match.py @@ -19,7 +19,6 @@ def test_prob_rr_match_dedupe(test_helpers, dialect): {"unique_id": 6, "first_name": "Jane", "surname": "Taylor"}, ] ) - df = helper.convert_frame(df) settings = { "link_type": "dedupe_only", @@ -71,8 +70,6 @@ def test_prob_rr_match_link_only(test_helpers, dialect): {"unique_id": 4, "first_name": "Alice", "surname": "Williams"}, ] ) - df_1 = helper.convert_frame(df_1) - df_2 = helper.convert_frame(df_2) settings = { "link_type": "link_only", @@ -114,8 +111,6 @@ def test_prob_rr_match_link_and_dedupe(test_helpers, dialect): {"unique_id": 3, "first_name": "Jane", "surname": "Taylor"}, ] ) - df_1 = helper.convert_frame(df_1) - df_2 = helper.convert_frame(df_2) settings = { "link_type": "link_and_dedupe", @@ -181,10 +176,6 @@ def test_prob_rr_match_link_only_multitable(test_helpers, dialect): map(lambda df: df.assign(city="Brighton"), (df_1, df_2, df_3, df_4)) ) - df_1 = helper.convert_frame(df_1) - df_2 = helper.convert_frame(df_2) - df_3 = helper.convert_frame(df_3) - df_4 = helper.convert_frame(df_4) dfs = [df_1, df_2, df_3, df_4] settings = { @@ -259,10 +250,6 @@ def test_prob_rr_match_link_and_dedupe_multitable(test_helpers, dialect): map(lambda df: df.assign(city="Brighton"), (df_1, df_2, df_3, df_4)) ) - df_1 = helper.convert_frame(df_1) - df_2 = helper.convert_frame(df_2) - df_3 = helper.convert_frame(df_3) - df_4 = helper.convert_frame(df_4) dfs = [df_1, df_2, df_3, df_4] settings = { @@ -340,7 +327,6 @@ def check_range(p): }, ] ) - df = helper.convert_frame(df) settings = { "link_type": "dedupe_only", diff --git a/tests/test_find_new_matches.py b/tests/test_find_new_matches.py index 48f207448f..dca24d9c36 100644 --- a/tests/test_find_new_matches.py +++ b/tests/test_find_new_matches.py @@ -106,4 +106,4 @@ def test_matches_work(test_helpers, dialect, fake_1000, tmp_path): ) matches = matches.as_record_dict() - assert len(matches) == 1 + assert len(matches) == 2 diff --git a/tests/test_full_example_deterministic_link.py b/tests/test_full_example_deterministic_link.py index 28bbe812fc..7052a9d26e 100644 --- a/tests/test_full_example_deterministic_link.py +++ b/tests/test_full_example_deterministic_link.py @@ -13,7 +13,6 @@ def test_deterministic_link_full_example(dialect, tmp_path, test_helpers): helper = test_helpers[dialect] df = pd.read_csv("./tests/datasets/fake_1000_from_splink_demos.csv") - df = helper.convert_frame(df) br_for_predict = [ "l.first_name = r.first_name and l.surname = r.surname and l.dob = r.dob", diff --git a/tests/test_graph_metrics.py b/tests/test_graph_metrics.py index fbfc64a8a3..38c9dd5253 100644 --- a/tests/test_graph_metrics.py +++ b/tests/test_graph_metrics.py @@ -231,15 +231,11 @@ def test_metrics(dialect, test_helpers): # pass in dummy frame to linker linker = helper.linker_with_registration( - helper.convert_frame(df_1), + df_1, {"link_type": "dedupe_only"}, ) - df_predict = linker.table_management.register_table( - helper.convert_frame(df_e), "predict" - ) - df_clustered = linker.table_management.register_table( - helper.convert_frame(df_c), "clusters" - ) + df_predict = linker.table_management.register_table(df_e, "predict") + df_clustered = linker.table_management.register_table(df_c, "clusters") cm = linker.clustering.compute_graph_metrics( df_predict, df_clustered, threshold_match_probability=0.95 @@ -351,15 +347,11 @@ def test_is_bridge(dialect, test_helpers): + [{"cluster_id": 3, "unique_id": i} for i in range(11, 18 + 1)] ) linker = helper.linker_with_registration( - helper.convert_frame(df_1), + df_1, {"link_type": "dedupe_only"}, ) - df_predict = linker.table_management.register_table( - helper.convert_frame(df_e), "br_predict" - ) - df_clustered = linker.table_management.register_table( - helper.convert_frame(df_c), "br_clusters" - ) + df_predict = linker.table_management.register_table(df_e, "br_predict") + df_clustered = linker.table_management.register_table(df_c, "br_clusters") # linker.debug_mode = True cm = linker.clustering.compute_graph_metrics( diff --git a/tests/test_km_distance_level.py b/tests/test_km_distance_level.py index 7e02f01446..25b5a7ac20 100644 --- a/tests/test_km_distance_level.py +++ b/tests/test_km_distance_level.py @@ -122,8 +122,6 @@ def test_km_distance_levels(dialect, test_helpers): settings_cll = {"link_type": "dedupe_only", "comparisons": [km_diff]} - df = helper.convert_frame(df) - linker = helper.linker_with_registration(df, settings_cl) cl_df_e = linker.inference.predict().as_pandas_dataframe() linker = helper.linker_with_registration(df, settings_cll) diff --git a/tests/test_regex_param.py b/tests/test_regex_param.py index d553a0d350..952df021c5 100644 --- a/tests/test_regex_param.py +++ b/tests/test_regex_param.py @@ -128,8 +128,7 @@ def test_regex(dialect, test_helpers, level_set, record_pairs_gamma): comparison_name = level_set["output_column_name"] - df = helper.convert_frame(df_pandas) - linker = helper.linker_with_registration(df, settings) + linker = helper.linker_with_registration(df_pandas, settings) linker_output = linker.inference.predict().as_pandas_dataframe() diff --git a/tests/test_score_missing_edges.py b/tests/test_score_missing_edges.py index 15d434d0a5..65400aca87 100644 --- a/tests/test_score_missing_edges.py +++ b/tests/test_score_missing_edges.py @@ -21,8 +21,6 @@ def test_score_missing_edges(test_helpers, dialect, link_type, copies_of_df): helper = test_helpers[dialect] - df = helper.convert_frame(df_pd) - settings = SettingsCreator( link_type=link_type, comparisons=[ @@ -38,7 +36,7 @@ def test_score_missing_edges(test_helpers, dialect, link_type, copies_of_df): retain_intermediate_calculation_columns=True, ) - linker_input = df if copies_of_df == 1 else [df for _ in range(copies_of_df)] + linker_input = df_pd if copies_of_df == 1 else [df_pd for _ in range(copies_of_df)] linker = helper.linker_with_registration(linker_input, settings) df_predict = linker.inference.predict() @@ -64,8 +62,6 @@ def test_score_missing_edges(test_helpers, dialect, link_type, copies_of_df): def test_score_missing_edges_all_edges(test_helpers, dialect, link_type, copies_of_df): helper = test_helpers[dialect] - df = helper.convert_frame(df_pd) - settings = SettingsCreator( link_type=link_type, comparisons=[ @@ -81,7 +77,7 @@ def test_score_missing_edges_all_edges(test_helpers, dialect, link_type, copies_ retain_intermediate_calculation_columns=True, ) - linker_input = df if copies_of_df == 1 else [df for _ in range(copies_of_df)] + linker_input = df_pd if copies_of_df == 1 else [df_pd for _ in range(copies_of_df)] linker = helper.linker_with_registration(linker_input, settings) df_predict = linker.inference.predict() @@ -128,11 +124,11 @@ def test_score_missing_edges_changed_column_names(test_helpers, dialect, link_ty source_dataset_column_name="sds", ) if link_type == "dedupe_only": - linker_input = helper.convert_frame(df) + linker_input = df else: df_2 = df.copy() df_2["sds"] = "frame_2" - linker_input = [helper.convert_frame(df), helper.convert_frame(df_2)] + linker_input = [df, df_2] linker = helper.linker_with_registration(linker_input, settings) df_predict = linker.inference.predict()