Repository: TileDB-Inc/TileDB-Py
Pull request #2046 (2024-08-21 09:52:54)

Make `from_pandas` respect `column_types` for index dimensions
Historically, `from_pandas` did not consider the `column_types` argument for index columns, only for regular columns. This PR resolves that issue.

---

[sc-19316]

johnkerl: Just a quick note -- TileDB attrs can handle UTF-8 or ASCII; dims, ASCII only but not UTF-8 at this time. (And per @davisp it would be a significant project to support these.) Just a heads-up in case a user requests UTF-8 on a dim.

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 254f9b00..ae25be63 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -62,6 +62,7 @@ test = [
     "psutil",
     "pyarrow",
     "pandas",
+    "dask[distributed]",
 ]
 
 [project.urls]
diff --git a/tiledb/dataframe_.py b/tiledb/dataframe_.py
index 9543d91a..489bb10f 100644
--- a/tiledb/dataframe_.py
+++ b/tiledb/dataframe_.py
@@ -420,7 +420,9 @@ def _sparse_from_dtypes(dtypes, sparse=None):
     return sparse if sparse is not None else False
 
 
-def create_dims(df, index_dims, tile=None, full_domain=False, filters=None):
+def create_dims(
+    df, index_dims, column_infos, tile=None, full_domain=False, filters=None
+):
     check_dataframe_deps()
     import pandas as pd
 
@@ -445,7 +447,10 @@ def create_dims(df, index_dims, tile=None, full_domain=False, filters=None):
         else:
             raise ValueError(f"Unknown column or index named {name!r}")
 
-        dtype = ColumnInfo.from_values(values).dtype
+        if name in column_infos:
+            dtype = column_infos[name].dtype
+        else:
+            dtype = ColumnInfo.from_values(values).dtype
         internal_dtype = dtype
 
         if name == "__tiledb_rows" and isinstance(index, pd.RangeIndex):
