Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
15 changes: 12 additions & 3 deletions bindings/c/src/table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -112,8 +112,7 @@ pub unsafe extern "C" fn paimon_read_builder_free(rb: *mut paimon_read_builder)
///
/// The `columns` parameter is a null-terminated array of null-terminated C strings.
/// Output order follows the caller-specified order. Unknown or duplicate names
/// cause `paimon_read_builder_new_read()` to fail; an empty list is a valid
/// zero-column projection.
/// are validated immediately; an empty list is a valid zero-column projection.
///
/// # Safety
/// `rb` must be a valid pointer from `paimon_table_new_read_builder`, or null (returns error).
Expand Down Expand Up @@ -149,6 +148,11 @@ pub unsafe extern "C" fn paimon_read_builder_with_projection(
ptr = ptr.add(1);
}

let col_refs: Vec<&str> = col_names.iter().map(String::as_str).collect();
if let Err(e) = state.table.new_read_builder().with_projection(&col_refs) {
return paimon_error::from_paimon(e);
}

state.projected_columns = Some(col_names);
std::ptr::null_mut()
}
Expand Down Expand Up @@ -229,7 +233,12 @@ pub unsafe extern "C" fn paimon_read_builder_new_read(
// Apply projection if set
if let Some(ref columns) = state.projected_columns {
let col_refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
rb_rust.with_projection(&col_refs);
if let Err(e) = rb_rust.with_projection(&col_refs) {
return paimon_result_new_read {
read: std::ptr::null_mut(),
error: paimon_error::from_paimon(e),
};
}
}

// Apply filter if set
Expand Down
9 changes: 5 additions & 4 deletions bindings/python/src/read.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,17 +69,18 @@ fn apply_read_config(
projection: &Option<Vec<String>>,
limit: Option<usize>,
filter: &Option<Predicate>,
) {
) -> PyResult<()> {
if let Some(projection) = projection {
let cols: Vec<&str> = projection.iter().map(String::as_str).collect();
builder.with_projection(&cols);
builder.with_projection(&cols).map_err(to_py_err)?;
}
if let Some(limit) = limit {
builder.with_limit(limit);
}
if let Some(filter) = filter {
builder.with_filter(filter.clone());
}
Ok(())
}

/// Extract a sequence of Python `Split` objects into core `DataSplit`s. Accepts
Expand Down Expand Up @@ -220,7 +221,7 @@ impl PyTableScan {
let splits = py.detach(|| {
rt.block_on(async {
let mut builder = self.table.new_read_builder();
apply_read_config(&mut builder, &self.projection, self.limit, &self.filter);
apply_read_config(&mut builder, &self.projection, self.limit, &self.filter)?;
let plan = builder.new_scan().plan().await.map_err(to_py_err)?;
Ok::<_, PyErr>(plan.splits().to_vec())
})
Expand All @@ -246,7 +247,7 @@ impl PyTableRead {
let batches = py.detach(|| {
rt.block_on(async {
let mut builder = self.table.new_read_builder();
apply_read_config(&mut builder, &self.projection, self.limit, &self.filter);
apply_read_config(&mut builder, &self.projection, self.limit, &self.filter)?;
// Validate config (e.g. projection) before the empty-splits fast
// path so an invalid projection fails consistently regardless of
// how many splits are passed.
Expand Down
3 changes: 2 additions & 1 deletion crates/integration_tests/tests/append_tables.rs
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,8 @@ async fn test_unpartitioned_projection() {

// Read with projection
let mut rb = table.new_read_builder();
rb.with_projection(&["value"]);
rb.with_projection(&["value"])
.expect("Projection should succeed");
let plan = rb.new_scan().plan().await.unwrap();
let read = rb.new_read().unwrap();
let result: Vec<RecordBatch> = read
Expand Down
33 changes: 21 additions & 12 deletions crates/integration_tests/tests/read_tables.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,9 @@ async fn scan_and_read<C: Catalog + ?Sized>(

let mut read_builder = table.new_read_builder();
if let Some(cols) = projection {
read_builder.with_projection(cols);
read_builder
.with_projection(cols)
.expect("Invalid projection");
}
let scan = read_builder.new_scan();
let plan = scan.plan().await.expect("Failed to plan scan");
Expand Down Expand Up @@ -107,7 +109,9 @@ async fn scan_and_read_with_projection_and_filter(
) -> (Plan, Vec<RecordBatch>) {
let mut read_builder = table.new_read_builder();
if let Some(cols) = projection {
read_builder.with_projection(cols);
read_builder
.with_projection(cols)
.expect("Invalid projection");
}
read_builder.with_filter(filter);
let scan = read_builder.new_scan();
Expand Down Expand Up @@ -486,7 +490,9 @@ async fn test_read_projection_empty() {
let table = get_table_from_catalog(&catalog, "simple_log_table").await;

let mut read_builder = table.new_read_builder();
read_builder.with_projection(&[]);
read_builder
.with_projection(&[])
.expect("Empty projection should succeed");
let read = read_builder
.new_read()
.expect("Empty projection should succeed");
Expand Down Expand Up @@ -528,9 +534,8 @@ async fn test_read_projection_unknown_column() {
let table = get_table_from_catalog(&catalog, "simple_log_table").await;

let mut read_builder = table.new_read_builder();
read_builder.with_projection(&["id", "nonexistent_column"]);
let err = read_builder
.new_read()
.with_projection(&["id", "nonexistent_column"])
.expect_err("Unknown columns should fail");

assert!(
Expand All @@ -551,9 +556,8 @@ async fn test_read_projection_all_invalid() {
let table = get_table_from_catalog(&catalog, "simple_log_table").await;

let mut read_builder = table.new_read_builder();
read_builder.with_projection(&["nonexistent_a", "nonexistent_b"]);
let err = read_builder
.new_read()
.with_projection(&["nonexistent_a", "nonexistent_b"])
.expect_err("All-invalid projection should fail");

assert!(
Expand All @@ -574,9 +578,8 @@ async fn test_read_projection_duplicate_column() {
let table = get_table_from_catalog(&catalog, "simple_log_table").await;

let mut read_builder = table.new_read_builder();
read_builder.with_projection(&["id", "id"]);
let err = read_builder
.new_read()
.with_projection(&["id", "id"])
.expect_err("Duplicate projection should fail");

assert!(
Expand Down Expand Up @@ -3164,7 +3167,9 @@ async fn test_read_data_evolution_table_with_row_id_projection() {

// Project _ROW_ID along with regular columns
let mut read_builder = table.new_read_builder();
read_builder.with_projection(&["_ROW_ID", "id", "name"]);
read_builder
.with_projection(&["_ROW_ID", "id", "name"])
.expect("Row ID projection should succeed");
let scan = read_builder.new_scan();
let plan = scan.plan().await.expect("Failed to plan scan");

Expand Down Expand Up @@ -3233,7 +3238,9 @@ async fn test_read_data_evolution_table_only_row_id_with_row_ranges() {
// Project only _ROW_ID with a partial row range
let mid = min_row_id + (max_row_id - min_row_id) / 2;
let mut read_builder = table.new_read_builder();
read_builder.with_projection(&["_ROW_ID"]);
read_builder
.with_projection(&["_ROW_ID"])
.expect("Row ID projection should succeed");
read_builder.with_row_ranges(vec![RowRange::new(min_row_id, mid)]);
let scan = read_builder.new_scan();
let plan = scan.plan().await.expect("plan");
Expand Down Expand Up @@ -3267,7 +3274,9 @@ async fn test_read_data_evolution_mixed_format_row_id_projection() {
let table = get_table_from_catalog(&catalog, "data_evolution_mixed_format_add_column").await;

let mut read_builder = table.new_read_builder();
read_builder.with_projection(&["_ROW_ID", "id"]);
read_builder
.with_projection(&["_ROW_ID", "id"])
.expect("Row ID projection should succeed");
let scan = read_builder.new_scan();
let plan = scan.plan().await.expect("Failed to plan scan");

Expand Down
14 changes: 9 additions & 5 deletions crates/integrations/datafusion/src/lateral_vector_search.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,9 +72,10 @@ impl QueryPlanner for PaimonQueryPlanner {
logical_plan: &LogicalPlan,
session_state: &SessionState,
) -> DFResult<Arc<dyn ExecutionPlan>> {
let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
LateralVectorSearchExtensionPlanner,
)]);
let planner = DefaultPhysicalPlanner::with_extension_planners(vec![
Arc::new(LateralVectorSearchExtensionPlanner),
Arc::new(crate::variant_pushdown::VariantExtractionExtensionPlanner),
]);
planner
.create_physical_plan(logical_plan, session_state)
.await
Expand All @@ -91,8 +92,10 @@ impl RewriteLateralVectorSearch {
}

pub(crate) fn optimizer_rules() -> Vec<Arc<dyn OptimizerRule + Send + Sync>> {
let mut rules: Vec<Arc<dyn OptimizerRule + Send + Sync>> =
vec![Arc::new(RewriteLateralVectorSearch::new())];
let mut rules: Vec<Arc<dyn OptimizerRule + Send + Sync>> = vec![
Arc::new(crate::variant_pushdown::RewriteVariantExtractions::new()),
Arc::new(RewriteLateralVectorSearch::new()),
];
rules.extend(Optimizer::default().rules);
rules
}
Expand Down Expand Up @@ -611,6 +614,7 @@ async fn read_target_rows(
let mut read_builder = table.new_read_builder();
read_builder
.with_projection(&projection_refs)
.map_err(to_datafusion_error)?
.with_row_ranges(row_ranges);
let plan = read_builder
.new_scan()
Expand Down
1 change: 1 addition & 0 deletions crates/integrations/datafusion/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ mod table;
mod table_function_args;
mod update;
mod variant_functions;
mod variant_pushdown;
mod vector_search;

use std::collections::HashMap;
Expand Down
54 changes: 38 additions & 16 deletions crates/integrations/datafusion/src/physical_plan/scan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties};
use futures::{StreamExt, TryStreamExt};
use paimon::spec::Predicate;
use paimon::spec::{DataField, Predicate};
use paimon::table::{ScanTrace, Table};
use paimon::DataSplit;

Expand All @@ -41,8 +41,8 @@ use crate::error::to_datafusion_error;
#[derive(Debug)]
pub struct PaimonTableScan {
table: Table,
/// Projected column names (if None, reads all columns).
projected_columns: Option<Vec<String>>,
/// Full Paimon read type for nested or connector-defined projections.
read_type: Vec<DataField>,
/// Filter translated from DataFusion expressions and reused during execute()
/// so reader-side pruning reaches the actual read path.
pushed_predicate: Option<Predicate>,
Expand All @@ -58,19 +58,22 @@ pub struct PaimonTableScan {
filter_exact: bool,
/// Metadata-pruning trace captured during eager scan planning.
scan_trace: Option<ScanTrace>,
/// Human-readable Variant extraction summary for explain output.
pushed_variants: Option<String>,
}

impl PaimonTableScan {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
schema: ArrowSchemaRef,
table: Table,
projected_columns: Option<Vec<String>>,
read_type: Vec<DataField>,
pushed_predicate: Option<Predicate>,
planned_partitions: Vec<Arc<[DataSplit]>>,
limit: Option<usize>,
filter_exact: bool,
scan_trace: Option<ScanTrace>,
pushed_variants: Option<String>,
) -> Self {
let plan_properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(schema.clone()),
Expand All @@ -80,13 +83,14 @@ impl PaimonTableScan {
));
Self {
table,
projected_columns,
read_type,
pushed_predicate,
planned_partitions,
plan_properties,
limit,
filter_exact,
scan_trace,
pushed_variants,
}
}

Expand Down Expand Up @@ -143,16 +147,13 @@ impl ExecutionPlan for PaimonTableScan {

let table = self.table.clone();
let schema = self.schema();
let projected_columns = self.projected_columns.clone();
let read_type = self.read_type.clone();
let pushed_predicate = self.pushed_predicate.clone();

let fut = async move {
let mut read_builder = table.new_read_builder();

if let Some(ref columns) = projected_columns {
let col_refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
read_builder.with_projection(&col_refs);
}
read_builder.with_read_type(read_type);
if let Some(filter) = pushed_predicate {
read_builder.with_filter(filter);
}
Expand Down Expand Up @@ -236,9 +237,12 @@ impl DisplayAs for PaimonTableScan {
self.planned_partitions.len()
)?;

if let Some(ref columns) = self.projected_columns {
write!(f, ", projection=[{}]", columns.join(", "))?;
}
let columns = self
.read_type
.iter()
.map(|field| field.name())
.collect::<Vec<_>>();
write!(f, ", projection=[{}]", columns.join(", "))?;
if let Some(ref predicate) = self.pushed_predicate {
write!(f, ", predicate={predicate}")?;
}
Expand All @@ -248,6 +252,9 @@ impl DisplayAs for PaimonTableScan {
if let Some(ref trace) = self.scan_trace {
write!(f, ", trace={trace}")?;
}
if let Some(ref pushed_variants) = self.pushed_variants {
write!(f, ", PushedVariants=[{pushed_variants}]")?;
}
Ok(())
}
}
Expand Down Expand Up @@ -282,18 +289,27 @@ mod tests {
)]))
}

fn test_read_type() -> Vec<DataField> {
vec![DataField::new(
0,
"id".to_string(),
DataType::Int(IntType::new()),
)]
}

#[test]
fn test_partition_count_empty_plan() {
let schema = test_schema();
let scan = PaimonTableScan::new(
schema,
dummy_table(),
None,
test_read_type(),
None,
vec![Arc::from(Vec::new())],
None,
false,
None,
None,
);
assert_eq!(scan.properties().output_partitioning().partition_count(), 1);
}
Expand All @@ -309,12 +325,13 @@ mod tests {
let scan = PaimonTableScan::new(
schema,
dummy_table(),
None,
test_read_type(),
None,
planned_partitions,
None,
false,
None,
None,
);
assert_eq!(scan.properties().output_partitioning().partition_count(), 3);
}
Expand Down Expand Up @@ -387,12 +404,17 @@ mod tests {
let scan = PaimonTableScan::new(
schema,
table,
Some(vec!["id".to_string()]),
vec![DataField::new(
0,
"id".to_string(),
DataType::Int(IntType::new()),
)],
Some(pushed_predicate),
vec![Arc::from(vec![split])],
None,
false,
None,
None,
);

let ctx = SessionContext::new();
Expand Down
Loading
Loading