From f7f8e976b78b31a0a3bcf330a82f4ef9d17e9c76 Mon Sep 17 00:00:00 2001
From: Greg Hogan <gregjhogan@gmail.com>
Date: Mon, 20 Apr 2020 16:25:43 -0700
Subject: [PATCH] ColumnStoreReader support for intermediate keys of flat
 dictionaries

old-commit-hash: cf131ee8e92f477992ed087a214c832c76ee4714
---
 common/column_store.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/common/column_store.py b/common/column_store.py
index 6f162db1b1..cb72076eea 100644
--- a/common/column_store.py
+++ b/common/column_store.py
@@ -36,7 +36,12 @@ class ColumnStoreReader():
     if self._is_dict:
       if not self._np_data:
         self._np_data = self._load(os.path.join(self._path, 'columnstore'))
-      return self._np_data.keys()
+      # return top level keys by matching on prefix and splitting on '/'
+      return {
+        k[len(self._prefix):].split('/')[0]
+        for k in self._np_data.keys()
+        if k.startswith(self._prefix)
+      }
 
     return list(self._keys)