Skip to content

Commit ab145c8

Browse files
alambalamb
authored
Move timestamp related tests out of context.rs and into sql integration test (#1696)
* Move some tests out of context.rs and into sql * Move support test out of context.rs and into sql tests * Fixup tests and make them compile
1 parent ed1de63 commit ab145c8

5 files changed

Lines changed: 247 additions & 258 deletions

File tree

‎datafusion/src/execution/context.rs‎

Lines changed: 0 additions & 165 deletions
Original file line numberDiff line numberDiff line change
@@ -2265,121 +2265,6 @@ mod tests {
22652265
Ok(())
22662266
}
22672267

2268-
#[tokio::test]
2269-
async fn aggregate_timestamps_sum() -> Result<()> {
2270-
let tmp_dir = TempDir::new()?;
2271-
let mut ctx = create_ctx(&tmp_dir, 1).await?;
2272-
ctx.register_table("t", test::table_with_timestamps())
2273-
.unwrap();
2274-
2275-
let results = plan_and_collect(
2276-
&mut ctx,
2277-
"SELECT sum(nanos), sum(micros), sum(millis), sum(secs) FROM t",
2278-
)
2279-
.await
2280-
.unwrap_err();
2281-
2282-
assert_eq!(results.to_string(), "Error during planning: The function Sum does not support inputs of type Timestamp(Nanosecond, None).");
2283-
2284-
Ok(())
2285-
}
2286-
2287-
#[tokio::test]
2288-
async fn aggregate_timestamps_count() -> Result<()> {
2289-
let tmp_dir = TempDir::new()?;
2290-
let mut ctx = create_ctx(&tmp_dir, 1).await?;
2291-
ctx.register_table("t", test::table_with_timestamps())
2292-
.unwrap();
2293-
2294-
let results = plan_and_collect(
2295-
&mut ctx,
2296-
"SELECT count(nanos), count(micros), count(millis), count(secs) FROM t",
2297-
)
2298-
.await
2299-
.unwrap();
2300-
2301-
let expected = vec![
2302-
"+----------------+-----------------+-----------------+---------------+",
2303-
"| COUNT(t.nanos) | COUNT(t.micros) | COUNT(t.millis) | COUNT(t.secs) |",
2304-
"+----------------+-----------------+-----------------+---------------+",
2305-
"| 3 | 3 | 3 | 3 |",
2306-
"+----------------+-----------------+-----------------+---------------+",
2307-
];
2308-
assert_batches_sorted_eq!(expected, &results);
2309-
2310-
Ok(())
2311-
}
2312-
2313-
#[tokio::test]
2314-
async fn aggregate_timestamps_min() -> Result<()> {
2315-
let tmp_dir = TempDir::new()?;
2316-
let mut ctx = create_ctx(&tmp_dir, 1).await?;
2317-
ctx.register_table("t", test::table_with_timestamps())
2318-
.unwrap();
2319-
2320-
let results = plan_and_collect(
2321-
&mut ctx,
2322-
"SELECT min(nanos), min(micros), min(millis), min(secs) FROM t",
2323-
)
2324-
.await
2325-
.unwrap();
2326-
2327-
let expected = vec![
2328-
"+----------------------------+----------------------------+-------------------------+---------------------+",
2329-
"| MIN(t.nanos) | MIN(t.micros) | MIN(t.millis) | MIN(t.secs) |",
2330-
"+----------------------------+----------------------------+-------------------------+---------------------+",
2331-
"| 2011-12-13 11:13:10.123450 | 2011-12-13 11:13:10.123450 | 2011-12-13 11:13:10.123 | 2011-12-13 11:13:10 |",
2332-
"+----------------------------+----------------------------+-------------------------+---------------------+",
2333-
];
2334-
assert_batches_sorted_eq!(expected, &results);
2335-
2336-
Ok(())
2337-
}
2338-
2339-
#[tokio::test]
2340-
async fn aggregate_timestamps_max() -> Result<()> {
2341-
let tmp_dir = TempDir::new()?;
2342-
let mut ctx = create_ctx(&tmp_dir, 1).await?;
2343-
ctx.register_table("t", test::table_with_timestamps())
2344-
.unwrap();
2345-
2346-
let results = plan_and_collect(
2347-
&mut ctx,
2348-
"SELECT max(nanos), max(micros), max(millis), max(secs) FROM t",
2349-
)
2350-
.await
2351-
.unwrap();
2352-
2353-
let expected = vec![
2354-
"+-------------------------+-------------------------+-------------------------+---------------------+",
2355-
"| MAX(t.nanos) | MAX(t.micros) | MAX(t.millis) | MAX(t.secs) |",
2356-
"+-------------------------+-------------------------+-------------------------+---------------------+",
2357-
"| 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10 |",
2358-
"+-------------------------+-------------------------+-------------------------+---------------------+",
2359-
];
2360-
assert_batches_sorted_eq!(expected, &results);
2361-
2362-
Ok(())
2363-
}
2364-
2365-
#[tokio::test]
2366-
async fn aggregate_timestamps_avg() -> Result<()> {
2367-
let tmp_dir = TempDir::new()?;
2368-
let mut ctx = create_ctx(&tmp_dir, 1).await?;
2369-
ctx.register_table("t", test::table_with_timestamps())
2370-
.unwrap();
2371-
2372-
let results = plan_and_collect(
2373-
&mut ctx,
2374-
"SELECT avg(nanos), avg(micros), avg(millis), avg(secs) FROM t",
2375-
)
2376-
.await
2377-
.unwrap_err();
2378-
2379-
assert_eq!(results.to_string(), "Error during planning: The function Avg does not support inputs of type Timestamp(Nanosecond, None).");
2380-
Ok(())
2381-
}
2382-
23832268
#[tokio::test]
23842269
async fn aggregate_avg_add() -> Result<()> {
23852270
let results = execute(
@@ -2418,56 +2303,6 @@ mod tests {
24182303
Ok(())
24192304
}
24202305

2421-
#[tokio::test]
2422-
async fn join_timestamp() -> Result<()> {
2423-
let tmp_dir = TempDir::new()?;
2424-
let mut ctx = create_ctx(&tmp_dir, 1).await?;
2425-
ctx.register_table("t", test::table_with_timestamps())
2426-
.unwrap();
2427-
2428-
let expected = vec![
2429-
"+-------------------------------+----------------------------+-------------------------+---------------------+-------+-------------------------------+----------------------------+-------------------------+---------------------+-------+",
2430-
"| nanos | micros | millis | secs | name | nanos | micros | millis | secs | name |",
2431-
"+-------------------------------+----------------------------+-------------------------+---------------------+-------+-------------------------------+----------------------------+-------------------------+---------------------+-------+",
2432-
"| 2011-12-13 11:13:10.123450 | 2011-12-13 11:13:10.123450 | 2011-12-13 11:13:10.123 | 2011-12-13 11:13:10 | Row 1 | 2011-12-13 11:13:10.123450 | 2011-12-13 11:13:10.123450 | 2011-12-13 11:13:10.123 | 2011-12-13 11:13:10 | Row 1 |",
2433-
"| 2018-11-13 17:11:10.011375885 | 2018-11-13 17:11:10.011375 | 2018-11-13 17:11:10.011 | 2018-11-13 17:11:10 | Row 0 | 2018-11-13 17:11:10.011375885 | 2018-11-13 17:11:10.011375 | 2018-11-13 17:11:10.011 | 2018-11-13 17:11:10 | Row 0 |",
2434-
"| 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10 | Row 3 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10 | Row 3 |",
2435-
"+-------------------------------+----------------------------+-------------------------+---------------------+-------+-------------------------------+----------------------------+-------------------------+---------------------+-------+",
2436-
];
2437-
2438-
let results = plan_and_collect(
2439-
&mut ctx,
2440-
"SELECT * FROM t as t1 \
2441-
JOIN (SELECT * FROM t) as t2 \
2442-
ON t1.nanos = t2.nanos",
2443-
)
2444-
.await
2445-
.unwrap();
2446-
assert_batches_sorted_eq!(expected, &results);
2447-
2448-
let results = plan_and_collect(
2449-
&mut ctx,
2450-
"SELECT * FROM t as t1 \
2451-
JOIN (SELECT * FROM t) as t2 \
2452-
ON t1.micros = t2.micros",
2453-
)
2454-
.await
2455-
.unwrap();
2456-
assert_batches_sorted_eq!(expected, &results);
2457-
2458-
let results = plan_and_collect(
2459-
&mut ctx,
2460-
"SELECT * FROM t as t1 \
2461-
JOIN (SELECT * FROM t) as t2 \
2462-
ON t1.millis = t2.millis",
2463-
)
2464-
.await
2465-
.unwrap();
2466-
assert_batches_sorted_eq!(expected, &results);
2467-
2468-
Ok(())
2469-
}
2470-
24712306
#[tokio::test]
24722307
async fn count_basic() -> Result<()> {
24732308
let results = execute("SELECT COUNT(c1), COUNT(c2) FROM test", 1).await?;

‎datafusion/src/test/mod.rs‎

Lines changed: 1 addition & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,7 @@ use crate::datasource::{MemTable, PartitionedFile, TableProvider};
2222
use crate::error::Result;
2323
use crate::from_slice::FromSlice;
2424
use crate::logical_plan::{LogicalPlan, LogicalPlanBuilder};
25-
use array::{
26-
Array, ArrayRef, StringArray, TimestampMicrosecondArray, TimestampMillisecondArray,
27-
TimestampNanosecondArray, TimestampSecondArray,
28-
};
25+
use array::{Array, ArrayRef};
2926
use arrow::array::{self, DecimalBuilder, Int32Array};
3027
use arrow::datatypes::{DataType, Field, Schema};
3128
use arrow::record_batch::RecordBatch;
@@ -185,14 +182,6 @@ pub fn make_partition(sz: i32) -> RecordBatch {
185182
RecordBatch::try_new(schema, vec![arr]).unwrap()
186183
}
187184

188-
/// Return a new table provider containing all of the supported timestamp types
189-
pub fn table_with_timestamps() -> Arc<dyn TableProvider> {
190-
let batch = make_timestamps();
191-
let schema = batch.schema();
192-
let partitions = vec![vec![batch]];
193-
Arc::new(MemTable::try_new(schema, partitions).unwrap())
194-
}
195-
196185
/// Return a new table which provide this decimal column
197186
pub fn table_with_decimal() -> Arc<dyn TableProvider> {
198187
let batch_decimal = make_decimal();
@@ -214,85 +203,6 @@ fn make_decimal() -> RecordBatch {
214203
RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap()
215204
}
216205

217-
/// Return record batch with all of the supported timestamp types
218-
/// values
219-
///
220-
/// Columns are named:
221-
/// "nanos" --> TimestampNanosecondArray
222-
/// "micros" --> TimestampMicrosecondArray
223-
/// "millis" --> TimestampMillisecondArray
224-
/// "secs" --> TimestampSecondArray
225-
/// "names" --> StringArray
226-
pub fn make_timestamps() -> RecordBatch {
227-
let ts_strings = vec![
228-
Some("2018-11-13T17:11:10.011375885995"),
229-
Some("2011-12-13T11:13:10.12345"),
230-
None,
231-
Some("2021-1-1T05:11:10.432"),
232-
];
233-
234-
let ts_nanos = ts_strings
235-
.into_iter()
236-
.map(|t| {
237-
t.map(|t| {
238-
t.parse::<chrono::NaiveDateTime>()
239-
.unwrap()
240-
.timestamp_nanos()
241-
})
242-
})
243-
.collect::<Vec<_>>();
244-
245-
let ts_micros = ts_nanos
246-
.iter()
247-
.map(|t| t.as_ref().map(|ts_nanos| ts_nanos / 1000))
248-
.collect::<Vec<_>>();
249-
250-
let ts_millis = ts_nanos
251-
.iter()
252-
.map(|t| t.as_ref().map(|ts_nanos| ts_nanos / 1000000))
253-
.collect::<Vec<_>>();
254-
255-
let ts_secs = ts_nanos
256-
.iter()
257-
.map(|t| t.as_ref().map(|ts_nanos| ts_nanos / 1000000000))
258-
.collect::<Vec<_>>();
259-
260-
let names = ts_nanos
261-
.iter()
262-
.enumerate()
263-
.map(|(i, _)| format!("Row {}", i))
264-
.collect::<Vec<_>>();
265-
266-
let arr_nanos = TimestampNanosecondArray::from_opt_vec(ts_nanos, None);
267-
let arr_micros = TimestampMicrosecondArray::from_opt_vec(ts_micros, None);
268-
let arr_millis = TimestampMillisecondArray::from_opt_vec(ts_millis, None);
269-
let arr_secs = TimestampSecondArray::from_opt_vec(ts_secs, None);
270-
271-
let names = names.iter().map(|s| s.as_str()).collect::<Vec<_>>();
272-
let arr_names = StringArray::from(names);
273-
274-
let schema = Schema::new(vec![
275-
Field::new("nanos", arr_nanos.data_type().clone(), true),
276-
Field::new("micros", arr_micros.data_type().clone(), true),
277-
Field::new("millis", arr_millis.data_type().clone(), true),
278-
Field::new("secs", arr_secs.data_type().clone(), true),
279-
Field::new("name", arr_names.data_type().clone(), true),
280-
]);
281-
let schema = Arc::new(schema);
282-
283-
RecordBatch::try_new(
284-
schema,
285-
vec![
286-
Arc::new(arr_nanos),
287-
Arc::new(arr_micros),
288-
Arc::new(arr_millis),
289-
Arc::new(arr_secs),
290-
Arc::new(arr_names),
291-
],
292-
)
293-
.unwrap()
294-
}
295-
296206
/// Asserts that given future is pending.
297207
pub fn assert_is_pending<'a, T>(fut: &mut Pin<Box<dyn Future<Output = T> + Send + 'a>>) {
298208
let waker = futures::task::noop_waker();

‎datafusion/tests/sql/aggregates.rs‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -473,3 +473,105 @@ async fn csv_query_array_agg_distinct() -> Result<()> {
473473

474474
Ok(())
475475
}
476+
477+
#[tokio::test]
478+
async fn aggregate_timestamps_sum() -> Result<()> {
479+
let mut ctx = ExecutionContext::new();
480+
ctx.register_table("t", table_with_timestamps()).unwrap();
481+
482+
let results = plan_and_collect(
483+
&mut ctx,
484+
"SELECT sum(nanos), sum(micros), sum(millis), sum(secs) FROM t",
485+
)
486+
.await
487+
.unwrap_err();
488+
489+
assert_eq!(results.to_string(), "Error during planning: The function Sum does not support inputs of type Timestamp(Nanosecond, None).");
490+
491+
Ok(())
492+
}
493+
494+
#[tokio::test]
495+
async fn aggregate_timestamps_count() -> Result<()> {
496+
let mut ctx = ExecutionContext::new();
497+
ctx.register_table("t", table_with_timestamps()).unwrap();
498+
499+
let results = execute_to_batches(
500+
&mut ctx,
501+
"SELECT count(nanos), count(micros), count(millis), count(secs) FROM t",
502+
)
503+
.await;
504+
505+
let expected = vec![
506+
"+----------------+-----------------+-----------------+---------------+",
507+
"| COUNT(t.nanos) | COUNT(t.micros) | COUNT(t.millis) | COUNT(t.secs) |",
508+
"+----------------+-----------------+-----------------+---------------+",
509+
"| 3 | 3 | 3 | 3 |",
510+
"+----------------+-----------------+-----------------+---------------+",
511+
];
512+
assert_batches_sorted_eq!(expected, &results);
513+
514+
Ok(())
515+
}
516+
517+
#[tokio::test]
518+
async fn aggregate_timestamps_min() -> Result<()> {
519+
let mut ctx = ExecutionContext::new();
520+
ctx.register_table("t", table_with_timestamps()).unwrap();
521+
522+
let results = execute_to_batches(
523+
&mut ctx,
524+
"SELECT min(nanos), min(micros), min(millis), min(secs) FROM t",
525+
)
526+
.await;
527+
528+
let expected = vec![
529+
"+----------------------------+----------------------------+-------------------------+---------------------+",
530+
"| MIN(t.nanos) | MIN(t.micros) | MIN(t.millis) | MIN(t.secs) |",
531+
"+----------------------------+----------------------------+-------------------------+---------------------+",
532+
"| 2011-12-13 11:13:10.123450 | 2011-12-13 11:13:10.123450 | 2011-12-13 11:13:10.123 | 2011-12-13 11:13:10 |",
533+
"+----------------------------+----------------------------+-------------------------+---------------------+",
534+
];
535+
assert_batches_sorted_eq!(expected, &results);
536+
537+
Ok(())
538+
}
539+
540+
#[tokio::test]
541+
async fn aggregate_timestamps_max() -> Result<()> {
542+
let mut ctx = ExecutionContext::new();
543+
ctx.register_table("t", table_with_timestamps()).unwrap();
544+
545+
let results = execute_to_batches(
546+
&mut ctx,
547+
"SELECT max(nanos), max(micros), max(millis), max(secs) FROM t",
548+
)
549+
.await;
550+
551+
let expected = vec![
552+
"+-------------------------+-------------------------+-------------------------+---------------------+",
553+
"| MAX(t.nanos) | MAX(t.micros) | MAX(t.millis) | MAX(t.secs) |",
554+
"+-------------------------+-------------------------+-------------------------+---------------------+",
555+
"| 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10.432 | 2021-01-01 05:11:10 |",
556+
"+-------------------------+-------------------------+-------------------------+---------------------+",
557+
];
558+
assert_batches_sorted_eq!(expected, &results);
559+
560+
Ok(())
561+
}
562+
563+
#[tokio::test]
564+
async fn aggregate_timestamps_avg() -> Result<()> {
565+
let mut ctx = ExecutionContext::new();
566+
ctx.register_table("t", table_with_timestamps()).unwrap();
567+
568+
let results = plan_and_collect(
569+
&mut ctx,
570+
"SELECT avg(nanos), avg(micros), avg(millis), avg(secs) FROM t",
571+
)
572+
.await
573+
.unwrap_err();
574+
575+
assert_eq!(results.to_string(), "Error during planning: The function Avg does not support inputs of type Timestamp(Nanosecond, None).");
576+
Ok(())
577+
}

0 commit comments

Comments
 (0)