Skip to content

Commit

Permalink
Add Handle suffix to DataChunk and LogicalType
Browse files Browse the repository at this point in the history
  • Loading branch information
Swoorup committed Jul 20, 2024
1 parent 8e6359b commit b05c6f4
Show file tree
Hide file tree
Showing 11 changed files with 122 additions and 118 deletions.
10 changes: 5 additions & 5 deletions crates/duckdb/examples/hello-ext/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ extern crate duckdb_loadable_macros;
extern crate libduckdb_sys;

use duckdb::{
core::{DataChunk, Inserter, LogicalType, LogicalTypeId},
core::{DataChunkHandle, Inserter, LogicalTypeHandle, LogicalTypeId},
vtab::{BindInfo, Free, FunctionInfo, InitInfo, VTab},
Connection, Result,
};
Expand Down Expand Up @@ -44,7 +44,7 @@ impl VTab for HelloVTab {
type BindData = HelloBindData;

unsafe fn bind(bind: &BindInfo, data: *mut HelloBindData) -> Result<(), Box<dyn std::error::Error>> {
bind.add_result_column("column0", LogicalType::new(LogicalTypeId::Varchar));
bind.add_result_column("column0", LogicalTypeHandle::from(LogicalTypeId::Varchar));
let param = bind.get_parameter(0).to_string();
unsafe {
(*data).name = CString::new(param).unwrap().into_raw();
Expand All @@ -59,7 +59,7 @@ impl VTab for HelloVTab {
Ok(())
}

unsafe fn func(func: &FunctionInfo, output: &mut DataChunk) -> Result<(), Box<dyn std::error::Error>> {
unsafe fn func(func: &FunctionInfo, output: &mut DataChunkHandle) -> Result<(), Box<dyn std::error::Error>> {
let init_info = func.get_init_data::<HelloInitData>();
let bind_info = func.get_bind_data::<HelloBindData>();

Expand All @@ -80,8 +80,8 @@ impl VTab for HelloVTab {
Ok(())
}

fn parameters() -> Option<Vec<LogicalType>> {
Some(vec![LogicalType::new(LogicalTypeId::Varchar)])
fn parameters() -> Option<Vec<LogicalTypeHandle>> {
Some(vec![LogicalTypeHandle::from(LogicalTypeId::Varchar)])
}
}

Expand Down
6 changes: 3 additions & 3 deletions crates/duckdb/src/appender/arrow.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use super::{ffi, Appender, Result};
use crate::{
core::{DataChunk, LogicalType},
core::{DataChunkHandle, LogicalTypeHandle},
error::result_from_duckdb_appender,
vtab::{record_batch_to_duckdb_data_chunk, to_duckdb_logical_type},
Error,
Expand Down Expand Up @@ -29,14 +29,14 @@ impl Appender<'_> {
#[inline]
pub fn append_record_batch(&mut self, record_batch: RecordBatch) -> Result<()> {
let schema = record_batch.schema();
let mut logical_type: Vec<LogicalType> = vec![];
let mut logical_type: Vec<LogicalTypeHandle> = vec![];
for field in schema.fields() {
let logical_t = to_duckdb_logical_type(field.data_type())
.map_err(|_op| Error::ArrowTypeToDuckdbType(field.to_string(), field.data_type().clone()))?;
logical_type.push(logical_t);
}

let mut data_chunk = DataChunk::new(&logical_type);
let mut data_chunk = DataChunkHandle::new(&logical_type);
record_batch_to_duckdb_data_chunk(&record_batch, &mut data_chunk).map_err(|_op| Error::AppendError)?;

let rc = unsafe { duckdb_append_data_chunk(self.app, data_chunk.get_ptr()) };
Expand Down
62 changes: 30 additions & 32 deletions crates/duckdb/src/core/data_chunk.rs
Original file line number Diff line number Diff line change
@@ -1,28 +1,41 @@
use super::{
logical_type::LogicalType,
logical_type::LogicalTypeHandle,
vector::{ArrayVector, FlatVector, ListVector, StructVector},
};
use crate::ffi::{
duckdb_create_data_chunk, duckdb_data_chunk, duckdb_data_chunk_get_column_count, duckdb_data_chunk_get_size,
duckdb_data_chunk_get_vector, duckdb_data_chunk_set_size, duckdb_destroy_data_chunk,
};

/// DataChunk in DuckDB.
pub struct DataChunk {
/// Handle to the DataChunk in DuckDB.
pub struct DataChunkHandle {
/// Pointer to the DataChunk in duckdb C API.
ptr: duckdb_data_chunk,

/// Whether this [DataChunk] own the [DataChunk::ptr].
/// Whether this [DataChunkHandle] own the [DataChunk::ptr].
owned: bool,
}

impl DataChunk {
/// Create a new [DataChunk] with the given [LogicalType]s.
pub fn new(logical_types: &[LogicalType]) -> Self {
impl Drop for DataChunkHandle {
fn drop(&mut self) {
if self.owned && !self.ptr.is_null() {
unsafe { duckdb_destroy_data_chunk(&mut self.ptr) }
self.ptr = std::ptr::null_mut();
}
}
}

impl DataChunkHandle {
pub(crate) unsafe fn new_unowned(ptr: duckdb_data_chunk) -> Self {
Self { ptr, owned: false }
}

/// Create a new [DataChunkHandle] with the given [LogicalTypeHandle]s.
pub fn new(logical_types: &[LogicalTypeHandle]) -> Self {
let num_columns = logical_types.len();
let mut c_types = logical_types.iter().map(|t| t.ptr).collect::<Vec<_>>();
let ptr = unsafe { duckdb_create_data_chunk(c_types.as_mut_ptr(), num_columns as u64) };
DataChunk { ptr, owned: true }
DataChunkHandle { ptr, owned: true }
}

/// Get the vector at the specific column index: `idx`.
Expand Down Expand Up @@ -50,49 +63,34 @@ impl DataChunk {
unsafe { duckdb_data_chunk_set_size(self.ptr, new_len as u64) };
}

/// Get the length / the number of rows in this [DataChunk].
/// Get the length / the number of rows in this [DataChunkHandle].
pub fn len(&self) -> usize {
unsafe { duckdb_data_chunk_get_size(self.ptr) as usize }
}

/// Check whether this [DataChunk] is empty.
/// Check whether this [DataChunkHandle] is empty.
pub fn is_empty(&self) -> bool {
self.len() == 0
}

/// Get the number of columns in this [DataChunk].
/// Get the number of columns in this [DataChunkHandle].
pub fn num_columns(&self) -> usize {
unsafe { duckdb_data_chunk_get_column_count(self.ptr) as usize }
}

/// Get the ptr of duckdb_data_chunk in this [DataChunk].
/// Get the ptr of duckdb_data_chunk in this [DataChunkHandle].
pub fn get_ptr(&self) -> duckdb_data_chunk {
self.ptr
}
}

impl From<duckdb_data_chunk> for DataChunk {
fn from(ptr: duckdb_data_chunk) -> Self {
Self { ptr, owned: false }
}
}

impl Drop for DataChunk {
fn drop(&mut self) {
if self.owned && !self.ptr.is_null() {
unsafe { duckdb_destroy_data_chunk(&mut self.ptr) }
self.ptr = std::ptr::null_mut();
}
}
}

#[cfg(test)]
mod test {
use super::{super::logical_type::LogicalTypeId, *};

#[test]
fn test_data_chunk_construction() {
let dc = DataChunk::new(&[LogicalType::new(LogicalTypeId::Integer)]);
let dc = DataChunkHandle::new(&[LogicalTypeHandle::from(LogicalTypeId::Integer)]);

assert_eq!(dc.num_columns(), 1);

Expand All @@ -101,7 +99,7 @@ mod test {

#[test]
fn test_vector() {
let datachunk = DataChunk::new(&[LogicalType::new(LogicalTypeId::Bigint)]);
let datachunk = DataChunkHandle::new(&[LogicalTypeHandle::from(LogicalTypeId::Bigint)]);
let mut vector = datachunk.flat_vector(0);
let data = vector.as_mut_slice::<i64>();

Expand All @@ -110,11 +108,11 @@ mod test {

#[test]
fn test_logi() {
let key = LogicalType::new(LogicalTypeId::Varchar);
let key = LogicalTypeHandle::from(LogicalTypeId::Varchar);

let value = LogicalType::new(LogicalTypeId::UTinyint);
let value = LogicalTypeHandle::from(LogicalTypeId::UTinyint);

let map = LogicalType::map(&key, &value);
let map = LogicalTypeHandle::map(&key, &value);

assert_eq!(map.id(), LogicalTypeId::Map);

Expand Down
54 changes: 27 additions & 27 deletions crates/duckdb/src/core/logical_type.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,11 +110,11 @@ impl From<u32> for LogicalTypeId {

/// DuckDB Logical Type.
/// <https://duckdb.org/docs/sql/data_types/overview>
pub struct LogicalType {
pub struct LogicalTypeHandle {
pub(crate) ptr: duckdb_logical_type,
}

impl Debug for LogicalType {
impl Debug for LogicalTypeHandle {
/// Debug implementation for LogicalType
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
let id = self.id();
Expand All @@ -134,7 +134,7 @@ impl Debug for LogicalType {
}
}

impl Drop for LogicalType {
impl Drop for LogicalTypeHandle {
/// Drop implementation for LogicalType
fn drop(&mut self) {
if !self.ptr.is_null() {
Expand All @@ -147,25 +147,25 @@ impl Drop for LogicalType {
}
}

impl From<duckdb_logical_type> for LogicalType {
/// Wrap a DuckDB logical type from C API
fn from(ptr: duckdb_logical_type) -> Self {
Self { ptr }
}
}

impl LogicalType {
/// Create a new [LogicalType] from [LogicalTypeId]
pub fn new(id: LogicalTypeId) -> Self {
impl From<LogicalTypeId> for LogicalTypeHandle {
/// Create a new [LogicalTypeHandle] from [LogicalTypeId]
fn from(id: LogicalTypeId) -> Self {
unsafe {
Self {
ptr: duckdb_create_logical_type(id as u32),
}
}
}
}

impl LogicalTypeHandle {
/// Create a DuckDB logical type from C API
pub(crate) unsafe fn new(ptr: duckdb_logical_type) -> Self {
Self { ptr }
}

/// Creates a map type from its child type.
pub fn map(key: &LogicalType, value: &LogicalType) -> Self {
pub fn map(key: &LogicalTypeHandle, value: &LogicalTypeHandle) -> Self {
unsafe {
Self {
ptr: duckdb_create_map_type(key.ptr, value.ptr),
Expand All @@ -174,7 +174,7 @@ impl LogicalType {
}

/// Creates a list type from its child type.
pub fn list(child_type: &LogicalType) -> Self {
pub fn list(child_type: &LogicalTypeHandle) -> Self {
unsafe {
Self {
ptr: duckdb_create_list_type(child_type.ptr),
Expand All @@ -183,7 +183,7 @@ impl LogicalType {
}

/// Creates an array type from its child type.
pub fn array(child_type: &LogicalType, array_size: u64) -> Self {
pub fn array(child_type: &LogicalTypeHandle, array_size: u64) -> Self {
unsafe {
Self {
ptr: duckdb_create_array_type(child_type.ptr, array_size),
Expand Down Expand Up @@ -213,7 +213,7 @@ impl LogicalType {
}

/// Make a `LogicalType` for `struct`
pub fn struct_type(fields: &[(&str, LogicalType)]) -> Self {
pub fn struct_type(fields: &[(&str, LogicalTypeHandle)]) -> Self {
let keys: Vec<CString> = fields.iter().map(|f| CString::new(f.0).unwrap()).collect();
let values: Vec<duckdb_logical_type> = fields.iter().map(|it| it.1.ptr).collect();
let name_ptrs = keys.iter().map(|it| it.as_ptr()).collect::<Vec<*const c_char>>();
Expand All @@ -230,7 +230,7 @@ impl LogicalType {
}

/// Make a `LogicalType` for `union`
pub fn union_type(fields: &[(&str, LogicalType)]) -> Self {
pub fn union_type(fields: &[(&str, LogicalTypeHandle)]) -> Self {
let keys: Vec<CString> = fields.iter().map(|f| CString::new(f.0).unwrap()).collect();
let values: Vec<duckdb_logical_type> = fields.iter().map(|it| it.1.ptr).collect();
let name_ptrs = keys.iter().map(|it| it.as_ptr()).collect::<Vec<*const c_char>>();
Expand Down Expand Up @@ -287,18 +287,18 @@ impl LogicalType {
_ => panic!("not a struct or union"),
}
};
Self::from(c_logical_type)
unsafe { Self::new(c_logical_type) }
}
}

#[cfg(test)]
mod test {
use crate::core::{LogicalType, LogicalTypeId};
use crate::core::{LogicalTypeHandle, LogicalTypeId};

#[test]
fn test_struct() {
let fields = &[("hello", LogicalType::new(crate::core::LogicalTypeId::Boolean))];
let typ = LogicalType::struct_type(fields);
let fields = &[("hello", LogicalTypeHandle::from(crate::core::LogicalTypeId::Boolean))];
let typ = LogicalTypeHandle::struct_type(fields);

assert_eq!(typ.num_children(), 1);
assert_eq!(typ.child_name(0), "hello");
Expand All @@ -307,7 +307,7 @@ mod test {

#[test]
fn test_decimal() {
let typ = LogicalType::decimal(10, 2);
let typ = LogicalTypeHandle::decimal(10, 2);

assert_eq!(typ.id(), crate::core::LogicalTypeId::Decimal);
assert_eq!(typ.decimal_width(), 10);
Expand All @@ -316,7 +316,7 @@ mod test {

#[test]
fn test_decimal_methods() {
let typ = LogicalType::new(crate::core::LogicalTypeId::Varchar);
let typ = LogicalTypeHandle::from(crate::core::LogicalTypeId::Varchar);

assert_eq!(typ.decimal_width(), 0);
assert_eq!(typ.decimal_scale(), 0);
Expand All @@ -325,10 +325,10 @@ mod test {
#[test]
fn test_union_type() {
let fields = &[
("hello", LogicalType::new(LogicalTypeId::Boolean)),
("world", LogicalType::new(LogicalTypeId::Integer)),
("hello", LogicalTypeHandle::from(LogicalTypeId::Boolean)),
("world", LogicalTypeHandle::from(LogicalTypeId::Integer)),
];
let typ = LogicalType::union_type(fields);
let typ = LogicalTypeHandle::union_type(fields);

assert_eq!(typ.num_children(), 2);

Expand Down
4 changes: 2 additions & 2 deletions crates/duckdb/src/core/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,6 @@ mod data_chunk;
mod logical_type;
mod vector;

pub use data_chunk::DataChunk;
pub use logical_type::{LogicalType, LogicalTypeId};
pub use data_chunk::DataChunkHandle;
pub use logical_type::{LogicalTypeHandle, LogicalTypeId};
pub use vector::*;
16 changes: 8 additions & 8 deletions crates/duckdb/src/core/vector.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
use std::{any::Any, ffi::CString, slice};

use libduckdb_sys::{duckdb_array_type_array_size, duckdb_array_vector_get_child};
use libduckdb_sys::{duckdb_array_type_array_size, duckdb_array_vector_get_child, DuckDbString};

use super::LogicalType;
use super::LogicalTypeHandle;
use crate::ffi::{
duckdb_list_entry, duckdb_list_vector_get_child, duckdb_list_vector_get_size, duckdb_list_vector_reserve,
duckdb_list_vector_set_size, duckdb_struct_type_child_count, duckdb_struct_type_child_name,
Expand Down Expand Up @@ -71,8 +71,8 @@ impl FlatVector {
}

/// Returns the logical type of the vector
pub fn logical_type(&self) -> LogicalType {
LogicalType::from(unsafe { duckdb_vector_get_column_type(self.ptr) })
pub fn logical_type(&self) -> LogicalTypeHandle {
unsafe { LogicalTypeHandle::new(duckdb_vector_get_column_type(self.ptr)) }
}

/// Set row as null
Expand Down Expand Up @@ -202,8 +202,8 @@ impl From<duckdb_vector> for ArrayVector {

impl ArrayVector {
/// Get the logical type of this ArrayVector.
pub fn logical_type(&self) -> LogicalType {
LogicalType::from(unsafe { duckdb_vector_get_column_type(self.ptr) })
pub fn logical_type(&self) -> LogicalTypeHandle {
unsafe { LogicalTypeHandle::new(duckdb_vector_get_column_type(self.ptr)) }
}

/// Returns the size of the array type.
Expand Down Expand Up @@ -259,8 +259,8 @@ impl StructVector {
}

/// Get the logical type of this struct vector.
pub fn logical_type(&self) -> LogicalType {
LogicalType::from(unsafe { duckdb_vector_get_column_type(self.ptr) })
pub fn logical_type(&self) -> LogicalTypeHandle {
unsafe { LogicalTypeHandle::new(duckdb_vector_get_column_type(self.ptr)) }
}

/// Get the name of the child by idx.
Expand Down
Loading

0 comments on commit b05c6f4

Please sign in to comment.