diff --git a/rust/datafusion/src/execution/context.rs b/rust/datafusion/src/execution/context.rs index a22b5d3c1e7..50b3a7a1add 100644 --- a/rust/datafusion/src/execution/context.rs +++ b/rust/datafusion/src/execution/context.rs @@ -1721,6 +1721,46 @@ mod tests { Ok(()) } + #[tokio::test] + async fn limit() -> Result<()> { + let tmp_dir = TempDir::new()?; + let mut ctx = create_ctx(&tmp_dir, 1)?; + ctx.register_table("t", table_with_sequence(1, 1000).unwrap()) + .unwrap(); + + let results = + plan_and_collect(&mut ctx, "SELECT i FROM t ORDER BY i DESC limit 3") + .await + .unwrap(); + + let expected = vec![ + "+------+", "| i |", "+------+", "| 1000 |", "| 999 |", "| 998 |", + "+------+", + ]; + + assert_batches_eq!(expected, &results); + + let results = plan_and_collect(&mut ctx, "SELECT i FROM t ORDER BY i limit 3") + .await + .unwrap(); + + let expected = vec![ + "+---+", "| i |", "+---+", "| 1 |", "| 2 |", "| 3 |", "+---+", + ]; + + assert_batches_eq!(expected, &results); + + let results = plan_and_collect(&mut ctx, "SELECT i FROM t limit 3") + .await + .unwrap(); + + // the actual rows are not guaranteed, so only check the count (should be 3) + let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, 3); + + Ok(()) + } + #[tokio::test] async fn case_sensitive_identifiers_functions() { let mut ctx = ExecutionContext::new();