feat(proto): support expressions in LIMIT skip and fetch

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!(

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<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,
 }
 

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 {

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__,
                 })
             }
         }

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<LogicalPlanNode>>,
-    /// 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<LogicalExprNode>>,
+    /// 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<LogicalExprNode>>,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct SelectionExecNode {

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<Expr> =
@@ -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,
                         },
                     ))),
                 })