Skip to content

Commit ee8b7c2

Browse files
SubhamSinghalSubhamSinghal
authored
Merge branch 'main' into window-topn-dense-rank
2 parents 7bac79e + 66c3840 commit ee8b7c2

7 files changed

Lines changed: 476 additions & 45 deletions

File tree

‎datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,5 +18,6 @@
1818
//! Validates query's actual memory usage is consistent with the specified memory
1919
//! limit.
2020
21+
mod smj_mem_validation;
2122
mod sort_mem_validation;
2223
mod utils;

‎datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! Memory-limit validation tests for sort-merge join queries.
19+
//!
20+
//! These tests run in separate processes to accurately measure memory usage.
21+
22+
use datafusion::prelude::SessionConfig;
23+
24+
use crate::memory_limit::memory_limit_validation::utils;
25+
26+
/// Ensures the planner selected a sort-merge join.
27+
const SMJ_OPERATOR_NAME: &str = "SortMergeJoinExec";
28+
29+
/// Configure a two-partition sort-merge join and reduce the sort reservation so
30+
/// the join can spill under the tested memory limits.
31+
fn smj_session_config() -> SessionConfig {
32+
SessionConfig::new()
33+
.with_target_partitions(2)
34+
.with_sort_spill_reservation_bytes(1024 * 1024)
35+
.set_bool("datafusion.optimizer.prefer_hash_join", false)
36+
}
37+
38+
/// Build a join with one large buffered key group and scalar output.
39+
fn smj_sum_query(series_len: usize) -> String {
40+
format!(
41+
"SELECT sum(rr.v) FROM generate_series(0, 0) AS l(k) \
42+
JOIN (SELECT i % 1 AS k, i AS v FROM generate_series(1, {series_len}) AS r(i)) rr \
43+
ON l.k = rr.k"
44+
)
45+
}
46+
47+
#[test]
48+
fn smj_with_mem_limit_1_runner() {
49+
utils::spawn_test_process("smj_mem_validation", "smj_with_mem_limit_1");
50+
}
51+
52+
#[test]
53+
fn smj_with_mem_limit_2_runner() {
54+
utils::spawn_test_process("smj_mem_validation", "smj_with_mem_limit_2");
55+
}
56+
57+
#[test]
58+
fn smj_no_mem_limit_runner() {
59+
utils::spawn_test_process("smj_mem_validation", "smj_no_mem_limit");
60+
}
61+
62+
/// Verify a 40 MB pool forces spilling within the RSS allowance.
63+
#[tokio::test]
64+
async fn smj_with_mem_limit_1() {
65+
utils::validate_query_with_memory_limits_and_config(
66+
40_000_000 * 4,
67+
Some(40_000_000),
68+
&smj_sum_query(5_000_000),
69+
&smj_sum_query(500_000),
70+
smj_session_config(),
71+
Some(SMJ_OPERATOR_NAME),
72+
Some(true),
73+
)
74+
.await;
75+
}
76+
77+
/// Verify a 16 MB pool forces spilling. The 5M join keys (~40 MB) stay resident
78+
/// independently of the pool limit, so this case needs a larger RSS allowance.
79+
#[tokio::test]
80+
async fn smj_with_mem_limit_2() {
81+
utils::validate_query_with_memory_limits_and_config(
82+
16_000_000 * 12,
83+
Some(16_000_000),
84+
&smj_sum_query(5_000_000),
85+
&smj_sum_query(500_000),
86+
smj_session_config(),
87+
Some(SMJ_OPERATOR_NAME),
88+
Some(true),
89+
)
90+
.await;
91+
}
92+
93+
#[tokio::test]
94+
async fn smj_no_mem_limit() {
95+
utils::validate_query_with_memory_limits_and_config(
96+
40_000_000 * 5,
97+
None,
98+
&smj_sum_query(5_000_000),
99+
&smj_sum_query(500_000),
100+
smj_session_config(),
101+
Some(SMJ_OPERATOR_NAME),
102+
Some(false),
103+
)
104+
.await;
105+
}

‎datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs‎

Lines changed: 10 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
//! This file is organized as:
2222
//! - Test runners that spawn individual test processes
2323
//! - Test cases that contain the actual validation logic
24-
use std::{process::Command, str};
2524
2625
use crate::memory_limit::memory_limit_validation::utils;
2726

@@ -32,67 +31,40 @@ use crate::memory_limit::memory_limit_validation::utils;
3231

3332
#[test]
3433
fn memory_limit_validation_runner_works_runner() {
35-
spawn_test_process("memory_limit_validation_runner_works");
34+
utils::spawn_test_process(
35+
"sort_mem_validation",
36+
"memory_limit_validation_runner_works",
37+
);
3638
}
3739

3840
#[test]
3941
fn sort_no_mem_limit_runner() {
40-
spawn_test_process("sort_no_mem_limit");
42+
utils::spawn_test_process("sort_mem_validation", "sort_no_mem_limit");
4143
}
4244

