Skip to content

Commit 9a195f6

Browse files
committed
fix(substrait): Handle scan counts outside the signed range
1 parent 0bbad2d commit 9a195f6

3 files changed

Lines changed: 151 additions & 9 deletions

File tree

‎datafusion/substrait/src/logical_plan/producer/rel/fetch_rel.rs‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
use crate::logical_plan::producer::SubstraitProducer;
1919
use datafusion::common::DFSchema;
20-
use datafusion::logical_expr::Limit;
20+
use datafusion::logical_expr::{FetchType, Limit, LogicalPlan, SkipType};
2121
use std::sync::Arc;
2222
use substrait::proto::rel::RelType;
2323
use substrait::proto::{FetchRel, Rel, fetch_rel};
@@ -26,7 +26,22 @@ pub fn from_limit(
2626
producer: &mut impl SubstraitProducer,
2727
limit: &Limit,
2828
) -> datafusion::common::Result<Box<Rel>> {
29-
let input = producer.handle_plan(limit.input.as_ref())?;
29+
// Pushing down LIMIT + OFFSET can produce a scan fetch above i64::MAX.
30+
// It is redundant when this Limit requests no rows beyond that bound.
31+
// Remove only that redundant fetch; retain the scan's exact offset.
32+
let mut input_plan = None;
33+
if let LogicalPlan::TableScan(scan) = limit.input.as_ref()
34+
&& let Some(scan_fetch) = scan.fetch
35+
&& i64::try_from(scan_fetch).is_err()
36+
&& let FetchType::Literal(Some(fetch)) = limit.get_fetch_type()?
37+
&& let SkipType::Literal(skip) = limit.get_skip_type()?
38+
&& skip.checked_add(fetch).is_some_and(|end| end <= scan_fetch)
39+
{
40+
let mut scan = scan.clone();
41+
scan.fetch = None;
42+
input_plan = Some(LogicalPlan::TableScan(scan));
43+
}
44+
let input = producer.handle_plan(input_plan.as_ref().unwrap_or(&limit.input))?;
3045
let empty_schema = Arc::new(DFSchema::empty());
3146

3247
let offset_mode = limit

‎datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
// under the License.
1717

1818
use crate::logical_plan::producer::{SubstraitProducer, to_substrait_named_struct};
19-
use datafusion::common::{DFSchema, ToDFSchema, substrait_datafusion_err};
19+
use datafusion::common::{
20+
DFSchema, ToDFSchema, not_impl_datafusion_err, substrait_datafusion_err,
21+
};
2022
use datafusion::logical_expr::utils::conjunction;
2123
use datafusion::logical_expr::{EmptyRelation, Expr, TableScan, Values, lit};
2224
use datafusion::scalar::ScalarValue;
@@ -61,6 +63,12 @@ pub fn from_table_scan(
6163
producer: &mut impl SubstraitProducer,
6264
scan: &TableScan,
6365
) -> datafusion::common::Result<Box<Rel>> {
66+
let skip = scan.skip.map(i64::try_from).transpose().map_err(|_| {
67+
not_impl_datafusion_err!("TableScan offset exceeds the supported Int64 range")
68+
})?;
69+
let fetch = scan.fetch.map(i64::try_from).transpose().map_err(|_| {
70+
not_impl_datafusion_err!("TableScan fetch exceeds the supported Int64 range")
71+
})?;
6472
let projection = scan.projection.as_ref().map(|p| {
6573
p.iter()
6674
.map(|i| StructItem {
@@ -116,15 +124,13 @@ pub fn from_table_scan(
116124
return Ok(read);
117125
}
118126
let empty_schema = Arc::new(DFSchema::empty());
119-
let offset_mode = scan
120-
.skip
121-
.map(|skip| producer.handle_expr(&lit(skip as i64), &empty_schema))
127+
let offset_mode = skip
128+
.map(|skip| producer.handle_expr(&lit(skip), &empty_schema))
122129
.transpose()?
123130
.map(Box::new)
124131
.map(fetch_rel::OffsetMode::OffsetExpr);
125-
let count_mode = scan
126-
.fetch
127-
.map(|fetch| producer.handle_expr(&lit(fetch as i64), &empty_schema))
132+
let count_mode = fetch
133+
.map(|fetch| producer.handle_expr(&lit(fetch), &empty_schema))
128134
.transpose()?
129135
.map(Box::new)
130136
.map(fetch_rel::CountMode::CountExpr);

‎datafusion/substrait/tests/cases/roundtrip_logical_plan.rs‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,127 @@ async fn roundtrip_table_scan_offset_and_fetch() -> Result<()> {
398398
Ok(())
399399
}
400400

401+
// Limit expressions use Int64, but limit pushdown can produce a larger usize fetch.
402+
#[tokio::test]
403+
#[cfg(target_pointer_width = "64")]
404+
async fn roundtrip_table_scan_fetch_above_i64_max() -> Result<()> {
405+
use datafusion::arrow::array::Int64Array;
406+
use datafusion::arrow::record_batch::RecordBatch;
407+
use datafusion::dataframe::DataFrame;
408+
409+
let ctx = SessionContext::new();
410+
ctx.register_batch(
411+
"t",
412+
RecordBatch::try_from_iter(vec![(
413+
"x",
414+
Arc::new(Int64Array::from(vec![0, 1, 2])) as ArrayRef,
415+
)])?,
416+
)?;
417+
let plan = ctx
418+
.sql("SELECT * FROM t LIMIT 9223372036854775807 OFFSET 1")
419+
.await?
420+
.into_optimized_plan()?;
421+
let LogicalPlan::Limit(limit) = &plan else {
422+
panic!("expected Limit")
423+
};
424+
let LogicalPlan::TableScan(scan) = limit.input.as_ref() else {
425+
panic!("expected TableScan")
426+
};
427+
assert_eq!(scan.fetch, Some(i64::MAX as usize + 1));
428+
let proto = to_substrait_plan(&plan, &ctx.state())?;
429+
let restored = from_substrait_plan(&ctx.state(), &proto).await?;
430+
for plan in [plan, restored] {
431+
let batches = DataFrame::new(ctx.state(), plan).collect().await?;
432+
datafusion::assert_batches_eq!(
433+
["+---+", "| x |", "+---+", "| 1 |", "| 2 |", "+---+"],
434+
&batches
435+
);
436+
}
437+
Ok(())
438+
}
439+
440+
#[tokio::test]
441+
#[cfg(target_pointer_width = "64")]
442+
async fn roundtrip_table_scan_signed_range() -> Result<()> {
443+
use datafusion::arrow::array::Int64Array;
444+
use datafusion::dataframe::DataFrame;
445+
use datafusion::datasource::provider_as_source;
446+
use datafusion::logical_expr::TableScanBuilder;
447+
448+
let ctx = SessionContext::new();
449+
let temp = tempfile::tempdir()?;
450+
let path = temp.path().join("data.csv");
451+
std::fs::write(&path, "x\n0\n1\n2\n")?;
452+
ctx.register_csv("t", path.to_str().unwrap(), CsvReadOptions::new())
453+
.await?;
454+
let source = provider_as_source(ctx.table_provider("t").await?);
455+
let max = i64::MAX as usize;
456+
for (skip, fetch, parent, succeeds) in [
457+
(None, Some(max), None, true),
458+
(Some(max), None, None, true),
459+
(None, Some(max + 1), None, false),
460+
(None, Some(usize::MAX), None, false),
461+
(Some(max + 1), None, None, false),
462+
(Some(usize::MAX), None, None, false),
463+
// The parent covers the scan bound exactly, or imposes a tighter bound.
464+
(None, Some(max + 1), Some((1, Some(max))), true),
465+
(None, Some(usize::MAX), Some((1, Some(2))), true),
466+
(Some(1), Some(max + 1), Some((0, Some(2))), true),
467+
// These parents do not make the scan fetch redundant.
468+
(None, Some(max + 1), Some((max, Some(max))), false),
469+
(None, Some(max + 1), Some((1, None)), false),
470+
] {
471+
let scan = TableScanBuilder::new("t", source.clone())
472+
.with_skip(skip)
473+
.with_fetch(fetch)
474+
.build()?;
475+
let mut builder = LogicalPlanBuilder::from(LogicalPlan::TableScan(scan));
476+
if let Some((skip, fetch)) = parent {
477+
builder = builder.limit(skip, fetch)?;
478+
}
479+
let plan = builder.build()?;
480+
let proto = to_substrait_plan(&plan, &ctx.state());
481+
if !succeeds {
482+
let error = proto.unwrap_err().to_string();
483+
let field = if skip.is_some() { "offset" } else { "fetch" };
484+
assert!(
485+
error.contains(&format!(
486+
"TableScan {field} exceeds the supported Int64 range"
487+
)),
488+
"{error}"
489+
);
490+
continue;
491+
}
492+
let proto = proto?;
493+
let restored = from_substrait_plan(&ctx.state(), &proto).await?;
494+
let expected = if skip == Some(max) {
495+
vec![]
496+
} else if parent.is_some() {
497+
vec![1, 2]
498+
} else {
499+
vec![0, 1, 2]
500+
};
501+
for plan in [plan, restored] {
502+
let batches = DataFrame::new(ctx.state(), plan).collect().await?;
503+
let actual = batches
504+
.iter()
505+
.flat_map(|batch| {
506+
batch
507+
.column(0)
508+
.as_any()
509+
.downcast_ref::<Int64Array>()
510+
.unwrap()
511+
.values()
512+
.iter()
513+
.copied()
514+
})
515+
.collect::<Vec<_>>();
516+
assert_eq!(actual, expected);
517+
}
518+
}
519+
Ok(())
520+
}
521+
401522
#[tokio::test]
402523
async fn select_with_limit_offset() -> Result<()> {
403524
roundtrip("SELECT * FROM data LIMIT 200 OFFSET 10").await?;

0 commit comments

Comments
 (0)