From a820e520ea250d458e651b5b27ca26112228bd77 Mon Sep 17 00:00:00 2001 From: Anton Karpenko Date: Wed, 15 Jul 2026 11:40:41 +0300 Subject: [PATCH] feat: Add logical plan serialization with placeholders in limit Now the logical plan can be de/serialized if there are placeholders in the limit. --- datafusion/core/src/physical_planner.rs | 24 ++++++++++++- datafusion/expr/src/logical_plan/plan.rs | 3 +- datafusion/proto/proto/datafusion.proto | 11 +++--- datafusion/proto/src/generated/pbjson.rs | 28 ++++++--------- datafusion/proto/src/generated/prost.rs | 14 ++++---- datafusion/proto/src/logical_plan/mod.rs | 45 ++++++++++++++---------- 6 files changed, 76 insertions(+), 49 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index b041e75a92583..29eeef9e305e4 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, SessionState}; 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, Partitioning as LogicalPartitioning, PlanType, Repartition, @@ -85,12 +86,14 @@ use datafusion_expr::expr::{ }; use datafusion_expr::expr_rewriter::unnormalize_cols; use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary; +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::{AggregateExprBuilder, AggregateFunctionExpr}; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::{ @@ -1202,6 +1205,25 @@ impl DefaultPhysicalPlanner { LogicalPlan::Subquery(_) => todo!(), LogicalPlan::SubqueryAlias(_) => children.one()?, LogicalPlan::Limit(limit) => { + // Try to evaluate skip and fetch expressions. + let simplifier = ExprSimplifier::new(SimplifyContext::default()); + + 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 dc4daf1ab7532..9c39cdaed0c9a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -3281,7 +3281,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, } diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 15c4538a3943a..7410e705e5091 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -349,11 +349,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/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index d851632276291..cf7d9bc001fbc 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -11072,25 +11072,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() } @@ -11170,24 +11166,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/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index 5c37aa76e6028..471364c4a20c7 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -554,12 +554,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, tag = "4")] + pub skip: ::core::option::Option, + /// Maximum number of rows to fetch; + /// If it is Literal and negative means no limit + #[prost(message, optional, tag = "5")] + pub fetch: ::core::option::Option, } #[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 a5d74d7f49fae..cd19cb7bcf61e 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -53,9 +53,7 @@ use datafusion_datasource_json::file_format::{ }; #[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; -use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, RecursiveQuery, SkipType, TableSource, Unnest, -}; +use datafusion_expr::{AggregateUDF, DmlStatement, RecursiveQuery, TableSource, Unnest}; use datafusion_expr::{ DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, Statement, WindowUDF, dml, @@ -801,15 +799,24 @@ 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.as_ref() { + Some(expr) => { + Some(from_proto::parse_expr(expr, ctx, extension_codec)?) + } + None => None, + }; + + let fetch = match limit.fetch.as_ref() { + 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 = @@ -1483,23 +1490,23 @@ 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(serialize_expr(expr.as_ref(), 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(serialize_expr(expr.as_ref(), 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, }, ))), })