Repository: vaexio/vaex

Issue:
Concatenate dataframe with manipulated column
Dear vaex team,

I would like to concatenate dataframes with a column that was converted to datetime64. However, I cannot access the column after concatenating. 

`import vaex`
`import numpy as np`
`x = np.arange(5) `
`y = ['2020-01-23','2020-01-24','2020-01-25','2020-01-26','2020-01-27']`
`df = vaex.from_arrays(x=x, y=y)`

Manipulation, e.g.
`df['y'] = df['y'].astype('datetime64[ms]')`
or
`df['y'] = df['y']+df['y']`

Concatenation:
`concatenated = df.concat(df)`

The original dataframe is fine:
`df.get_column_names`
`df['y']`

while the concatenated dataframe produces errors:
`concatenated.get_column_names`
KeyError: "Unknown variables or column: 'str_cat(__y, __y)'"
`concatenated['y']`
NameError: name '__y' is not defined

Is there a problem in how I use the concatenation or is it a bug?
Kind regards,
Maya


Discussion:
Hi @MayaStr 

Thank you for reporting this, and providing such a nice and reproducible example. We will try to attend to this.

Cheers, 
Jovan.

Resolving patch:
diff --git a/packages/vaex-core/vaex/column.py b/packages/vaex-core/vaex/column.py
--- a/packages/vaex-core/vaex/column.py
+++ b/packages/vaex-core/vaex/column.py
@@ -1,5 +1,6 @@
 import logging
 import os
+import warnings
 
 import six
 import numpy as np
@@ -145,13 +146,15 @@ def __init__(self, expressions, dtype=None):
                         except:
                             self.fill_value = expressions[0].values.fill_value
 
-            # np.datetime64 and find_common_type don't mix very well
             any_strings = any([dtype == str_type for dtype in dtypes])
             if any_strings:
                 self.dtype = str_type
             else:
+                # np.datetime64/timedelta64 and find_common_type don't mix very well
                 if all([dtype.type == np.datetime64 for dtype in dtypes]):
                     self.dtype = dtypes[0]
+                elif all([dtype.type == np.timedelta64 for dtype in dtypes]):
+                    self.dtype = dtypes[0]
                 else:
                     if all([dtype == dtypes[0] for dtype in dtypes]):  # find common types doesn't always behave well
                         self.dtype = dtypes[0]
@@ -183,21 +186,24 @@ def __len__(self):
 
     def trim(self, i1, i2):
         start, stop = i1, i2
-        i = 0
-        offset = 0
+        i = 0  # the current expression
+        offset = 0  # and the offset the current expression has wrt the dataframe
+        # find the first expression we overlap with
         while start >= offset + len(self.expressions[i].df):
             offset += len(self.expressions[i].df)
             i += 1
-        if stop <= offset + len(self.expressions[i].df):  # single df
-            expressions = [self.expressions[i][start-offset:stop-offset+1]]
-        else:  # multile expressions
-            expressions = [self.expressions[i][start-offset:]]
-            while stop >= offset + len(self.expressions[i].df):
+        expressions = [self.expressions[i][start-offset:]]
+        offset += len(self.expressions[i].df)
+        if stop > offset:
+            # we need more expression
+            i += 1
+            # keep adding complete expression till we find the it to be the last
+            while stop > offset + len(self.expressions[i].df):
                 offset += len(self.expressions[i].df)
-                i += 1
                 expressions.append(self.expressions[i])
-            # if stop <= offset + len(self.expressions[i])
-            expressions.append(self.expressions[i][start-offset:stop-offset+1])
+                i += 1
+            if stop > offset:  # add the tail part
+                expressions.append(self.expressions[i][:stop-offset])
         return ColumnConcatenatedLazy(expressions, self.dtype)
 
     def __getitem__(self, slice):
diff --git a/packages/vaex-core/vaex/dataframe.py b/packages/vaex-core/vaex/dataframe.py
--- a/packages/vaex-core/vaex/dataframe.py
+++ b/packages/vaex-core/vaex/dataframe.py
@@ -2945,7 +2945,10 @@ def add_column(self, name, f_or_array, dtype=None):
                 self._dtypes_override[valid_name] = dtype
             else:
                 if isinstance(ar, np.ndarray) and ar.dtype.kind == 'O':
