test[gpu]: cover binary export edge cases

Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>

Author: 0ax1

vortex-cuda/benches/arrow_binary_cuda.rs

@@ -26,7 +26,6 @@ use vortex::array::validity::Validity;
 use vortex::buffer::Buffer;
 use vortex::buffer::ByteBuffer;
 use vortex::dtype::DType;
-use vortex::dtype::Nullability;
 use vortex::error::VortexExpect;
 use vortex::error::VortexResult;
 use vortex::session::VortexSession;
@@ -45,6 +44,15 @@ async fn binary_on_device(
     views: Buffer<BinaryView>,
     buffers: Arc<[ByteBuffer]>,
     ctx: &mut CudaExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    binary_on_device_with_validity(views, buffers, Validity::NonNullable, ctx).await
+}
+
+async fn binary_on_device_with_validity(
+    views: Buffer<BinaryView>,
+    buffers: Arc<[ByteBuffer]>,
+    validity: Validity,
+    ctx: &mut CudaExecutionCtx,
 ) -> VortexResult<ArrayRef> {
     let views = ctx
         .ensure_on_device(BufferHandle::new_host(views.into_byte_buffer()))
@@ -60,8 +68,8 @@ async fn binary_on_device(
     Ok(VarBinViewArray::new_handle(
         views,
         device_buffers.into(),
-        DType::Binary(Nullability::NonNullable),
-        Validity::NonNullable,
+        DType::Binary(validity.nullability()),
+        validity,
     )
     .into_array())
 }
@@ -85,6 +93,14 @@ async fn out_of_line_binary(len: usize, ctx: &mut CudaExecutionCtx) -> VortexRes
     binary_on_device(views, Arc::from([values]), ctx).await
 }
 
+async fn sliced_validity_binary(len: usize, ctx: &mut CudaExecutionCtx) -> VortexResult<ArrayRef> {
+    let views =
+        Buffer::from_iter((0..len).map(|idx| BinaryView::make_view(&idx.to_le_bytes(), 0, 0)));
+    let validity = Validity::from_iter((0..=len).map(|idx| idx % 3 != 0)).slice(1..len + 1)?;
+
+    binary_on_device_with_validity(views, Arc::from([]), validity, ctx).await
+}
+
 unsafe fn release_arrow_device_array(array: &mut ArrowDeviceArray) {
     unsafe {
         if let Some(release) = array.array.release {
@@ -154,6 +170,35 @@ fn benchmark_arrow_binary_export(c: &mut Criterion) {
                 });
             },
         );
+
+        group.throughput(Throughput::Bytes(
+            (len * (size_of::<BinaryView>() + 8) + len.div_ceil(8)) as u64,
+        ));
+        group.bench_with_input(
+            BenchmarkId::new("cuda/arrow_binary/sliced_validity", len_label),
+            &len,
+            |b, &len| {
+                b.iter_custom(|iters| {
+                    let timed = TimedLaunchStrategy::default();
+                    let timer = timed.timer();
+
+                    let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+                        .vortex_expect("failed to create execution context")
+                        .with_launch_strategy(Arc::new(timed));
+                    let array = block_on(sliced_validity_binary(len, &mut cuda_ctx))
+                        .vortex_expect("failed to create binary fixture");
+
+                    for _ in 0..iters {
+                        let mut exported =
+                            block_on(array.clone().export_device_array(&mut cuda_ctx))
+                                .vortex_expect("failed to export device array");
+                        unsafe { release_arrow_device_array(&mut exported) };
+                    }
+
+                    Duration::from_nanos(timer.load(Ordering::Relaxed))
+                });
+            },
+        );
     }
 
     group.finish();

vortex-cuda/src/arrow/canonical.rs

@@ -1965,6 +1965,58 @@ mod tests {
         Ok(())
     }
 
+    #[crate::test]
+    async fn test_export_binary_empty_and_all_null() -> VortexResult<()> {
+        let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+            .vortex_expect("failed to create execution context");
+
+        let empty = VarBinViewArray::from_iter_nullable_bin(std::iter::empty::<Option<&[u8]>>())
+            .into_array();
+        let mut exported = empty.export_device_array_with_schema(&mut ctx).await?;
+        let field = Field::try_from(&exported.schema)?;
+        assert_eq!(field, Field::new("", DataType::Binary, true));
+        assert_binary_layout(&exported.array.array, 0, 0, &[0], b"")?;
+        assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA);
+        unsafe { release_exported_array(&raw mut exported.array.array) };
+
+        let all_null =
+            VarBinViewArray::from_iter_nullable_bin([None::<&[u8]>, None::<&[u8]>]).into_array();
+        let mut exported = all_null.export_device_array_with_schema(&mut ctx).await?;
+        let field = Field::try_from(&exported.schema)?;
+        assert_eq!(field, Field::new("", DataType::Binary, true));
+        assert_binary_layout(&exported.array.array, 2, 2, &[0, 0, 0], b"")?;
+        assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA);
+        unsafe { release_exported_array(&raw mut exported.array.array) };
+
+        Ok(())
+    }
+
+    #[crate::test]
+    async fn test_export_binary_invalid_data_buffer_ref_errors() -> VortexResult<()> {
+        let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+            .vortex_expect("failed to create execution context");
+
+        let view = BinaryView::make_view(b"this references a missing data buffer", 0, 0);
+        let array = VarBinViewArray::new_handle(
+            BufferHandle::new_host(Buffer::from_iter([view]).into_byte_buffer()),
+            Arc::from([]),
+            DType::Binary(Nullability::NonNullable),
+            Validity::NonNullable,
+        )
+        .into_array();
+
+        let err = array
+            .export_device_array_with_schema(&mut ctx)
+            .await
+            .expect_err("missing binary data buffer should fail");
+        assert!(
+            err.to_string()
+                .contains("a view references an invalid data buffer")
+        );
+
+        Ok(())
+    }
+
     #[crate::test]
     async fn test_export_list() -> VortexResult<()> {
         let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty())