Preserve columns in isna, notna

rapidsai · Apr 12, 2024 · aba0c7c · aba0c7c
1 parent 227ea9b
commit aba0c7c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/python/cudf/cudf/core/frame.py b/python/cudf/cudf/core/frame.py
@@ -1120,7 +1120,9 @@ def isna(self):
         array([False, False,  True,  True, False, False])
         """
         data_columns = (col.isnull() for col in self._columns)
-        return self._from_data_like_self(zip(self._column_names, data_columns))
+        return self._from_data_like_self(
+            self._data._from_columns_like_self(data_columns)
+        )
 
     # Alias for isna
     isnull = isna
@@ -1199,7 +1201,9 @@ def notna(self):
         array([ True,  True, False, False,  True,  True])
         """
         data_columns = (col.notnull() for col in self._columns)
-        return self._from_data_like_self(zip(self._column_names, data_columns))
+        return self._from_data_like_self(
+            self._data._from_columns_like_self(data_columns)
+        )
 
     # Alias for notna
     notnull = notna

diff --git a/python/cudf/cudf/tests/test_dataframe.py b/python/cudf/cudf/tests/test_dataframe.py
@@ -10964,6 +10964,8 @@ def test_squeeze(axis, data):
     [
         lambda df: df.where(df < 2, 2),
         lambda df: df.nans_to_nulls(),
+        lambda df: df.isna(),
+        lambda df: df.notna(),
         lambda df: abs(df),
         lambda df: -df,
         lambda df: ~df,