Fixed dist_gen.zipf_choice sampled ordering (#53)

Co-authored-by: Tim Band <t.b@ucl>

Author: tim-band

datafaker/generators.py

@@ -667,20 +667,20 @@ def __init__(
         self._fit = fit_from_buckets(counts, estimated_counts)
         if suppress_count == 0:
             if sample_count is None:
-                self._query = f"SELECT {column_name} AS value FROM {table_name} GROUP BY value ORDER BY COUNT({column_name}) DESC"
+                self._query = f"SELECT {column_name} AS value FROM {table_name} WHERE {column_name} IS NOT NULL GROUP BY value ORDER BY COUNT({column_name}) DESC"
                 self._comment = f"All the values that appear in column {column_name} of table {table_name}"
                 self._annotation = None
             else:
-                self._query = f"SELECT value FROM (SELECT {column_name} AS value FROM {table_name} ORDER BY RANDOM() LIMIT {sample_count}) AS _inner GROUP BY value"
+                self._query = f"SELECT value FROM (SELECT {column_name} AS value FROM {table_name} WHERE {column_name} IS NOT NULL ORDER BY RANDOM() LIMIT {sample_count}) AS _inner GROUP BY value ORDER BY COUNT(value) DESC"
                 self._comment = f"The values that appear in column {column_name} of a random sample of {sample_count} rows of table {table_name}"
                 self._annotation = "sampled"
         else:
             if sample_count is None:
-                self._query = f"SELECT value FROM (SELECT {column_name} AS value, COUNT({column_name}) AS count FROM {table_name} GROUP BY value) AS _inner WHERE {suppress_count} < count"
+                self._query = f"SELECT value FROM (SELECT {column_name} AS value, COUNT({column_name}) AS count FROM {table_name} WHERE {column_name} IS NOT NULL GROUP BY value ORDER BY count DESC) AS _inner WHERE {suppress_count} < count"
                 self._comment = f"All the values that appear in column {column_name} of table {table_name} more than {suppress_count} times"
                 self._annotation = "suppressed"
             else:
-                self._query = f"SELECT value FROM (SELECT value, COUNT(value) AS count FROM (SELECT {column_name} AS value FROM {table_name} ORDER BY RANDOM() LIMIT {sample_count}) AS _inner GROUP BY value) AS _inner WHERE {suppress_count} < count"
+                self._query = f"SELECT value FROM (SELECT value, COUNT(value) AS count FROM (SELECT {column_name} AS value FROM {table_name} WHERE {column_name} IS NOT NULL ORDER BY RANDOM() LIMIT {sample_count}) AS _inner GROUP BY value ORDER BY count DESC) AS _inner WHERE {suppress_count} < count"
                 self._comment = f"The values that appear more than {suppress_count} times in column {column_name}, out of a random sample of {sample_count} rows of table {table_name}"
                 self._annotation = "sampled and suppressed"
     def nominal_kwargs(self):

tests/examples/instrument.sql

@@ -99,3 +99,7 @@ INSERT INTO public.signature_model VALUES (3, 'Veleno', 2, 2);
 INSERT INTO public.signature_model VALUES (4, 'Grifter', NULL, NULL);
 INSERT INTO public.signature_model VALUES (5, 'Proton', 3, 1);
 INSERT INTO public.signature_model VALUES (6, 'Isabelle', NULL, 3);
+INSERT INTO public.signature_model VALUES (7, 'Natalia', NULL, 1);
+INSERT INTO public.signature_model VALUES (8, 'Ray', 2, NULL);
+INSERT INTO public.signature_model VALUES (9, 'DamageControl', 2, 3);
+INSERT INTO public.signature_model VALUES (10, 'Lucy', 3, 1);

tests/test_interactive.py

@@ -481,7 +481,7 @@ def test_set_generator_choice(self):
             self.assertEqual(gc.config["src-stats"][0]["name"], f"auto__{TABLE}__{COLUMN}")
             self.assertEqual(
                 gc.config["src-stats"][0]["query"],
-                f"SELECT {COLUMN} AS value FROM {TABLE} GROUP BY value ORDER BY COUNT({COLUMN}) DESC",
+                f"SELECT {COLUMN} AS value FROM {TABLE} WHERE {COLUMN} IS NOT NULL GROUP BY value ORDER BY COUNT({COLUMN}) DESC",
             )
 
     def test_old_generators_remain(self):
@@ -834,7 +834,7 @@ def test_create_with_missingness(self):
 
 
 class GeneratorTests(GeneratesDBTestCase):
-    """ Testing configure-missing with generation. """
+    """ Testing configure-generators with generation. """
     dump_file_path = "instrument.sql"
     database_name = "instrument"
     schema_name = "public"
@@ -883,3 +883,25 @@ def test_set_null(self):
                 self.assertEqual(row.name, "")
                 self.assertIsNone(row.based_on)
             self.assertEqual(count, 3)
+
+    def test_dist_gen_sampled_produces_ordered_src_stats(self):
+        """ Tests that choosing a sampled choice generator produces ordered src stats """
+        with self._get_cmd({}) as gc:
+            gc.do_next("signature_model.player_id")
+            gc.do_set("dist_gen.zipf_choice [sampled]")
+            gc.do_next("signature_model.based_on")
+            gc.do_set("dist_gen.zipf_choice [sampled]")
+            gc.do_quit("")
+            config = gc.config
+            self.set_configuration(config)
+            src_stats = self.get_src_stats(config)
+        player_ids = [
+            s["value"]
+            for s in src_stats["auto__signature_model__player_id"]["results"]
+        ]
+        self.assertListEqual(player_ids, [2, 3, 1])
+        based_ons = [
+            s["value"]
+            for s in src_stats["auto__signature_model__based_on"]["results"]
+        ]
+        self.assertListEqual(based_ons, [1, 3, 2])

tests/utils.py

@@ -177,6 +177,7 @@ def get_src_stats(self, config) -> dict[str, any]:
         (self.stats_fd, self.stats_file_path) = mkstemp(".yaml", "src_stats_", text=True)
         with os.fdopen(self.stats_fd, "w", encoding="utf-8") as stats_fh:
             stats_fh.write(yaml.dump(src_stats))
+        return src_stats
 
     def create_generators(self, config) -> None:
         """ ``create-generators`` with ``src-stats.yaml`` and the rest, producing ``df.py`` """