From b6d6669582c35552a4aa757dbad2263c49ce4266 Mon Sep 17 00:00:00 2001 From: Sammy Sidhu Date: Thu, 14 Dec 2023 20:25:31 -0800 Subject: [PATCH] add more scan operator tests --- .../src/optimization/rules/push_down_limit.rs | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/src/daft-plan/src/optimization/rules/push_down_limit.rs b/src/daft-plan/src/optimization/rules/push_down_limit.rs index 6cf3704dce..a679fef3a0 100644 --- a/src/daft-plan/src/optimization/rules/push_down_limit.rs +++ b/src/daft-plan/src/optimization/rules/push_down_limit.rs @@ -127,7 +127,10 @@ mod tests { rules::PushDownLimit, Optimizer, }, - test::{dummy_scan_node, dummy_scan_node_with_pushdowns}, + test::{ + dummy_scan_node, dummy_scan_node_with_pushdowns, + dummy_scan_operator_node_with_pushdowns, + }, LogicalPlan, PartitionScheme, }; @@ -259,6 +262,27 @@ mod tests { Ok(()) } + /// Tests that Limit does push into external Source with existing larger limit. + /// + /// Limit-Source[existing_limit] -> Source[new_limit] + #[test] + fn limit_does_push_into_scan_operator_if_larger_limit() -> DaftResult<()> { + let plan = dummy_scan_operator_node_with_pushdowns( + vec![ + Field::new("a", DataType::Int64), + Field::new("b", DataType::Utf8), + ], + Pushdowns::default().with_limit(Some(10)), + ) + .limit(5, false)? + .build(); + let expected = "\ + Limit: 5\ + \n Source: Operator = AnonymousScanOperator: File paths=[/foo], Format-specific config = Json(JsonSourceConfig { buffer_size: None, chunk_size: None }), Storage config = Native(NativeStorageConfig { io_config: None, multithreaded_io: true }), File schema = a (Int64), b (Utf8), Partitioning keys = [], Limit pushdown = 5, Output schema = a (Int64), b (Utf8)"; + assert_optimized_plan_eq(plan, expected)?; + Ok(()) + } + /// Tests that Limit does not push into in-memory Source. #[test] #[cfg(feature = "python")]