Bugfix: Fix column subsetting bug in numpy deserializer

arctic/serialization/numpy_arrays.py

@@ -145,7 +145,7 @@ def objify(self, doc, columns=None, as_df=True):
         """
         cols = columns or doc[METADATA][COLUMNS]
         data = {}
-        valid_columns = doc[METADATA][LENGTHS]
+        valid_columns = set(cols).intersection(doc[METADATA][LENGTHS])
         missing_columns = set(cols).difference(valid_columns)
         for col in valid_columns:
             d = decompress(doc[DATA][doc[METADATA][LENGTHS][col][0]: doc[METADATA][LENGTHS][col][1] + 1])
@@ -158,11 +158,12 @@ def objify(self, doc, columns=None, as_df=True):
                 d = ma.masked_array(d, mask)
             data[col] = d
 
-        for col in missing_columns:
-            # if there is missing data in a chunk, we can default to NaN and
-            empty = np.empty(len(d))
-            empty[:] = np.nan
-            data[col] = empty
+        if data:
+            for col in missing_columns:
+                # if there is missing data in a chunk, we can default to NaN and
+                empty = np.empty(len(d))
+                empty[:] = np.nan
+                data[col] = empty
 
         if as_df:
             return pd.DataFrame(data)
@@ -218,15 +219,17 @@ def deserialize(self, data, columns=None):
         meta = data[0][METADATA] if isinstance(data, list) else data[METADATA]
         index = INDEX in meta
 
+        if not isinstance(data, list):
+            data = [data]
+
         if columns:
             if index:
                 columns = columns[:]
                 columns.extend(meta[INDEX])
             if len(columns) > len(set(columns)):
                 raise Exception("Duplicate columns specified, cannot de-serialize")
-
-        if not isinstance(data, list):
-            data = [data]
+        else:
+            columns = [col for doc in data for col in doc[METADATA][COLUMNS]]
 
         df = defaultdict(list)
 

tests/unit/serialization/test_numpy_arrays.py

@@ -21,6 +21,29 @@ def test_with_strings():
     assert_frame_equal(f.objify(f.docify(df)), df)
 
 
+def test_frame_converter_with_all_valid_column_subset():
+    f = FrameConverter()
+    df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)),
+                      columns=list('ABCD'))
+
+    assert_frame_equal(f.objify(f.docify(df), columns=['A', 'B']), df[['A', 'B']])
+
+
+def test_frame_converter_with_some_invalid_column_subset():
+    f = FrameConverter()
+    df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)),
+                      columns=list('ABCD'))
+    expected = pd.DataFrame({'A': df['A'], 'N': np.nan})
+    assert_frame_equal(f.objify(f.docify(df), columns=['A', 'N']), expected[['A', 'N']])
+
+
+def test_frame_converter_with_no_valid_column_subset():
+    f = FrameConverter()
+    df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)),
+                      columns=list('ABCD'))
+    assert f.objify(f.docify(df), columns=['N']).empty
+
+
 def test_with_objects_raises():
     class Example(object):
         def __init__(self, data):
@@ -53,6 +76,40 @@ def test_with_index():
     assert_frame_equal(df, n.deserialize(a))
 
 
+def test_invalid_column_subset_with_index():
+    df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)),
+                      columns=list('ABCD'))
+    df = df.set_index(['A'])
+    n = FrametoArraySerializer()
+    a = n.serialize(df)
+    expected = pd.DataFrame({'N': np.nan}, index=df.index)
+    assert_frame_equal(expected, n.deserialize(a, columns=['N']))
+
+
+@pytest.mark.parametrize('index', [True, False])
+@pytest.mark.parametrize('columns', [None, ['B', 'D'], ['D', 'B']])
+def test_multiple_data_input_different_columns(columns, index):
+    df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)),
+                      columns=list('ABCD'))
+    if index:
+        df = df.set_index(['A'])
+    n = FrametoArraySerializer()
+    a = n.serialize(df[['B']])
+    b = n.serialize(df[['D']])
+    expected = df[['B']].append(df[['D']], ignore_index=not index)
+    assert_frame_equal(expected, n.deserialize([a, b], columns=columns))
+
+
+@pytest.mark.parametrize('column', [['B'], ['D']])
+def test_multiple_data_input_with_no_index_and_invalid_column(column):
+    df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)),
+                      columns=list('ABCD'))
+    n = FrametoArraySerializer()
+    a = n.serialize(df[['B']])
+    b = n.serialize(df[['D']])
+    assert_frame_equal(df[column], n.deserialize([a, b], columns=column))
+
+
 def test_with_nans():
     df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)),
                       columns=list('ABCD'))