import keyword
import re
import unittest
from collections import namedtuple

import numpy as np
import pandas as pd
import pytest
from numpy.testing import assert_array_almost_equal, assert_array_equal

import riptable as rt
from riptable import (
    Categorical,
    Dataset,
    Date,
    DateTimeNano,
    FastArray,
    NumpyCharTypes,
    Struct,
    TimeSpan,
    utcnow,
)
from riptable.rt_enum import (
    FILTERED_LONG_NAME,
    INVALID_DICT,
    CategoryMode,
    TypeRegister,
)
from riptable.rt_numpy import arange, isnan, logical, tile
from riptable.rt_utils import describe

from .test_utils import get_all_categorical_data


def almost_eq(ds, arr, places=5, rtol=1e-5, atol=1e-5, equal_nan=True):

    epsilon = 10 ** (-places)
    rtol = epsilon

    if ds.shape != arr.shape:
        return False
    for _j, _col in enumerate(ds):
        _dsc = ds[_col]
        _arrc = arr[:, _j]
        if np.issubdtype(_dsc.dtype, np.floating):
            # N.B.: Do not compare the elements that both numpy and riptable give nan or inf.
            #       In the case of modulo by zero, numpy gives nan but riptable gives inf.
            finite_mask = ~((~np.isfinite(_dsc)) & (~np.isfinite(_arrc)))
            if not np.allclose(_dsc[finite_mask], _arrc[finite_mask], rtol, atol, equal_nan):
                return False
        else:
            try:
                if not (_dsc == _arrc).all():
                    return False
            except TypeError:
                return False
    return True


# def almost_eq(ds, arr, places=5):
#    epsilon = 10 ** (-places)
#    rval = True
#    if ds.shape != arr.shape:
#        return False
#    for _j, _col in enumerate(ds):
#        _dsc = ds[_col]
#        _arrc = arr[:, _j]
#        for v1, v2 in zip(_dsc, _arrc.astype(_dsc.dtype.type)):
#            if isinstance(v1, str):
#                rval &= (v1 == v2)
#            else:
#                if not (np.isfinite(v1) and np.isfinite(v2)):
#                    continue
#                rval &= (abs(v1 - v2) < epsilon)
#    return rval


