Skip to content

perf(exec): parallelize lance exec scan across fragments #236

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ringfa11 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(exec): parallelize lance exec scan across fragments #236

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 78 additions & 12 deletions rust/ffi/exec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use datafusion::prelude::SessionContext;
use datafusion_common::{DataFusionError, Result as DFResult};
use datafusion_expr::Expr;
use futures::StreamExt;
use lance_table::format::Fragment;

use crate::error::{clear_last_error, set_last_error, ErrorCode};
use crate::exec_ir::{agg_ir_to_df_expr, output_type_to_arrow, parse_exec_ir_v1};
Expand All @@ -25,6 +26,7 @@ struct LanceExecPartition {
schema: Arc<Schema>,
projection: Arc<[String]>,
filter: Option<Expr>,
fragments: Option<Vec<Fragment>>,
}

impl PartitionStream for LanceExecPartition {
Expand Down Expand Up @@ -55,6 +57,7 @@ impl PartitionStream for LanceExecPartition {
let dataset = self.dataset.clone();
let projection = self.projection.clone();
let filter = self.filter.clone();
let fragments = self.fragments.clone();

builder.spawn_on(
async move {
Expand All @@ -64,6 +67,9 @@ impl PartitionStream for LanceExecPartition {
if let Some(filter) = &filter {
scan.filter_expr(filter.clone());
}
if let Some(fragments) = &fragments {
scan.with_fragments(fragments.clone());
}
scan.scan_in_order(false);

let mut stream = scan
Expand Down Expand Up @@ -100,7 +106,7 @@ impl PartitionStream for LanceExecPartition {
#[derive(Debug)]
struct LanceExecTableProvider {
schema: Arc<Schema>,
partition: Arc<LanceExecPartition>,
partitions: Vec<Arc<LanceExecPartition>>,
}

#[async_trait::async_trait]
Expand All @@ -120,9 +126,14 @@ impl TableProvider for LanceExecTableProvider {
_filters: &[Expr],
limit: Option<usize>,
) -> DFResult<Arc<dyn datafusion::physical_plan::ExecutionPlan>> {
let partitions: Vec<Arc<dyn PartitionStream>> = self
.partitions
.iter()
.map(|p| p.clone() as Arc<dyn PartitionStream>)
.collect();
Ok(Arc::new(StreamingTableExec::try_new(
self.schema(),
vec![self.partition.clone()],
partitions,
projection,
vec![],
false,
Expand All @@ -131,6 +142,42 @@ impl TableProvider for LanceExecTableProvider {
}
}

fn split_fragments(
dataset: Arc<lance::Dataset>,
schema: Arc<Schema>,
projection: Arc<[String]>,
filter: Option<Expr>,
target_partitions: usize,
) -> Vec<Arc<LanceExecPartition>> {
let all_fragments = dataset.fragments().as_ref().clone();
let n = target_partitions.min(all_fragments.len());
if n <= 1 {
return vec![Arc::new(LanceExecPartition {
dataset: dataset.clone(),
schema: schema.clone(),
projection: projection.clone(),
filter: filter.clone(),
fragments: None,
})];
}
let mut buckets: Vec<Vec<Fragment>> = vec![Vec::new(); n];
for (i, frag) in all_fragments.iter().enumerate() {
buckets[i % n].push(frag.clone());
}
buckets
.into_iter()
.map(|fragments| {
Arc::new(LanceExecPartition {
dataset: dataset.clone(),
schema: schema.clone(),
projection: projection.clone(),
filter: filter.clone(),
fragments: Some(fragments),
})
})
.collect()
}

fn projected_schema(handle: &DatasetHandle, projection: &[String]) -> Result<Arc<Schema>, String> {
let base_schema = handle.arrow_schema.as_ref();
let mut fields = Vec::with_capacity(projection.len());
Expand All @@ -148,7 +195,11 @@ fn projected_schema(handle: &DatasetHandle, projection: &[String]) -> Result<Arc
Ok(Arc::new(Schema::new(fields)))
}

async fn build_exec_df(handle: &DatasetHandle, exec_ir: &[u8]) -> Result<DataFrame, String> {
async fn build_exec_df(
handle: &DatasetHandle,
exec_ir: &[u8],
threads: usize,
) -> Result<DataFrame, String> {
let exec_ir = parse_exec_ir_v1(exec_ir).map_err(|e| format!("exec_ir parse: {e}"))?;

let filter = if exec_ir.filter_ir.is_empty() {
Expand All @@ -163,19 +214,25 @@ async fn build_exec_df(handle: &DatasetHandle, exec_ir: &[u8]) -> Result<DataFra
let projection: Arc<[String]> = exec_ir.scan_projection.clone().into();
let schema = projected_schema(handle, projection.as_ref())?;

let partition = Arc::new(LanceExecPartition {
dataset: handle.dataset.clone(),
schema: schema.clone(),
let ctx = SessionContext::new();
let target_partitions = if threads > 0 {
threads
} else {
ctx.copied_config().target_partitions()
};
let partitions = split_fragments(
handle.dataset.clone(),
schema.clone(),
projection,
filter,
});
target_partitions,
);

let provider = LanceExecTableProvider {
schema: schema.clone(),
partition,
partitions,
};

let ctx = SessionContext::new();
ctx.register_table("t", Arc::new(provider))
.map_err(|e| e.to_string())?;
let df = ctx.table("t").await.map_err(|e| e.to_string())?;
Expand Down Expand Up @@ -269,7 +326,7 @@ fn get_exec_schema_inner(
let bytes = unsafe { slice_from_ptr(exec_ir, exec_ir_len, "exec_ir")? };

let schema_res = runtime::block_on(async {
let df = build_exec_df(handle, bytes).await?;
let df = build_exec_df(handle, bytes, 0).await?;
let plan = df.create_physical_plan().await.map_err(|e| e.to_string())?;
Ok::<_, String>(plan.schema())
})
Expand All @@ -280,13 +337,20 @@ fn get_exec_schema_inner(
Ok(Arc::new(Schema::new(schema.fields().clone())))
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Default)]
pub struct LanceExecContext {
pub threads: u32,
}

#[no_mangle]
pub unsafe extern "C" fn lance_create_dataset_exec_stream_ir(
dataset: *mut c_void,
exec_ir: *const u8,
exec_ir_len: usize,
exec_ctx: *const LanceExecContext,
) -> *mut c_void {
match create_dataset_exec_stream_ir_inner(dataset, exec_ir, exec_ir_len) {
match create_dataset_exec_stream_ir_inner(dataset, exec_ir, exec_ir_len, exec_ctx) {
Ok(stream) => {
clear_last_error();
Box::into_raw(Box::new(stream)) as *mut c_void
Expand All @@ -302,12 +366,14 @@ fn create_dataset_exec_stream_ir_inner(
dataset: *mut c_void,
exec_ir: *const u8,
exec_ir_len: usize,
exec_ctx: *const LanceExecContext,
) -> FfiResult<StreamHandle> {
let handle = unsafe { dataset_handle(dataset)? };
let bytes = unsafe { slice_from_ptr(exec_ir, exec_ir_len, "exec_ir")? };
let threads = unsafe { exec_ctx.as_ref() }.map(|c| c.threads as usize).unwrap_or(0);

let stream_res = runtime::block_on(async {
let df = build_exec_df(handle, bytes).await?;
let df = build_exec_df(handle, bytes, threads).await?;
df.execute_stream().await.map_err(|e| e.to_string())
})
.map_err(|e| FfiError::new(ErrorCode::Exec, format!("runtime: {e}")))?;
Expand Down
7 changes: 6 additions & 1 deletion src/include/lance_ffi.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,10 @@ typedef struct LanceDebugCounters {
uint64_t commit_count;
} LanceDebugCounters;

typedef struct LanceExecContext {
uint32_t threads;
} LanceExecContext;

void *lance_create_session(uint64_t index_cache_size_bytes,
uint64_t metadata_cache_size_bytes);
void lance_close_session(void *session);
Expand Down Expand Up @@ -92,7 +96,8 @@ void lance_close_stream(void *stream);
void *lance_get_exec_schema(void *dataset, const uint8_t *exec_ir,
size_t exec_ir_len);
void *lance_create_dataset_exec_stream_ir(void *dataset, const uint8_t *exec_ir,
size_t exec_ir_len);
size_t exec_ir_len,
const LanceExecContext *exec_ctx);

int32_t lance_last_error_code();
const char *lance_last_error_message();
Expand Down
5 changes: 4 additions & 1 deletion src/lance_scan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3184,12 +3184,15 @@ LanceExecLocalInit(ExecutionContext &context, TableFunctionInitInput &input,
auto result =
make_uniq<LanceExecLocalState>(std::move(chunk), context.client);
result->global_state = &global;
LanceExecContext exec_ctx;
exec_ctx.threads =
NumericCast<uint32_t>(DBConfig::GetConfig(context.client).options.maximum_threads);
result->stream = lance_create_dataset_exec_stream_ir(
bind_data.dataset,
bind_data.exec_ir.empty()
? nullptr
: reinterpret_cast<const uint8_t *>(bind_data.exec_ir.data()),
bind_data.exec_ir.size());
bind_data.exec_ir.size(), &exec_ctx);
if (!result->stream) {
throw IOException("Failed to create Lance exec stream" +
LanceFormatErrorSuffix());
Expand Down