4345
#[test]
4446
fn sort_with_mem_limit_1_runner() {
45-
spawn_test_process("sort_with_mem_limit_1");
47+
utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_1");
4648
}
4749

4850
#[test]
4951
fn sort_with_mem_limit_2_runner() {
50-
spawn_test_process("sort_with_mem_limit_2");
52+
utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2");
5153
}
5254

5355
#[test]
5456
fn sort_with_mem_limit_3_runner() {
55-
spawn_test_process("sort_with_mem_limit_3");
57+
utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_3");
5658
}
5759

5860
#[test]
5961
fn sort_with_mem_limit_2_cols_1_runner() {
60-
spawn_test_process("sort_with_mem_limit_2_cols_1");
62+
utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2_cols_1");
6163
}
6264

6365
#[test]
6466
fn sort_with_mem_limit_2_cols_2_runner() {
65-
spawn_test_process("sort_with_mem_limit_2_cols_2");
66-
}
67-
68-
/// Helper function that executes a test in a separate process with the required
69-
/// environment variable set. Re-invokes the current test binary directly,
70-
/// avoiding cargo overhead and recompilation.
71-
fn spawn_test_process(test: &str) {
72-
let test_path =
73-
format!("memory_limit::memory_limit_validation::sort_mem_validation::{test}");
74-
75-
let exe = std::env::current_exe().expect("Failed to get test binary path");
76-
77-
let output = Command::new(exe)
78-
.arg(&test_path)
79-
.arg("--exact")
80-
.arg("--nocapture")
81-
.env("DATAFUSION_TEST_MEM_LIMIT_VALIDATION", "1")
82-
.output()
83-
.expect("Failed to execute test command");
84-
85-
let stdout = str::from_utf8(&output.stdout).unwrap_or("");
86-
let stderr = str::from_utf8(&output.stderr).unwrap_or("");
87-
88-
assert!(
89-
output.status.success(),
90-
"Test '{}' failed with status: {}\nstdout:\n{}\nstderr:\n{}",
91-
test,
92-
output.status,
93-
stdout,
94-
stderr
95-
);
67+
utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2_cols_2");
9668
}
9769

9870
// ===========================================================================

‎datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs‎

Lines changed: 101 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,14 @@
1616
// under the License.
1717

1818
use datafusion_common_runtime::SpawnedTask;
19+
use std::process::Command;
20+
use std::str;
1921
use std::sync::Arc;
2022
use std::sync::atomic::{AtomicUsize, Ordering};
2123
use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System};
2224
use tokio::time::{Duration, interval};
2325

26+
use datafusion::physical_plan::{ExecutionPlan, collect, displayable};
2427
use datafusion::prelude::{SessionConfig, SessionContext};
2528
use datafusion_common::human_readable_size;
2629
use datafusion_execution::{memory_pool::FairSpillPool, runtime_env::RuntimeEnvBuilder};
@@ -98,6 +101,42 @@ where
98101
(result, peak_rss)
99102
}
100103

