Skip to content

Commit 0592a4f

Browse files
committed
Complete fixed-size binary language bindings
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent ec1abcc commit 0592a4f

6 files changed

Lines changed: 45 additions & 1 deletion

File tree

lang/cpp/include/vortex/dtype.hpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ enum class DataTypeVariant {
3737
Decimal = DTYPE_DECIMAL,
3838
// Nested fixed-size list
3939
FixedSizeList = DTYPE_FIXED_SIZE_LIST,
40+
// Fixed-size binary data
41+
FixedSizeBinary = DTYPE_FIXED_SIZE_BINARY,
4042
};
4143

4244
// Primitive type
@@ -95,6 +97,7 @@ class DataType {
9597
DataType list_element() const;
9698
DataType fixed_size_list_element() const;
9799
uint32_t fixed_size_list_size() const;
100+
uint32_t fixed_size_binary_size() const;
98101

99102
private:
100103
friend struct detail::Access;
@@ -135,6 +138,7 @@ DataType float32(bool nullable = false);
135138
DataType float64(bool nullable = false);
136139
DataType utf8(bool nullable = false);
137140
DataType binary(bool nullable = false);
141+
DataType fixed_size_binary(uint32_t byte_width, bool nullable = false);
138142
DataType decimal(uint8_t precision, int8_t scale, bool nullable = false);
139143
DataType list(DataType element, bool nullable = false);
140144
DataType fixed_size_list(DataType element, uint32_t size, bool nullable = false);

lang/cpp/src/dtype.cpp

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,10 @@ uint32_t DataType::fixed_size_list_size() const {
124124
return vx_dtype_fixed_size_list_size(handle_.get());
125125
}
126126

127+
uint32_t DataType::fixed_size_binary_size() const {
128+
return vx_dtype_fixed_size_binary_size(handle_.get());
129+
}
130+
127131
namespace dtype {
128132

129133
DataType null() {
@@ -174,6 +178,9 @@ DataType utf8(bool nullable) {
174178
DataType binary(bool nullable) {
175179
return Access::adopt<DataType>(vx_dtype_new_binary(nullable));
176180
}
181+
DataType fixed_size_binary(uint32_t byte_width, bool nullable) {
182+
return Access::adopt<DataType>(vx_dtype_new_fixed_size_binary(byte_width, nullable));
183+
}
177184
DataType decimal(uint8_t precision, int8_t scale, bool nullable) {
178185
return Access::adopt<DataType>(vx_dtype_new_decimal(precision, scale, nullable));
179186
}

lang/cpp/tests/dtype.cpp

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,13 @@ TEST_CASE("Decimal dtype", "[dtype]") {
2626
REQUIRE_THROWS_AS(d.list_element(), VortexException);
2727
}
2828

29+
TEST_CASE("Fixed-size binary dtype", "[dtype]") {
30+
auto d = dtype::fixed_size_binary(16, dtype::Nullable);
31+
REQUIRE(d.variant() == DataTypeVariant::FixedSizeBinary);
32+
REQUIRE(d.fixed_size_binary_size() == 16);
33+
REQUIRE(d.nullable());
34+
}
35+
2936
TEST_CASE("copy dtype", "[dtype]") {
3037
auto d = dtype::int32(true);
3138
DataType d2 = d;

vortex-ffi/cinclude/vortex.h

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,10 @@ typedef enum {
111111
* Nested fixed-size list type.
112112
*/
113113
DTYPE_FIXED_SIZE_LIST = 9,
114+
/**
115+
* Fixed-size binary data.
116+
*/
117+
DTYPE_FIXED_SIZE_BINARY = 10,
114118
} vx_dtype_variant;
115119

116120
/**
@@ -1000,6 +1004,11 @@ const vx_dtype *vx_dtype_new_utf8(bool is_nullable);
10001004
*/
10011005
const vx_dtype *vx_dtype_new_binary(bool is_nullable);
10021006

1007+
/**
1008+
* Create a new fixed-size binary data type.
1009+
*/
1010+
const vx_dtype *vx_dtype_new_fixed_size_binary(uint32_t byte_width, bool is_nullable);
1011+
10031012
/**
10041013
* Create a new list data type.
10051014
*
@@ -1075,6 +1084,11 @@ const vx_dtype *vx_dtype_fixed_size_list_element(const vx_dtype *dtype);
10751084
*/
10761085
uint32_t vx_dtype_fixed_size_list_size(const vx_dtype *dtype);
10771086

1087+
/**
1088+
* Returns the byte width of a fixed-size binary type.
1089+
*/
1090+
uint32_t vx_dtype_fixed_size_binary_size(const vx_dtype *dtype);
1091+
10781092
/**
10791093
* Checks if the type is time.
10801094
*/

vortex-ffi/examples/dtype.c

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,10 @@ void print_decimal_dtype(const vx_dtype *dtype) {
8686
printf("decimal(precision=%u, scale=%d)", precision, scale);
8787
}
8888

89+
void print_fixed_size_binary_dtype(const vx_dtype *dtype) {
90+
printf("fixed binary(size=%u)", vx_dtype_fixed_size_binary_size(dtype));
91+
}
92+
8993
void print_dtype(const vx_dtype *dtype) {
9094
switch (vx_dtype_get_variant(dtype)) {
9195
case DTYPE_NULL:
@@ -118,6 +122,9 @@ void print_dtype(const vx_dtype *dtype) {
118122
case DTYPE_DECIMAL:
119123
print_decimal_dtype(dtype);
120124
break;
125+
case DTYPE_FIXED_SIZE_BINARY:
126+
print_fixed_size_binary_dtype(dtype);
127+
break;
121128
}
122129
printf("%c\n", vx_dtype_is_nullable(dtype) ? '?' : ' ');
123130
}

vortex-python/test/test_array.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,12 @@ def test_varbin_array_round_trip():
2525
assert arr.to_arrow_array() == a
2626

2727

28+
def test_fixed_size_binary_array_round_trip():
29+
a = pa.array([b"abc", None, b"def"], type=pa.binary(3))
30+
arr = vortex.array(a)
31+
assert arr.to_arrow_array() == a
32+
33+
2834
def test_varbin_array_take():
2935
a = vortex.array(pa.array(["a", "b", "c", "d"], type=pa.string_view()))
3036
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():
5056
[
5157
pa.duration("us"),
5258
pa.month_day_nano_interval(),
53-
pa.binary(3),
5459
],
5560
)
5661
def test_unsupported_arrow_type_raises_value_error(arrow_type: pa.DataType):

0 commit comments

Comments
 (0)