Cast enum values to enum dtype (#1854)

* For empty enumerations, if we do not explicitly cast, it is automatically casted to dtype of float64

Author: nguyenv

tiledb/enumeration.py

@@ -90,11 +90,11 @@ def values(self) -> NDArray:
         :rtype: NDArray
         """
         if self.dtype.kind == "U":
-            return np.array(super().str_values())
+            return np.array(super().str_values(), dtype=np.str_)
         elif self.dtype.kind == "S":
             return np.array(super().str_values(), dtype=np.bytes_)
         else:
-            return super().values()
+            return np.array(super().values(), dtype=self.dtype)
 
     def extend(self, values: Sequence[Any]) -> Enumeration:
         """Add additional values to the enumeration.

tiledb/tests/test_enumeration.py

@@ -116,3 +116,43 @@ def test_array_schema_enumeration_nullable(self, sparse, pass_df):
             expected_validity = [False, False, True, False, False]
             assert_array_equal(A[:]["a"].mask, expected_validity)
             assert_array_equal(A.df[:]["a"].isna(), expected_validity)
+
+    @pytest.mark.parametrize(
+        "dtype, values",
+        [
+            (np.int8, np.array([1, 2, 3], np.int8)),
+            (np.uint8, np.array([1, 2, 3], np.uint8)),
+            (np.int16, np.array([1, 2, 3], np.int16)),
+            (np.uint16, np.array([1, 2, 3], np.uint16)),
+            (np.int32, np.array([1, 2, 3], np.int32)),
+            (np.uint32, np.array([1, 2, 3], np.uint32)),
+            (np.int64, np.array([1, 2, 3], np.int64)),
+            (np.uint64, np.array([1, 2, 3], np.uint64)),
+            (np.dtype("S"), np.array(["a", "b", "c"], np.dtype("S"))),
+            (np.dtype("U"), np.array(["a", "b", "c"], np.dtype("U"))),
+        ],
+    )
+    def test_enum_dtypes(self, dtype, values):
+        # create empty
+        enmr = tiledb.Enumeration("e", False, dtype=dtype)
+        if dtype in (np.dtype("S"), np.dtype("U")):
+            assert enmr.dtype.kind == enmr.values().dtype.kind == dtype.kind
+        else:
+            assert enmr.dtype == enmr.values().dtype == dtype
+            assert_array_equal(enmr.values(), [])
+
+        # then extend with values
+        enmr = enmr.extend(values)
+        if dtype in (np.dtype("S"), np.dtype("U")):
+            assert enmr.dtype.kind == enmr.values().dtype.kind == dtype.kind
+        else:
+            assert enmr.dtype == enmr.values().dtype == dtype
+            assert_array_equal(enmr.values(), values)
+
+        # create with values
+        enmr = tiledb.Enumeration("e", False, values=values)
+        if dtype in (np.dtype("S"), np.dtype("U")):
+            assert enmr.dtype.kind == enmr.values().dtype.kind == dtype.kind
+        else:
+            assert enmr.dtype == enmr.values().dtype == dtype
+            assert_array_equal(enmr.values(), values)

tiledb/tests/test_schema_evolution.py

@@ -174,7 +174,12 @@ def test_schema_evolution_with_enmr(tmp_path):
 
 @pytest.mark.parametrize(
     "type,data",
-    (("int", [0]), ("bool", [True, False]), ("str", ["abc", "defghi", "jk"])),
+    (
+        ("int", [0]),
+        ("bool", [True, False]),
+        ("str", ["abc", "defghi", "jk"]),
+        ("bytes", [b"abc", b"defghi", b"jk"]),
+    ),
 )
 def test_schema_evolution_extend_enmr(tmp_path, type, data):
     uri = str(tmp_path)