Correctly set the validity buffer when castin dictionary to values

Author: XanthosXanthopoulos

apis/python/tests/test_dataframe.py

@@ -2784,6 +2784,52 @@ def test_extend_enumerations(tmp_path):
             assert (readback_df[c] == written_df[c]).all()
 
 
+def test_write_dictionary_to_non_enum_column(tmp_path):
+    written_df = pd.DataFrame(
+        {
+            "soma_joinid": pd.Series([0, 1, 2, 3, 4, 5], dtype=np.int64),
+            "str": pd.Series(["A", "B", "A", "B", "B", None], dtype="category"),
+            "byte": pd.Series([b"A", b"B", b"A", b"B", b"B", None], dtype="category"),
+            "bool": pd.Series([True, False, True, False, False, None], dtype="category"),
+            "int64": pd.Series([0, 1, 2, 0, 1, None], dtype="Int64").astype("category"),
+            "uint64": pd.Series([0, 1, 2, 0, 1, None], dtype="UInt64").astype("category"),
+            "int32": pd.Series([0, 1, 2, 0, 1, None], dtype="Int32").astype("category"),
+            "uint32": pd.Series([0, 1, 2, 0, 1, None], dtype="UInt32").astype("category"),
+            "int16": pd.Series([0, 1, 2, 0, 1, None], dtype="Int16").astype("category"),
+            "uint16": pd.Series([0, 1, 2, 0, 1, None], dtype="UInt16").astype("category"),
+            "int8": pd.Series([0, 1, 2, 0, 1, None], dtype="Int8").astype("category"),
+            "uint8": pd.Series([0, 1, 2, 0, 1, None], dtype="UInt8").astype("category"),
+            "float32": pd.Series([0, 1.1, 2.1, 0, 1.1, None], dtype="Float32").astype("category"),
+            "float64": pd.Series([0, 1.1, 2.1, 0, 1.1, None], dtype="Float64").astype("category"),
+        },
+    )
+
+    schema = pa.schema([
+        pa.field("soma_joinid", pa.int64()),
+        pa.field("str", pa.large_string(), nullable=True),
+        pa.field("byte", pa.large_binary(), nullable=True),
+        pa.field("bool", pa.bool_(), nullable=True),
+        pa.field("int64", pa.int64(), nullable=True),
+        pa.field("uint64", pa.uint64(), nullable=True),
+        pa.field("int32", pa.int32(), nullable=True),
+        pa.field("uint32", pa.uint32(), nullable=True),
+        pa.field("int16", pa.int16(), nullable=True),
+        pa.field("uint16", pa.uint16(), nullable=True),
+        pa.field("int8", pa.int8(), nullable=True),
+        pa.field("uint8", pa.uint8(), nullable=True),
+        pa.field("float32", pa.float32(), nullable=True),
+        pa.field("float64", pa.float64(), nullable=True),
+    ])
+
+    with soma.DataFrame.create(str(tmp_path), schema=schema, domain=[[0, 9]]) as soma_dataframe:
+        tbl = pa.Table.from_pandas(written_df, preserve_index=False)
+        soma_dataframe.write(tbl)
+
+    with soma.open(str(tmp_path)) as soma_dataframe:
+        readback_tbl = soma_dataframe.read().concat()
+        assert tbl.to_pylist() == readback_tbl.to_pylist()
+
+
 def test_multiple_writes_with_str_enums(tmp_path):
     uri = tmp_path.as_posix()
 

libtiledbsoma/src/soma/managed_query.cc

@@ -681,12 +681,12 @@ void ManagedQuery::_cast_dictionary_values(ArrowSchema* schema, ArrowArray* arra
                 "values");
     }
 
-    setup_write_column(
-        schema->name,
-        array->length,
-        std::move(data_buffer),
-        (uint64_t*)nullptr,
-        nullptr);  // validities are set by index column
+    std::unique_ptr<std::uint8_t[]> validity = nullptr;
+    if (schema->flags & ARROW_FLAG_NULLABLE) {
+        validity = _cast_validity_buffer_ptr(array);
+    }
+
+    setup_write_column(schema->name, array->length, std::move(data_buffer), (uint64_t*)nullptr, std::move(validity));
 }
 
 template <>
@@ -774,12 +774,13 @@ void ManagedQuery::_cast_dictionary_values<std::string>(ArrowSchema* schema, Arr
         std::tie(data_buffer, offset_buffer) = extract_values.operator()<uint32_t>(schema, array);
     }
 
+    std::unique_ptr<std::uint8_t[]> validity = nullptr;
+    if (schema->flags & ARROW_FLAG_NULLABLE) {
+        validity = _cast_validity_buffer_ptr(array);
+    }
+
     setup_write_column(
-        schema->name,
-        array->length,
-        std::move(data_buffer),
-        std::move(offset_buffer),
-        nullptr);  // validities are set by index column
+        schema->name, array->length, std::move(data_buffer), std::move(offset_buffer), std::move(validity));
 }
 
 template <>
@@ -829,12 +830,12 @@ void ManagedQuery::_cast_dictionary_values<bool>(ArrowSchema* schema, ArrowArray
                 "values");
     }
 
-    setup_write_column(
-        schema->name,
-        array->length,
-        std::move(data_buffer),
-        (uint64_t*)nullptr,
-        nullptr);  // validities are set by index column
+    std::unique_ptr<std::uint8_t[]> validity = nullptr;
+    if (schema->flags & ARROW_FLAG_NULLABLE) {
+        validity = _cast_validity_buffer_ptr(array);
+    }
+
+    setup_write_column(schema->name, array->length, std::move(data_buffer), (uint64_t*)nullptr, std::move(validity));
 }
 
 template <>