Skip to content

Commit 848119f

Browse files
adriangbclaude
andcommitted
Add preferred_ordering field to TableScan
This commit adds a new optional field `preferred_ordering` to the `TableScan` logical plan node to support sort pushdown optimizations. Changes include: - Add `preferred_ordering: Option<Vec<SortExpr>>` field to `TableScan` struct - Add `try_new_with_preferred_ordering` constructor method - Update all `TableScan` constructors throughout the codebase to include the new field - Update `Debug`, `PartialEq`, `Hash`, and `PartialOrd` implementations - Update pattern matching in optimizer and other modules The preferred_ordering field is currently not used by any optimization rules but provides the foundation for future sort pushdown implementations. This is part 2 of 2 PRs split from apache#17273 as requested in apache#17273 (comment) 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <[email protected]>
1 parent 1d9e138 commit 848119f

File tree

5 files changed

+32
-1
lines changed

5 files changed

+32
-1
lines changed

datafusion/expr/src/logical_plan/plan.rs

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2552,6 +2552,8 @@ pub struct TableScan {
25522552
pub filters: Vec<Expr>,
25532553
/// Optional number of rows to read
25542554
pub fetch: Option<usize>,
2555+
/// Optional preferred ordering for the scan
2556+
pub preferred_ordering: Option<Vec<SortExpr>>,
25552557
}
25562558

25572559
impl Debug for TableScan {
@@ -2563,6 +2565,7 @@ impl Debug for TableScan {
25632565
.field("projected_schema", &self.projected_schema)
25642566
.field("filters", &self.filters)
25652567
.field("fetch", &self.fetch)
2568+
.field("preferred_ordering", &self.preferred_ordering)
25662569
.finish_non_exhaustive()
25672570
}
25682571
}
@@ -2574,6 +2577,7 @@ impl PartialEq for TableScan {
25742577
&& self.projected_schema == other.projected_schema
25752578
&& self.filters == other.filters
25762579
&& self.fetch == other.fetch
2580+
&& self.preferred_ordering == other.preferred_ordering
25772581
}
25782582
}
25792583

@@ -2593,18 +2597,22 @@ impl PartialOrd for TableScan {
25932597
pub filters: &'a Vec<Expr>,
25942598
/// Optional number of rows to read
25952599
pub fetch: &'a Option<usize>,
2600+
/// Optional preferred ordering for the scan
2601+
pub preferred_ordering: &'a Option<Vec<SortExpr>>,
25962602
}
25972603
let comparable_self = ComparableTableScan {
25982604
table_name: &self.table_name,
25992605
projection: &self.projection,
26002606
filters: &self.filters,
26012607
fetch: &self.fetch,
2608+
preferred_ordering: &self.preferred_ordering,
26022609
};
26032610
let comparable_other = ComparableTableScan {
26042611
table_name: &other.table_name,
26052612
projection: &other.projection,
26062613
filters: &other.filters,
26072614
fetch: &other.fetch,
2615+
preferred_ordering: &other.preferred_ordering,
26082616
};
26092617
comparable_self.partial_cmp(&comparable_other)
26102618
}
@@ -2617,6 +2625,7 @@ impl Hash for TableScan {
26172625
self.projected_schema.hash(state);
26182626
self.filters.hash(state);
26192627
self.fetch.hash(state);
2628+
self.preferred_ordering.hash(state);
26202629
}
26212630
}
26222631

@@ -2670,8 +2679,22 @@ impl TableScan {
26702679
projected_schema,
26712680
filters,
26722681
fetch,
2682+
preferred_ordering: None,
26732683
})
26742684
}
2685+
2686+
pub fn try_new_with_preferred_ordering(
2687+
table_name: impl Into<TableReference>,
2688+
table_source: Arc<dyn TableSource>,
2689+
projection: Option<Vec<usize>>,
2690+
filters: Vec<Expr>,
2691+
fetch: Option<usize>,
2692+
preferred_ordering: Option<Vec<SortExpr>>,
2693+
) -> Result<Self> {
2694+
let mut table_scan = Self::try_new(table_name, table_source, projection, filters, fetch)?;
2695+
table_scan.preferred_ordering = preferred_ordering;
2696+
Ok(table_scan)
2697+
}
26752698
}
26762699

26772700
// Repartition the plan based on a partitioning scheme.
@@ -4896,6 +4919,7 @@ mod tests {
48964919
projected_schema: Arc::clone(&schema),
48974920
filters: vec![],
48984921
fetch: None,
4922+
preferred_ordering: None,
48994923
}));
49004924
let col = schema.field_names()[0].clone();
49014925

@@ -4926,6 +4950,7 @@ mod tests {
49264950
projected_schema: Arc::clone(&unique_schema),
49274951
filters: vec![],
49284952
fetch: None,
4953+
preferred_ordering: None,
49294954
}));
49304955
let col = schema.field_names()[0].clone();
49314956

datafusion/expr/src/logical_plan/tree_node.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -599,6 +599,7 @@ impl LogicalPlan {
599599
projected_schema,
600600
filters,
601601
fetch,
602+
preferred_ordering,
602603
}) => filters.map_elements(f)?.update_data(|filters| {
603604
LogicalPlan::TableScan(TableScan {
604605
table_name,
@@ -607,6 +608,7 @@ impl LogicalPlan {
607608
projected_schema,
608609
filters,
609610
fetch,
611+
preferred_ordering,
610612
})
611613
}),
612614
LogicalPlan::Distinct(Distinct::On(DistinctOn {

datafusion/optimizer/src/optimize_projections/mod.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,7 @@ fn optimize_projections(
242242
filters,
243243
fetch,
244244
projected_schema: _,
245+
preferred_ordering,
245246
} = table_scan;
246247

247248
// Get indices referred to in the original (schema with all fields)
@@ -250,12 +251,13 @@ fn optimize_projections(
250251
Some(projection) => indices.into_mapped_indices(|idx| projection[idx]),
251252
None => indices.into_inner(),
252253
};
253-
return TableScan::try_new(
254+
return TableScan::try_new_with_preferred_ordering(
254255
table_name,
255256
source,
256257
Some(projection),
257258
filters,
258259
fetch,
260+
preferred_ordering,
259261
)
260262
.map(LogicalPlan::TableScan)
261263
.map(Transformed::yes);

datafusion/optimizer/src/push_down_filter.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3055,6 +3055,7 @@ mod tests {
30553055
projection,
30563056
source: Arc::new(test_provider),
30573057
fetch: None,
3058+
preferred_ordering: None,
30583059
});
30593060

30603061
Ok(LogicalPlanBuilder::from(table_scan))

datafusion/proto/src/logical_plan/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,7 @@ fn from_table_source(
271271
projected_schema,
272272
filters: vec![],
273273
fetch: None,
274+
preferred_ordering: None,
274275
});
275276

276277
LogicalPlanNode::try_from_logical_plan(&r, extension_codec)

0 commit comments

Comments
 (0)