udf-chain-physical-fusion-Daft-AtomGit

Msrc/daft-logical-plan/src/logical_plan.rs+8-2

	@@ -529,10 +529,16 @@ impl LogicalPlan {
529	Self::UDFProject(UDFProject {	529	Self::UDFProject(UDFProject {
530	expr,	530	expr,
531	passthrough_columns,	531	passthrough_columns,
		532	+ udf_properties,
532	..	533	..
533	}) => Self::UDFProject(	534	}) => Self::UDFProject(
534	- UDFProject::try_new(input.clone(), expr.clone(), passthrough_columns.clone())	535	+ UDFProject::try_new_fused(
535	- .unwrap(),	536	+ input.clone(),
		537	+ expr.clone(),
		538	+ passthrough_columns.clone(),
		539	+ udf_properties.clone(),
		540	+ )
		541	+ .unwrap(),
536	),	542	),
537	Self::Filter(Filter { predicate, .. }) => {	543	Self::Filter(Filter { predicate, .. }) => {
538	Self::Filter(Filter::try_new(input.clone(), predicate.clone()).unwrap())	544	Self::Filter(Filter::try_new(input.clone(), predicate.clone()).unwrap())

Msrc/daft-logical-plan/src/ops/udf.rs+28-0

	@@ -63,6 +63,34 @@ impl UDFProject {
63	})	63	})
64	}	64	}
65		65
		66	+ /// Construct a UDF project whose expression contains a compatible chain
		67	+ /// of Python UDFs. The physical-fusion optimizer has already checked that
		68	+ /// all UDFs share the supplied scheduling and error-handling properties.
		69	+ pub(crate) fn try_new_fused(
		70	+ input: Arc<LogicalPlan>,
		71	+ expr: ExprRef,
		72	+ passthrough_columns: Vec<ExprRef>,
		73	+ udf_properties: UDFProperties,
		74	+ ) -> Result<Self> {
		75	+ let output_field = expr.to_field(&input.schema())?;
		76	+ let fields = passthrough_columns
		77	+ .iter()
		78	+ .map(\|e\| e.to_field(&input.schema()))
		79	+ .chain(std::iter::once(Ok(output_field)))
		80	+ .collect::<DaftResult<Vec<_>>>()?;
		81	+
		82	+ Ok(Self {
		83	+ plan_id: None,
		84	+ node_id: None,
		85	+ input,
		86	+ expr,
		87	+ udf_properties,
		88	+ passthrough_columns,
		89	+ projected_schema: Arc::new(Schema::new(fields)),
		90	+ stats_state: StatsState::NotMaterialized,
		91	+ })
		92	+ }
		93	+
