Skip to content

Commit a36f41e

Browse files
paleolimbotpaleolimbot
committed
maybe prepare + schema
1 parent 50d6d3f commit a36f41e

4 files changed

Lines changed: 121 additions & 9 deletions

File tree

‎r/adbcdrivermanager/NAMESPACE‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,13 @@ S3method("[[",adbc_async_task)
88
S3method("[[",adbc_error)
99
S3method("[[",adbc_xptr)
1010
S3method("[[<-",adbc_xptr)
11-
S3method(adbc_async_task_cancel,adbc_async_execute_query)
11+
S3method(adbc_async_task_cancel,adbc_async_statement_cancellable)
1212
S3method(adbc_async_task_cancel,default)
1313
S3method(adbc_async_task_result,adbc_async_execute_query)
14+
S3method(adbc_async_task_result,adbc_async_prepare)
1415
S3method(adbc_async_task_result,adbc_async_sleep)
1516
S3method(adbc_async_task_result,adbc_async_statement_stream_get_next)
17+
S3method(adbc_async_task_result,adbc_async_statement_stream_schema)
1618
S3method(adbc_connection_init,adbc_database_log)
1719
S3method(adbc_connection_init,adbc_database_monkey)
1820
S3method(adbc_connection_init,default)

‎r/adbcdrivermanager/R/async.R‎

Lines changed: 67 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -134,8 +134,37 @@ adbc_async_task_result.adbc_async_sleep <- function(task) {
134134
task$user_data$duration_ms
135135
}
136136

