diff --git a/buildinginitializer/buildinginitializer.py b/buildinginitializer/buildinginitializer.py
index 7333ce4acc24c4a0edba6719e489b9a39f8a73da..2cd0287bdf0b9c15728eae9e39d8dcb4b1196ea5 100644
--- a/buildinginitializer/buildinginitializer.py
+++ b/buildinginitializer/buildinginitializer.py
@@ -68,6 +68,10 @@ class BuildingInitializer:
         self.building_id_column = building_id_column
         self.building_geometry_column = building_geometry_column
 
+        self.outputs = set()
+        for rule in self.rule_handler:
+            self.outputs.update(rule.outputs)
+
         if source_db is None:
             raise ValueError("No source database object is found.")
         if building_db is None:
@@ -251,12 +255,8 @@ class BuildingInitializer:
                 Connected cursor to the OSM Replication database.
         """
 
-        # Define all unique column names
-        columns = set()
-        for _, args in result.items():
-            for key in args.keys():
-                columns.add(key)
-        columns = list(columns)
+        # The rule outputs define the columns that should be upserted.
+        columns = list(self.outputs)
 
         # Format values so that they can be uploaded as a batch on a PostGIS database.
         values = []
@@ -270,15 +270,15 @@ class BuildingInitializer:
         columns = ",".join(columns)
         # Format values
         values = ",".join(values)
-
         # Add all relevant columns to the database. If the building ID already exists, update
         # the attributes.
         sql_statement = f"""
-                INSERT INTO {self.building_table} ({self.building_id_column}, {columns})
-                VALUES {values}
-                ON CONFLICT ({self.building_id_column}) DO UPDATE
-                SET {on_conflict}
-                """
+            INSERT INTO {self.building_table} ({self.building_id_column}, {columns})
+            VALUES {values}
+            ON CONFLICT ({self.building_id_column}) DO UPDATE
+            SET {on_conflict}
+            """
+
         cursor.execute(sql_statement)
 
     @staticmethod
@@ -438,7 +438,7 @@ class BuildingInitializer:
             SELECT ol.key, ol.type, buildings.id AS building_id
             FROM osm_lands AS ol
             INNER JOIN buildings ON ST_Intersects(buildings.geom, ol.geometry)
-            WHERE buildings.geom && ol.geometry;
+            WHERE buildings.geom && ol.geometry
             """
         cursor.execute(sql_statement)
         for key, value, member in cursor:
diff --git a/tests/test_buildinginitializer.py b/tests/test_buildinginitializer.py
index d7b95f8091afc0c34143101635fb4b6ed46baa14..06101fec0addbc5f7cc22252e93ef541491c5b71 100644
--- a/tests/test_buildinginitializer.py
+++ b/tests/test_buildinginitializer.py
@@ -87,17 +87,20 @@ def test_upsert_building_information(source_db, building_db):
     """
 
     building_initializer = OSMInitializer({}, source_db=source_db, building_db=building_db)
+    building_initializer.outputs = {"occupancy", "storeys", "relation_id"}
     building_information = {
         173039036: {"occupancy": "TEST-1", "storeys": 10, "relation_id": None},
         352224937: {"storeys": 8, "relation_id": 5, "occupancy": "TEST-2"},
     }
     building_db.connect()
     building_initializer.upsert(building_information, building_db.cursor)
-    building_db.cursor.execute("""
+    building_db.cursor.execute(
+        """
         SELECT osm_id, occupancy, storeys, relation_id
         FROM obm_buildings
         WHERE osm_id IN (173039036,352224937)
-        """)
+        """
+    )
     for building_id, occupancy, storeys, relation_id in building_db.cursor:
         if building_id == 173039036:
             assert occupancy == "TEST-1"
@@ -121,11 +124,13 @@ def test_run(source_db, building_db, rules):
     building_ids = [352224937, 173039036, 27139591]
     building_initializer.run(building_ids=building_ids, batch_size=4)
     building_db.connect()
-    building_db.cursor.execute("""
+    building_db.cursor.execute(
+        """
         SELECT osm_id, storeys
         FROM obm_buildings
         WHERE osm_id IN (173039036,352224937)
-        """)
+        """
+    )
     for building_id, storeys in building_db.cursor:
         if building_id == 173039036:
             assert storeys == 4