From 915a7c4195d2e42a52c8f61d8c42f3f779ddfbf9 Mon Sep 17 00:00:00 2001
From: Greg Hogan <gregjhogan@gmail.com>
Date: Mon, 20 Apr 2020 16:56:38 -0700
Subject: [PATCH] ColumnStoreWriter.add_dict() support for multiple types

old-commit-hash: e23f1fcf0ac89c1e9d3ee6c4a64cf49745b919a4
---
 common/column_store.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/common/column_store.py b/common/column_store.py
index cb72076eea..14d7accc3d 100644
--- a/common/column_store.py
+++ b/common/column_store.py
@@ -158,7 +158,7 @@ class ColumnStoreWriter():
     # TODO(mgraczyk): This implementation will need to change if we add zip or compression.
     return ColumnStoreWriter(os.path.join(self._path, group_name))
 
-  def add_dict(self, data, dtype=None, compression=False, overwrite=False):
+  def add_dict(self, data, dtypes=None, compression=False, overwrite=False):
     # default name exists to have backward compatibility with equivalent directory structure
     npy_path = os.path.join(self._path, "columnstore")
     mkdirs_exists_ok(os.path.dirname(npy_path))
@@ -166,6 +166,7 @@ class ColumnStoreWriter():
     flat_dict = dict()
     _flatten_dict(flat_dict, "", data)
     for k, v in flat_dict.items():
+      dtype = dtypes[k] if dtypes is not None and k in dtypes else None
       flat_dict[k] = np.array(v, copy=False, dtype=dtype)
 
     if overwrite: