Return DataFrames from _extract_tuple so ReferentialIntegrity scores instead of erroring

sdmetrics/reports/multi_table/_properties/base.py

@@ -56,9 +56,13 @@ def _get_num_iterations(self, metadata):
     def _extract_tuple(data, relation):
         parent_data = data[relation['parent_table_name']]
         child_data = data[relation['child_table_name']]
+        primary_key = relation['parent_primary_key']
+        foreign_key = relation['child_foreign_key']
+        primary_key = [primary_key] if not isinstance(primary_key, list) else primary_key
+        foreign_key = [foreign_key] if not isinstance(foreign_key, list) else foreign_key
         return (
-            parent_data[relation['parent_primary_key']],
-            child_data[relation['child_foreign_key']],
+            parent_data[primary_key],
+            child_data[foreign_key],
         )
 
     def _compute_average(self):

tests/unit/reports/multi_table/_properties/test_base.py

@@ -120,7 +120,8 @@ def test__extract_tuple(self):
         real_columns = base_property._extract_tuple(real_data, relation)
 
         # Assert
-        assert real_columns == (real_data['users']['user_id'], real_data['sessions']['user_id'])
+        pd.testing.assert_frame_equal(real_columns[0], real_data['users'][['user_id']])
+        pd.testing.assert_frame_equal(real_columns[1], real_data['sessions'][['user_id']])
 
     def test__generate_details_property(self):
         """Test the ``_generate_details`` method."""

tests/unit/reports/multi_table/_properties/test_relationship_validity.py

@@ -92,7 +92,31 @@ def test__extract_tuple(self, real_data_fixture):
         real_columns = relationship_validity._extract_tuple(real_data, relation)
 
         # Assert
-        assert real_columns == (real_data['users']['user_id'], real_data['sessions']['user_id'])
+        pd.testing.assert_frame_equal(real_columns[0], real_data['users'][['user_id']])
+        pd.testing.assert_frame_equal(real_columns[1], real_data['sessions'][['user_id']])
+
+    def test__generate_details_referential_integrity_score(
+        self, real_data_fixture, synthetic_data_fixture, metadata_fixture
+    ):
+        """Test that ``ReferentialIntegrity`` returns a score instead of erroring.
+
+        ``_extract_tuple`` used to pass the key columns as ``pd.Series``, which made
+        ``ReferentialIntegrity`` fail with ``AttributeError: 'Series' object has no
+        attribute 'columns'`` and record a ``nan`` score.
+        """
+        # Setup
+        relationship_validity = RelationshipValidity()
+
+        # Run
+        relationship_validity._generate_details(
+            real_data_fixture, synthetic_data_fixture, metadata_fixture
+        )
+
+        # Assert
+        details = relationship_validity.details
+        ri_row = details[details['Metric'] == 'ReferentialIntegrity'].iloc[0]
+        assert ri_row['Error'] is None
+        assert not np.isnan(ri_row['Score'])
 
     def test__get_num_iteration(self):
         """Test the ``_get_num_iterations`` method."""