@@ -659,6 +664,7 @@ def _create_array(uri, df, sparse, full_domain, index_dims, column_infos, tiledb
     dims, dim_metadata = create_dims(
         df,
         index_dims,
+        column_infos,
         full_domain=full_domain,
         tile=tiledb_args.get("tile"),
         filters=tiledb_args.get("dim_filters", True),
diff --git a/tiledb/libtiledb.pyx b/tiledb/libtiledb.pyx
index dd871839..0a20b112 100644
--- a/tiledb/libtiledb.pyx
+++ b/tiledb/libtiledb.pyx
@@ -2158,7 +2158,7 @@ cdef class Query(object):
 
     def label_index(self, labels):
         """Apply Array.label_index with query parameters."""
-        from .multirange_indexer import LabelIndexer
+        from .multirange_indexing import LabelIndexer
         return LabelIndexer(self.array, tuple(labels), query=self)
 
     @property


Test patch:
diff --git a/tiledb/tests/test_dask.py b/tiledb/tests/test_dask.py
index 6ca8da76..4d275190 100644
--- a/tiledb/tests/test_dask.py
+++ b/tiledb/tests/test_dask.py
@@ -10,10 +10,10 @@ import tiledb
 from .common import DiskTestCase
 
 # Skip this test if dask is unavailable
-da = pytest.importorskip("dask.array")
+da_array = pytest.importorskip("dask.array")
+da_distributed = pytest.importorskip("dask.distributed")
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="does not run on windows")
 class TestDaskSupport(DiskTestCase):
     def test_dask_from_numpy_1d(self):
         uri = self.path("np_1attr")
@@ -22,10 +22,10 @@ class TestDaskSupport(DiskTestCase):
         T.close()
 
         with tiledb.open(uri) as T:
-            D = da.from_tiledb(T)
+            D = da_array.from_tiledb(T)
             np.testing.assert_array_equal(D, A)
 
-        D2 = da.from_tiledb(uri)
+        D2 = da_array.from_tiledb(uri)
         np.testing.assert_array_equal(D2, A)
         self.assertAlmostEqual(
             np.mean(A), D2.mean().compute(scheduler="single-threaded")
@@ -43,18 +43,6 @@ class TestDaskSupport(DiskTestCase):
         tiledb.DenseArray.create(uri, schema)
 
     @pytest.mark.filterwarnings("ignore:There is no current event loop")
-    @pytest.mark.filterwarnings(
-        # In Python 3.7 on POSIX systems, Hurricane outputs a warning message
-        # that "make_current is deprecated." This should be fixed by Dask in
-        # future releases.
-        "ignore:make_current is deprecated"
-    )
-    @pytest.mark.skipif(
-        condition=(
-            sys.version_info >= (3, 11) and (datetime.now() < datetime(2023, 1, 6))
-        ),
-        reason="https://github.com/dask/distributed/issues/6785",
-    )
     def test_dask_multiattr_2d(self):
         uri = self.path("multiattr")
 
@@ -66,32 +54,30 @@ class TestDaskSupport(DiskTestCase):
             T[:] = {"attr1": ar1, "attr2": ar2}
         with tiledb.DenseArray(uri, mode="r", attr="attr2") as T:
             # basic round-trip from dask.array
-            D = da.from_tiledb(T, attribute="attr2")
+            D = da_array.from_tiledb(T, attribute="attr2")
             np.testing.assert_array_equal(ar2, np.array(D))
 
         # smoke-test computation
         # note: re-init from_tiledb each time, or else dask just uses the cached materialization
-        D = da.from_tiledb(uri, attribute="attr2")
+        D = da_array.from_tiledb(uri, attribute="attr2")
         self.assertAlmostEqual(np.mean(ar2), D.mean().compute(scheduler="threads"))
-        D = da.from_tiledb(uri, attribute="attr2")
+        D = da_array.from_tiledb(uri, attribute="attr2")
         self.assertAlmostEqual(
             np.mean(ar2), D.mean().compute(scheduler="single-threaded")
         )
-        D = da.from_tiledb(uri, attribute="attr2")
+        D = da_array.from_tiledb(uri, attribute="attr2")
         self.assertAlmostEqual(np.mean(ar2), D.mean().compute(scheduler="processes"))
 
         # test dask.distributed
-        from dask.distributed import Client
-
-        D = da.from_tiledb(uri, attribute="attr2")
-        with Client():
+        D = da_array.from_tiledb(uri, attribute="attr2")
+        with da_distributed.Client():
             np.testing.assert_approx_equal(D.mean().compute(), np.mean(ar2))
 
     def test_dask_write(self):
         uri = self.path("dask_w")
-        D = da.random.random(10, 10)
+        D = da_array.random.random(10, 10)
         D.to_tiledb(uri)
-        DT = da.from_tiledb(uri)
+        DT = da_array.from_tiledb(uri)
         np.testing.assert_array_equal(D, DT)
 
     def test_dask_overlap_blocks(self):
@@ -101,10 +87,10 @@ class TestDaskSupport(DiskTestCase):
         T.close()
 
         with tiledb.open(uri) as T:
-            D = da.from_tiledb(T)
+            D = da_array.from_tiledb(T)
             np.testing.assert_array_equal(D, A)
 
-        D2 = da.from_tiledb(uri)
+        D2 = da_array.from_tiledb(uri)
         np.testing.assert_array_equal(D2, A)
 
         D3 = D2.map_overlap(
@@ -133,7 +119,7 @@ class TestDaskSupport(DiskTestCase):
         with tiledb.open(uri, "w", attr="TDB_VALUES") as T:
             T[:] = A
 
-        D2 = da.from_tiledb(uri, attribute="TDB_VALUES")
+        D2 = da_array.from_tiledb(uri, attribute="TDB_VALUES")
 
         D3 = D2.map_overlap(
             lambda x: x + 1, depth={0: 0, 1: 1, 2: 1}, dtype=D2.dtype, boundary="none"
@@ -160,7 +146,7 @@ class TestDaskSupport(DiskTestCase):
         with tiledb.open(uri, "w", attr="TDB_VALUES") as D1:
             D1[:] = A
 
-        D2 = da.from_tiledb(uri, attribute="TDB_VALUES")
+        D2 = da_array.from_tiledb(uri, attribute="TDB_VALUES")
 
         D3 = D2.map_blocks(lambda x: x + 1, dtype=D2.dtype).compute(
             scheduler="processes"
@@ -168,11 +154,6 @@ class TestDaskSupport(DiskTestCase):
         np.testing.assert_array_equal(D2 + 1, D3)
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="does not run on windows")
-@pytest.mark.skipif(
-    sys.version_info[:2] == (3, 8),
-    reason="Fails on Python 3.8 due to dask worker restarts",
-)
 def test_sc33742_dask_array_object_dtype_conversion():
     # This test verifies that an array can be converted to buffer after serialization
     # through several dask.distributed compute steps. The original source of the issue
@@ -182,7 +163,6 @@ def test_sc33742_dask_array_object_dtype_conversion():
 
     import dask
     import numpy as np
-    from dask.distributed import Client, LocalCluster
 
     @dask.delayed
     def get_data():
@@ -213,7 +193,9 @@ def test_sc33742_dask_array_object_dtype_conversion():
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
         global client
-        client = Client(LocalCluster(scheduler_port=9786, dashboard_address=9787))
+        client = da_distributed.Client(
+            da_distributed.LocalCluster(scheduler_port=9786, dashboard_address=9787)
+        )
 
         w = []
 
diff --git a/tiledb/tests/test_dimension_label.py b/tiledb/tests/test_dimension_label.py
index 0daeab2e..bb7c98e5 100644
--- a/tiledb/tests/test_dimension_label.py
+++ b/tiledb/tests/test_dimension_label.py
@@ -1,3 +1,5 @@
+from collections import OrderedDict
+
 import numpy as np
 import pytest
 
@@ -416,3 +418,68 @@ class DimensionLabelTestCase(DiskTestCase):
                             result["value"], attr_data[:, index:]
                         )
                     np.testing.assert_array_equal(result[label_name], label_index)
+
+    @pytest.mark.skipif(
+        tiledb.libtiledb.version()[0] == 2 and tiledb.libtiledb.version()[1] < 15,
+        reason="dimension labels requires libtiledb version 2.15 or greater",
+    )
+    def test_dimension_label_on_query(self):
+        uri = self.path("query_label_index")
+
+        dim1 = tiledb.Dim("d1", domain=(1, 4))
+        dim2 = tiledb.Dim("d2", domain=(1, 3))
+        dom = tiledb.Domain(dim1, dim2)
+        att = tiledb.Attr("a1", dtype=np.int64)
+        dim_labels = {
+            0: {"l1": dim1.create_label_schema("decreasing", np.int64)},
+            1: {
+                "l2": dim2.create_label_schema("increasing", np.int64),
+                "l3": dim2.create_label_schema("increasing", np.float64),
+            },
+        }
+        schema = tiledb.ArraySchema(domain=dom, attrs=(att,), dim_labels=dim_labels)
+        tiledb.Array.create(uri, schema)
+
+        a1_data = np.reshape(np.arange(1, 13), (4, 3))
+        l1_data = np.arange(4, 0, -1)
+        l2_data = np.arange(-1, 2)
+        l3_data = np.linspace(0, 1.0, 3)
+
+        with tiledb.open(uri, "w") as A:
+            A[:] = {"a1": a1_data, "l1": l1_data, "l2": l2_data, "l3": l3_data}
+
+        with tiledb.open(uri, "r") as A:
+            np.testing.assert_equal(
+                A.query().label_index(["l1"])[3:4],
+                OrderedDict(
+                    {"l1": np.array([4, 3]), "a1": np.array([[1, 2, 3], [4, 5, 6]])}
+                ),
+            )
+            np.testing.assert_equal(
+                A.query().label_index(["l1", "l3"])[2, 0.5:1.0],
+                OrderedDict(
+                    {
+                        "l3": np.array([0.5, 1.0]),
+                        "l1": np.array([2]),
+                        "a1": np.array([[8, 9]]),
+                    }
+                ),
+            )
+            np.testing.assert_equal(
+                A.query().label_index(["l2"])[:, -1:0],
+                OrderedDict(
+                    {
+                        "l2": np.array([-1, 0]),
+                        "a1": np.array([[1, 2], [4, 5], [7, 8], [10, 11]]),
+                    },
+                ),
+            )
+            np.testing.assert_equal(
+                A.query().label_index(["l3"])[:, 0.5:1.0],
+                OrderedDict(
+                    {
+                        "l3": np.array([0.5, 1.0]),
+                        "a1": np.array([[2, 3], [5, 6], [8, 9], [11, 12]]),
+                    },
+                ),
+            )
diff --git a/tiledb/tests/test_pandas_dataframe.py b/tiledb/tests/test_pandas_dataframe.py
index 51c02d74..72a22f00 100644
--- a/tiledb/tests/test_pandas_dataframe.py
+++ b/tiledb/tests/test_pandas_dataframe.py
@@ -127,7 +127,7 @@ def make_dataframe_categorical():
     return df
 
 
-class TestColumnInfo:
+class TestColumnInfo(DiskTestCase):
     def assertColumnInfo(self, info, info_dtype, info_repr=None, info_nullable=False):
         assert isinstance(info.dtype, np.dtype)
         assert info.dtype == info_dtype
@@ -250,6 +250,96 @@ class TestColumnInfo:
                     # check that the column name is included in the error message
                     assert "supported (column foo)" in str(exc.value)
 
+    def test_apply_dtype_index_ingest(self):
+        uri = self.path("index_dtype_default_dtype")
+        tiledb.from_pandas(
+            uri,
+            pd.DataFrame({"a": np.arange(0, 20), "b": np.arange(20, 40)}),
+            sparse=True,
+            index_dims=["a"],
+        )
+        with tiledb.open(uri) as A:
+            if sys.platform == "win32" and sys.version_info[:2] == (3, 8):
+                assert A.schema.domain.dim(0).dtype == np.int32
+            else:
+                assert A.schema.domain.dim(0).dtype == np.int64
+
+        uri = self.path("index_dtype_casted_dtype")
+        tiledb.from_pandas(
+            uri,
+            pd.DataFrame({"a": np.arange(0, 20), "b": np.arange(20, 40)}),
+            sparse=True,
+            index_dims=["a"],
+            column_types={"a": np.uint8},
+        )
+        with tiledb.open(uri) as A:
+            assert A.schema.domain.dim(0).dtype == np.uint8
+
+        # multiple index dims
+        uri = self.path("index_dtype_default_dtype_multi")
+        tiledb.from_pandas(
+            uri,
+            pd.DataFrame(
+                {
+                    "a": np.random.random_sample(20),
+                    "b": [str(uuid.uuid4()) for _ in range(20)],
+                }
+            ),
+            sparse=True,
+            index_dims=["a", "b"],
+        )
+        with tiledb.open(uri) as A:
+            assert A.schema.domain.dim(0).dtype == np.float64
+            assert A.schema.domain.dim(1).dtype == np.bytes_
+
+        uri = self.path("index_dtype_casted_dtype_multi")
+        tiledb.from_pandas(
+            uri,
+            pd.DataFrame(
+                {
+                    "a": np.random.random_sample(20),
+                    "b": [str(uuid.uuid4()) for _ in range(20)],
+                }
+            ),
+            sparse=True,
+            index_dims=["a", "b"],
+            column_types={"a": np.float32, "b": np.bytes_},
+        )
+        with tiledb.open(uri) as A:
+            assert A.schema.domain.dim(0).dtype == np.float32
+            assert A.schema.domain.dim(1).dtype == np.bytes_
+
+    def test_apply_dtype_index_schema_only(self):
+        uri = self.path("index_dtype_casted_dtype")
+        tiledb.from_pandas(
+            uri,
+            pd.DataFrame({"a": np.arange(0, 20), "b": np.arange(20, 40)}),
+            sparse=True,
+            index_dims=["a"],
+            column_types={"a": np.uint8},
+            mode="schema_only",
+        )
+        with tiledb.open(uri) as A:
+            assert A.schema.domain.dim(0).dtype == np.uint8
+
+        uri = self.path("index_dtype_casted_dtype_multi")
+        tiledb.from_pandas(
+            uri,
+            pd.DataFrame(
+                {
+                    "a": np.random.random_sample(20),
+                    "b": [str(uuid.uuid4()) for _ in range(20)],
+                }
+            ),
+            sparse=True,
+            index_dims=["a", "b"],
+            column_types={"a": np.float32, "b": np.bytes_},
+            mode="schema_only",
+        )
+        with tiledb.open(uri) as A:
+            assert A.schema.domain.dim(0).dtype == np.float32
+            assert A.schema.domain.dim(1).dtype == np.bytes_
+
 
 class TestDimType:
     @pytest.mark.parametrize(