ARROW-12214: [Rust][DataFusion] Add tests for limit
I was debugging another issue (not a bug in DataFusion I don't think) but noticed there wasn't any coverage for LIMIT in exec.rs, so I figured I would add some.
(well really I was writing a test to trigger what I thought was a bug in DataFusion -- lol)
Closes #9897 from alamb/limit_fix
Authored-by: Andrew Lamb <andrew@nerdnetworks.org>
Signed-off-by: Andrew Lamb <andrew@nerdnetworks.org>
diff --git a/rust/datafusion/src/execution/context.rs b/rust/datafusion/src/execution/context.rs
index a22b5d3..50b3a7a 100644
--- a/rust/datafusion/src/execution/context.rs
+++ b/rust/datafusion/src/execution/context.rs
@@ -1722,6 +1722,46 @@
}
#[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();
ctx.register_table("t", table_with_sequence(1, 1).unwrap())
