Complete fixed-size binary language bindings

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

lang/cpp/include/vortex/dtype.hpp

@@ -37,6 +37,8 @@ enum class DataTypeVariant {
     Decimal = DTYPE_DECIMAL,
     // Nested fixed-size list
     FixedSizeList = DTYPE_FIXED_SIZE_LIST,
+    // Fixed-size binary data
+    FixedSizeBinary = DTYPE_FIXED_SIZE_BINARY,
 };
 
 // Primitive type
@@ -95,6 +97,7 @@ class DataType {
     DataType list_element() const;
     DataType fixed_size_list_element() const;
     uint32_t fixed_size_list_size() const;
+    uint32_t fixed_size_binary_size() const;
 
 private:
     friend struct detail::Access;
@@ -135,6 +138,7 @@ DataType float32(bool nullable = false);
 DataType float64(bool nullable = false);
 DataType utf8(bool nullable = false);
 DataType binary(bool nullable = false);
+DataType fixed_size_binary(uint32_t byte_width, bool nullable = false);
 DataType decimal(uint8_t precision, int8_t scale, bool nullable = false);
 DataType list(DataType element, bool nullable = false);
 DataType fixed_size_list(DataType element, uint32_t size, bool nullable = false);

lang/cpp/src/dtype.cpp

@@ -124,6 +124,10 @@ uint32_t DataType::fixed_size_list_size() const {
     return vx_dtype_fixed_size_list_size(handle_.get());
 }
 
+uint32_t DataType::fixed_size_binary_size() const {
+    return vx_dtype_fixed_size_binary_size(handle_.get());
+}
+
 namespace dtype {
 
 DataType null() {
@@ -174,6 +178,9 @@ DataType utf8(bool nullable) {
 DataType binary(bool nullable) {
     return Access::adopt<DataType>(vx_dtype_new_binary(nullable));
 }
+DataType fixed_size_binary(uint32_t byte_width, bool nullable) {
+    return Access::adopt<DataType>(vx_dtype_new_fixed_size_binary(byte_width, nullable));
+}
 DataType decimal(uint8_t precision, int8_t scale, bool nullable) {
     return Access::adopt<DataType>(vx_dtype_new_decimal(precision, scale, nullable));
 }

lang/cpp/tests/dtype.cpp

@@ -26,6 +26,13 @@ TEST_CASE("Decimal dtype", "[dtype]") {
     REQUIRE_THROWS_AS(d.list_element(), VortexException);
 }
 
+TEST_CASE("Fixed-size binary dtype", "[dtype]") {
+    auto d = dtype::fixed_size_binary(16, dtype::Nullable);
+    REQUIRE(d.variant() == DataTypeVariant::FixedSizeBinary);
+    REQUIRE(d.fixed_size_binary_size() == 16);
+    REQUIRE(d.nullable());
+}
+
 TEST_CASE("copy dtype", "[dtype]") {
     auto d = dtype::int32(true);
     DataType d2 = d;

vortex-ffi/cinclude/vortex.h

@@ -111,6 +111,10 @@ typedef enum {
      * Nested fixed-size list type.
      */
     DTYPE_FIXED_SIZE_LIST = 9,
+    /**
+     * Fixed-size binary data.
+     */
+    DTYPE_FIXED_SIZE_BINARY = 10,
 } vx_dtype_variant;
 
 /**
@@ -1021,6 +1025,11 @@ const vx_dtype *vx_dtype_new_utf8(bool is_nullable);
  */
 const vx_dtype *vx_dtype_new_binary(bool is_nullable);
 
+/**
+ * Create a new fixed-size binary data type.
+ */
+const vx_dtype *vx_dtype_new_fixed_size_binary(uint32_t byte_width, bool is_nullable);
+
 /**
  * Create a new list data type.
  *
@@ -1096,6 +1105,11 @@ const vx_dtype *vx_dtype_fixed_size_list_element(const vx_dtype *dtype);
  */
 uint32_t vx_dtype_fixed_size_list_size(const vx_dtype *dtype);
 
+/**
+ * Returns the byte width of a fixed-size binary type.
+ */
+uint32_t vx_dtype_fixed_size_binary_size(const vx_dtype *dtype);
+
 /**
  * Checks if the type is time.
  */

vortex-ffi/examples/dtype.c

@@ -86,6 +86,10 @@ void print_decimal_dtype(const vx_dtype *dtype) {
     printf("decimal(precision=%u, scale=%d)", precision, scale);
 }
 
+void print_fixed_size_binary_dtype(const vx_dtype *dtype) {
+    printf("fixed binary(size=%u)", vx_dtype_fixed_size_binary_size(dtype));
+}
+
 void print_dtype(const vx_dtype *dtype) {
     switch (vx_dtype_get_variant(dtype)) {
     case DTYPE_NULL:
@@ -118,6 +122,9 @@ void print_dtype(const vx_dtype *dtype) {
     case DTYPE_DECIMAL:
         print_decimal_dtype(dtype);
         break;
+    case DTYPE_FIXED_SIZE_BINARY:
+        print_fixed_size_binary_dtype(dtype);
+        break;
     }
     printf("%c\n", vx_dtype_is_nullable(dtype) ? '?' : ' ');
 }

vortex-python/test/test_array.py

@@ -25,6 +25,12 @@ def test_varbin_array_round_trip():
     assert arr.to_arrow_array() == a
 
 
+def test_fixed_size_binary_array_round_trip():
+    a = pa.array([b"abc", None, b"def"], type=pa.binary(3))
+    arr = vortex.array(a)
+    assert arr.to_arrow_array() == a
+
+
 def test_varbin_array_take():
     a = vortex.array(pa.array(["a", "b", "c", "d"], type=pa.string_view()))
     assert a.take(vortex.array(pa.array([0, 2]))).to_arrow_array() == pa.array(
@@ -50,7 +56,6 @@ def test_scalar_at_out_of_bounds():
     [
         pa.duration("us"),
         pa.month_day_nano_interval(),
-        pa.binary(3),
     ],
 )
 def test_unsupported_arrow_type_raises_value_error(arrow_type: pa.DataType):