-                    types = list({type(k) for k in ar if np.all(k == k) and k is not None})
+                    ar_data = ar
+                    if np.ma.isMaskedArray(ar):
+                        ar_data = ar.data
+                    types = list({type(k) for k in ar_data if np.all(k == k) and k is not None})
                     if len(types) == 1 and issubclass(types[0], six.string_types):
                         self._dtypes_override[valid_name] = str_type
                     if len(types) == 0:  # can only be if all nan right?
@@ -5833,30 +5836,28 @@ class DataFrameConcatenated(DataFrameLocal):
 
     def __init__(self, dfs, name=None):
         super(DataFrameConcatenated, self).__init__(None, None, [])
-        self.dfs = dfs
+        # to reduce complexity, we 'extract' the dataframes (i.e. remove filter)
+        self.dfs = dfs = [df.extract() for df in dfs]
         self.name = name or "-".join(df.name for df in self.dfs)
         self.path = "-".join(df.path for df in self.dfs)
         first, tail = dfs[0], dfs[1:]
-        for df in dfs:
-            assert df.filtered is False, "we don't support filtering for concatenated DataFrames"
-        for column_name in first.get_column_names(virtual=False):
-            if all([column_name in df.get_column_names(virtual=False) for df in tail]):
+        for column_name in first.get_column_names(virtual=False, hidden=True):
+            if all([column_name in df.get_column_names(virtual=False, hidden=True) for df in tail]):
                 self.column_names.append(column_name)
         self.columns = {}
-        for column_name in self.get_column_names(virtual=False):
+        for column_name in self.get_column_names(virtual=False, hidden=True):
             self.columns[column_name] = ColumnConcatenatedLazy([df[column_name] for df in dfs])
             self._save_assign_expression(column_name)
 
         for name in list(first.virtual_columns.keys()):
             if all([first.virtual_columns[name] == df.virtual_columns.get(name, None) for df in tail]):
-                self.virtual_columns[name] = first.virtual_columns[name]
+                self.add_virtual_column(name, first.virtual_columns[name])
                 self.column_names.append(name)
             else:
                 self.columns[name] = ColumnConcatenatedLazy([df[name] for df in dfs])
                 self.column_names.append(name)
             self._save_assign_expression(name)
 
-
         for df in dfs[:1]:
             for name, value in list(df.variables.items()):
                 if name not in self.variables:
@@ -5868,6 +5869,7 @@ def __init__(self, dfs, name=None):
         self._index_end = self._length_unfiltered
 
     def is_masked(self, column):
+        column = _ensure_string_from_expression(column)
         if column in self.columns:
             return self.columns[column].is_masked
         else:


Test patch:
diff --git a/tests/common.py b/tests/common.py
--- a/tests/common.py
+++ b/tests/common.py
@@ -116,18 +116,32 @@ def ds_trimmed():
     ds.set_active_range(2, 12)
     return ds.trim()
 
-@pytest.fixture(params=['ds_filtered', 'ds_half', 'ds_trimmed', 'ds_remote'])
-def ds(request, ds_filtered, ds_half, ds_trimmed, ds_remote):
-    named = dict(ds_filtered=ds_filtered, ds_half=ds_half, ds_trimmed=ds_trimmed, ds_remote=ds_remote)
+@pytest.fixture()
+def df_trimmed(ds_trimmed):
+    return ds_trimmed
+
+@pytest.fixture()
+def df_concat(ds_trimmed):
+    df = ds_trimmed
+    df1 = df[:2]   # length 2
+    df2 = df[2:3]  # length 1
+    df3 = df[3:7]  # length 4
+    df4 = df[7:10]  # length 3
+    return vaex.concat([df1, df2, df3, df4])
+
+
+@pytest.fixture(params=['ds_filtered', 'ds_half', 'ds_trimmed', 'ds_remote', 'df_concat'])
+def ds(request, ds_filtered, ds_half, ds_trimmed, ds_remote, df_concat):
+    named = dict(ds_filtered=ds_filtered, ds_half=ds_half, ds_trimmed=ds_trimmed, ds_remote=ds_remote, df_concat=df_concat)
     return named[request.param]
 
 @pytest.fixture
 def df(ds):
     return ds
 
