Skip to content

Commit abbfa58

Browse files
committed
async duckdb scan
Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>
1 parent 49f98c3 commit abbfa58

5 files changed

Lines changed: 164 additions & 44 deletions

File tree

‎vortex-duckdb/cpp/table_function.cpp‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
#include "duckdb/function/table_function.hpp"
1717
#include "duckdb/main/capi/capi_internal.hpp"
1818
#include "duckdb/main/connection.hpp"
19+
#include "duckdb/parallel/async_result.hpp"
1920
#include "duckdb/parser/parsed_data/create_table_function_info.hpp"
2021
#include "duckdb/planner/operator/logical_get.hpp"
2122

@@ -279,16 +280,31 @@ init_local(ExecutionContext &, TableFunctionInitInput &input, GlobalTableFunctio
279280
return make_uniq<CTableLocalData>(std::move(cdata));
280281
}
281282

283+
struct VortexWaitTask final : AsyncTask {
284+
explicit VortexWaitTask(void *ffi_local) : ffi_local(ffi_local) {
285+
}
286+
void Execute() override {
287+
duckdb_table_function_wait_and_fetch(ffi_local);
288+
}
289+
void *ffi_local;
290+
};
291+
282292
void function(ClientContext &, TableFunctionInput &input, DataChunk &output) {
283293
void *const ffi_global = input.global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
284294
void *const ffi_local = input.local_state->Cast<CTableLocalData>().ffi_data->DataPtr();
285295

286296
duckdb_data_chunk chunk = reinterpret_cast<duckdb_data_chunk>(&output);
287297
duckdb_vx_error error_out = nullptr;
288-
duckdb_table_function_scan(ffi_global, ffi_local, chunk, &error_out);
298+
const bool blocked_on_io = duckdb_table_function_scan(ffi_global, ffi_local, chunk, &error_out);
299+
289300
if (error_out) {
290301
throw InvalidInputException(IntoErrString(error_out));
291302
}
303+
if (blocked_on_io) {
304+
vector<unique_ptr<AsyncTask>> tasks(1);
305+
tasks[0] = make_uniq<VortexWaitTask>(ffi_local);
306+
input.async_result = AsyncResult(std::move(tasks));
307+
}
292308
}
293309

294310
using FilterVec = vector<unique_ptr<Expression>>;

‎vortex-duckdb/include/vortex.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,11 +48,13 @@ bool duckdb_table_function_pushdown_projection_aggregates(void *bind_data,
4848
duckdb_vx_error *error_out);
4949

