From 7d464abcdbb2b5018e7c8c24d76ac520b36780cd Mon Sep 17 00:00:00 2001 From: Anton Karpenko Date: Fri, 7 Aug 2026 10:46:37 +0300 Subject: [PATCH] feat(proto): support expressions in LIMIT skip and fetch Serialize and deserialize LIMIT expressions in logical plan protobufs, and simplify them before physical planning so evaluable expressions can be used. --- datafusion/core/src/physical_planner.rs | 25 ++++++++++- datafusion/expr/src/logical_plan/plan.rs | 6 ++- .../proto-models/proto/datafusion.proto | 11 +++-- .../proto-models/src/generated/pbjson.rs | 28 +++++------- .../proto-models/src/generated/prost.rs | 14 +++--- datafusion/proto/src/logical_plan/mod.rs | 44 +++++++++++-------- 6 files changed, 78 insertions(+), 50 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index da8e0f2f574d7..0ddfb550c2510 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -29,7 +29,8 @@ use crate::error::{DataFusionError, Result}; use crate::execution::context::ExecutionProps; use crate::logical_expr::utils::generate_sort_key; use crate::logical_expr::{ - Aggregate, EmptyRelation, Join, Projection, Sort, TableScan, Unnest, Values, Window, + Aggregate, EmptyRelation, Join, Limit, Projection, Sort, TableScan, Unnest, Values, + Window, }; use crate::logical_expr::{Expr, LogicalPlan, PlanType, Repartition}; use crate::physical_expr::{ @@ -88,12 +89,14 @@ use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessar use datafusion_expr::physical_planning_context::{ PhysicalPlanningContext, ScalarSubqueryResults, SubqueryIndex, }; +use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{expr_to_columns, split_conjunction}; use datafusion_expr::{ Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp, }; +use datafusion_optimizer::simplify_expressions::ExprSimplifier; use datafusion_physical_expr::aggregate::{ AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder, }; @@ -1216,6 +1219,26 @@ impl DefaultPhysicalPlanner { } LogicalPlan::SubqueryAlias(_) => children.one()?, LogicalPlan::Limit(limit) => { + // Try to evaluate skip and fetch expressions. + let context = SimplifyContext::builder().build(); + let simplifier = ExprSimplifier::new(context); + + let skip = match &limit.skip { + Some(expr) => Some(Box::new(simplifier.simplify(*expr.clone())?)), + None => None, + }; + + let fetch = match &limit.fetch { + Some(expr) => Some(Box::new(simplifier.simplify(*expr.clone())?)), + None => None, + }; + + let limit = Limit { + input: Arc::clone(&limit.input), + skip, + fetch, + }; + let input = children.one()?; let SkipType::Literal(skip) = limit.get_skip_type()? else { return not_impl_err!( diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 9ac27b46a78e6..62f04e41ba510 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -3674,7 +3674,8 @@ pub struct Limit { pub enum SkipType { /// The skip expression is a literal value. Literal(usize), - /// Currently only supports expressions that can be folded into constants. + //// Currently supports all expressions that can be evaluated. + /// UnsupportedExpr means that the expression is not considered by the analyzer/optimizer. UnsupportedExpr, } @@ -3683,7 +3684,8 @@ pub enum FetchType { /// The fetch expression is a literal value. /// `Literal(None)` means the fetch expression is not provided. Literal(Option), - /// Currently only supports expressions that can be folded into constants. + //// Currently supports all expressions that can be evaluated. + /// UnsupportedExpr means that the expression is not considered by the analyzer/optimizer. UnsupportedExpr, } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index cbc41a7c5713e..8db00a1617032 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -438,11 +438,14 @@ message CrossJoinNode { } message LimitNode { + reserved 2, 3; LogicalPlanNode input = 1; - // The number of rows to skip before fetch; non-positive means don't skip any - int64 skip = 2; - // Maximum number of rows to fetch; negative means no limit - int64 fetch = 3; + // The number of rows to skip before fetch; + // If it is Literal and non-positive means don't skip any + LogicalExprNode skip = 4; + // Maximum number of rows to fetch; + // If it is Literal and negative means no limit + LogicalExprNode fetch = 5; } message SelectionExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 7f9b9eddc5ff5..fdea1d813afbe 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -11840,25 +11840,21 @@ impl serde::Serialize for LimitNode { if self.input.is_some() { len += 1; } - if self.skip != 0 { + if self.skip.is_some() { len += 1; } - if self.fetch != 0 { + if self.fetch.is_some() { len += 1; } let mut struct_ser = serializer.serialize_struct("datafusion.LimitNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; } - if self.skip != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("skip", ToString::to_string(&self.skip).as_str())?; + if let Some(v) = self.skip.as_ref() { + struct_ser.serialize_field("skip", v)?; } - if self.fetch != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("fetch", ToString::to_string(&self.fetch).as_str())?; + if let Some(v) = self.fetch.as_ref() { + struct_ser.serialize_field("fetch", v)?; } struct_ser.end() } @@ -11938,24 +11934,20 @@ impl<'de> serde::Deserialize<'de> for LimitNode { if skip__.is_some() { return Err(serde::de::Error::duplicate_field("skip")); } - skip__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; + skip__ = map_.next_value()?; } GeneratedField::Fetch => { if fetch__.is_some() { return Err(serde::de::Error::duplicate_field("fetch")); } - fetch__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; + fetch__ = map_.next_value()?; } } } Ok(LimitNode { input: input__, - skip: skip__.unwrap_or_default(), - fetch: fetch__.unwrap_or_default(), + skip: skip__, + fetch: fetch__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index f7633483080f1..5ca22295fcc5d 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -746,12 +746,14 @@ pub struct CrossJoinNode { pub struct LimitNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - /// The number of rows to skip before fetch; non-positive means don't skip any - #[prost(int64, tag = "2")] - pub skip: i64, - /// Maximum number of rows to fetch; negative means no limit - #[prost(int64, tag = "3")] - pub fetch: i64, + /// The number of rows to skip before fetch; + /// If it is Literal and non-positive means don't skip any + #[prost(message, optional, boxed, tag = "4")] + pub skip: ::core::option::Option<::prost::alloc::boxed::Box>, + /// Maximum number of rows to fetch; + /// If it is Literal and negative means no limit + #[prost(message, optional, boxed, tag = "5")] + pub fetch: ::core::option::Option<::prost::alloc::boxed::Box>, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct SelectionExecNode { diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 653ae9ab05355..b98ee562dd822 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -60,8 +60,8 @@ use datafusion_datasource_json::file_format::{ use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; use datafusion_expr::dml::InsertOp; use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning, - RecursiveQuery, SkipType, TableSource, Unnest, WriteOp, + AggregateUDF, DmlStatement, HigherOrderUDF, RangePartitioning, RecursiveQuery, + TableSource, Unnest, WriteOp, }; use datafusion_expr::{ DistinctOn, DropView, Expr, JoinConstraint, LogicalPlan, LogicalPlanBuilder, @@ -961,15 +961,22 @@ impl AsLogicalPlan for LogicalPlanNode { LogicalPlanType::Limit(limit) => { let input: LogicalPlan = into_logical_plan!(limit.input, ctx, extension_codec)?; - let skip = limit.skip.max(0) as usize; - - let fetch = if limit.fetch < 0 { - None - } else { - Some(limit.fetch as usize) + let skip = match &limit.skip { + Some(expr) => { + Some(from_proto::parse_expr(expr, ctx, extension_codec)?) + } + None => None, + }; + let fetch = match &limit.fetch { + Some(expr) => { + Some(from_proto::parse_expr(expr, ctx, extension_codec)?) + } + None => None, }; - LogicalPlanBuilder::from(input).limit(skip, fetch)?.build() + LogicalPlanBuilder::from(input) + .limit_by_expr(skip, fetch)? + .build() } LogicalPlanType::Join(join) => { let left_keys: Vec = @@ -1715,23 +1722,22 @@ impl AsLogicalPlan for LogicalPlanNode { limit.input.as_ref(), extension_codec, )?; - let SkipType::Literal(skip) = limit.get_skip_type()? else { - return Err(proto_error( - "LogicalPlan::Limit only supports literal skip values", - )); + let skip = match &limit.skip { + Some(expr) => Some(Box::new(serialize_expr(expr, extension_codec)?)), + None => None, }; - let FetchType::Literal(fetch) = limit.get_fetch_type()? else { - return Err(proto_error( - "LogicalPlan::Limit only supports literal fetch values", - )); + + let fetch = match &limit.fetch { + Some(expr) => Some(Box::new(serialize_expr(expr, extension_codec)?)), + None => None, }; Ok(LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::Limit(Box::new( protobuf::LimitNode { input: Some(Box::new(input)), - skip: skip as i64, - fetch: fetch.unwrap_or(i64::MAX as usize) as i64, + skip, + fetch, }, ))), })