-@pytest.fixture(params=['ds_filtered', 'ds_half', 'ds_trimmed'])
-def ds_local(request, ds_filtered, ds_half, ds_trimmed):
-    named = dict(ds_filtered=ds_filtered, ds_half=ds_half, ds_trimmed=ds_trimmed)
+@pytest.fixture(params=['ds_filtered', 'ds_half', 'ds_trimmed', 'df_concat'])
+def ds_local(request, ds_filtered, ds_half, ds_trimmed, df_concat):
+    named = dict(ds_filtered=ds_filtered, ds_half=ds_half, ds_trimmed=ds_trimmed, df_concat=df_concat)
     return named[request.param]
 
 @pytest.fixture
@@ -190,7 +204,7 @@ def create_base_ds():
     obj_mask = np.array([False] * 9 + [True])
     obj = nm.copy().astype('object')
     obj[2:12] = np.ma.MaskedArray(data=obj_data, mask=obj_mask, dtype='object')
-    dataset.add_column("obj", obj)
+    dataset.add_column("obj", obj, dtype=np.dtype('O'))
 
     booleans = np.ones(21, dtype=np.bool)
     booleans[[4, 6, 8, 14, 16, 19]] = False
diff --git a/tests/concat_test.py b/tests/concat_test.py
--- a/tests/concat_test.py
+++ b/tests/concat_test.py
@@ -1,6 +1,7 @@
+import pytest
 import vaex
 import numpy as np
-from functools import reduce
+
 
 def test_concat():
     x1, y1, z1 = np.arange(3), np.arange(3,0,-1), np.arange(10,13)
@@ -38,6 +39,7 @@ def test_concat():
     assert (np.array(dd2.evaluate('z')) == np.array(ds.evaluate('z'))).all()
     assert (np.array(dd2.evaluate('w')) == np.array(ds.evaluate('w'))).all()
 
+
 def test_concat_unequals_virtual_columns():
     ds1 = vaex.from_scalars(x=1, y=2)
     ds2 = vaex.from_scalars(x=2, y=3)
@@ -51,6 +53,7 @@ def test_concat_unequals_virtual_columns():
     assert ds.w.tolist() == [1+2, 2+3]
     assert ds.z.tolist() == [1+2, 2*3]
 
+
 def test_concat_arrow_strings():
     df1 = vaex.from_arrays(x=vaex.string_column(['aap', 'noot', 'mies']))
     df2 = vaex.from_arrays(x=vaex.string_column(['a', 'b', 'c']))
@@ -58,6 +61,7 @@ def test_concat_arrow_strings():
     assert df.dtype('x') == df1.dtype('x')
     assert df.x.tolist() == ['aap', 'noot', 'mies', 'a', 'b', 'c']
 
+
 def test_concat_mixed_types():
     x1 = np.zeros(3) + np.nan
     x2 = vaex.string_column(['hi', 'there'])
@@ -70,3 +74,33 @@ def test_concat_mixed_types():
     assert df[2:4].x.tolist() == [None, 'hi']
     assert df[3:4].x.tolist() == ['hi']
     assert df[3:5].x.tolist() == ['hi', 'there']
+
+
+def test_dtypes(df_concat):
+    assert df_concat.timedelta.dtype.kind == 'm'
+    assert df_concat.datetime.dtype.kind == 'M'
+
+
+def test_hidden(df_trimmed):
+    xlist = (df_trimmed.x + 1).tolist()
+    df = df_trimmed.copy()
+    df['x'] = df.x + 1
+    dfc = df.concat(df)  # make sure concat copies hidden columns
+    assert dfc.x.tolist() == xlist + xlist
+
+
+def test_concat_filtered(df_trimmed):
+    df = df_trimmed
+    df1 = df[df.x < 5]
+    df2 = df[df.x >= 5]
+    dfc = df1.concat(df2)
+    assert dfc.x.tolist() == df.x.tolist()
+
+
+@pytest.mark.parametrize("i1", list(range(0, 6)))
+@pytest.mark.parametrize("length", list(range(1, 6)))
+def test_sliced_concat(i1, length, df_concat):
+    i2 = i1 + length
+    x = df_concat.x.tolist()
+    df = df_concat[i1:i2]
+    assert df.x.tolist() == x[i1:i2]