5050
extern
51-
void duckdb_table_function_scan(void *global_init_data,
51+
bool duckdb_table_function_scan(void *global_init_data,
5252
void *local_init_data,
5353
duckdb_data_chunk output,
5454
duckdb_vx_error *error_out);
5555

56+
extern void duckdb_table_function_wait_and_fetch(void *local_init_data);
57+
5658
extern bool duckdb_table_function_pushdown_expression(duckdb_vx_expr expr);
5759

5860
extern

‎vortex-duckdb/src/ffi.rs‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ use crate::duckdb::TableInitInput;
3030
use crate::duckdb::try_or;
3131
use crate::duckdb::try_or_null;
3232
use crate::table_function::Cardinality;
33+
use crate::table_function::ScanResult;
3334
use crate::table_function::TableFunctionBind;
3435
use crate::table_function::TableFunctionGlobal;
3536
use crate::table_function::TableFunctionLocal;
@@ -45,6 +46,7 @@ use crate::table_function::scan;
4546
use crate::table_function::statistics;
4647
use crate::table_function::table_scan_progress;
4748
use crate::table_function::to_string;
49+
use crate::table_function::wait_and_fetch;
4850

4951
#[unsafe(no_mangle)]
5052
unsafe extern "C-unwind" fn duckdb_table_function_to_string(
@@ -148,27 +150,34 @@ unsafe extern "C-unwind" fn duckdb_table_function_scan(
148150
local_init_data: *mut c_void,
149151
output: cpp::duckdb_data_chunk,
150152
error_out: *mut cpp::duckdb_vx_error,
151-
) {
153+
) -> bool {
152154
let global_init_data = unsafe { global_init_data.cast::<TableFunctionGlobal>().as_ref() }
153155
.vortex_expect("global_init_data null pointer");
154156
let local_init_data = unsafe { local_init_data.cast::<TableFunctionLocal>().as_mut() }
155157
.vortex_expect("local_init_data null pointer");
156158
let data_chunk = unsafe { DataChunk::borrow_mut(output) };
157159

158160
match scan(local_init_data, global_init_data, data_chunk) {
159-
Ok(()) => {
160-
// The data chunk is already filled by the function.
161-
// No need to do anything here.
161+
Ok(result) => result == ScanResult::BlockedOnIO,
162+
Err(e) => {
163+
unsafe {
164+
error_out.write(cpp::duckdb_vx_error_create(
165+
e.to_string().as_ptr().cast(),
166+
e.to_string().len(),
167+
));
168+
}
169+
false
162170
}
163-
Err(e) => unsafe {
164-
error_out.write(cpp::duckdb_vx_error_create(
165-
e.to_string().as_ptr().cast(),
166-
e.to_string().len(),
167-
));
168-
},
169171
}
170172
}
171173

174+
#[unsafe(no_mangle)]
175+
unsafe extern "C-unwind" fn duckdb_table_function_wait_and_fetch(local_init_data: *mut c_void) {
176+
let local_init_data = unsafe { local_init_data.cast::<TableFunctionLocal>().as_mut() }
177+
.vortex_expect("local_init_data null pointer");
178+
wait_and_fetch(local_init_data);
179+
}
180+
172181
#[unsafe(no_mangle)]
173182
pub unsafe extern "C-unwind" fn duckdb_table_function_pushdown_expression(
174183
expr: cpp::duckdb_vx_expr,

‎vortex-duckdb/src/table_function.rs‎

Lines changed: 105 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ use vortex::file::v2::FileStatsLayoutReader;
4141
use vortex::io::kanal_ext::KanalExt as _;
4242
use vortex::io::runtime::BlockingRuntime as _;
4343
use vortex::io::runtime::current::ThreadSafeIterator;
44+
use vortex::io::runtime::current::TryRecv;
4445
use vortex::layout::scan::multi::MultiLayoutChild;
4546
use vortex::layout::scan::multi::MultiLayoutDataSource;
4647
use vortex::metrics::tracing::get_global_labels;
@@ -131,6 +132,7 @@ impl<'a> TableInitInput<'a> {
131132

132133
type ScanItem = VortexResult<(ArrayRef, Arc<ConversionCache>)>;
133134
type DataSourceIterator = ThreadSafeIterator<ScanItem>;
135+
type ScanItemFuture = BoxFuture<'static, Option<ScanItem>>;
134136

135137
pub struct TableFunctionGlobal {
136138
iterator: DataSourceIterator,
@@ -158,8 +160,10 @@ pub struct TableFunctionLocal {
158160
exporter: Option<ArrayExporter>,
159161
partition_index: u64,
160162
file_index: usize,
161-
// Aggregate scan accumulated partials. Empty for non-aggregate scan
163+
/// Aggregate scan accumulated partials. Empty for non-aggregate scan
162164
partials: Vec<Box<dyn DynAccumulator>>,
165+
/// In progress next item
166+
pending: Option<ScanItemFuture>,
163167
}
164168

165169
pub struct PartitionData {
@@ -422,6 +426,7 @@ pub fn init_local(
422426
partition_index: 0,
423427
file_index: 0,
424428
partials,
429+
pending: None,
425430
}
426431
}
427432

@@ -443,11 +448,43 @@ fn convert_result(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Struc
443448
})
444449
}
445450

451+
/// Write aggregation output row to "chunk"
452+
fn finalize_aggregate(
453+
global_state: &TableFunctionGlobal,
454+
chunk: &mut DataChunkRef,
455+
) -> VortexResult<()> {
456+
// 0 means we're the last thread, u64::MAX means output is written.
457+
// is_err() means CAS didn't succeed
458+
if global_state
459+
.pending
460+
.compare_exchange(0, u64::MAX, Ordering::AcqRel, Ordering::Relaxed)
461+
.is_err()
462+
{
463+
return Ok(());
464+
}
465+
466+
let mut accumulators = global_state.partials.lock();
467+
let row_count = global_state.row_count.load(Ordering::Acquire) as i64;
468+
let mut accum_iter = accumulators.iter_mut();
469+
for (idx, aggregate) in global_state.aggregates.iter().enumerate() {
470+
let value = match aggregate {
471+
ColumnAggregate::Real { .. } => {
472+
let accum = accum_iter.next().vortex_expect("partial for real agg");
473+
Value::try_from(accum.finish()?)?
474+
}
475+
ColumnAggregate::CountStar => Value::from(row_count),
476+
};
477+
chunk.get_vector_mut(idx).reference_value(&value);
478+
}
479+
chunk.set_len(1);
480+
Ok(())
481+
}
482+
446483
fn scan_aggregate(
447484
local_state: &mut TableFunctionLocal,
448485
global_state: &TableFunctionGlobal,
449486
chunk: &mut DataChunkRef,
450-
) -> VortexResult<()> {
487+
) -> VortexResult<ScanResult> {
451488
let aggregates_len = global_state.aggregates.len();
452489
// seen[k] = output column for requested column k.
453490
// If min(x), max(x), avg(y) are requested, seen = { 0: 0, 1: 1}
@@ -468,32 +505,13 @@ fn scan_aggregate(
468505

469506
let mut ctx = SESSION.create_execution_ctx();
470507
loop {
471-
let Some(result) = local_state.iterator.next() else {
472-
// 0 means we're the last thread, u64::MAX means output is written.
473-
// is_err() means CAS didn't succeed
474-
if global_state
475-
.pending
476-
.compare_exchange(0, u64::MAX, Ordering::AcqRel, Ordering::Relaxed)
477-
.is_err()
478-
{
479-
return Ok(());
508+
let result = match poll_next_item(local_state) {
509+
NextItem::Ready(item) => item,
510+
NextItem::BlockedOnIO => return Ok(ScanResult::BlockedOnIO),
511+
NextItem::None => {
512+
finalize_aggregate(global_state, chunk)?;
513+
return Ok(ScanResult::Exported);
480514
}
481-
482-
let mut accumulators = global_state.partials.lock();
483-
let row_count = global_state.row_count.load(Ordering::Acquire) as i64;
484-
let mut accum_iter = accumulators.iter_mut();
485-
for (idx, aggregate) in global_state.aggregates.iter().enumerate() {
486-
let value = match aggregate {
487-
ColumnAggregate::Real { .. } => {
488-
let accum = accum_iter.next().vortex_expect("partial for real agg");
489-
Value::try_from(accum.finish()?)?
490-
}
491-
ColumnAggregate::CountStar => Value::from(row_count),
492-
};
493-
chunk.get_vector_mut(idx).reference_value(&value);
494-
}
495-
chunk.set_len(1);
496-
return Ok(());
497515
};
498516
let array = convert_result(result?.0, &mut ctx)?;
499517

@@ -517,22 +535,77 @@ fn scan_aggregate(
517535
}
518536
}
519537

538+
enum NextItem {
539+
Ready(ScanItem),
540+
None,
541+
BlockedOnIO,
542+
}
543+
544+
/// Poll next item without blocking caller thread
545+
fn poll_next_item(local_state: &mut TableFunctionLocal) -> NextItem {
546+
if let Some(future) = local_state.pending.as_mut() {
547+
let mut ctx = Context::from_waker(std::task::Waker::noop());
548+
return match future.poll_unpin(&mut ctx) {
549+
Poll::Ready(item) => {
550+
local_state.pending = None;
551+
if let Some(item) = item {
552+
NextItem::Ready(item)
553+
} else {
554+
NextItem::None
555+
}
556+
}
557+
Poll::Pending => NextItem::BlockedOnIO,
558+
};
559+
}
560+
561+
match local_state.iterator.try_recv() {
562+
TryRecv::Item(item) => NextItem::Ready(item),
563+
TryRecv::Closed => NextItem::None,
564+
TryRecv::Empty => {
565+
let rx = local_state.iterator.receiver();
566+
local_state.pending = Some(Box::pin(async move { rx.recv().await.ok() }));
567+
NextItem::BlockedOnIO
568+
}
569+
}
570+
}
571+
572+
/// Finish array receive started by scan()
573+
pub fn wait_and_fetch(local_state: &mut TableFunctionLocal) {
574+
if let Some(mut future) = local_state.pending.take() {
575+
// Duckdb calls this on a background task thread pool which is separate
576+
// from table function processing pool. This call therefore doesn't
577+
// block the query worker
578+
let item = RUNTIME.block_on(&mut future);
579+
// When this assignment is happening, thread related to "local_state"
580+
// is suspended. It will be resumed only when wait_and_fetch returns
581+
// so there's no race condition
582+
local_state.pending = Some(Box::pin(std::future::ready(item)));
583+
};
584+
}
585+
586+
#[derive(PartialEq, Eq)]
587+
pub enum ScanResult {
588+
Exported,
589+
BlockedOnIO,
590+
}
591+
520592
pub fn scan(
521593
local_state: &mut TableFunctionLocal,
522594
global_state: &TableFunctionGlobal,
523595
chunk: &mut DataChunkRef,
524-
) -> VortexResult<()> {
596+
) -> VortexResult<ScanResult> {
525597
if !local_state.partials.is_empty() {
526598
return scan_aggregate(local_state, global_state, chunk);
527599
}
528600

529601
loop {
530602
if local_state.exporter.is_none() {
531603
let mut ctx = SESSION.create_execution_ctx();
532-
let Some(result) = local_state.iterator.next() else {
533-
return Ok(());
604+
let (array_result, conversion_cache) = match poll_next_item(local_state) {
605+
NextItem::Ready(item) => item?,
606+
NextItem::None => return Ok(ScanResult::Exported),
607+
NextItem::BlockedOnIO => return Ok(ScanResult::BlockedOnIO),
534608
};
535-
let (array_result, conversion_cache) = result?;
536609
local_state.file_index = conversion_cache.file_index;
537610
let array_result = convert_result(array_result, &mut ctx)?;
538611

@@ -576,7 +649,7 @@ pub fn scan(
576649
.reference_value(&Value::from(local_state.file_index as u64));
577650
}
578651

579-
Ok(())
652+
Ok(ScanResult::Exported)
580653
}
581654

582655
/// Scan progress as a percentage (0.0–100.0).

‎vortex-io/src/runtime/current.rs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,26 @@ pub struct ThreadSafeIterator<T> {
139139
driver: Arc<Mutex<Option<smol::Task<()>>>>,
140140
}
141141

142+
pub enum TryRecv<T> {
143+
Item(T),
144+
Empty,
145+
Closed,
146+
}
147+
148+
impl<T> ThreadSafeIterator<T> {
149+
pub fn receiver(&self) -> kanal::AsyncReceiver<T> {
150+
self.results.clone()
151+
}
152+
153+
pub fn try_recv(&self) -> TryRecv<T> {
154+
match self.results.try_recv() {
155+
Ok(Some(v)) => TryRecv::Item(v),
156+
Ok(None) => TryRecv::Empty,
157+
Err(_) => TryRecv::Closed,
158+
}
159+
}
160+
}
161+
142162
// Manual clone implementation since `T` does not need to be `Clone`.
143163
impl<T> Clone for ThreadSafeIterator<T> {
144164
fn clone(&self) -> Self {

0 commit comments

Comments
 (0)