1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

use std::{ffi::CStr, ops::Deref, sync::Arc};

use arrow::{
    array::{RecordBatch, StructArray},
    error::ArrowError,
    ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema},
};

use crate::{connection::ConnectionHandle, ffi, statement::PreparedStatementHandle};

#[derive(Debug)]
pub struct ArrowResultHandle {
    pub handle: ffi::duckdb_arrow,
    _parent: ArrowResultParent,
}

#[derive(Debug)]
pub enum ArrowResultParent {
    Connection(Arc<ConnectionHandle>),
    Statement(Arc<PreparedStatementHandle>),
}

impl ArrowResultHandle {
    /// # Safety
    /// Takes ownership
    pub unsafe fn from_raw_connection(
        handle: ffi::duckdb_arrow,
        connection: Arc<ConnectionHandle>,
    ) -> Self {
        Self {
            handle,
            _parent: ArrowResultParent::Connection(connection),
        }
    }
    /// # Safety
    /// Takes ownership
    pub unsafe fn from_raw_statement(
        handle: ffi::duckdb_arrow,
        statement: Arc<PreparedStatementHandle>,
    ) -> Self {
        Self {
            handle: handle,
            _parent: ArrowResultParent::Statement(statement),
        }
    }
    /// # Safety
    /// Does not check if there is actually an error
    pub unsafe fn error(&self) -> String {
        let err = ffi::duckdb_query_arrow_error(self.handle);
        CStr::from_ptr(err).to_string_lossy().into_owned()
    }
    pub fn column_count(&self) -> u64 {
        unsafe { ffi::duckdb_arrow_column_count(self.handle) }
    }
    pub fn row_count(&self) -> u64 {
        unsafe { ffi::duckdb_arrow_row_count(self.handle) }
    }
    pub fn rows_changed(&self) -> u64 {
        unsafe { ffi::duckdb_arrow_rows_changed(self.handle) }
    }

    /// # Safety
    /// The result must be consumed before calling this again
    pub unsafe fn query_array(&self) -> Result<Result<RecordBatch, ArrowError>, String> {
        let mut out_schema = FFI_ArrowSchema::empty();
        if unsafe {
            ffi::duckdb_query_arrow_schema(
                self.handle,
                &mut std::ptr::addr_of_mut!(out_schema) as *mut _ as *mut ffi::duckdb_arrow_schema,
            )
        } != ffi::DuckDBSuccess
        {
            return Err("duckdb_query_arrow_schema()".to_owned());
        }
        let mut out_array = FFI_ArrowArray::empty();
        if unsafe {
            ffi::duckdb_query_arrow_array(
                self.handle,
                &mut std::ptr::addr_of_mut!(out_array) as *mut _ as *mut ffi::duckdb_arrow_array,
            )
        } != ffi::DuckDBSuccess
        {
            return Err("duckdb_query_arrow_array()".to_owned());
        }
        let arr = from_ffi(out_array, &out_schema)
            .map(StructArray::from)
            .map(RecordBatch::from);
        Ok(arr)
    }
}

impl Deref for ArrowResultHandle {
    type Target = ffi::duckdb_arrow;

    fn deref(&self) -> &Self::Target {
        &self.handle
    }
}

impl Drop for ArrowResultHandle {
    fn drop(&mut self) {
        unsafe { ffi::duckdb_destroy_arrow(&mut self.handle) }
    }
}