From 85004078bc60d664614a9581cecc5ac3ed0280af Mon Sep 17 00:00:00 2001
From: Michael Wang <michaelwang0905@icloud.com>
Date: Mon, 7 Mar 2022 14:59:32 -0800
Subject: [PATCH] Add test for memory test utils

---
 python/cudf/cudf/testing/_utils.py     |  2 +
 python/cudf/cudf/tests/test_testing.py | 62 +++++++++++++++++++++++++-
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/python/cudf/cudf/testing/_utils.py b/python/cudf/cudf/testing/_utils.py
index ab6e3466b25..f6b5e0f3ccc 100644
--- a/python/cudf/cudf/testing/_utils.py
+++ b/python/cudf/cudf/testing/_utils.py
@@ -344,6 +344,8 @@ def assert_column_memory_eq(
     """
     assert lhs.base_data_ptr == rhs.base_data_ptr
     assert lhs.base_mask_ptr == rhs.base_mask_ptr
+    assert lhs.base_size == rhs.base_size
+    assert lhs.offset == rhs.offset
     assert lhs.size == rhs.size
     assert len(lhs.base_children) == len(rhs.base_children)
     for lhs_child, rhs_child in zip(lhs.base_children, rhs.base_children):
diff --git a/python/cudf/cudf/tests/test_testing.py b/python/cudf/cudf/tests/test_testing.py
index 0b27c562d75..4dc4d86d94c 100644
--- a/python/cudf/cudf/tests/test_testing.py
+++ b/python/cudf/cudf/tests/test_testing.py
@@ -1,7 +1,8 @@
-# Copyright (c) 2020, NVIDIA CORPORATION.
+# Copyright (c) 2020-2022, NVIDIA CORPORATION.
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 import pytest
 
 import cudf
@@ -11,10 +12,28 @@
     assert_index_equal,
     assert_series_equal,
 )
-from cudf.testing._utils import NUMERIC_TYPES, OTHER_TYPES, assert_eq
+from cudf.testing._utils import (
+    NUMERIC_TYPES,
+    OTHER_TYPES,
+    assert_column_memory_eq,
+    assert_column_memory_ne,
+    assert_eq,
+)
 from cudf.testing.testing import assert_column_equal
 
 
+@pytest.fixture(
+    params=[
+        pa.array([*range(10)]),
+        pa.array(["hello", "world", "rapids", "AI"]),
+        pa.array([[1, 2, 3], [4, 5], [6], [], [7]]),
+        pa.array([{"f0": "hello", "f1": 42}, {"f0": "world", "f1": 3}]),
+    ]
+)
+def arrow_arrays(request):
+    return request.param
+
+
 @pytest.mark.parametrize("rdata", [[1, 2, 5], [1, 2, 6], [1, 2, 5, 6]])
 @pytest.mark.parametrize("exact", ["equiv", True, False])
 @pytest.mark.parametrize("check_names", [True, False])
@@ -369,3 +388,42 @@ def test_basic_scalar_equality(left, right):
 def test_basic_scalar_inequality(left, right):
     with pytest.raises(AssertionError, match=r".*not (almost )?equal.*"):
         assert_eq(left, right)
+
+
+def test_assert_column_memory_basic(arrow_arrays):
+    left = cudf.core.column.ColumnBase.from_arrow(arrow_arrays)
+    right = cudf.core.column.ColumnBase.from_arrow(arrow_arrays)
+
+    with pytest.raises(AssertionError):
+        assert_column_memory_eq(left, right)
+    assert_column_memory_ne(left, right)
+
+
+def test_assert_column_memory_slice(arrow_arrays):
+    col = cudf.core.column.ColumnBase.from_arrow(arrow_arrays)
+    left = col[0:1]
+    right = col[1:2]
+
+    with pytest.raises(AssertionError):
+        assert_column_memory_eq(left, right)
+    assert_column_memory_ne(left, right)
+
+    with pytest.raises(AssertionError):
+        assert_column_memory_eq(left, col)
+    assert_column_memory_ne(left, col)
+
+    with pytest.raises(AssertionError):
+        assert_column_memory_eq(right, col)
+    assert_column_memory_ne(right, col)
+
+
+def test_assert_column_memory_basic_same(arrow_arrays):
+    data = cudf.core.column.ColumnBase.from_arrow(arrow_arrays)
+    buf = cudf.core.buffer.Buffer(data=data.base_data, owner=data)
+
+    left = cudf.core.column.build_column(buf, dtype=np.int32)
+    right = cudf.core.column.build_column(buf, dtype=np.int32)
+
+    assert_column_memory_eq(left, right)
+    with pytest.raises(AssertionError):
+        assert_column_memory_ne(left, right)