104+
/// Helper function that executes a test in a separate process with the required
105+
/// environment variable set. Re-invokes the current test binary directly,
106+
/// avoiding cargo overhead and recompilation.
107+
pub fn spawn_test_process(module: &str, test: &str) {
108+
let test_path = format!("memory_limit::memory_limit_validation::{module}::{test}");
109+
let exe = std::env::current_exe().expect("Failed to get test binary path");
110+
let output = Command::new(exe)
111+
.arg(&test_path)
112+
.arg("--exact")
113+
.arg("--nocapture")
114+
.env("DATAFUSION_TEST_MEM_LIMIT_VALIDATION", "1")
115+
.output()
116+
.expect("Failed to execute test command");
117+
118+
let stdout = str::from_utf8(&output.stdout).unwrap_or("");
119+
let stderr = str::from_utf8(&output.stderr).unwrap_or("");
120+
assert!(
121+
output.status.success(),
122+
"Test '{test}' failed with status: {}\nstdout:\n{stdout}\nstderr:\n{stderr}",
123+
output.status,
124+
);
125+
}
126+
127+
fn operator_spill_count(plan: &dyn ExecutionPlan, operator_name: &str) -> usize {
128+
let own = if plan.name() == operator_name {
129+
plan.metrics().and_then(|m| m.spill_count()).unwrap_or(0)
130+
} else {
131+
0
132+
};
133+
own + plan
134+
.children()
135+
.into_iter()
136+
.map(|child| operator_spill_count(child.as_ref(), operator_name))
137+
.sum::<usize>()
138+
}
139+
101140
/// Query runner that validates the memory usage of the query.
102141
///
103142
/// Note this function is supposed to run in a separate process for accurate memory
@@ -132,6 +171,30 @@ pub async fn validate_query_with_memory_limits(
132171
mem_limit_bytes: Option<i64>,
133172
query: &str,
134173
baseline_query: &str,
174+
) {
175+
let session_config = SessionConfig::new().with_target_partitions(4); // Make sure the configuration is the same if test is running on different machines
176+
validate_query_with_memory_limits_and_config(
177+
expected_mem_bytes,
178+
mem_limit_bytes,
179+
query,
180+
baseline_query,
181+
session_config,
182+
None,
183+
None,
184+
)
185+
.await;
186+
}
187+
188+
/// Validate memory usage with a custom session configuration and optional
189+
/// operator and spill assertions.
190+
pub async fn validate_query_with_memory_limits_and_config(
191+
expected_mem_bytes: i64,
192+
mem_limit_bytes: Option<i64>,
193+
query: &str,
194+
baseline_query: &str,
195+
session_config: SessionConfig,
196+
expected_operator_name: Option<&str>,
197+
expected_operator_spill: Option<bool>,
135198
) {
136199
if std::env::var("DATAFUSION_TEST_MEM_LIMIT_VALIDATION").is_err() {
137200
println!("Skipping test because DATAFUSION_TEST_MEM_LIMIT_VALIDATION is not set");
@@ -151,18 +214,50 @@ pub async fn validate_query_with_memory_limits(
151214
None => runtime_builder.build_arc().unwrap(),
152215
};
153216

154-
let session_config = SessionConfig::new().with_target_partitions(4); // Make sure the configuration is the same if test is running on different machines
155-
156217
let ctx = SessionContext::new_with_config_rt(session_config, runtime);
157218

158219
let df = ctx.sql(query).await.unwrap();
220+
let physical_plan = df.create_physical_plan().await.unwrap();
221+
222+
if let Some(expected) = expected_operator_name {
223+
let plan_display = displayable(physical_plan.as_ref()).indent(true).to_string();
224+
assert!(
225+
plan_display.contains(expected),
226+
"expected physical plan to contain `{expected}`, but got:\n{plan_display}",
227+
);
228+
}
229+
159230
// Run a query with 10% data to estimate the constant overhead
160-
let df_small = ctx.sql(baseline_query).await.unwrap();
231+
let baseline_plan = ctx
232+
.sql(baseline_query)
233+
.await
234+
.unwrap()
235+
.create_physical_plan()
236+
.await
237+
.unwrap();
238+
let baseline_task_ctx = ctx.task_ctx();
239+
let (_, baseline_max_rss) = measure_max_rss(|| async move {
240+
collect(baseline_plan, baseline_task_ctx).await.unwrap()
241+
})
242+
.await;
161243

162-
let (_, baseline_max_rss) =
163-
measure_max_rss(|| async { df_small.collect().await.unwrap() }).await;
244+
let execution_plan = Arc::clone(&physical_plan);
245+
let execution_task_ctx = ctx.task_ctx();
246+
let (_, max_rss) = measure_max_rss(|| async move {
247+
collect(execution_plan, execution_task_ctx).await.unwrap()
248+
})
249+
.await;
164250

165-
let (_, max_rss) = measure_max_rss(|| async { df.collect().await.unwrap() }).await;
251+
if let (Some(operator), Some(expect_spill)) =
252+
(expected_operator_name, expected_operator_spill)
253+
{
254+
let spill_count = operator_spill_count(physical_plan.as_ref(), operator);
255+
assert_eq!(
256+
spill_count > 0,
257+
expect_spill,
258+
"unexpected spill_count={spill_count} for {operator}",
259+
);
260+
}
166261

167262
println!(
168263
"Memory before: {}, Memory after: {}",

‎datafusion/functions/src/datetime/to_date.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ impl ToDateFunc {
100100
args,
101101
|s, format| {
102102
string_to_timestamp_millis_formatted(s, format)
103-
.map(|n| n / (24 * 60 * 60 * 1_000))
103+
.map(|n| n.div_euclid(24 * 60 * 60 * 1_000))
104104
.and_then(|v| {
105105
v.try_into().map_err(|_| {
106106
internal_datafusion_err!("Unable to cast to Date32 for converting from i64 to i32 failed")

‎datafusion/sqllogictest/test_files/datetime/dates.slt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,12 @@ SELECT to_date('01-14-2023 01:01:30+05:30', '%q', '%d-%m-%Y %H/%M/%S', '%+', '%m
139139
----
140140
2023-01-13
141141

142+
# Formatted pre-epoch datetimes retain their calendar date
143+
query D
144+
SELECT to_date('1969-12-31 12:00:00', '%Y-%m-%d %H:%M:%S');
145+
----
146+
1969-12-31
147+
142148
statement error DataFusion error: Execution error: to_date function unsupported data type at index 1: List
143149
SELECT to_date('2022-08-03T14:38:50+05:30', make_array('%s', '%q', '%d-%m-%Y %H:%M:%S%#z', '%+'));
144150

0 commit comments

Comments
 (0)