137+
#' @export
138+
adbc_async_task_cancel.adbc_async_statement_cancellable <- function(task) {
139+
adbc_statement_cancel(task$user_data$statement)
140+
TRUE
141+
}
142+
143+
adbc_statement_prepare_async <- function(statement) {
144+
task <- adbc_async_task(
145+
c("adbc_async_prepare", "adbc_async_statement_cancellable")
146+
)
147+
148+
user_data <- task$user_data
149+
user_data$statement <- statement
150+
.Call(RAdbcAsyncTaskLaunchPrepare, task, statement)
151+
152+
task
153+
}
154+
155+
#' @export
156+
adbc_async_task_result.adbc_async_prepare <- function(task) {
157+
if (!identical(task$return_code, 0L)) {
158+
stop_for_error(task$return_code, task$error_xptr)
159+
}
160+
161+
task$user_data$statement
162+
}
163+
137164
adbc_statement_execute_query_async <- function(statement, stream = NULL) {
138-
task <- adbc_async_task("adbc_async_execute_query")
165+
task <- adbc_async_task(
166+
c("adbc_async_execute_query", "adbc_async_statement_cancellable")
167+
)
139168

140169
user_data <- task$user_data
141170
user_data$statement <- statement
@@ -151,12 +180,6 @@ adbc_statement_execute_query_async <- function(statement, stream = NULL) {
151180
task
152181
}
153182

154-
#' @export
155-
adbc_async_task_cancel.adbc_async_execute_query <- function(task) {
156-
adbc_statement_cancel(task$user_data$statement)
157-
TRUE
158-
}
159-
160183
#' @export
161184
adbc_async_task_result.adbc_async_execute_query <- function(task) {
162185
if (!identical(task$return_code, 0L)) {
@@ -170,8 +193,44 @@ adbc_async_task_result.adbc_async_execute_query <- function(task) {
170193
)
171194
}
172195

196+
adbc_statement_stream_get_schema_async <- function(statement, stream) {
197+
task <- adbc_async_task(
198+
c("adbc_async_statement_stream_get_next", "adbc_async_statement_cancellable")
199+
)
200+
201+
user_data <- task$user_data
202+
user_data$statement <- statement
203+
user_data$stream <- stream
204+
user_data$schema <- nanoarrow::nanoarrow_allocate_schema()
205+
206+
user_data$rows_affected <- .Call(
207+
RAdbcAsyncTaskLaunchStreamGetSchema,
208+
task,
209+
stream,
210+
user_data$schema
211+
)
212+
213+
task
214+
}
215+
216+
217+
#' @export
218+
adbc_async_task_result.adbc_async_statement_stream_schema <- function(task) {
219+
if (!identical(task$return_code, 0L)) {
220+
adbc_statement_release(task$user_data$statement)
221+
stop(task$user_data$stream$get_last_error())
222+
}
223+
224+
list(
225+
statement = task$user_data$statement,
226+
array = task$user_data$schema
227+
)
228+
}
229+
173230
adbc_statement_stream_get_next_async <- function(statement, stream) {
174-
task <- adbc_async_task("adbc_async_statement_stream_get_next")
231+
task <- adbc_async_task(
232+
c("adbc_async_statement_stream_get_next", "adbc_async_statement_cancellable")
233+
)
175234

176235
user_data <- task$user_data
177236
user_data$statement <- statement

‎r/adbcdrivermanager/src/async.cc‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
#include <deque>
2424
#include <future>
25+
#include <memory>
2526
#include <mutex>
2627
#include <string>
2728
#include <thread>
@@ -42,6 +43,17 @@ static inline void later_ensure_initialized() {
4243

4344
static void later_task_callback_wrapper(void* data);
4445

46+
struct ArrowArrayCustomDeleter {
47+
void operator()(ArrowArray* array) const {
48+
if (array->release != nullptr) {
49+
array->release(array);
50+
}
51+
delete array;
52+
}
53+
};
54+
55+
using UniqueArrowArrayPtr = std::unique_ptr<ArrowArray, ArrowArrayCustomDeleter>;
56+
4557
enum class RAdbcAsyncTaskStatus { NOT_STARTED, STARTED, READY };
4658

4759
struct RAdbcAsyncTask {
@@ -217,6 +229,22 @@ extern "C" SEXP RAdbcAsyncTaskLaunchSleep(SEXP task_xptr, SEXP duration_ms_sexp)
217229
return R_NilValue;
218230
}
219231

232+
extern "C" SEXP RAdbcAsyncTaskLaunchPrepare(SEXP task_xptr, SEXP statement_xptr) {
233+
auto task = adbc_from_xptr<RAdbcAsyncTask>(task_xptr);
234+
error_for_started_task(task);
235+
236+
auto statement = adbc_from_xptr<AdbcStatement>(statement_xptr);
237+
238+
task->result = std::async(std::launch::async, [task, statement] {
239+
*(task->return_code) = AdbcStatementPrepare(statement, task->return_error);
240+
task->ScheduleCallbackIfSet();
241+
});
242+
243+
task->status = RAdbcAsyncTaskStatus::STARTED;
244+
UNPROTECT(1);
245+
return R_NilValue;
246+
}
247+
220248
extern "C" SEXP RAdbcAsyncTaskLaunchExecuteQuery(SEXP task_xptr, SEXP statement_xptr,
221249
SEXP stream_xptr) {
222250
auto task = adbc_from_xptr<RAdbcAsyncTask>(task_xptr);
@@ -245,6 +273,23 @@ extern "C" SEXP RAdbcAsyncTaskLaunchExecuteQuery(SEXP task_xptr, SEXP statement_
245273
return rows_affected_sexp;
246274
}
247275

276+
extern "C" SEXP RAdbcAsyncTaskLaunchStreamGetSchema(SEXP task_xptr, SEXP stream_xptr,
277+
SEXP schema_xptr) {
278+
auto task = adbc_from_xptr<RAdbcAsyncTask>(task_xptr);
279+
error_for_started_task(task);
280+
281+
auto stream = adbc_from_xptr<ArrowArrayStream>(stream_xptr);
282+
auto schema = adbc_from_xptr<ArrowSchema>(schema_xptr);
283+
284+
task->result = std::async(std::launch::async, [task, stream, schema] {
285+
*(task->return_code) = stream->get_schema(stream, schema);
286+
task->ScheduleCallbackIfSet();
287+
});
288+
289+
task->status = RAdbcAsyncTaskStatus::STARTED;
290+
return R_NilValue;
291+
}
292+
248293
extern "C" SEXP RAdbcAsyncTaskLaunchStreamGetNext(SEXP task_xptr, SEXP stream_xptr,
249294
SEXP array_xptr) {
250295
auto task = adbc_from_xptr<RAdbcAsyncTask>(task_xptr);

‎r/adbcdrivermanager/src/init.c‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,11 @@ SEXP RAdbcAsyncTaskData(SEXP task_xptr);
2727
SEXP RAdbcAsyncTaskWaitFor(SEXP task_xptr, SEXP duration_ms_sexp);
2828
SEXP RAdbcAsyncTaskWait(SEXP task_xptr, SEXP resolution_ms_sexp);
2929
SEXP RAdbcAsyncTaskLaunchSleep(SEXP task_xptr, SEXP duration_ms_sexp);
30+
SEXP RAdbcAsyncTaskLaunchPrepare(SEXP task_xptr, SEXP statement_xptr);
3031
SEXP RAdbcAsyncTaskLaunchExecuteQuery(SEXP task_xptr, SEXP statement_xptr,
3132
SEXP stream_xptr);
33+
SEXP RAdbcAsyncTaskLaunchStreamGetSchema(SEXP task_xptr, SEXP stream_xptr,
34+
SEXP schema_xptr);
3235
SEXP RAdbcAsyncTaskLaunchStreamGetNext(SEXP task_xptr, SEXP stream_xptr, SEXP array_xptr);
3336
SEXP RAdbcVoidDriverInitFunc(void);
3437
SEXP RAdbcMonkeyDriverInitFunc(void);
@@ -118,7 +121,10 @@ static const R_CallMethodDef CallEntries[] = {
118121
{"RAdbcAsyncTaskWaitFor", (DL_FUNC)&RAdbcAsyncTaskWaitFor, 2},
119122
{"RAdbcAsyncTaskWait", (DL_FUNC)&RAdbcAsyncTaskWait, 2},
120123
{"RAdbcAsyncTaskLaunchSleep", (DL_FUNC)&RAdbcAsyncTaskLaunchSleep, 2},
124+
{"RAdbcAsyncTaskLaunchPrepare", (DL_FUNC)&RAdbcAsyncTaskLaunchPrepare, 2},
121125
{"RAdbcAsyncTaskLaunchExecuteQuery", (DL_FUNC)&RAdbcAsyncTaskLaunchExecuteQuery, 3},
126+
{"RAdbcAsyncTaskLaunchStreamGetSchema", (DL_FUNC)&RAdbcAsyncTaskLaunchStreamGetSchema,
127+
3},
122128
{"RAdbcAsyncTaskLaunchStreamGetNext", (DL_FUNC)&RAdbcAsyncTaskLaunchStreamGetNext, 3},
123129
{"RAdbcVoidDriverInitFunc", (DL_FUNC)&RAdbcVoidDriverInitFunc, 0},
124130
{"RAdbcMonkeyDriverInitFunc", (DL_FUNC)&RAdbcMonkeyDriverInitFunc, 0},

0 commit comments

Comments
 (0)