66	pub fn with_plan_id(mut self, plan_id: usize) -> Self {	94	pub fn with_plan_id(mut self, plan_id: usize) -> Self {
67	self.plan_id = Some(plan_id);	95	self.plan_id = Some(plan_id);
68	self	96	self

Msrc/daft-logical-plan/src/optimization/optimizer.rs+7-1

	@@ -9,7 +9,7 @@ use super::{
9	rules::{	9	rules::{
10	DetectMonotonicId, DropIntoBatches, DropRepartition, EliminateCrossJoin, EliminateOffsets,	10	DetectMonotonicId, DropIntoBatches, DropRepartition, EliminateCrossJoin, EliminateOffsets,
11	EliminateSubqueryAliasRule, EnrichWithStats, ExtractWindowFunction, FilterNullJoinKey,	11	EliminateSubqueryAliasRule, EnrichWithStats, ExtractWindowFunction, FilterNullJoinKey,
12	- LiftProjectFromAgg, MaterializeScans, OptimizerRule, PushDownAggregation,	12	+ FuseUDFProjects, LiftProjectFromAgg, MaterializeScans, OptimizerRule, PushDownAggregation,
13	PushDownAntiSemiJoin, PushDownFilter, PushDownJoinPredicate, PushDownLimit,	13	PushDownAntiSemiJoin, PushDownFilter, PushDownJoinPredicate, PushDownLimit,
14	PushDownProjection, PushDownShard, ReorderJoins, RewriteCountDistinct, RewriteOffset,	14	PushDownProjection, PushDownShard, ReorderJoins, RewriteCountDistinct, RewriteOffset,
15	ShardScans, SimplifyExpressionsRule, SimplifyNullFilteredJoin, SplitExplodeFromProject,	15	ShardScans, SimplifyExpressionsRule, SimplifyNullFilteredJoin, SplitExplodeFromProject,
	@@ -203,6 +203,12 @@ impl OptimizerBuilder {
203	vec![Box::new(PushDownProjection::new())],	203	vec![Box::new(PushDownProjection::new())],
204	RuleExecutionStrategy::FixedPoint(None),	204	RuleExecutionStrategy::FixedPoint(None),
205	),	205	),
		206	+ // Collapse compatible UDF chains after SplitUDFs has made their
		207	+ // task boundaries explicit.
		208	+ RuleBatch::new(
		209	+ vec![Box::new(FuseUDFProjects::new())],
		210	+ RuleExecutionStrategy::FixedPoint(None),
		211	+ ),
206	// --- Push down aggregations ---	212	// --- Push down aggregations ---
207	RuleBatch::new(	213	RuleBatch::new(
208	vec![Box::new(PushDownAggregation::new(	214	vec![Box::new(PushDownAggregation::new(

Asrc/daft-logical-plan/src/optimization/rules/fuse_udf_projects.rs+60-0

@@ -0,0 +1,60 @@
	1	+use std::{collections::HashMap, sync::Arc};
	2	+
	3	+use common_error::DaftResult;
	4	+use common_treenode::{Transformed, TreeNode};
	5	+use daft_dsl::optimization::replace_columns_with_expressions;
	6	+
	7	+use super::OptimizerRule;
	8	+use crate::{LogicalPlan, ops::UDFProject};
	9	+
	10	+/// Fuse adjacent, semantically compatible Python UDF projects into one
	11	+/// physical project. This removes an actor/task and materialization boundary
	12	+/// without changing the user-visible expression graph.
	13	+#[derive(Default, Debug)]
	14	+pub struct FuseUDFProjects {}
	15	+
	16	+impl FuseUDFProjects {
	17	+ pub fn new() -> Self {
	18	+ Self {}
	19	+ }
	20	+
	21	+ fn compatible(parent: &UDFProject, child: &UDFProject) -> bool {
	22	+ if !parent.passthrough_columns.is_empty() \|\| !child.passthrough_columns.is_empty() {
	23	+ return false;
	24	+ }
	25	+
	26	+ let mut parent_props = parent.udf_properties.clone();
	27	+ let mut child_props = child.udf_properties.clone();
	28	+ // Function names are diagnostic metadata, not execution semantics.
	29	+ parent_props.name.clear();
	30	+ child_props.name.clear();
	31	+ parent_props == child_props
	32	+ }
	33	+}
	34	+
	35	+impl OptimizerRule for FuseUDFProjects {
	36	+ fn try_optimize(&self, plan: Arc<LogicalPlan>) -> DaftResult<Transformed<Arc<LogicalPlan>>> {
	37	+ plan.transform_up(\|node\| {
	38	+ let LogicalPlan::UDFProject(parent) = node.as_ref() else {
	39	+ return Ok(Transformed::no(node));
	40	+ };
	41	+ let LogicalPlan::UDFProject(child) = parent.input.as_ref() else {
	42	+ return Ok(Transformed::no(node));
	43	+ };
	44	+ if !Self::compatible(parent, child) {
	45	+ return Ok(Transformed::no(node));
	46	+ }
	47	+
	48	+ let replacements = HashMap::from([(child.expr.name().to_string(), child.expr.clone())]);
	49	+ let fused_expr = replace_columns_with_expressions(parent.expr.clone(), &replacements);
	50	+ let mut fused = UDFProject::try_new_fused(
	51	+ child.input.clone(),
	52	+ fused_expr,
	53	+ vec![],
	54	+ parent.udf_properties.clone(),
	55	+ )?;
	56	+ fused.stats_state = parent.stats_state.clone();
	57	+ Ok(Transformed::yes(Arc::new(LogicalPlan::UDFProject(fused))))
	58	+ })
	59	+ }
	60	+}

Msrc/daft-logical-plan/src/optimization/rules/mod.rs+2-0

	@@ -7,6 +7,7 @@ mod eliminate_subquery_alias;
7	mod enrich_with_stats;	7	mod enrich_with_stats;
8	mod extract_window_function;	8	mod extract_window_function;
9	mod filter_null_join_key;	9	mod filter_null_join_key;
		10	+mod fuse_udf_projects;
10	mod granular_projections;	11	mod granular_projections;
11	mod lift_project_from_agg;	12	mod lift_project_from_agg;
12	mod materialize_scans;	13	mod materialize_scans;
	@@ -37,6 +38,7 @@ pub use eliminate_subquery_alias::EliminateSubqueryAliasRule;
37	pub use enrich_with_stats::EnrichWithStats;	38	pub use enrich_with_stats::EnrichWithStats;
38	pub use extract_window_function::ExtractWindowFunction;	39	pub use extract_window_function::ExtractWindowFunction;
39	pub use filter_null_join_key::FilterNullJoinKey;	40	pub use filter_null_join_key::FilterNullJoinKey;
		41	+pub use fuse_udf_projects::FuseUDFProjects;
40	pub use granular_projections::SplitGranularProjection;	42	pub use granular_projections::SplitGranularProjection;
41	pub use lift_project_from_agg::LiftProjectFromAgg;	43	pub use lift_project_from_agg::LiftProjectFromAgg;
42	pub use materialize_scans::MaterializeScans;	44	pub use materialize_scans::MaterializeScans;