2 files changed, 7 insertions, 12 deletions
diff --git a/src/sqlite.rs b/src/sqlite.rs
index c274b63..2185e66 100644
--- a/src/sqlite.rs
+++ b/src/sqlite.rs
@@ -1,18 +1,13 @@
 use rusqlite;
 
-use std::collections::HashMap;
-
-
-#[derive(Debug)]
-pub struct Zero;
+use std::collections::HashSet;
 
 
 pub fn get_column_names(
     dbconn: &rusqlite::Connection,
     table_name: &str,
-// TODO: Use a HashSet instead
-) -> Result<HashMap<String, Zero>, crate::Error> {
-    let mut column_names = HashMap::new();
+) -> Result<HashSet<String>, crate::Error> {
+    let mut column_names = HashSet::new();
 
     let mut stmt = dbconn.prepare(
         &format!(
@@ -32,7 +27,7 @@ pub fn get_column_names(
     for row_result in rows {
         let row = row_result?;
 
-        column_names.insert(row, Zero{});
+        column_names.insert(row);
     }
 
     Ok(column_names)
diff --git a/src/yaml.rs b/src/yaml.rs
index 9b0988c..af71779 100644
--- a/src/yaml.rs
+++ b/src/yaml.rs
@@ -1,7 +1,7 @@
 use rusqlite;
 use yaml_rust::yaml;
 
-use std::collections::HashMap;
+use std::collections::HashSet;
 
 
 mod sql;
@@ -18,7 +18,7 @@ pub fn extract(
     doc: &mut yaml::Yaml,
     tx: &rusqlite::Transaction,
     table_name: &str,
-    table_columns: &HashMap<String, crate::sqlite::Zero>,
+    table_columns: &HashSet<String>,
 ) -> Result<(), crate::Error> {
     match doc {
         yaml::Yaml::Array(ref mut array) => {
@@ -30,7 +30,7 @@ pub fn extract(
             use std::borrow::Cow;
 
             let keys: Vec<yaml::Yaml> = hash.keys().map(|k| k.clone()).collect();
-            let columns_as_yaml: Vec<yaml::Yaml> = table_columns.keys()
+            let columns_as_yaml: Vec<yaml::Yaml> = table_columns.iter()
                 .map(|c| yaml::Yaml::from_str(c))
                 .collect();