class TestDataset(unittest.TestCase):
    @staticmethod
    def get_basic_dataset(nrows=10):
        return Dataset({_k: list(range(_i * nrows, (_i + 1) * nrows)) for _i, _k in enumerate("abcdefghijklmnop")})

    @staticmethod
    def get_arith_dataset(include_strings=False, dtype=None):
        ds = Dataset(
            {
                "A": [0, 6, 9],
                "B": [1.2, 3.1, 9.6],
                "G": [-1.6, 2.7, 4.6],
                "C": [2.4, 6.2, 19.2],
            }
        )
        if dtype:
            for _k, _v in ds.items():
                ds[_k] = _v.astype(dtype)
        if include_strings:
            ds.S = ["string_00", "string_01", "string_02"]
            ds.U = ["ℙƴ☂ℌøἤ_00", "ℙƴ☂ℌøἤ_01", "ℙƴ☂ℌøἤ_02"]
        return ds

    def test_col_ctor_01(self):
        list1 = [(_k, [_i]) for _i, _k in enumerate("bac")]
        with self.assertRaises(TypeError):
            _ = Dataset({"a": (1,)})
        with self.assertRaises(TypeError):
            _ = Dataset(list1)
        with self.assertRaises(ValueError):
            _ = Dataset({"a": [1, 2], "b": 3})
        list2 = [Dataset({_k: [_i]}) for _i, _k in enumerate("bac")]
        with self.assertRaises(TypeError):  # for this use concat_rows() or concat_columns()
            _ = Dataset(list2)
        with self.assertRaises(TypeError):  # for this use concat_rows() or concat_columns()
            _ = Dataset([Struct(_d.asdict()) for _d in list2])
        with self.assertRaises(TypeError):  # this one will be supported later w/ a sep. method
            _ = Dataset([_d.asdict() for _d in list2])
        assert (Dataset(None) == Dataset({})).all(
            axis=None
        ), f"Dataset default argument for 'None' input value should create an empty Dataset."
        list3 = [{_k: [_i]} for _i, _k in enumerate("bac")]
        with self.assertRaises(TypeError):
            _ = Dataset(list3)
        dict1 = {_k: [_i] for _i, _k in enumerate("bac")}
        st = Struct(dict1)
        ds = Dataset(st)
        self.assertEqual(ds.shape, (1, 3))
        for _i, _k in enumerate("bac"):
            self.assertEqual(ds[0, _k], st[_k][0])
        ds1 = Dataset({"a": [1], "b": 2})
        ds2 = Dataset({"a": [[1]], "b": 2})
        self.assertTrue((ds1 == ds2).all(axis=None))
        self.assertEqual(ds2.shape, (1, 2))
        self.assertEqual(ds2.a.shape, (1,))
        self.assertEqual(ds2.b.shape, (1,))
        ds = Dataset()
        self.assertEqual(ds.shape, (0, 0))

    # def test_col_ctor_02(self):
    #    hold_WarnOnInvalidNames = Struct.WarnOnInvalidNames
    #    Struct.WarnOnInvalidNames = False
    #    tempsave = Struct.AllowAnyName
    #    Struct.AllowAnyName = False
    #    with self.assertRaises(ValueError):
    #        _ = Dataset({'a': [0], '0': [1]})
    #    with self.assertRaises(ValueError):
    #        _ = Dataset({'a': [0], 'a-b-c': [1]})
    #    # we auto capitalize now
    #    #for kwd in keyword.kwlist:
    #    #    _ = Dataset({'a': [0], kwd: [1]})
    #    #for memb in dir(Dataset):
    #    #    with self.assertRaises(ValueError):
    #    #        _ = Dataset({'a': [0], memb: [1]})
    #    Struct.WarnOnInvalidNames = hold_WarnOnInvalidNames
    #    Struct.AllowAnyName = tempsave

    def test_col_ctor_02(self):
        arr = arange(5)
        inv_keys = ["True", "False", "None", "size"]
        inv_dict = {k: arr for k in inv_keys}
        with self.assertWarns(UserWarning):
            ds = Dataset(inv_dict)
        inv_keys[-1] = "Size"

        self.assertTrue(bool(np.all(inv_keys == list(ds))))
        for k in inv_keys:
            self.assertTrue(bool(np.all(ds[k] == arr)))

    def test_col_ctor_03(self):
        tempsave = Struct.AllowAnyName
        Struct.AllowAnyName = False
        hold_WarnOnInvalidNames = Struct.WarnOnInvalidNames
        Struct.WarnOnInvalidNames = True
        with self.assertWarns(UserWarning):
            _ = Dataset({"a": [0], "0": [1]})
        with self.assertWarns(UserWarning):
            _ = Dataset({"a": [0], "a-b-c": [1]})
        for kwd in keyword.kwlist:
            with self.assertWarns(UserWarning):
                _ = Dataset({"a": [0], kwd: [1]})
        # note: dir has changed
        # for memb in dir(Dataset):
        #    with self.assertWarns(UserWarning, msg=memb):
        #        # We expect really stupid errors resulting from not avoiding invalid names.
        #        try:
        #            _ = Dataset({'a': [0], memb: [1]})
        #        except (TypeError, AttributeError):
        #            pass
        Struct.WarnOnInvalidNames = hold_WarnOnInvalidNames
        Struct.AllowAnyName = tempsave

    def test_col_ctor_04(self):
        hold_UseFastArray = Struct.UseFastArray
        dd = {"a": FastArray([0, 1]), "b": np.array([10, 11])}
        Struct.set_fast_array(True)
        ds1 = Dataset(dd)
        self.assertIsInstance(ds1.a, FastArray)
        self.assertIsInstance(ds1.b, FastArray)
        Struct.set_fast_array(False)
        ds2 = Dataset(dd)
        self.assertIsInstance(ds2.a, np.ndarray)
        self.assertIsInstance(ds2.b, np.ndarray)
        self.assertNotIsInstance(ds2.a, FastArray)
        self.assertNotIsInstance(ds2.b, FastArray)
        Struct.set_fast_array(hold_UseFastArray)

    def test_col_ctor_05(self):
        dd = {"a": list(range(4)), "b": list("abcd"), "c": list("μεαν")}
        ds = Dataset(dd)
        self.assertEqual(ds.shape, (4, 3))
        dd = {
            "a": np.array(list(range(4))),
            "b": np.array(list("abcd"), dtype=object),
            "c": np.array(list("μεαν"), dtype=object),
        }
        ds = Dataset(dd)
        self.assertEqual(ds.shape, (4, 3))

    def test_astype(self):
        ds = self.get_arith_dataset(include_strings=True)
        dtypes0 = ds.dtypes
        ds1 = ds.astype(np.int16)
        for _k in list("ABCG"):
            self.assertEqual(ds1[_k].dtype, np.dtype(np.int16))
        for _k in list("US"):
            self.assertEqual(ds1[_k].dtype, dtypes0[_k])
        ds.S = ds.A.astype(str)
        ds.S = ds.B.astype(str)
        # 9/28/2018 SJK: Dataset no longer flips Unicode arrays to Categorical, removed unicode column
        ds1 = ds.astype(np.float16, ignore_non_computable=True)
        for _k in list("ABCG"):
            self.assertEqual(ds1[_k].dtype, np.dtype(np.float16))

    def test_view_vs_copy_01(self):
        # Here, ensure that all standard indexing produces a view, not a copy;
        # the chief exception being "fancy row-indexing".
        # Also ensure that Dataset.copy(deep=True) does the correct thing!
        # First, numeric only.  Create two datasets, identical, but from totally different memory.
        ds1 = self.get_basic_dataset(nrows=1000)
        ds1[10:] = ds1[10:].astype(np.float32)
        ds2 = self.get_basic_dataset(nrows=1000)
        ds2[10:] = ds2[10:].astype(np.float32)
        self.assertTrue((ds1 == ds2).all(axis=None))
        ds1.a[5:100:5] *= -1
        self.assertTrue((ds1.a[5:100:5] == ds2.a[5:100:5] * -1).all())
        self.assertTrue((ds1[1:] == ds2[1:]).all(axis=None))
        ds1[5:100:5, ["b", "d", "e"]] += -2
        self.assertTrue((ds1[5:100:5, ["b", "d", "e"]] == ds2[5:100:5, ["b", "d", "e"]] + -2).all(axis=None))
        self.assertTrue((ds1.c == ds2.c).all())
        self.assertTrue((ds1[5:] == ds2[5:]).all(axis=None))
        ds1[5:100:5, 5:7] = 84
        self.assertTrue((ds1.f[5:100:5] == 84).all())
        self.assertTrue((ds1.g[5:100:5] == 84).all())
        self.assertTrue((ds1[7:] == ds2[7:]).all(axis=None))
        ds1[[2, 4, 79], ["h", "i"]] = 99
        self.assertTrue((ds1.h[[2, 4, 79]] == 99).any())
        self.assertTrue((ds1.i[[2, 4, 79]] == 99).any())
        self.assertTrue((ds1[9:] == ds2[9:]).all(axis=None))
        # Now, a copy:
        ds3 = ds1[5:100:5, ["j", "k"]].copy()
        ds3[:, :] = -2000
        # New is changed.
        self.assertTrue((ds3.j == -2000).all())
        self.assertTrue((ds3.k == -2000).all())
        # Orig. is unchanged.
        self.assertTrue((ds1.j[5:100:5] != -2000).all())
        self.assertTrue((ds1.k[5:100:5] != -2000).all())
        self.assertTrue((ds1[9:] == ds2[9:]).all(axis=None))

    def test_basic_interface(self):
        cols = ["a", "b", "c", "μεαν"]
        dict1 = {_k: [_i] for _i, _k in enumerate(cols)}
        ds1 = Dataset(dict1)
        self.assertEqual(list(ds1.keys()), cols)
        self.assertEqual(list(ds1.keys()), cols)
        self.assertEqual(list(ds1), cols)
        for _idx, (_k, _v) in enumerate(ds1.items()):
            self.assertEqual(ds1[_k][0], dict1[_k][0])
            self.assertEqual(cols[_idx][0], _k[0])
            self.assertEqual(ds1[_idx][0], dict1[_k][0])
            self.assertEqual(getattr(ds1, _k)[0], dict1[_k][0])
        self.assertEqual(list(ds1.keys()), cols)
        self.assertEqual(list(ds1), cols)
        self.assertEqual(list(reversed(ds1)), list(reversed(cols)))
        self.assertEqual(ds1.size, ds1.get_ncols() * ds1.get_nrows())
        self.assertEqual(ds1.shape[1], ds1.get_ncols())
        self.assertEqual(len(ds1), ds1.get_nrows())
        self.assertEqual(ds1.shape, (ds1.get_nrows(), ds1.get_ncols()))

    def test_col_map_02(self):
        st = Dataset({_k: list(range(_i * 10, (_i + 1) * 10)) for _i, _k in enumerate("abcdefghijklmnopqrst")})
        cmap = {
            "a": "x",  # 1-transition
            "b": "c",
            "c": "b",  # 2-cycle
            "f": "f",  # 1-cycle
            "g": "h",
            "h": "j",
            "j": "k",
            "k": "g",  # 4-cycle
            "m": "n",
            "n": "o",
            "o": "y",  # 3-transition
            "p": "q",
            "q": "z",
            "z": "r",
        }  # false 3-cycle, really a 2-transition
        st.col_map(cmap)
        for _k in "abcdefghijklmnopqrst":
            self.assertEqual(st[cmap.get(_k, _k)][1], (ord(_k) - ord("a")) * 10 + 1)

    def test_col_add_prefix(self):
        ds = Dataset({"col_" + str(i): arange(5) for i in range(5)})
        correct = ["new_" + col for col in list(ds.keys())]
        ds.col_add_prefix("new_")
        self.assertTrue(bool(np.all(ds.keys() == correct)))
        ds["new_new_col_1"] = arange(5)

        # test no overwrite
        correct = ["new_" + col for col in list(ds.keys())]
        ds.col_add_prefix("new_")
        self.assertTrue(bool(np.all(ds.keys() == correct)))

        # test restore labels
        ds = Dataset({"col_" + str(i): arange(5) for i in range(5)})
        ds = ds.gb(["col_1", "col_2"]).sum()
        correct = ["new_" + label for label in ds.label_get_names()]
        ds.col_add_prefix("new_")
        self.assertTrue(bool(np.all(ds.label_get_names() == correct)))

        # test restore sort
        ds = Dataset({"col_" + str(i): arange(5) for i in range(5)})
        ds.sort_view(["col_4", "col_3"])
        correct = ["new_" + label for label in ds._col_sortlist]
        ds.col_add_prefix("new_")
        self.assertTrue(bool(np.all(ds._col_sortlist == correct)))

    def test_np_keep_unicode(self):
        arr = np.random.choice(["aaaaa", "bbb", "ccc", "dd"], 30)
        self.assertTrue(arr.dtype.char == "U")
        ds = Dataset({"uni": arr}, unicode=True)
        self.assertTrue(ds.uni.dtype.char == "U")

    def test_swap_01(self):
        orig_cols = ["a", "b", "c", "d"]
        dict1 = {_k: [_i] for _i, _k in enumerate(orig_cols)}
        st = Dataset(dict1)
        for _i, _k in enumerate(orig_cols):
            self.assertEqual(st[_k][0], _i)
        st.a, st.b = st.b, st.a
        self.assertEqual(st["a"][0], 1)
        self.assertEqual(st["b"][0], 0)
        st.col_swap(["a", "b"], ["b", "a"])
        for _i, _k in enumerate(orig_cols):
            self.assertEqual(st[_k][0], _i)
        st.col_swap(["a", "b", "c"], ["c", "b", "a"])
        self.assertEqual(st["a"][0], 2)
        self.assertEqual(st["b"][0], 1)
        self.assertEqual(st["c"][0], 0)
        self.assertEqual(list(st.keys()), orig_cols)
        st = Dataset(dict1)
        st.col_swap(["a", "b", "c", "d"], ["c", "b", "a", "d"])
        self.assertEqual(st["a"][0], 2)
        self.assertEqual(st["b"][0], 1)
        self.assertEqual(st["c"][0], 0)
        self.assertEqual(st["d"][0], 3)
        self.assertEqual(list(st.keys()), orig_cols)

    def test_swap_02(self):
        ds = self.get_basic_dataset(nrows=4)
        for _i, _k in enumerate(ds):
            self.assertEqual(ds[_k][0], _i * 4)
        ds.a, ds.b = ds.b, ds.a
        self.assertEqual(ds["a"][0], 4)
        self.assertEqual(ds["b"][0], 0)
        ds = self.get_arith_dataset()
        # this type of set item does not work because a string mask returns a dataset
        # ds[['A', 'B']] = ds[['B', 'A']]
        a = ds.A
        ds.A = ds.B
        ds.B = a
        self.assertEqual(ds.A.tolist(), [1.2, 3.1, 9.6])
        self.assertEqual(ds.B.tolist(), [0, 6, 9])
        # this is also broken because getitem happens before setitem
        # ds[1:3, ['G', 'C']] = ds[1:3, ['C', 'G']]
        g = ds.G[1:3].copy()
        c = ds.C[1:3].copy()
        ds.G[1:3] = c
        ds.C[1:3] = g
        self.assertEqual(ds.G.tolist(), [-1.6, 6.2, 19.2])
        self.assertEqual(ds.C.tolist(), [2.4, 2.7, 4.6])
        # self.col_swap

    def test_add_remove_column(self):

        tempsave = Struct.AllowAnyName
        Struct.AllowAnyName = False

        ds = Dataset({"a": [0], "b": [1], "c": [2], "μεαν": [3]})
        ds.d = [4]
        self.assertEqual(list(ds.keys()), ["a", "b", "c", "μεαν", "d"])
        self.assertEqual(ds.d, [4])
        self.assertEqual(ds["d"], [4])

        with self.assertRaises(IndexError):
            ds["1"] = [5]
        with self.assertRaises(IndexError):
            ds["a b"] = [5]

        for kwd in keyword.kwlist:
            if kwd not in ["True", "False", "None"]:
                with self.assertRaises(IndexError):
                    ds[kwd] = [5]
        for memb in dir(Dataset):
            with self.assertRaises(IndexError):
                ds[memb] = [5]
        self.assertEqual(list(ds.keys()), ["a", "b", "c", "μεαν", "d"])
        del ds.b
        # SJK 11/27/2018 - we support this method of column removal now. hits the same path as del ds.c
        # with self.assertRaises(AttributeError):
        #    del ds['c']
        ds.col_remove("c")
        del ds.μεαν
        self.assertEqual(list(ds.keys()), ["a", "d"])
        ds.e = 6
        ds["f"] = 7
        ds.g = [8]
        ds.h = [[9]]
        self.assertEqual(ds.h.shape, (1,))
        # this is allowed now since a tuple is listlike
        # with self.assertRaises(TypeError):
        #    ds.i = (10, )
        with self.assertRaises(TypeError):
            ds.i = {10}
        with self.assertRaises(TypeError):
            ds.i = {"a": 1, "b": 2}
        ds.j = "a"
        ds.k = "μεαν"
        self.assertEqual(list(ds.keys()), list("adefghjk"))
        ds = Dataset({})
        ds.A = 1
        self.assertEqual(ds.shape, (1, 1))
        ds = Dataset({})
        ds.A = [1, 2, 3]
        self.assertEqual(ds.shape, (3, 1))
        ds = Dataset({})
        ds.A = np.array([1, 2, 3])
        self.assertEqual(ds.shape, (3, 1))
        with self.assertRaises(ValueError):
            ds.B = np.array([1, 2, 3, 4])

        Struct.AllowAnyName = tempsave

    def test_broadcast(self):
        ds = Dataset({"test": arange(10)})
        ds.x = [1]
        ds.y = (1,)
        self.assertTrue(ds.x[9] == 1)
        self.assertTrue(ds.y[9] == 1)

    def test_indexing_01(self):
        ds = Dataset({"a": [0], "b": [1], "c": [2], "μεαν": [3]})
        self.assertTrue((ds.a == [0]).all())
        self.assertTrue((ds.b == [1]).all())
        self.assertTrue((ds.c == [2]).all())
        self.assertTrue((ds["a"] == [0]).all())
        self.assertTrue((ds["b"] == [1]).all())
        self.assertTrue((ds["c"] == [2]).all())
        self.assertTrue((ds[np.str_("a")] == [0]).all())
        self.assertEqual(ds.μεαν, [3])
        self.assertEqual(ds["μεαν"], [3])
        ds["μεαν"] = [3]
        self.assertTrue((ds[b"c"] == [2]).all())
        self.assertTrue((ds[np.str_("c")] == [2]).all())
        with self.assertRaises(KeyError):
            _ = ds["q"]
        with self.assertRaises(KeyError):
            _ = ds["1"]
        with self.assertRaises(KeyError):
            _ = ds["a b"]
        with self.assertRaises(TypeError):
            _ = ds[0.0]
        with self.assertRaises(TypeError):
            _ = ds[:, 0.0]
        with self.assertRaises(AttributeError):
            _ = ds.q
        for kwd in keyword.kwlist:
            with self.assertRaises(KeyError):
                _ = ds[kwd]
        for memb in dir(Dataset):
            with self.assertRaises(KeyError):
                _ = ds[memb]

    def test_indexing_02(self):
        # rows
        ds = self.get_basic_dataset()
        self.assertEqual(ds[5, :].a, 5)
        rows = [6, 6, 1, 5, 8]
        self.assertTrue((ds[rows, :].a == rows).all())
        sel = np.zeros(ds.shape[0], dtype=bool)
        sel[rows] = True
        srows = [1, 5, 6, 8]
        with self.assertRaises(IndexError):
            _ = ds[[]]
        with self.assertRaises(IndexError):
            _ = ds[
                5,
            ]
        with self.assertRaises(TypeError):
            _ = ds[0.0, :]
        with self.assertRaises(IndexError):
            _ = ds[["1"], :]
        with self.assertRaises(IndexError):
            _ = ds["1", :]
        with self.assertRaises(TypeError):
            _ = ds[[complex(1)], :]
        with self.assertRaises(TypeError):
            _ = ds[complex(1), :]
        with self.assertRaises(IndexError):
            _ = ds[
                rows,
            ]
        with self.assertRaises(IndexError):
            _ = ds[
                sel,
            ]
        with self.assertRaises(IndexError):
            _ = ds[ds.get_nrows(), :]
        with self.assertRaises(TypeError):
            _ = ds[None, :]
        self.assertTrue((ds[sel, :].a == srows).all())
        self.assertTrue((ds[sel.tolist(), :].a == srows).all())

    def test_indexing_03(self):
        # cols
        ds = self.get_basic_dataset()
        colb = list(range(10, 20))
        self.assertTrue((ds["b"] == colb).all())
        self.assertTrue((ds[:, "b"] == colb).all())
        cols = ["b", "c", "d", "g"]
        row0d = {"b": 10, "c": 20, "d": 30, "g": 60}
        self.assertEqual(dict(ds[:, cols][0, :].asdict()), row0d)  # temp hack
        ncols = np.array(cols)
        self.assertEqual(dict(ds[:, ncols][0, :].asdict()), row0d)  # temp hack
        self.assertIsInstance(ds.a, FastArray)
        self.assertIsInstance(ds[0], FastArray)
        self.assertIsInstance(ds[[0]], Dataset)
        self.assertEqual(list(ds[["a", "b"]].keys()), ["a", "b"])
        self.assertEqual(list(ds[[b"a", b"b"]].keys()), ["a", "b"])
        self.assertEqual(list(ds[[b"a", "b"]].keys()), ["a", "b"])
        self.assertEqual(list(ds[["a", b"b"]].keys()), ["a", "b"])
        self.assertEqual(list(ds[["a", np.str_("b")]].keys()), ["a", "b"])
        dupcols = ds[:, ["b", "c", "d", "c", "g"]]
        # no longer raise IndexError, check that column removed
        self.assertEqual(dupcols.shape[1], 4)
        with self.assertRaises(KeyError):
            _ = ds["q"]
        with self.assertRaises(AttributeError):
            _ = ds.q

    def test_indexing_04(self):
        # cols
        # There is a behavioral change: don't let magic occur.  Index on single 'value' -> columns.
        # Since integer, list(int), logical and slice indexing are allowed for columns as well,
        # the following are all actually allowed COLUMN operations, but they do not give what is
        # perhaps expected!
        ds = self.get_basic_dataset()
        cols = [6, 6, 1, 5, 8]
        dupcols = ds[cols]  # there are duplicates in rows
        self.assertEqual(dupcols.shape[1], 4)
        sel = np.zeros(ds.shape[1], dtype=bool)
        sel[cols] = True
        sel_cnames = [_k for _i, _k in enumerate(list(ds.keys())) if sel[_i]]
        scols = [6, 1, 5, 8]
        scols_cnames = list(map(list(ds.keys()).__getitem__, scols))
        self.assertEqual(list(ds[scols].keys()), scols_cnames)
        self.assertEqual(list(ds[[4, 5]].keys()), ["e", "f"])
        with self.assertRaises(TypeError):
            _ = ds[:, None]
        with self.assertRaises(TypeError):
            _ = ds[None]
        self.assertEqual(list(ds[sel].keys()), sel_cnames)
        self.assertEqual(list(ds[sel.tolist()].keys()), sel_cnames)

    def test_indexing_05(self):
        # rows & cols
        ds = self.get_basic_dataset()
        self.assertEqual(ds[5, "a"], 5)
        rows = [6, 6, 1, 5, 8]
        self.assertTrue((ds[rows, "a"] == rows).all())
        sel = np.zeros(ds.shape[0], dtype=bool)
        sel[rows] = True
        srows = [1, 5, 6, 8]
        self.assertTrue((ds[sel, "a"] == srows).all())
        self.assertTrue((ds[sel, ["a", "c"]].a == srows).all())
        self.assertTrue((ds[sel, ["a", "c"]].c == [_i + 20 for _i in srows]).all())
        self.assertIsInstance(ds[5, "a"], np.integer)
        with self.assertRaises(TypeError):
            _ = ds[0.0, 0.0]
        with self.assertRaises(TypeError):
            _ = ds[None, None]

    def test_indexing_06(self):
        # slicing
        ds = self.get_basic_dataset()
        nrows, ncols = ds.shape
        self.assertEqual(ds[1:-1, :].shape, (nrows - 2, ncols))
        self.assertEqual(ds[:, ::2].shape, (nrows, ncols // 2))
        self.assertEqual(ds[::2].shape, (nrows, ncols // 2))
        self.assertEqual(ds[4:-3, 1::3].shape, (nrows - 7, ncols // 3))
        self.assertEqual(sum(ds[4:-3, 1::3].sum().tolist()[0]), 1125)
        self.assertEqual(sum(ds[1:-3:2, 1:8:3].sum().tolist()[0]), 387)

    def test_indexing_07(self):
        # get/set do matching and sane things, singles
        ds = self.get_basic_dataset()
        ds["newcol"] = -1
        self.assertTrue((ds["newcol"] == FastArray([-1] * ds.get_nrows())).all())
        ds["another"] = -2
        self.assertTrue((ds["another"] == FastArray([-2] * ds.get_nrows())).all())
        with self.assertRaises(NotImplementedError):
            ds[3, "good_and_bad"] = -3
        with self.assertRaises(IndexError):
            ds[ds.get_nrows(), :] = -999
        # now accepted
        ds[["another", "yet_another"]] = -3
        ds = self.get_basic_dataset()
        o2c = ds[2, "c"]
        ds.a = -1
        self.assertTrue((ds["a"] == FastArray([-1] * ds.get_nrows())).all())
        ds.a = range(-100, -100 + ds.shape[0])
        self.assertTrue((ds["a"] == FastArray(range(-100, -100 + ds.shape[0]))).all())
        ds.b[1] = -2
        self.assertEqual(ds[1, "b"], -2)
        o2c = ds[2, "c"]
        ds[2, :].c = -3
        self.assertNotEqual(-3, o2c)
        self.assertEqual(ds[2, "c"], o2c)
        o3r = ds[3, :]
        ds[3, :] = -4
        self.assertTrue(
            (ds[3, :] != o3r).all(axis=None),
            msg="Setting a row should modify in place!",
        )
        ds[4, "e"] = -5
        self.assertEqual(ds[4, "e"], -5)
        self.assertEqual(ds.e[4], -5)
        self.assertEqual(ds[4, :].e, -5)
        ds = self.get_basic_dataset()
        ds["a"] = list(range(-100, -100 + ds.shape[0]))
        self.assertTrue((ds.a == FastArray(range(-100, -100 + ds.shape[0]))).all())
        ds[1, :] = range(-100, -100 + ds.shape[1])
        self.assertEqual(ds[1, :].tolist(), [list(range(-100, -100 + ds.shape[1]))])

    def test_indexing_08(self):
        # get/set do matching and sane things, multiples
        ds = self.get_basic_dataset()
        ds[["a", "b"]] = -1
        self.assertTrue((ds["a"] == FastArray([-1] * ds.get_nrows())).all())
        self.assertTrue((ds["b"] == FastArray([-1] * ds.get_nrows())).all())
        ds = self.get_basic_dataset()
        with self.assertRaises(ValueError):
            ds[["a", "b"]] = list(range(-100, -90))
        with self.assertRaises(ValueError):
            ds[["a", "b"]] = [list(range(-100, -90))] * 2
        other1 = Dataset(
            {f"i{_i:02d}": list(range(_i * ds.get_nrows(), _i * ds.get_nrows() + ds.get_nrows())) for _i in range(2)}
        )
        # ds[['g', 'k']] = other1
        ds["g"] = other1[0]
        ds["k"] = other1[1]
        for _i, _k in enumerate("gk"):
            self.assertTrue((ds[_k] == FastArray(range(ds.get_nrows() * _i, ds.get_nrows() * (_i + 1)))).all())
        ds = self.get_basic_dataset()
        other2 = Dataset({f"i{_i:02d}": list(range(_i * 3, _i * 3 + 3)) for _i in range(ds.shape[1])})
        for _i, _k in enumerate("ab"):
            self.assertTrue((ds[_k] == FastArray(range(ds.get_nrows() * _i, ds.get_nrows() * (_i + 1)))).all())

        # ds[[5, 7, 8], :] = other2
        for i in range(ds.shape[1]):
            ds[i][[5, 7, 8]] = other2[i]
        ds = self.get_basic_dataset()
        other3 = Dataset({"A": [-1, -2, -3], "B": [-4, -5, -6]})
        with self.assertRaises(ValueError):
            ds[:, ["g", "k"]] = other3
        with self.assertRaises(IndexError):
            ds[[5, 7, 8], :] = other3
        for i, k in enumerate("gk"):
            ds[k][[5, 7, 8]] = other3[i]
        # ds[[5, 7, 8], ['g', 'k']] = other3
        self.assertTrue((ds.g == [60, 61, 62, 63, 64, -1, 66, -2, -3, 69]).all())
        self.assertTrue((ds.h == [70, 71, 72, 73, 74, 75, 76, 77, 78, 79]).all())
        self.assertTrue((ds.k == [100, 101, 102, 103, 104, -4, 106, -5, -6, 109]).all())

    def test_indexing_08_previous(self):
        # get/set do matching and sane things, multiples
        ds = self.get_basic_dataset()
        ds[["a", "b"]] = -1
        self.assertTrue((ds["a"] == FastArray([-1] * ds.get_nrows())).all())
        self.assertTrue((ds["b"] == FastArray([-1] * ds.get_nrows())).all())
        ds = self.get_basic_dataset()
        with self.assertRaises(ValueError):
            ds[["a", "b"]] = list(range(-100, -90))
        with self.assertRaises(ValueError):
            ds[["a", "b"]] = [list(range(-100, -90))] * 2
        other1 = Dataset(
            {f"i{_i:02d}": list(range(_i * ds.get_nrows(), _i * ds.get_nrows() + ds.get_nrows())) for _i in range(2)}
        )
        ds[["g", "k"]] = other1
        for _i, _k in enumerate("gk"):
            self.assertTrue((ds[_k] == FastArray(range(ds.get_nrows() * _i, ds.get_nrows() * (_i + 1)))).all())
        ds = self.get_basic_dataset()
        other2 = Dataset({f"i{_i:02d}": list(range(_i * 3, _i * 3 + 3)) for _i in range(ds.shape[1])})
        for _i, _k in enumerate("ab"):
            self.assertTrue((ds[_k] == FastArray(range(ds.get_nrows() * _i, ds.get_nrows() * (_i + 1)))).all())
        ds[[5, 7, 8], :] = other2
        ds = self.get_basic_dataset()
        other3 = Dataset({"A": [-1, -2, -3], "B": [-4, -5, -6]})
        with self.assertRaises(ValueError):
            ds[:, ["g", "k"]] = other3
        with self.assertRaises(IndexError):
            ds[[5, 7, 8], :] = other3
        ds[[5, 7, 8], ["g", "k"]] = other3
        self.assertTrue((ds.g == [60, 61, 62, 63, 64, -1, 66, -2, -3, 69]).all())
        self.assertTrue((ds.h == [70, 71, 72, 73, 74, 75, 76, 77, 78, 79]).all())
        self.assertTrue((ds.k == [100, 101, 102, 103, 104, -4, 106, -5, -6, 109]).all())

    def test_indexing_09(self):
        # get/set with row and column filters
        ds1 = Dataset({"a": [1, 2, 3], "b": [4, 5, 6]})
        ds2 = Dataset({"a": [7, 8, 9], "b": [10, 11, 12]})
        filt = [False, True, False]
        ds1[filt, :] = ds2[filt, :]
        self.assertEqual(ds1.a[1], 8)
        self.assertEqual(ds1.b[1], 11)
        ds3 = Dataset({"a": [13], "b": [14]})
        ds1[filt, :] = ds3
        self.assertEqual(ds1.a[1], 13)
        self.assertEqual(ds1.b[1], 14)
        filt = [True, False, False]
        ds3 = ds1[filt, :]
        self.assertEqual(ds3.a[0], 1)
        self.assertEqual(ds3.b[0], 4)
        ds1 = ds2
        self.assertTrue((ds1.a == [7, 8, 9]).all())
        self.assertTrue((ds1.b == [10, 11, 12]).all())
        with self.assertRaises(ValueError):
            ds1[:, ["a", "b"]] = FastArray([1, 2, 3])

    # def test_row_indexing(self):
    #    strings = FastArray(['c', 'a', 'b', 'b', 'a', 'b', 'c', 'b', 'a', 'a', 'c', 'c', 'b', 'c', 'a'])
    #    int1    = FastArray([2, 1, 2, 1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1])
    #    int2    = None

    def test_concat_rows(self):
        N = 5
        dset1, dset2, _ = self._concat_test_data(N)
        # print('dset1\n',dset1)
        # print(dset2)
        dset = Dataset.concat_rows([dset1, dset2])
        nr1 = dset1.get_nrows()
        nr2 = dset2.get_nrows()
        ncols = len(set(dset1) | set(dset2))
        self.assertEqual(dset.shape, (nr1 + nr2, ncols))
        for _k, _v in dset.items():
            if _k in dset1 and _k in dset2:
                self.assertTrue((_v == np.hstack((dset1[_k], dset2[_k]))).all())
            elif _k in dset1:
                self.assertTrue((_v[:nr1] == dset1[_k]).all())
            else:
                self.assertTrue((_v[nr1:] == dset2[_k]).all())

    def test_copy_names(self):

        ds = Dataset({"col_" + str(i): arange(5) for i in range(5)})
        ds2 = ds[:2, :]
        for k, v in ds.items():
            cp = ds2[k]
            self.assertEqual(v.get_name(), cp.get_name())

    def test_concat_rows_with_categoricals(self):
        # String labels
        ds1 = Dataset({"Ticker": ["a", "a", "c", "b", "b"], "Price": 100 + np.arange(5)})
        ds2 = Dataset({"Ticker": ["b", "b", "d", "d", "c"], "Price": 105 + np.arange(5)})
        ds1.Ticker = Categorical(ds1.Ticker)
        ds2.Ticker = Categorical(ds2.Ticker)
        ds = Dataset.concat_rows([ds1, ds2])
        self.assertTrue((ds.Ticker.as_string_array == ["a", "a", "c", "b", "b", "b", "b", "d", "d", "c"]).all())
        self.assertTrue((ds.Price == 100 + np.arange(10)).all())
        # Numeric labels
        ds1 = Dataset({"Ticker": [1, 1, 3, 2, 2], "Price": 100 + np.arange(5)})
        ds2 = Dataset({"Ticker": [2, 2, 4, 4, 3], "Price": 105 + np.arange(5)})
        ds1.Ticker = Categorical(ds1.Ticker, [1, 2, 3, 4, 5], from_matlab=True)
        ds2.Ticker = Categorical(ds2.Ticker, [1, 2, 3, 4, 5], from_matlab=True)
        ds = Dataset.concat_rows([ds1, ds2])
        self.assertTrue((ds.Ticker == [1, 1, 3, 2, 2, 2, 2, 4, 4, 3]).all())
        self.assertTrue((ds.Price == 100 + np.arange(10)).all())
        N = 16
        dset1 = Dataset(dict(A=np.arange(N), C=np.ones(N), B=N * ["c"]))
        dset2 = Dataset(dict(A=np.arange(N, 2 * N, 1), D=np.zeros(N), B=N * ["d"]))
        test = ["test1", "test2"]
        dset2["D"] = Categorical(dset2["D"], categories=test, from_matlab=True)
        ds3 = Dataset.concat_rows([dset1, dset2])
        self.assertTrue(all(ds3.D.view(FastArray) == np.zeros(N * 2)))
        self.assertTrue(all(ds3.D.expand_array == [FILTERED_LONG_NAME] * N * 2))

    def test_concat_datetime(self):
        dtn = TypeRegister.DateTimeNano([1541239200000000000, 1541325600000000000], from_tz="NYC", to_tz="NYC")
        dts = dtn.hour_span
        ds1 = Dataset({"dtn": dtn, "dts": dts})
        ds2 = Dataset({"dtn": dtn, "dts": dts})
        ds3 = Dataset.concat_rows([ds1, ds2])
        self.assertTrue(isinstance(ds3.dtn, TypeRegister.DateTimeNano))
        self.assertTrue(isinstance(ds3.dts, TypeRegister.TimeSpan))

    def test_concat_columns(self):
        N = 5
        dset1, dset2, dsets = self._concat_test_data(N)
        for do_copy in (True, False):
            # unequal lengths
            with self.assertRaises(ValueError):
                _ = Dataset.concat_columns([dset1, dset2.head(2)], do_copy)
            with self.assertRaises(ValueError):
                _ = Dataset.concat_columns(dsets, do_copy, on_duplicate="invalid")
                # duplicate columns
            with self.assertRaises(KeyError):
                _ = Dataset.concat_columns(dsets, do_copy)
            # first and last
            with self.assertWarns(UserWarning):
                dset = Dataset.concat_columns(dsets, do_copy, on_duplicate="first")
            self.assertTrue((dset.A == dset1.A).all())
            dset = Dataset.concat_columns(dsets, do_copy, on_duplicate="last", on_mismatch="ignore")
            self.assertTrue((dset.A == dset2.A).all())
            # shape
            with self.assertRaises(RuntimeError):
                dset = Dataset.concat_columns(dsets, do_copy, on_duplicate="first", on_mismatch="raise")
            dset = Dataset.concat_columns(dsets, do_copy, on_duplicate="first", on_mismatch="ignore")
            self.assertEqual(dset.shape, (N, 4))
            self.assertEqual(list(dset.keys()), ["A", "B", "C", "D"])
        for do_copy, test in ((True, self.assertNotEqual), (False, self.assertEqual)):
            # copy vs. view
            dset1, dset2, dsets = self._concat_test_data(N)
            with self.assertWarns(UserWarning):
                dset = Dataset.concat_columns(dsets, do_copy, on_duplicate="first")
            dset.A[0] = -1
            test(dset1.A[0], dset.A[0])
            dset1, dset2, dsets = self._concat_test_data(N)
            dset = Dataset.concat_columns([dset1], do_copy)
            dset.A[0] = -1
            test(dset1.A[0], dset.A[0])
            dset = Dataset.concat_columns([dset1, dset1], do_copy, on_duplicate="first")
            self.assertTrue((dset == dset1).all(axis=None))

    def _concat_test_data(self, N):
        dset1 = Dataset(dict(A=np.arange(N), B=np.ones(N), C=N * ["c"]))
        dset2 = Dataset(dict(A=np.arange(N, 2 * N, 1), B=np.zeros(N), D=N * ["d"]))
        dsets = [dset1, dset2]
        return dset1, dset2, dsets

    def test_from_tagged_rows(self):
        ds0 = Dataset({"a": [1, 2, 3], "b": [11, 12, 13]})
        ds1 = Dataset.from_tagged_rows([{"a": 1, "b": 11}, {"a": 2, "b": 12}, {"a": 3, "b": 13}])
        ds2 = Dataset.from_tagged_rows(Struct(_r) for _r in [{"a": 1, "b": 11}, {"a": 2, "b": 12}, {"a": 3, "b": 13}])
        ds3 = Dataset.from_tagged_rows(Dataset(_r) for _r in [{"a": 1, "b": 11}, {"a": 2, "b": 12}, {"a": 3, "b": 13}])
        Row = namedtuple("Row", ["a", "b"])
        ds4 = Dataset.from_tagged_rows(Row(**_r) for _r in [{"a": 1, "b": 11}, {"a": 2, "b": 12}, {"a": 3, "b": 13}])
        self.assertIsInstance(ds1, Dataset)
        self.assertTrue((ds1 == ds0).all(axis=None))
        self.assertTrue((ds2 == ds0).all(axis=None))
        self.assertTrue((ds3 == ds0).all(axis=None))
        self.assertTrue((ds4 == ds0).all(axis=None))
        with self.assertRaises(NotImplementedError):
            _ = Dataset.from_tagged_rows(
                [
                    {"a": 1, "b": 11},
                    {"a": 2, "b": 12},
                    {"a": 3},
                    {"b": 14},
                    {"c": 25},
                    {"a": 6, "b": 16, "c": 26},
                ]
            )

    def test_from_rows(self):
        ds0 = Dataset({"a": [1, 2, 3], "b": [11, 12, 13]})
        ds1 = Dataset.from_rows([[1, 11], [2, 12], [3, 13]], ["a", "b"])
        ds2 = Dataset.from_rows(((1, 11), (2, 12), (3, 13)), ["a", "b"])
        ds3 = Dataset.from_rows(np.array([[1, 11], [2, 12], [3, 13]]), ["a", "b"])
        self.assertIsInstance(ds1, Dataset)
        self.assertTrue((ds1 == ds0).all(axis=None))
        self.assertTrue((ds2 == ds0).all(axis=None))
        self.assertTrue((ds3 == ds0).all(axis=None))
        with self.assertRaises(ValueError):
            _ = Dataset.from_rows([[1, 11], [2, 12, 0], [3, 13]], ["a", "b"])
        with self.assertRaises(ValueError):
            _ = Dataset.from_rows(
                [
                    [1, 11],
                    [
                        2,
                    ],
                    [3, 13],
                ],
                ["a", "b"],
            )
        with self.assertRaises(TypeError):
            _ = Dataset.from_rows([{"a": 1, "b": 11}, {"a": 2, "b": 12}], ["a", "b"])

    def test_from_jagged_rows(self):
        ds1 = Dataset.from_jagged_rows([["A", "B", "C"], ["D"], ["E", "F", "GGG"], "H"])
        # from_jagged will add the letter 'C' in front now by default
        nd = INVALID_DICT[ds1["C1"].dtype.num]
        ds0 = Dataset(
            {
                "C0": ["A", "D", "E", "H"],
                "C1": ["B", nd, "F", nd],
                "C2": ["C", nd, "GGG", nd],
            }
        )
        self.assertIsInstance(ds1, Dataset)
        self.assertTrue((ds1 == ds0).all(axis=None))

    # 10/16/2018 - SJK removed these tests because they use NodataType, which has been removed
    # def test_from_jagged_dict(self):
    #    d = {'name': ['bob', 'mary', 'sue', 'john'],
    #        'letters': [['A', 'B', 'C'], ['D'], ['E', 'F', 'G'], 'H']}
    #    nd = Categorical.NODATA.cast(str)
    #    ds0 = Dataset({'name': ['bob', 'mary', 'sue', 'john'],
    #                    'letters0': ['A', 'D', 'E', 'H'], 'letters1': ['B', nd, 'F', nd],
    #                    'letters2': ['C', nd, 'G', nd]})
    #    ds1 = Dataset.from_jagged_dict(d)
    #    self.assertIsInstance(ds1, Dataset)
    #    self.assertTrue((ds1 == ds0).all(axis=None))
    #    with self.assertRaises(ValueError):
    #        _ = Dataset.from_jagged_dict({'name': ['bob', 'mary', 'sue'],
    #            'letters': [['A', 'B', 'C'], ['D'], ['E', 'F', 'G'], 'H']})
    #    ds2 = Dataset({'name': ['bob', 'bob', 'bob', 'mary', 'sue', 'sue', 'sue', 'john'],
    #                    'letters': ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']})
    #    ds3 = Dataset.from_jagged_dict(d, stacked=True)
    #    self.assertTrue((ds2 == ds3).all(axis=None))

    def test_melt(self):
        ds = Dataset({"A": ["a", "b", "c"], "B": [1, 3, 5], "C": [2, 4, 6]})
        tm0 = Dataset({"A": ["a", "b", "c"], "variable": ["B", "B", "B"], "value": [1, 3, 5]})
        tm0t = ds.melt(id_vars=["A"], value_vars=["B"])
        self.assertTrue((tm0 == tm0t).all(axis=None))
        tm1 = Dataset(
            {
                "A": ["a", "b", "c", "a", "b", "c"],
                "variable": ["B", "B", "B", "C", "C", "C"],
                "value": [1, 3, 5, 2, 4, 6],
            }
        )
        tm1t = ds.melt(id_vars=["A"], value_vars=["B", "C"])
        self.assertTrue((tm1 == tm1t).all(axis=None))
        tm2 = Dataset({"A": ["a", "b", "c"], "myVarname": ["B", "B", "B"], "myValname": [1, 3, 5]})
        tm2t = ds.melt(
            id_vars=["A"],
            value_vars=["B"],
            var_name="myVarname",
            value_name="myValname",
        )
        self.assertTrue((tm2 == tm2t).all(axis=None))

    def test_melt_datetimenano(self):
        ds = Dataset(
            {
                "A": ["a", "b", "c"],
                "t": rt.DateTimeNano(["20200925 10:00", "20200926 10:00", "20200927 10:00"], from_tz="NYC"),
            }
        )
        ds["t1"] = ds["t"] + 1e9 * 60
        ds["t2"] = ds["t"] + 1e9 * 60 * 2
        tm0 = Dataset(
            {
                "A": ["a", "b", "c", "a", "b", "c", "a", "b", "c"],
                "variable": ["t", "t", "t", "t1", "t1", "t1", "t2", "t2", "t2"],
                "value": rt.DateTimeNano(
                    [
                        "20200925 10:00",
                        "20200926 10:00",
                        "20200927 10:00",
                        "20200925 10:01",
                        "20200926 10:01",
                        "20200927 10:01",
                        "20200925 10:02",
                        "20200926 10:02",
                        "20200927 10:02",
                    ],
                    from_tz="NYC",
                ),
            }
        )
        tm0t = ds.melt(id_vars=["A"])
        result = tm0 == tm0t
        self.assertTrue(result.all(axis=None))

    def test_pivot(self):
        pass

    def test_tolist(self):
        ds = self.get_basic_dataset()
        self.assertEqual(
            ds[4, :].tolist(),
            [[4, 14, 24, 34, 44, 54, 64, 74, 84, 94, 104, 114, 124, 134, 144, 154]],
        )
        self.assertEqual(
            ds[[4, 5, 8], :].tolist(),
            [
                [4, 14, 24, 34, 44, 54, 64, 74, 84, 94, 104, 114, 124, 134, 144, 154],
                [5, 15, 25, 35, 45, 55, 65, 75, 85, 95, 105, 115, 125, 135, 145, 155],
                [8, 18, 28, 38, 48, 58, 68, 78, 88, 98, 108, 118, 128, 138, 148, 158],
            ],
        )

    def test_asrows(self):
        ds = self.get_basic_dataset()
        self.assertIsInstance(next(ds.asrows()), Dataset)
        self.assertIsInstance(next(ds.asrows(as_type="array")), np.ndarray)
        _it = ds.asrows()
        for _i in range(ds.get_nrows()):
            self.assertEqual(next(_it).tolist(), ds[_i, :].tolist())
        _it = ds.asrows()
        for _i in range(ds.get_nrows()):
            self.assertEqual(next(_it).tolist(), ds[_i, :].tolist())

    def test_aggregators_01(self):
        ds = self.get_basic_dataset()
        self.assertEqual(
            ds.sum().tolist()[0],
            [
                45,
                145,
                245,
                345,
                445,
                545,
                645,
                745,
                845,
                945,
                1045,
                1145,
                1245,
                1345,
                1445,
                1545,
            ],
        )
        self.assertTrue(all((ds.sum(axis=0) == ds.sum()).all().tolist()[0]))
        self.assertEqual(
            ds.sum(axis=1).tolist(),
            [1200, 1216, 1232, 1248, 1264, 1280, 1296, 1312, 1328, 1344],
        )
        self.assertEqual(ds.sum(axis=None), 12720)
        self.assertEqual(
            ds.mean().tolist()[0],
            [
                4.5,
                14.5,
                24.5,
                34.5,
                44.5,
                54.5,
                64.5,
                74.5,
                84.5,
                94.5,
                104.5,
                114.5,
                124.5,
                134.5,
                144.5,
                154.5,
            ],
        )
        for method in "min max sum mean std var median count".split():
            func = getattr(ds, method)
            self.assertIsInstance(func(), Struct)
            self.assertNotIsInstance(func(as_dataset=False), Dataset)
            self.assertIsInstance(func(as_dataset=True), Dataset)
            for axislist, atype in (
                ([0, "c", "C", "col", "COL", "column", "COLUMN"], Struct),
                ([1, "r", "R", "row", "ROW"], np.ndarray),
                ([None, "all", "ALL"], (int, float, np.integer, np.floating)),
            ):
                for axis in axislist:
                    self.assertIsInstance(func(axis=axis), atype)
            with self.assertRaises(NotImplementedError):
                _ = func(axis=3)
            with self.assertRaises(NotImplementedError):
                _ = func(axis="EVERYTHING")

    def test_aggregators_02(self):
        ds = self.get_basic_dataset()[:, :3]
        dvals = b"abcdefghij"
        evals = b"ABCDEFGHIJ"
        ds.d = list(dvals.decode())
        ds.e = list(evals.decode())
        self.assertEqual(ds.count().tolist()[0], [ds.get_nrows()] * ds.get_ncols())
        self.assertEqual(ds.count(fill_value=None).tolist()[0], [ds.get_nrows()] * 3)
        self.assertEqual(ds.max().tolist()[0], [9, 19, 29, b"j", b"J"])
        self.assertEqual(ds.max(fill_value={}).tolist()[0], [9, 19, 29])
        self.assertEqual(ds.max(fill_value=b"".join).tolist()[0], [9, 19, 29, dvals, evals])
        self.assertEqual(ds.max(fill_value="banana").tolist()[0], [9, 19, 29, b"banana", b"banana"])
        self.assertEqual(ds.max(fill_value=0).tolist()[0], [9, 19, 29, 0, 0])
        self.assertEqual(
            ds.max(fill_value={"d": None, "e": b"".join}).tolist()[0],
            [9, 19, 29, evals],
        )
        self.assertEqual(
            ds.max(fill_value={"d": "apple", "e": "orange"}).tolist()[0],
            [9, 19, 29, b"apple", b"orange"],
        )

    def test_aggregators_03(self):
        ds1 = self.get_basic_dataset()
        ds1.S = list("ABCDEFGHIJ")
        ds2 = self.get_basic_dataset()
        for axis in (0, 1, None):
            self.assertEqual(ds1.sum(axis=axis).tolist(), ds2.sum(axis=axis).tolist())
        self.assertEqual(
            ds1.sum(axis=0, fill_value=0)[list(ds2)].tolist(),
            ds2.sum(axis=0, fill_value=0).tolist(),
        )
        self.assertEqual(ds1.sum(axis=0, fill_value=0).S, 0)
        self.assertEqual(
            ds1.sum(axis=1, fill_value=0).tolist(),
            ds2.sum(axis=1, fill_value=0).tolist(),
        )
        self.assertEqual(
            [_x - 1 for _x in ds1.sum(axis=1, fill_value=1).tolist()],
            ds2.sum(axis=1, fill_value=1).tolist(),
        )
        self.assertEqual(ds1.sum(axis=None, fill_value=0), ds2.sum(axis=None, fill_value=0))
        self.assertEqual(ds1.sum(axis=None, fill_value=1) - 1, ds2.sum(axis=None, fill_value=1))
        self.assertEqual(ds1.sum(axis=None, fill_value=len) - 10, ds2.sum(axis=None, fill_value=1))

    def test_any(self):
        ds1 = Dataset(
            {
                "A": [True, False, False],
                "B": [True, True, False],
                "C": [False, False, False],
            }
        )
        self.assertTrue((ds1.any(as_dataset=False) == Struct({"A": True, "B": True, "C": False})).all())
        self.assertTrue((ds1.any(axis=0, as_dataset=False) == Struct({"A": True, "B": True, "C": False})).all())
        self.assertTrue((ds1.any(as_dataset=True) == Dataset({"A": True, "B": True, "C": False})).all(axis=None))
        self.assertTrue(
            (ds1.any(axis=0, as_dataset=True) == Dataset({"A": True, "B": True, "C": False})).all(axis=None)
        )
        self.assertTrue(np.all(ds1.any(axis=1) == [True, True, False]))
        self.assertTrue(ds1.any(axis=None))
        self.assertFalse(Dataset({}).any(axis=None))  # there does not exist one which is true
        ds1.S = ["X", "Y", ""]
        self.assertTrue(
            (ds1.any(axis=0, as_dataset=False) == Struct({"A": True, "B": True, "C": False, "S": True})).all()
        )

    def test_all(self):
        ds1 = Dataset(
            {
                "A": [True, False, False],
                "B": [True, True, True],
                "C": [False, False, False],
            }
        )
        self.assertTrue((ds1.all(as_dataset=False) == Struct({"A": False, "B": True, "C": False})).all())
        self.assertTrue((ds1.all(axis=0, as_dataset=False) == Struct({"A": False, "B": True, "C": False})).all())
        self.assertTrue((ds1.all(as_dataset=True) == Dataset({"A": False, "B": True, "C": False})).all(axis=None))
        self.assertTrue(
            (ds1.all(axis=0, as_dataset=True) == Dataset({"A": False, "B": True, "C": False})).all(axis=None)
        )
        self.assertTrue(np.all(ds1.all(axis=1) == [False, False, False]))
        self.assertTrue(np.all(ds1[[0, 1], ["A", "B"]].all(axis=1) == [True, False]))
        self.assertFalse(ds1.all(axis=None))
        self.assertTrue(ds1[0, ["A", "B"]].all(axis=None))
        self.assertTrue(Dataset({}).all(axis=None))  # all that exist are true
        ds1.S = ["X", "Y", ""]
        self.assertTrue(
            (ds1.all(axis=0, as_dataset=False) == Struct({"A": False, "B": True, "C": False, "S": False})).all()
        )

    def test_comparison_01(self):
        ds1 = Dataset({_k: list(range(_i * 3, (_i + 1) * 3)) for _i, _k in enumerate("AB")})
        ds2 = Dataset({_k: list(range(_i * 3, (_i + 1) * 3)) for _i, _k in enumerate("AB")})
        ds2[0, "A"] += 1
        ds2[1, "B"] += 1
        ds3 = Dataset({_k: list(range(_i * 3, (_i + 1) * 3)) for _i, _k in enumerate("AC")})
        self.assertTrue((ds1 == ds1).all(axis=None))
        self.assertTrue((ds1 == ds1).any(axis=None))
        self.assertFalse((ds1 == ds2).all(axis=None))
        self.assertTrue(((ds1 == ds2) == Dataset({"A": [False, True, True], "B": [True, False, True]})).all(axis=None))
        self.assertTrue((ds1 <= ds2).all(axis=None))
        self.assertTrue(((ds1 < ds2) == Dataset({"A": [True, False, False], "B": [False, True, False]})).all(axis=None))
        self.assertEqual(list((ds2 == ds3).keys()), ["A", "B", "C"])
        self.assertEqual(list((ds3 == ds2).keys()), ["A", "C", "B"])
        for comp in "eq ne lt le gt ge".split():
            self.assertIsInstance((getattr(ds1, f"__{comp}__")(ds2)).all(axis=None), bool)

    def test_comparison_02(self):
        ds1 = Dataset({"A": [5.5], "B": [20]})
        ds2 = Dataset({"A": [5.5], "B": ["20"]})
        ds3 = Dataset({"A": [6.5], "B": [30]})
        self.assertTrue(((ds1 < ds3) == Dataset({"A": [True], "B": [True]})).all(axis=None))
        with self.assertRaises(TypeError):
            _ = ds2 < ds3
        for comp in "eq ne lt le gt ge".split():
            self.assertIsInstance((getattr(ds1, f"__{comp}__")(ds3)).all(axis=None), bool)

    def test_arith_ops_01(self):
        with self.assertWarns(RuntimeWarning):
            for op in "+= -= *= /= //= %= **=".split():  # @= **=
                for val in (
                    10,
                    5.5,
                    1,
                    0.5,
                    0,
                    -0.5,
                    -1,
                    -5.5 - 10,
                    [1, 2.0, 3],
                    np.array([1, 2.0, 3]),
                    FastArray([1, 2.0, 3]),
                ):
                    ds = self.get_arith_dataset(dtype=np.float32)
                    origds = ds.copy()
                    arr = np.array(ds.tolist())
                    origarr = arr.copy()
                    val2 = np.reshape(val, (3, -1)) if isinstance(val, (list, np.ndarray)) else val
                    self.assertTrue(almost_eq(ds, arr, places=5), msg=f"Failure for base case.")
                    exec(f"ds {op} val; arr {op} val2")
                    self.assertTrue(
                        almost_eq(ds, arr, places=5),
                        msg=f"Failure for ds {op} {val}\n origds:{origds}\n {origarr}\nds:{ds}\narr:{arr}",
                    )
                    exec(f"ds {op} val2; arr {op} val2")
                    self.assertTrue(
                        almost_eq(ds, arr, places=5),
                        msg=f"Failure for ds {op} {val2}\n origds:{origds}\n {origarr}\nds:{ds}\narr:{arr}",
                    )
            for op in "<<= >>= &= ^= |=".split():
                for val in (
                    10,
                    5,
                    1,
                    0,
                    [1, 2, 3],
                    np.array([1, 2, 3]),
                    FastArray([1, 2, 3]),
                ):
                    ds = self.get_arith_dataset(dtype=np.int32)
                    arr = np.array(ds.tolist())
                    val2 = np.reshape(val, (3, -1)) if isinstance(val, (list, np.ndarray)) else val
                    self.assertTrue(almost_eq(ds, arr, places=5), msg=f"Failure for base case.")
                    exec(f"ds {op} val; arr {op} val2")
                    self.assertTrue(almost_eq(ds, arr, places=5), msg=f"Failure for ds {op} {val}")
                    exec(f"ds {op} val2; arr {op} val2")
                    self.assertTrue(almost_eq(ds, arr, places=5), msg=f"Failure for ds {op} {val2}")

    def test_arith_ops_02(self):
        with self.assertWarns(RuntimeWarning):
            ds = self.get_arith_dataset(dtype=np.float32)
            arr = np.array(ds.tolist())
            self.assertTrue(almost_eq(ds, arr, places=5), msg=f"Failure for base case.")
            for op in "+ - * / // % **".split():  # @
                for val in (
                    10,
                    5.5,
                    1,
                    0.5,
                    0,
                    -0.5,
                    -1,
                    -5.5 - 10,
                    [1, 2.0, 3],
                    np.array([1, 2.0, 3]),
                    FastArray([1, 2.0, 3]),
                ):
                    val2 = np.reshape(val, (3, -1)) if isinstance(val, (list, np.ndarray)) else val
                    ds0 = eval(f"ds {op} val")  # ds {op} val2 fails for op in {'%', '**'}
                    arr0 = eval(f"arr {op} val2")
                    self.assertTrue(
                        almost_eq(ds0, arr0, places=5),
                        msg=f"Failure for ds {op} {val}\n",
                    )
                    ds1 = eval(f"ds {op} val2")  # ds {op} val2 fails for op in {'%', '**'}
                    self.assertTrue(
                        almost_eq(ds1, arr0, places=5),
                        msg=f"Failure for ds {op} {val2}",
                    )
            ds = self.get_arith_dataset(dtype=np.int32)
            arr = np.array(ds.tolist())
            self.assertTrue(almost_eq(ds, arr, places=5), msg=f"Failure for base case.")
            for op in "<< >> & ^ |".split():  # @
                for val in (
                    10,
                    5,
                    1,
                    0,
                    [1, 2, 3],
                    np.array([1, 2, 3]),
                    FastArray([1, 2, 3]),
                ):
                    val2 = np.reshape(val, (3, -1)) if isinstance(val, (list, np.ndarray)) else val
                    ds0 = eval(f"ds {op} val")  # ds {op} val2 fails for op in {'<<', '>>'}
                    arr0 = eval(f"arr {op} val2")
                    self.assertTrue(almost_eq(ds0, arr0, places=5), msg=f"Failure for ds {op} {val}")
                    ds1 = eval(f"ds {op} val2")  # ds {op} val2 fails for op in {'<<', '>>'}
                    self.assertTrue(
                        almost_eq(ds1, arr0, places=5),
                        msg=f"Failure for ds {op} {val2}",
                    )

    def test_arith_ops_03(self):
        ds = self.get_arith_dataset()
        arr = np.array(ds.tolist())
        self.assertTrue((-ds).tolist() == (-arr).tolist(), msg=f"Failure for -ds")
        self.assertTrue((+ds).tolist() == (+arr).tolist(), msg=f"Failure for +ds")
        self.assertTrue(abs(ds).tolist() == abs(arr).tolist(), msg=f"Failure for abs(ds)")
        self.assertTrue(ds.abs().tolist() == abs(ds).tolist(), msg=f"Failure for ds.abs()")
        ds = self.get_arith_dataset(dtype=np.int32)
        arr = np.array(ds.tolist())
        self.assertTrue((~ds).tolist() == (~arr).tolist(), msg=f"Failure for ~ds")

    def test_arith_ops_04(self):
        ds0 = self.get_arith_dataset(include_strings=False)
        ds1 = self.get_arith_dataset(include_strings=True)
        arith_cnames = list(ds0.keys())
        str_cnames = [_c for _c in ds1 if _c not in ds0]
        orig_strs_cols = ds1[str_cnames].copy(deep=True)
        self.assertTrue((ds0.abs() == ds1.abs()[arith_cnames]).all(axis=None))
        self.assertTrue((ds1.abs()[str_cnames] == orig_strs_cols).all(axis=None))
        self.assertTrue(((ds0 + 10) == (ds1 + 10)[arith_cnames]).all(axis=None))
        self.assertTrue(((ds1 + 10)[str_cnames] == orig_strs_cols).all(axis=None))
        ds0 += 5
        ds1 += 5
        self.assertTrue((ds0 == ds1[arith_cnames]).all(axis=None))
        self.assertTrue((ds1[str_cnames] == orig_strs_cols).all(axis=None))

    def test_arith_ops_05(self):
        ds1 = self.get_arith_dataset(include_strings=True)
        for wrong in (
            None,
            [0, [1, 2], 3],
            np.ones(ds1.shape),
            np.ones((1, ds1.get_ncols())),
        ):
            with self.assertRaises((ValueError, KeyError, TypeError)):
                _ = ds1 + wrong
        st1 = Struct({"A": 1, "B": 2, "C1": 3, "G": np.arange(1, 4)})
        ds2 = ds1 + st1
        self.assertEqual((ds1.A + 1).tolist(), ds2.A.tolist())
        self.assertEqual((ds1.B + 2).tolist(), ds2.B.tolist())
        self.assertEqual(ds1.C.tolist(), ds2.C.tolist())
        self.assertEqual((ds1.G + [1, 2, 3]).tolist(), ds2.G.tolist())
        self.assertEqual((ds1.S).tolist(), ds2.S.tolist())
        self.assertEqual((ds1.U).tolist(), ds2.U.tolist())
        ds3 = 2 * ds1
        ds4 = ds1 + ds3
        self.assertEqual((ds1.A * 3).tolist(), ds4.A.tolist())
        self.assertEqual((ds1.B * 3).tolist(), ds4.B.tolist())
        self.assertEqual((ds1.C * 3).tolist(), ds4.C.tolist())
        self.assertEqual((ds1.G * 3).tolist(), ds4.G.tolist())
        self.assertEqual((ds1.S).tolist(), ds4.S.tolist())
        self.assertEqual((ds1.U).tolist(), ds4.U.tolist())
        ds1 += st1
        self.assertTrue((ds1 == ds2).all(axis=None))

    def test_arith_ops_06(self):
        ds1 = self.get_arith_dataset(dtype=np.int32)
        self.assertTrue(((ds1 * 3) == (3 * ds1)).all(axis=None))
        self.assertTrue(((ds1 + 3) == (3 + ds1)).all(axis=None))
        self.assertTrue(((-ds1 + 3) == (3 - ds1)).all(axis=None))
        self.assertTrue(((ds1 | 3) == (3 | ds1)).all(axis=None))
        self.assertTrue(((ds1 & 3) == (3 & ds1)).all(axis=None))
        self.assertTrue(((ds1 ^ 3) == (3 ^ ds1)).all(axis=None))

    def test_arith_ops_07(self):
        for dtype in (np.int32, np.float32):
            # import pdb; pdb.set_trace()
            rhs = self.get_arith_dataset(dtype=dtype, include_strings=True)
            rhs[0, "A"] = 1  # no need for div-by-zero warnings in this test
            prec = 5
            for lhs in (3, 5.5):  # used in eval
                for op in "/ // % **".split():  # ** should be last as it modifies rhs and prec.
                    if op == "**":
                        rhs = rhs.abs()
                        rhs[2, :] /= 2  # don't get too big!
                        prec = 1
                    res = eval(f"lhs {op} rhs")
                    for _cn, _col in res.items():
                        if _col.dtype.char in NumpyCharTypes.Noncomputable or isinstance(_col, Categorical):
                            self.assertTrue((_col == rhs[_cn]).all())
                            continue
                        for _r in range(len(_col)):
                            elem = rhs[_r, _cn]  # used in eval
                            self.assertAlmostEqual(
                                _col[_r],
                                eval(f"lhs {op} elem"),
                                places=prec,
                                msg=f"operation {op} elem {elem}",
                            )

    def test_arith_ops_08(self):
        ds1 = self.get_basic_dataset(nrows=100).astype(np.float32)
        ds1.S = ds1.a.astype(str)
        ds2 = ds1.copy()
        ds1[2, :] /= 2
        self.assertTrue((ds1[2, :] == ds2[2, :] / 2).all(axis=None))
        ds1[3, ["e", "f"]] /= 3
        self.assertTrue((ds1[3, ["e", "f"]] == ds2[3, ["e", "f"]] / 3).all(axis=None))
        ds1[4:10:2, :] /= 4
        self.assertTrue((ds1[4:10:2, :] == ds2[4:10:2, :] / 4).all(axis=None))
        ds1[12:16:2, 10:12] /= 5
        self.assertTrue((ds1[12:16:2, 10:12] == ds2[12:16:2, 10:12] / 5).all(axis=None))
        self.assertTrue((ds1 == ds2)[1:9:2, :4].all(axis=None))
        self.assertTrue((ds1 == ds2)[1:9:2, 6:].all(axis=None))

    def test_dataset_objects(self):
        ds = Dataset({"float_obj": np.array([1.0, 2.0, 3.0], dtype=object)})
        self.assertTrue(ds.float_obj.dtype.char in NumpyCharTypes.AllFloat)

        # mixed object will always default to flip to string now SJK 3/7/2019
        ds = Dataset({"mixed_object": np.array([np.nan, "str", 1], dtype=object)})
        self.assertTrue(ds.mixed_object.dtype.char == "S")

        ds = Dataset({"mixed_string_start": np.array(["str", np.nan, 1], dtype=object)})

    def test_sample(self):
        arr = tile(FastArray([10, 20]), 100)
        ds = Dataset({"arr": arr})
        f = logical(arange(200) % 2)
        s = ds.sample(N=50, filter=f)
        self.assertEqual(s._nrows, 50)
        self.assertTrue(bool(np.all(s.arr == 20)))

    def test_dataset_pandas(self):
        import pandas as pd

        df = pd.DataFrame({"A": ["a", "b", "c", "a"]})
        df["B"] = df["A"].astype("category")
        ds = Dataset({"A": df.A, "B": df.B})
        self.assertIsInstance(ds.B, TypeRegister.Categorical)
        self.assertTrue((df.A == ds.A.astype(str)).all())
        self.assertTrue((df.B == ds.B.as_string_array.astype(str)).all())

    def _test_output(self):
        ds = self.get_basic_dataset()[:4, :3]
        self.assertEqual(
            str(ds),
            """#   a    b    c
-   -   --   --
0   0   10   20
1   1   11   21
2   2   12   22
3   3   13   23""",
        )
        self.assertTrue(
            re.match(
                r"\[4 rows x 3 columns\]\s+total\s+bytes:\s+[\d.]+\s+[KMGTPEZY]?B",
                ds._last_row_stats(),
            )
        )

    def test_head_tail(self):
        ds = self.get_basic_dataset(30)
        self.assertEqual(ds.head().shape, (20, ds.get_ncols()))
        self.assertEqual(ds.tail().shape, (20, ds.get_ncols()))
        self.assertEqual(ds.head(n=5).shape, (5, ds.get_ncols()))
        self.assertEqual(ds.tail(n=25).shape, (25, ds.get_ncols()))

    def test_drop_duplicates(self):
        import pandas as pd

        ds = Dataset(
            {
                "strcol": np.random.choice(["a", "b", "c", "d"], 15),
                "intcol": np.random.randint(0, 3, 15),
                "rand": np.random.rand(15),
            }
        )
        df = pd.DataFrame(ds.asdict())

        rt_drop_first = ds.drop_duplicates(["strcol", "intcol"], keep="first")
        pd_drop_first = df.drop_duplicates(["strcol", "intcol"], keep="first")

        for colname in ds:
            rtcol = rt_drop_first[colname]
            pdcol = pd_drop_first[colname]

            self.assertTrue(bool(np.all(rtcol == pdcol.values)))

    def test_drop_duplicates_empty_dataset(self):
        empty_ds = Dataset({"a": [], "b": []})
        new_ds = empty_ds.drop_duplicates("a")
        self.assertIsNot(new_ds, empty_ds)
        self_ds = empty_ds.drop_duplicates("a", inplace=True)
        self.assertIs(self_ds, empty_ds)

    def test_describe(self):
        ds = self.get_basic_dataset()
        # TODO NW Should be String, not FastArray
        ds.b = FastArray(ds.b.astype(str))
        labels = describe(None)
        col_a_res = np.array([10.0, 10.0, 0.0, 4.5, 3.028, 0.0, 0.9, 2.25, 4.5, 6.75, 8.1, 9.0, 4.5])
        col_c_res = col_a_res + 20.0
        col_c_res[:2] = col_a_res[:2]

        # TJD: test assumes intimate of describe display order, needs to be fixed
        col_c_res[2] = col_a_res[2]
        col_c_res[4] = col_a_res[4]

        def norm(v):
            return np.sqrt(v.dot(v))

        #
        res = ds.describe()
        self.assertIsInstance(res, Dataset)
        self.assertEqual(res.shape, (len(labels), ds.get_ncols()))  # one added, one dropped
        self.assertIsInstance(res.Stats, FastArray)
        # TODO NW Decode not needed if Stats were a String and not a FastArray
        self.assertEqual([s.decode("unicode_escape") for s in res.Stats.tolist()], labels)
        self.assertAlmostEqual(norm(res.a - col_a_res), 0, places=3)
        self.assertNotIn("b", res)
        # print("***",res.c)
        # print("***",col_c_res)
        self.assertAlmostEqual(norm(res.c - col_c_res), 0, places=3)
        #
        # TODO NW When using Strings, fill_value should be Categorical.NODATA
        res = ds.describe(fill_value=INVALID_DICT[ds.b.dtype.num])
        self.assertIsInstance(res, Dataset)
        self.assertEqual(res.shape, (len(labels), ds.get_ncols() + 1))  # one added, none dropped
        self.assertIsInstance(res.Stats, FastArray)
        # TODO NW Decode not needed if Stats were a String and not a FastArray
        self.assertEqual([s.decode("unicode_escape") for s in res.Stats.tolist()], labels)
        self.assertAlmostEqual(norm(res.a - col_a_res), 0, places=3)
        # TODO NW Decode not needed if String used above for ds.b
        # TODO NW When using Strings, comparison value on right should be Categorical.NODATA
        # TJD: strings are not in describe
        self.assertEqual(res.b.tolist(), [INVALID_DICT[ds.b.dtype.num]] * len(labels))
        self.assertAlmostEqual(norm(res.c - col_c_res), 0, places=3)

        ds = Dataset({"arr": arange(10.0)})
        ds.arr = np.nan
        result = ds.describe().arr
        self.assertEqual(result[0], 10)
        self.assertEqual(result[1], 0)
        self.assertTrue(bool(np.all(isnan(result[2:]))))

    # Regression test for RIP-442 - Error displaying dataset with multiple 'key' columns with zero rows
    def test_repr_multikey_columns_with_zero_rows(self):
        ds = Dataset({"K1": Categorical(["A", "A", "A"]), "K2": Categorical(["B", "B", "B"]), "V": [1, 2, 3]})
        red_filtered = ds.filter(ds.V == 0)  # zero row Dataset
        # apply a reduction
        red = ds.cat(["K1", "K2"]).first(ds.V)

        # filter resulting in zero rows
        red_filtered = red.filter(red.V == 0)
        assert isinstance(repr(red_filtered), str)

    def test_footer_get_values(self):
        ds = Dataset({"col_" + str(i): arange(5) for i in range(5)})
        nofooters = ds.footer_get_values()
        self.assertEqual(nofooters, {})

        ds.footer_set_values("sum", {"col_1": 10})
        with self.assertWarns(UserWarning):
            missing = ds.footer_get_values("mean")
            self.assertEqual(missing, {})

        allcols = ds.footer_get_values()
        self.assertEqual(len(allcols), 1)
        self.assertEqual(len(allcols["sum"]), len(ds))
        self.assertTrue(allcols["sum"][0] is None)
        self.assertEqual(allcols["sum"][1], 10)

        allcols_fill = ds.footer_get_values(fill_value="INV")
        self.assertEqual(allcols_fill["sum"][0], "INV")

        somecols = ds.footer_get_values(columns="col_0")
        self.assertTrue(somecols["sum"][0] is None)

        somecols = ds.footer_get_values(columns=["col_0", "col_1"])
        self.assertEqual(somecols["sum"][1], 10)

    def test_footer_set_values(self):
        ds = Dataset({"col_" + str(i): arange(5) for i in range(5)})
        ds.footer_set_values("sum", {"col_1": 10})
        footerdict = ds.col_get_attribute("col_1", "Footer")
        self.assertEqual(footerdict["sum"], 10)

        ds.footer_set_values("mean", {"col_1": 2.0})
        footerdict = ds.col_get_attribute("col_1", "Footer")
        self.assertEqual(footerdict["mean"], 2.0)

        footerdict = ds.col_get_attribute("col_0", "Footer")
        self.assertTrue(footerdict is None)

    def test_footer_remove(self):
        ds = Dataset({"col_" + str(i): arange(5) for i in range(5)})
        ds.footer_set_values("row1", dict(zip(ds, ["footer" + str(i) for i in range(len(ds))])))
        with self.assertWarns(UserWarning):
            ds.footer_remove("row2")

        ds.footer_remove("row1", ["col_3", "col_1"])
        footer3 = ds.col_get_attribute("col_3", "Footer")
        self.assertEqual(len(footer3), 0)
        footer1 = ds.col_get_attribute("col_1", "Footer")
        self.assertEqual(len(footer1), 0)
        footer2 = ds.col_get_attribute("col_2", "Footer")
        self.assertEqual(footer2["row1"], "footer2")

        with self.assertRaises(IndexError):
            ds.footer_remove("row1", "col123")

        ds.footer_remove()
        self.assertTrue(ds.footers is None)

    def test_mask_or_isnan(self):
        d = {
            "intcol": FastArray([-128, 2, 3, -128, 1, 2, 3], dtype=np.int8),
            "fltcol": FastArray([np.nan, 2, np.nan, 4, 1, 2, 3], dtype=np.float32),
        }
        correct = FastArray([True, False, True, True, False, False, False])
        ds = Dataset(d)
        mask = ds.mask_or_isnan()
        self.assertTrue(bool(np.all(mask == correct)))

    def test_output(self):
        ds = self.get_basic_dataset()[:4, :3]

        self.assertEqual(
            str(ds),
            """#   a    b    c
-   -   --   --
0   0   10   20
1   1   11   21
2   2   12   22
3   3   13   23""",
        )
        self.assertRegex(
            ds._last_row_stats(),
            r"\[4 rows x 3 columns\]\s+total\s+bytes:\s+[\d.]+\s+[KMGTPEZY]?B",
        )
        ds = self.get_arith_dataset(include_strings=True)
        self.assertEqual(
            str(ds),
            "#   A      B       G       C   S           U        \n-   -   ----   -----   -----   ---------   ---------\n0   0   1.20   -1.60    2.40   string_00   ℙƴ☂ℌøἤ_00\n1   6   3.10    2.70    6.20   string_01   ℙƴ☂ℌøἤ_01\n2   9   9.60    4.60   19.20   string_02   ℙƴ☂ℌøἤ_02",
        )
        self.assertRegex(
            ds._last_row_stats(),
            r"\[3 rows x 6 columns\]\s+total\s+bytes:\s+[\d.]+\s+[KMGTPEZY]?B",
        )

    def test_to_pandas(self):
        cols = ["c", "b", "a"]
        ds = Dataset(dict([("c", [4, 5]), ("b", [2, 3]), ("a", [0, 1])]))
        self.assertEqual(list(ds.keys()), cols)
        df = ds.to_pandas()
        self.assertEqual(list(df.columns), cols)
        self.assertIsInstance(df.a, pd.Series)
        self.assertIsInstance(df.b, pd.Series)
        self.assertIsInstance(df.c, pd.Series)
        self.assertTrue((df.a == [0, 1]).all())
        self.assertTrue((df.b == [2, 3]).all())
        self.assertTrue((df.c == [4, 5]).all())

    def test_from_pandas(self):
        cols = ["c", "b", "a"]
        df = pd.DataFrame({"a": [0, 1], "b": [2, 3], "c": [4, 5]}, columns=cols)
        ds = Dataset.from_pandas(df)
        self.assertEqual(list(ds.keys()), cols)
        check_type = FastArray if Dataset.UseFastArray else np.ndarray
        self.assertIsInstance(ds.a, check_type)
        self.assertIsInstance(ds.b, check_type)
        self.assertIsInstance(ds.c, check_type)
        self.assertTrue((ds.a == [0, 1]).all())
        self.assertTrue((ds.b == [2, 3]).all())
        self.assertTrue((ds.c == [4, 5]).all())

    def test_ds_df_roundtrip_with_categoricals_and_datetimes(self):
        from enum import IntEnum

        def _bytes_to_string(arr):
            if arr.dtype.char == "S":
                arr = arr.astype("U")
            return arr

        def _datetime_to_int(arr):
            return arr.apply(lambda x: x.strftime("%Y%m%d")).astype(int)

        # Create a DataFrame
        df = pd.DataFrame({"A": ["a", "b", "c", "a"]})
        df["B"] = df["A"].astype("category")
        codes = [0, 1, 0, 2]
        df["C"] = pd.Categorical.from_codes(codes, list("qwe"))
        df["D"] = pd.Categorical.from_codes(codes, np.arange(3) ** 2)
        df["E"] = pd.Categorical.from_codes(codes, np.arange(0, 0.3, 0.1))
        df["F"] = pd.Categorical.from_codes(codes, np.arange(3, dtype="float32") ** 2)
        df["G"] = pd.Series(pd.date_range("20190101", periods=4))
        for (key, tz) in [
            ("H", "UTC"),
            ("I", "GMT"),
            ("J", "US/Eastern"),
            ("K", "Europe/Dublin"),
        ]:
            df[key] = pd.Series(pd.date_range("20190101", periods=4, tz=tz))

        # Create a Dataset from the DataFrame and compare
        ds = Dataset.from_pandas(df)
        self.assertIsInstance(ds.B, TypeRegister.Categorical)
        self.assertIsInstance(ds.C, TypeRegister.Categorical)
        self.assertTrue((df.A == ds.A.astype(str)).all())
        for key in "BCDEF":
            self.assertTrue((df[key] == _bytes_to_string(ds[key].expand_array)).all())
        for (key, tz) in [
            ("G", "UTC"),
            ("H", "UTC"),
            ("I", "GMT"),
            ("J", "NYC"),
            ("K", "DUBLIN"),
        ]:
            self.assertTrue((_datetime_to_int(df[key]) == ds[key].yyyymmdd).all())
            self.assertTrue(ds[key]._timezone._from_tz == "UTC")
            self.assertTrue(ds[key]._timezone._to_tz == tz)

        # Make two of the categoricals into enum type
        class MyEnum(IntEnum):
            a = 0
            b = 1
            c = 2

        ds.B = TypeRegister.Categorical([0, 1, 2, 0], MyEnum)
        ds.C = TypeRegister.Categorical(codes, {"q": 0, "w": 1, "e": 2})

        # Create a DataFrame from the Dataset and compare with the original
        # DataFrame (ie. test the round-trip)
        df2 = ds.to_pandas()
        self.assertEqual(df2.B.dtype.name, "category")
        self.assertEqual(df2.C.dtype.name, "category")
        for key in df:
            # Timezone unaware pandas datetime should get converted to UTC in riptable
            if key == "G":
                self.assertTrue((pd.DatetimeIndex(df[key]).tz_localize("UTC") == df2[key]).all())
            else:
                self.assertTrue((df[key] == df2[key]).all())

    def test_to_pandas_categorical_mapping(self):
        country_map = {0: "USA", 2: "IRL", 4: "GBR", 8: "AUS", 16: "CHN", 32: "JPN"}
        k = "dict_cat"
        ds = Dataset({k: Categorical([0, 2, 32], categories=country_map)})
        df = ds.to_pandas()
        ds_from_pandas = Dataset.from_pandas(df)
        assert_array_equal(ds[k].as_singlekey().expand_array, ds_from_pandas[k].expand_array)
        assert_array_equal(ds[k].category_array, ds_from_pandas[k].category_array)

    def test_from_pandas_num_string(self):
        dates = ["20190101", "20190101"]
        df = pd.DataFrame({"a": dates})
        ds = Dataset.from_pandas(df)
        assert_array_equal(ds["a"], np.array(dates, dtype="S"))

    def test_from_pandas_obj_number(self):
        df = pd.DataFrame({"a": np.array([1, 2, np.nan], dtype="O")})
        ds = Dataset.from_pandas(df)
        assert_array_almost_equal(ds["a"], np.array([1.0, 2.0, np.nan]))

    def test_from_pandas_obj_unicode(self):
        df = pd.DataFrame({"a": np.array(["Okay", "Namé"], dtype="O")})
        ds = Dataset.from_pandas(df)
        assert_array_equal(ds["a"], np.array(["Okay", "Namé"]))

    def test_to_pandas_timespan(self):
        dt_list = [np.nan, 11096000000000, 86401000000000]
        ds = Dataset({"a": TimeSpan(dt_list)})
        df = ds.to_pandas()
        assert_array_equal(df["a"], pd.to_timedelta(dt_list))

    def test_from_pandas_timespan(self):
        dt_list = [np.nan, 11096000000000, 86401000000000]
        df = pd.DataFrame({"a": pd.to_timedelta(dt_list)})
        ds = Dataset.from_pandas(df)
        assert_array_equal(ds["a"], TimeSpan(dt_list))

    def test_to_pandas_nullable_int(self):
        for dtype_str, pd_dtype_str in zip(
            ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"],
            ["Int8", "Int16", "Int32", "Int64", "UInt8", "UInt16", "UInt32", "UInt64"],
        ):
            dtype = np.dtype(dtype_str)
            invalid_value = INVALID_DICT[dtype.num]
            ds = Dataset({"a": FastArray([1, 2, invalid_value], dtype=dtype)})
            # not all versions of pandas have .array
            if hasattr(pd, "array"):
                df = ds.to_pandas()
                expected_arr = pd.Series([1, 2, np.nan], dtype=pd_dtype_str)
                self.assertTrue(
                    df["a"].equals(expected_arr),
                    msg=f"actual:\n{df.a}\nexpected:\n{expected_arr}",
                )

    def test_to_pandas_keep_invalid_sentinel(self):
        for dtype_str in ["i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8"]:
            dtype = np.dtype(dtype_str)
            invalid_value = INVALID_DICT[dtype.num]
            ds = Dataset({"a": FastArray([1, 2, invalid_value], dtype=dtype)})
            df = ds.to_pandas(use_nullable=False)
            expected_arr = pd.Series([1, 2, invalid_value], dtype=dtype)
            self.assertTrue(
                df["a"].equals(expected_arr),
                msg=f"actual:\n{df.a}\nexpected:\n{expected_arr}",
            )

    def test_from_pandas_nullable_int(self):
        # not all versions of pandas have .array
        if hasattr(pd, "array"):
            for dtype_str, pd_dtype_str in zip(
                [
                    "int8",
                    "int16",
                    "int32",
                    "int64",
                    "uint8",
                    "uint16",
                    "uint32",
                    "uint64",
                ],
                [
                    "Int8",
                    "Int16",
                    "Int32",
                    "Int64",
                    "UInt8",
                    "UInt16",
                    "UInt32",
                    "UInt64",
                ],
            ):
                df = pd.DataFrame({"a": pd.array([1, 2, np.nan], dtype=pd_dtype_str)})
                ds = Dataset.from_pandas(df)
                expected_arr = FastArray([1, 2, 3], dtype=dtype_str)
                expected_arr[2] = INVALID_DICT[expected_arr.dtype.num]
                assert_array_equal(ds["a"], expected_arr)

    def test_from_pandas_no_index_zero_string_column(self):
        df = pd.DataFrame({"a": [b"a", b"b"]}, index=[1, 2])
        ds = Dataset.from_pandas(df)
        assert_array_equal(ds["a"], np.array([b"a", b"b"]))

    def test_to_pandas_dst_nonexistent(self):
        ds = Dataset({"a": DateTimeNano(["19900401 02:00:00"], from_tz="NYC")})
        df = ds.to_pandas()
        fmt_string = "%Y%m%d %H:%M:%S.%f"
        expected_arr = ds["a"].strftime(fmt_string)
        arr = df["a"].dt.strftime(fmt_string)
        assert_array_equal(arr, expected_arr)

    def test_as_pandas_df_warn(self):
        ds = Dataset({"a": [1, 2, 3]})
        with self.assertWarns(FutureWarning):
            df = ds.as_pandas_df()

    def test_imatrix(self):
        arrsize = 3
        ds = Dataset({"time": arange(arrsize * 1.0), "data": arange(arrsize)})
        im = ds.imatrix_make(dtype=np.int32)
        assert_array_equal(im[:, 0], arange(arrsize, dtype=np.int32))

        im = ds.imatrix_make(dtype=np.float32)
        assert_array_equal(im[:, 1], arange(arrsize, dtype=np.float32))

    def test_imatrix_isnan(self):
        ds = Dataset({"a": [np.nan, np.nan, np.nan], "b": arange(3)})
        x = ds.imatrix_make()
        assert_array_equal(x.isnan(), np.isnan(x))

        ds = Dataset({"a": [np.nan, np.nan, np.nan], "b": arange(3) + 1})
        x = ds.imatrix_make()
        result = ds.imatrix_make().fillna(0)

        assert_array_equal(x.isnan(), np.isnan(x))
        test = np.asfortranarray(
            [
                [
                    0,
                    1,
                ],
                [0, 2],
                [0, 3],
            ]
        )
        self.assertTrue(np.all(test == result))

    def test_dataset_footer_copy(self):
        ds = Dataset({"x1": ["A", "B"]})
        ds.footer_set_values("sum", {"x1": 3})
        ds.footer_remove()
        ds.copy()

    def test_equals(self):
        ds = Dataset({"a": [True, True, False, True], "b": [False, False, True, False]})
        ds2 = Dataset({"a": [True, True, True, True], "b": [False, False, False, False]})
        self.assertTrue(ds.equals(ds))
        self.assertTrue(ds.equals(ds2) == False)

    def test_pivot2(self):
        ds2 = Dataset(
            {
                "date": [20190101] * 4,
                "symbol": ["AAPL", "MSFT", "EBAY", "MSFT"],
                "test": ["n1", "n2", "n3", "n2"],
                "shares": [100, 200, 300, 200],
            }
        )
        with self.assertRaises(ValueError):
            ds2.pivot("test", columns="symbol", values="shares")

        ds = Dataset(
            {
                "foo": ["one", "one", "one", "two", "two", "two"],
                "bar": ["A", "B", "C", "A", "B", "C"],
                "baz": [1, 2, 3, 4, 5, 6],
                "zoo": ["x", "y", "z", "q", "w", "t"],
            }
        )

        ms = ds.pivot(labels="foo", columns="bar")
        newds = ds.pivot(labels="foo", columns="bar", values="baz")
        self.assertTrue(np.all(newds.A == FastArray([1, 4])))

    def test_allnames(self):
        previous = Struct.AllNames
        Struct.AllNames = True
        ds = Dataset({"a": arange(5), "_test": arange(5.0)})
        self.assertTrue(np.all(ds._test == arange(5.0)))
        ds["sum"] = 3
        self.assertTrue(np.all(ds["sum"] == FastArray([3, 3, 3, 3, 3])))
        Struct.AllNames = previous

    def test_scalar_expansion(self):
        ds = Dataset({"a": arange(5), "b": arange(5.0), "c": list("abcde")})
        x = utcnow(1)
        # test adding a scalar or array of 1 to dataset of row lwength 5
        ds.x = x
        y = x[0]
        ds.y = y
        self.assertTrue(np.all(ds.x == ds.y))
        z = x - utcnow(1)
        ds.z = -z
        ds.z1 = -z[0]
        self.assertTrue(np.all(ds.z == ds.z1))

    def test_pandas_objectarray(self):
        import pandas as pd

        df = pd.DataFrame({"test": np.arange(10).astype("U")})
        x = Dataset(df)
        self.assertTrue(x.test[1] == b"1")
        x = Dataset(df, unicode=True)
        self.assertTrue(x.test[1] == "1")

    def test_filter_mask_copy(self):
        ds = Dataset()
        # also check that len returns 0 instead of None
        self.assertTrue(len(ds) == 0)
        ds.MyStrKey = Categorical(list("ABCD"))
        ds.MyIntKey = [1, 2, 3, 4]
        ds.MyDate = Date([1, 2, 3, 4])
        ds.MyValue = 1, 2, 1, 2
        output = ds.filter(ds.MyValue == 1, inplace=False)
        self.assertIsNot(ds, output)
        self.assertEqual(ds.get_nrows(), 4)
        self.assertEqual(output.get_nrows(), 2)
        self.assertEqual(ds.get_ncols(), 4)
        self.assertEqual(output.get_ncols(), 4)
        # TODO: Assert column ordering is the same before/after the operation.
        result = output.cat(["MyStrKey", "MyDate"]).null()
        self.assertTrue(isinstance(result["MyDate"], Date))

    def test_filter_mask_inplace(self):
        ds = Dataset()
        # also check that len returns 0 instead of None
        self.assertTrue(len(ds) == 0)
        ds.MyStrKey = Categorical(list("ABCD"))
        ds.MyIntKey = [1, 2, 3, 4]
        ds.MyDate = Date([1, 2, 3, 4])
        ds.MyValue = 1, 2, 1, 2
        output = ds.filter(ds.MyValue == 1, inplace=True)
        self.assertIs(ds, output)
        self.assertEqual(ds.get_nrows(), 2)
        self.assertEqual(output.get_nrows(), 2)
        self.assertEqual(ds.get_ncols(), 4)
        self.assertEqual(output.get_ncols(), 4)
        # TODO: Assert column ordering is the same before/after the operation.
        result = output.cat(["MyStrKey", "MyDate"]).null()
        self.assertTrue(isinstance(result["MyDate"], Date))

    def test_filter_fancy_copy(self):
        ds = Dataset()
        # also check that len returns 0 instead of None
        self.assertTrue(len(ds) == 0)
        ds.MyStrKey = Categorical(list("ABCD"))
        ds.MyIntKey = [1, 2, 3, 4]
        ds.MyDate = Date([1, 2, 3, 4])
        ds.MyValue = 1, 2, 1, 2
        output = ds.filter([0, 2], inplace=False)
        self.assertIsNot(ds, output)
        self.assertEqual(ds.get_nrows(), 4)
        self.assertEqual(output.get_nrows(), 2)
        self.assertEqual(ds.get_ncols(), 4)
        self.assertEqual(output.get_ncols(), 4)
        # TODO: Assert column ordering is the same before/after the operation.
        result = output.cat(["MyStrKey", "MyDate"]).null()
        self.assertTrue(isinstance(result["MyDate"], Date))

    def test_filter_fancy_inplace(self):
        ds = Dataset()
        # also check that len returns 0 instead of None
        self.assertTrue(len(ds) == 0)
        ds.MyStrKey = Categorical(list("ABCD"))
        ds.MyIntKey = [1, 2, 3, 4]
        ds.MyDate = Date([1, 2, 3, 4])
        ds.MyValue = 1, 2, 1, 2
        output = ds.filter([0, 2], inplace=True)
        self.assertIs(ds, output)
        self.assertEqual(ds.get_nrows(), 2)
        self.assertEqual(output.get_nrows(), 2)
        self.assertEqual(ds.get_ncols(), 4)
        self.assertEqual(output.get_ncols(), 4)
        # TODO: Assert column ordering is the same before/after the operation.
        result = output.cat(["MyStrKey", "MyDate"]).null()
        self.assertTrue(isinstance(result["MyDate"], Date))

    def test_filter_scalar_copy(self):
        ds = Dataset()
        # also check that len returns 0 instead of None
        self.assertTrue(len(ds) == 0)
        ds.MyStrKey = Categorical(list("ABCD"))
        ds.MyIntKey = [1, 2, 3, 4]
        ds.MyDate = Date([1, 2, 3, 4])
        ds.MyValue = 1, 2, 1, 2
        output = ds.filter(2, inplace=False)
        self.assertIsNot(ds, output)
        self.assertEqual(ds.get_nrows(), 4)
        self.assertEqual(output.get_nrows(), 1)
        self.assertEqual(ds.get_ncols(), 4)
        self.assertEqual(output.get_ncols(), 4)
        # TODO: Assert column ordering is the same before/after the operation.
        result = output.cat(["MyStrKey", "MyDate"]).null()
        self.assertTrue(isinstance(result["MyDate"], Date))

    def test_filter_mask_inplace_norows(self):
        ds = Dataset()
        # also check that len returns 0 instead of None
        self.assertTrue(len(ds) == 0)
        ds.MyStrKey = Categorical(list("ABCD"))
        ds.MyIntKey = [1, 2, 3, 4]
        ds.MyDate = Date([1, 2, 3, 4])
        ds.MyValue = 1, 2, 1, 2
        output = ds.filter(ds.MyValue == 13, inplace=True)
        self.assertIs(ds, output)
        self.assertEqual(ds.get_nrows(), 0)
        self.assertEqual(output.get_nrows(), 0)
        self.assertEqual(ds.get_ncols(), 4)
        self.assertEqual(output.get_ncols(), 4)
        # TODO: Assert column ordering is the same before/after the operation.
        result = output.cat(["MyStrKey", "MyDate"]).null()
        self.assertTrue(isinstance(result["MyDate"], Date))

    def test_filter_mask_copy_norows(self):
        ds = Dataset()
        # also check that len returns 0 instead of None
        self.assertTrue(len(ds) == 0)
        ds.MyStrKey = Categorical(list("ABCD"))
        ds.MyIntKey = [1, 2, 3, 4]
        ds.MyDate = Date([1, 2, 3, 4])
        ds.MyValue = 1, 2, 1, 2
        output = ds.filter(ds.MyValue == 13, inplace=False)
        self.assertIsNot(ds, output)
        self.assertEqual(ds.get_nrows(), 4)
        self.assertEqual(output.get_nrows(), 0)
        self.assertEqual(ds.get_ncols(), 4)
        self.assertEqual(output.get_ncols(), 4)
        # TODO: Assert column ordering is the same before/after the operation.
        result = output.cat(["MyStrKey", "MyDate"]).null()
        self.assertTrue(isinstance(result["MyDate"], Date))

    def test_filter_mask_length_mismatch_raises(self):
        ds = Dataset()
        # also check that len returns 0 instead of None
        self.assertTrue(len(ds) == 0)
        ds.MyStrKey = Categorical(list("ABCD"))
        ds.MyIntKey = [1, 2, 3, 4]
        ds.MyDate = Date([1, 2, 3, 4])
        ds.MyValue = 1, 2, 1, 2
        with self.assertRaises(ValueError):
            _ = ds.filter((ds.MyValue == 13)[:3], inplace=False)

    def test_underscorename(self):
        ds = Dataset()
        ds.num = arange(10)
        with self.assertRaises(IndexError):
            ds["_num"] = arange(10)

    def test_overwrite_column_with_scalar(self) -> None:
        # Test that a Dataset column can be overwritten with a scalar,
        # and that scalar is broadcast to the correct size to match the Dataset.
        # This is a regression test for some behavior that appears to be broken by
        # https://github.com/rtosholdings/riptable/commit/811960b3de521e19a1945602fb5a8b2193845b1d
        ds = rt.Dataset(
            {
                "a": rt.full(20, 1.2345, dtype=np.float32),
                "b": rt.arange(20, dtype=np.uint64),
                "c": rt.FA([11, -13, -17, 19, 23]).tile(4),
            }
        )
        orig_rowcount = ds.get_nrows()

        # Overwrite each of the columns with a scalar value.
        scalars = {"a": np.int16(12345), "b": True, "c": "hello"}
        assert set(ds.keys()) == set(scalars.keys())

        for col_name in ds.keys():
            ds[col_name] = scalars[col_name]

        assert ds.get_nrows() == orig_rowcount

        # Make sure each of the columns has the expected type and value
        # given the scalar that was assigned to it.
        for col_name in ds.keys():
            col = ds[col_name]
            scalar_value = scalars[col_name]
            expected_dtype = (
                scalar_value.dtype if isinstance(scalar_value, np.generic) else np.min_scalar_type(scalar_value)
            )
            assert col.dtype == expected_dtype
            assert_array_equal(col, scalar_value)


@pytest.mark.parametrize("categorical", get_all_categorical_data())
def test_dataset_to_dataframe_roundtripping(categorical):
    k = "categorical"
    ds = Dataset({k: categorical})
    df = ds.to_pandas()
    ds_from_pandas = Dataset.from_pandas(df)
    assert_array_equal(ds[k].as_singlekey().expand_array, ds_from_pandas[k].expand_array)
    # add support for checking category arrays of multikey categoricals
    if categorical.category_mode != CategoryMode.MultiKey:
        assert_array_equal(ds[k].category_array, ds_from_pandas[k].category_array)


if __name__ == "__main__":
    tester = unittest.main()
