Refactor tree building (#173)

* Rename materialize_tree to materialize_table

* done and tested

Author: yoid2000

docs/development.md

@@ -11,7 +11,7 @@ Activate `poetry` environment: `poetry shell`. (can skip, then prepend `poetry r
 - Check: `flake8 . && mypy . && black --check . && isort . --check`
 
  To filter out acceptable lines:
-     `poetry run flake8 . | findstr /v "E501" | findstr /v "F405" | findstr /v "F403"`
+     `poetry run flake8 syndiffix | findstr /v "E501" | findstr /v "F405" | findstr /v "F403"`
 
 ### Creating a new release
 

syndiffix/clustering/common.py

@@ -24,6 +24,7 @@ def __repr__(self) -> str:
 
 
 # Owner, Stitch columns, Derived columns
+# The total set of columns is Stitch+Derived.
 DerivedCluster = tuple[StitchOwner, list[ColumnId], list[ColumnId]]
 
 
@@ -33,7 +34,9 @@ class Clusters:
     derived_clusters: list[DerivedCluster]
 
 
-TreeMaterializer = Callable[[Forest, list[ColumnId]], tuple[list[MicrodataRow], Combination]]
+TreeBuilder = Callable[[Forest, list[ColumnId]], None]
+
+TableMaterializer = Callable[[Forest, list[ColumnId]], tuple[list[MicrodataRow], Combination]]
 
 
 def microdata_row_to_row(microdata_row: MicrodataRow) -> Row:

syndiffix/clustering/stitching.py

@@ -360,29 +360,36 @@ def _do_patch(
 
 
 def _stitch(
-    materialize_tree: TreeMaterializer,
+    materialize_table: TableMaterializer,
     forest: Forest,
     metadata: StitchingMetadata,
     left: tuple[list[MicrodataRow], Combination],
     derived_cluster: DerivedCluster,
 ) -> tuple[list[MicrodataRow], Combination]:
     (_, stitch_columns, derived_columns) = derived_cluster
 
-    right = materialize_tree(forest, stitch_columns + derived_columns)
+    right = materialize_table(forest, stitch_columns + derived_columns)
 
     if len(stitch_columns) == 0:
         return _do_patch(forest.unsafe_rng, left, right)
     else:
         return _do_stitch(forest, metadata, left, right, derived_cluster)
 
 
+def build_forest(tree_builder: TreeBuilder, forest: Forest, clusters: Clusters) -> None:
+    tree_builder(forest, clusters.initial_cluster)
+    for derived_cluster in clusters.derived_clusters:
+        (_, stitch_columns, derived_columns) = derived_cluster
+        tree_builder(forest, stitch_columns + derived_columns)
+
+
 def build_table(
-    materialize_tree: TreeMaterializer, forest: Forest, metadata: StitchingMetadata, clusters: Clusters
+    materialize_table: TableMaterializer, forest: Forest, metadata: StitchingMetadata, clusters: Clusters
 ) -> tuple[list[Row], Combination]:
-    acc = materialize_tree(forest, clusters.initial_cluster)
+    acc = materialize_table(forest, clusters.initial_cluster)
 
     for derived_cluster in clusters.derived_clusters:
-        acc = _stitch(materialize_tree, forest, metadata, acc, derived_cluster)
+        acc = _stitch(materialize_table, forest, metadata, acc, derived_cluster)
 
     rows, columns = acc
     return [microdata_row_to_row(row) for row in rows], columns

syndiffix/synthesizer.py

@@ -9,7 +9,7 @@
 
 from .bucket import harvest
 from .clustering.common import MicrodataRow
-from .clustering.stitching import StitchingMetadata, build_table
+from .clustering.stitching import StitchingMetadata, build_forest, build_table
 from .clustering.strategy import ClusteringStrategy, DefaultClustering, MlClustering
 from .common import *
 from .counters import (
@@ -118,8 +118,19 @@ def __init__(
         for col_id, converter in enumerate(self.column_convertors):
             converter.analyze_tree(self.forest.get_tree((ColumnId(col_id),)))
 
+        def tree_builder(forest: Forest, columns: list[ColumnId]) -> None:
+            combination = tuple(sorted(columns))
+            # get_tree builds the trees (except 1dim, which have already been built)
+            _ = forest.get_tree(combination)
+
+        build_forest(
+            tree_builder,
+            self.forest,
+            self.clusters,
+        )
+
     def sample(self) -> pd.DataFrame:
-        def materialize_tree(forest: Forest, columns: list[ColumnId]) -> tuple[list[MicrodataRow], Combination]:
+        def materialize_table(forest: Forest, columns: list[ColumnId]) -> tuple[list[MicrodataRow], Combination]:
             combination = tuple(sorted(columns))
             tree = forest.get_tree(combination)
             buckets = harvest(tree, self.forest.derive_unsafe_rng())
@@ -134,7 +145,7 @@ def materialize_tree(forest: Forest, columns: list[ColumnId]) -> tuple[list[Micr
             )
 
         rows, root_combination = build_table(
-            materialize_tree,
+            materialize_table,
             self.forest,
             StitchingMetadata(self.column_is_integral, self.entropy_1dim),
             self.clusters,

tests/clustering/test_stitching.py

@@ -35,7 +35,7 @@ def test_stitching() -> None:
         (ColumnId(2), ColumnId(3)): _build_rows(col_c_right, col_d),
     }
 
-    def materialize_tree(_forest: Forest, columns: list[ColumnId]) -> tuple[list[MicrodataRow], Combination]:
+    def materialize_table(_forest: Forest, columns: list[ColumnId]) -> tuple[list[MicrodataRow], Combination]:
         combination = tuple(sorted(columns))
         return (microtables[combination], combination)
 
@@ -47,7 +47,7 @@ def materialize_tree(_forest: Forest, columns: list[ColumnId]) -> tuple[list[Mic
         ],
     )
 
-    rows, combination = build_table(materialize_tree, forest, _dummy_metadata(4), clusters)
+    rows, combination = build_table(materialize_table, forest, _dummy_metadata(4), clusters)
 
     assert combination == (0, 1, 2, 3)
     assert rows == [
@@ -67,7 +67,7 @@ def materialize_tree(_forest: Forest, columns: list[ColumnId]) -> tuple[list[Mic
 def test_empty_microtables() -> None:
     forest = load_forest("dummy.csv")
 
-    def materialize_tree(_forest: Forest, columns: list[ColumnId]) -> tuple[list[MicrodataRow], Combination]:
+    def materialize_table(_forest: Forest, columns: list[ColumnId]) -> tuple[list[MicrodataRow], Combination]:
         combination = tuple(sorted(columns))
         return ([], combination)
 
@@ -78,7 +78,7 @@ def materialize_tree(_forest: Forest, columns: list[ColumnId]) -> tuple[list[Mic
         ],
     )
 
-    rows, combination = build_table(materialize_tree, forest, _dummy_metadata(3), clusters)
+    rows, combination = build_table(materialize_table, forest, _dummy_metadata(3), clusters)
 
     assert combination == (0, 1, 2)
     assert rows == []
@@ -98,7 +98,7 @@ def test_patching() -> None:
         (ColumnId(3),): _build_rows(col_d),
     }
 
-    def materialize_tree(_forest: Forest, columns: list[ColumnId]) -> tuple[list[MicrodataRow], Combination]:
+    def materialize_table(_forest: Forest, columns: list[ColumnId]) -> tuple[list[MicrodataRow], Combination]:
         combination = tuple(sorted(columns))
         return (microtables[combination], combination)
 
@@ -110,7 +110,7 @@ def materialize_tree(_forest: Forest, columns: list[ColumnId]) -> tuple[list[Mic
         ],
     )
 
-    rows, combination = build_table(materialize_tree, forest, _dummy_metadata(4), clusters)
+    rows, combination = build_table(materialize_table, forest, _dummy_metadata(4), clusters)
 
     assert combination == (0, 1, 2, 3)
     assert rows == [

tests/test_synthesizer.py

@@ -306,3 +306,32 @@ def test_pid() -> None:
     # Check that every value in syn_data['c1'] begins with either 'a' or 'b'
     for value in syn_data["c1"]:
         assert value.startswith("a") or value.startswith("b"), f"Value '{value}' does not start with 'a' or 'b'"
+
+
+def test_tree_builder() -> None:
+    # Create a dataframe with three columns
+    np.random.seed(42)  # For reproducible tests
+    df = pd.DataFrame(
+        {
+            "col1": np.random.choice(["A", "B", "C"], size=100),
+            "col2": np.random.randint(0, 10, size=100),
+            "col3": np.random.uniform(0, 1, size=100),
+        }
+    )
+
+    # Create synthesizer (this should build all trees)
+    syn = Synthesizer(df)
+
+    # Import necessary types for combinations
+    from itertools import combinations
+
+    from syndiffix.common import ColumnId
+
+    # Test all possible combinations of columns (1, 2, and 3 columns)
+    column_indices = [ColumnId(0), ColumnId(1), ColumnId(2)]  # Indices for col1, col2, col3
+
+    for r in range(1, 4):  # 1, 2, and 3 columns
+        for combination in combinations(column_indices, r):
+            # Convert to tuple as expected by the tree cache
+            tree = syn.forest._tree_cache.get(combination)
+            assert tree is not None, f"Tree not found for combination {combination}"