diff --git a/src/algorithms/fd_algorithm.cpp b/src/algorithms/fd_algorithm.cpp
index d1997e7361..575f32583c 100644
--- a/src/algorithms/fd_algorithm.cpp
+++ b/src/algorithms/fd_algorithm.cpp
@@ -63,7 +63,14 @@ std::vector<Column const*> FDAlgorithm::GetKeys() const {
         }
     }
 
-    size_t const number_of_columns = input_table_->GetNumberOfColumns();
+    if (fds_count_per_col.empty()) return keys;
+    assert(fds_count_per_col.begin()->first->GetSchema() != nullptr);
+    size_t const number_of_columns = fds_count_per_col.begin()->first->GetSchema()->GetNumColumns();
+    assert(std::all_of(
+            ++fds_count_per_col.begin(), fds_count_per_col.end(),
+            [first_schema = fds_count_per_col.begin()->first->GetSchema()](auto const& kv_pair) {
+                return kv_pair.first->GetSchema() == first_schema;
+            }));
     for (auto const& [col, num]: fds_count_per_col) {
         if (num + 1 + cols_of_equal_values == number_of_columns) {
             keys.push_back(col);