diff --git a/include/lbug_rs.h b/include/lbug_rs.h index 7c4ae65..84b7fd7 100644 --- a/include/lbug_rs.h +++ b/include/lbug_rs.h @@ -1,7 +1,7 @@ #pragma once -#include #include +#include #include #include #include @@ -13,8 +13,8 @@ #include "c_api/lbug.h" #include "main/lbug.h" #else -#include #include +#include #endif namespace lbug_rs { @@ -40,7 +40,8 @@ struct TypeListBuilder { }; std::unique_ptr create_type_list(); -inline void type_list_insert(TypeListBuilder& list, std::unique_ptr type) { +inline void type_list_insert(TypeListBuilder& list, + std::unique_ptr type) { list.insert(std::move(type)); } @@ -143,7 +144,7 @@ std::unique_ptr connection_create_arrow_rel_table_csr( lbug::main::Connection& connection, std::string_view tableName, std::string_view srcTableName, std::string_view dstTableName, ArrowSchema indicesSchema, std::unique_ptr indicesArrays, ArrowSchema indptrSchema, - std::unique_ptr indptrArrays); + std::unique_ptr indptrArrays, std::string_view dstColName = "to"); std::unique_ptr connection_drop_arrow_table( lbug::main::Connection& connection, std::string_view tableName); inline std::unique_ptr connection_prepare( diff --git a/src/connection.rs b/src/connection.rs index 0bebd46..5a5a6bf 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -263,8 +263,8 @@ impl<'a> Connection<'a> { #[cfg(feature = "arrow")] /// Registers Arrow memory in CSR form as a relationship table. /// - /// The `indices_batches` schema must include a destination column named `to`. The - /// `indptr_batches` schema must contain at least one offset column. + /// The `indices_batches` schema must include a destination column named by `dst_col_name` + /// (defaults to `"to"`). The `indptr_batches` schema must contain at least one offset column. /// /// *Requires the `arrow` feature* pub fn create_arrow_rel_table_csr( @@ -274,6 +274,7 @@ impl<'a> Connection<'a> { indptr_batches: &[arrow::record_batch::RecordBatch], src_table_name: &str, dst_table_name: &str, + dst_col_name: &str, ) -> Result, Error> { let (indices_schema, indices_arrays) = export_arrow_batches(indices_batches)?; let (indptr_schema, indptr_arrays) = export_arrow_batches(indptr_batches)?; @@ -287,6 +288,7 @@ impl<'a> Connection<'a> { indices_arrays, indptr_schema, indptr_arrays, + ffi::StringView::new(dst_col_name), )?; Self::query_result_from_ffi(result) } @@ -606,7 +608,61 @@ Invalid input : expected rule oC_SingleQuery (line: 1, o indptr_schema, vec![Arc::new(UInt64Array::from(vec![0, 1, 2]))], )?; - conn.create_arrow_rel_table_csr("Knows", &[indices], &[indptr], "Person", "Person")?; + conn.create_arrow_rel_table_csr("Knows", &[indices], &[indptr], "Person", "Person", "to")?; + + let result = conn.query( + "MATCH (a:Person)-[r:Knows]->(b:Person) \ + RETURN a.id, r.weight, b.id ORDER BY a.id, b.id;", + )?; + assert_eq!(result.to_string(), "a.id|r.weight|b.id\n0|7|1\n1|9|0\n"); + temp_dir.close()?; + Ok(()) + } + + #[test] + #[cfg(feature = "arrow")] + fn test_create_arrow_rel_table_csr_custom_dst_col() -> Result<()> { + use arrow::array::{Int64Array, UInt64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use std::sync::Arc; + + let temp_dir = tempfile::tempdir()?; + let db = Database::new(temp_dir.path().join("test"), SYSTEM_CONFIG_FOR_TESTS)?; + let conn = Connection::new(&db)?; + let node_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let nodes = + RecordBatch::try_new(node_schema, vec![Arc::new(Int64Array::from(vec![0, 1]))])?; + conn.create_arrow_table("Person", &[nodes])?; + + let indices_schema = Arc::new(Schema::new(vec![ + Field::new("dest", DataType::UInt64, false), + Field::new("weight", DataType::Int64, false), + ])); + let indices = RecordBatch::try_new( + indices_schema, + vec![ + Arc::new(UInt64Array::from(vec![1, 0])), + Arc::new(Int64Array::from(vec![7, 9])), + ], + )?; + let indptr_schema = Arc::new(Schema::new(vec![Field::new( + "indptr", + DataType::UInt64, + false, + )])); + let indptr = RecordBatch::try_new( + indptr_schema, + vec![Arc::new(UInt64Array::from(vec![0, 1, 2]))], + )?; + conn.create_arrow_rel_table_csr( + "Knows", + &[indices], + &[indptr], + "Person", + "Person", + "dest", + )?; let result = conn.query( "MATCH (a:Person)-[r:Knows]->(b:Person) \ diff --git a/src/ffi/arrow.rs b/src/ffi/arrow.rs index 10107f8..564b8ac 100644 --- a/src/ffi/arrow.rs +++ b/src/ffi/arrow.rs @@ -93,7 +93,7 @@ pub(crate) mod ffi_arrow { arrays: UniquePtr, ) -> Result>>; - fn connection_create_arrow_rel_table_csr<'a, 'b, 'c, 'db>( + fn connection_create_arrow_rel_table_csr<'a, 'b, 'c, 'd, 'db>( connection: Pin<&mut Connection<'db>>, table_name: StringView<'a>, src_table_name: StringView<'b>, @@ -102,6 +102,7 @@ pub(crate) mod ffi_arrow { indices_arrays: UniquePtr, indptr_schema: ArrowSchema, indptr_arrays: UniquePtr, + dst_col_name: StringView<'d>, ) -> Result>>; fn connection_drop_arrow_table<'a, 'db>( diff --git a/src/lbug_rs.cpp b/src/lbug_rs.cpp index 014eca3..b3d7c7f 100644 --- a/src/lbug_rs.cpp +++ b/src/lbug_rs.cpp @@ -138,7 +138,8 @@ std::unique_ptr connection_create_arrow_rel_table( arrays->arrays.size(), &capiResult); if (capiResult._query_result == nullptr) { auto error = lbug_get_last_error(); - std::string message = error == nullptr ? "Failed to create Arrow relationship table" : error; + std::string message = + error == nullptr ? "Failed to create Arrow relationship table" : error; if (error != nullptr) { lbug_destroy_string(error); } @@ -152,16 +153,18 @@ std::unique_ptr connection_create_arrow_rel_table_csr( lbug::main::Connection& connection, std::string_view tableName, std::string_view srcTableName, std::string_view dstTableName, ArrowSchema indicesSchema, std::unique_ptr indicesArrays, ArrowSchema indptrSchema, - std::unique_ptr indptrArrays) { + std::unique_ptr indptrArrays, std::string_view dstColName) { lbug_connection capiConnection{&connection}; lbug_query_result capiResult{}; auto tableNameString = std::string(tableName); auto srcTableNameString = std::string(srcTableName); auto dstTableNameString = std::string(dstTableName); - (void)lbug_connection_create_arrow_rel_table_csr(&capiConnection, - tableNameString.c_str(), srcTableNameString.c_str(), dstTableNameString.c_str(), - &indicesSchema, indicesArrays->arrays.data(), indicesArrays->arrays.size(), &indptrSchema, - indptrArrays->arrays.data(), indptrArrays->arrays.size(), &capiResult); + auto dstColNameString = std::string(dstColName); + (void)lbug_connection_create_arrow_rel_table_csr(&capiConnection, tableNameString.c_str(), + srcTableNameString.c_str(), dstTableNameString.c_str(), &indicesSchema, + indicesArrays->arrays.data(), indicesArrays->arrays.size(), &indptrSchema, + indptrArrays->arrays.data(), indptrArrays->arrays.size(), dstColNameString.c_str(), + &capiResult); if (capiResult._query_result == nullptr) { auto error = lbug_get_last_error(); std::string message = @@ -179,8 +182,8 @@ std::unique_ptr connection_drop_arrow_table( lbug::main::Connection& connection, std::string_view tableName) { lbug_connection capiConnection{&connection}; lbug_query_result capiResult{}; - (void)lbug_connection_drop_arrow_table( - &capiConnection, std::string(tableName).c_str(), &capiResult); + (void)lbug_connection_drop_arrow_table(&capiConnection, std::string(tableName).c_str(), + &capiResult); if (capiResult._query_result == nullptr) { auto error = lbug_get_last_error(); std::string message = error == nullptr ? "Failed to drop Arrow table" : error;