velvia commented on a change in pull request #9970:
URL: https://github.com/apache/arrow/pull/9970#discussion_r611955641
##########
File path: rust/datafusion/src/execution/context.rs
##########
@@ -1403,6 +1403,121 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn aggregate_timestamps_sum() -> Result<()> {
+ let tmp_dir = TempDir::new()?;
+ let mut ctx = create_ctx(&tmp_dir, 1)?;
+ ctx.register_table("t", test::table_with_timestamps())
+ .unwrap();
+
+ let results = plan_and_collect(
+ &mut ctx,
+ "SELECT sum(nanos), sum(micros), sum(millis), sum(secs) FROM t",
+ )
+ .await
+ .unwrap_err();
+
+ assert_eq!(results.to_string(), "Error during planning: Coercion from
[Timestamp(Nanosecond, None)] to the signature Uniform(1, [Int8, Int16, Int32,
Int64, UInt8, UInt16, UInt32, UInt64, Float32, Float64]) failed.");
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn aggregate_timestamps_count() -> Result<()> {
+ let tmp_dir = TempDir::new()?;
+ let mut ctx = create_ctx(&tmp_dir, 1)?;
+ ctx.register_table("t", test::table_with_timestamps())
+ .unwrap();
+
+ let results = plan_and_collect(
+ &mut ctx,
+ "SELECT count(nanos), count(micros), count(millis), count(secs)
FROM t",
+ )
+ .await
+ .unwrap();
+
+ let expected = vec![
Review comment:
Testing string formatting seems kinda yucky and prone to fail when the
formatting changes. Is there a better way to test the output here?
##########
File path: rust/datafusion/src/test/mod.rs
##########
@@ -182,6 +185,93 @@ pub fn make_partition(sz: i32) -> RecordBatch {
RecordBatch::try_new(schema, vec![arr]).unwrap()
}
+/// Return a new table provider containing all of the supported timestamp types
+pub fn table_with_timestamps() -> Arc<dyn TableProvider> {
+ let batch = make_timestamps();
+ let schema = batch.schema();
+ let partitions = vec![vec![batch]];
+ Arc::new(MemTable::try_new(schema, partitions).unwrap())
+}
+
+/// Return record batch with all of the supported timestamp types
+/// values
+///
+/// Columns are named:
+/// "nanos" --> TimestampNanosecondArray
+/// "micros" --> TimestampMicrosecondArray
+/// "millis" --> TimestampMillisecondArray
+/// "secs" --> TimestampSecondArray
+/// "names" --> StringArray
+pub fn make_timestamps() -> RecordBatch {
Review comment:
👍 this will be really useful.
##########
File path: rust/datafusion/src/scalar.rs
##########
@@ -72,12 +77,14 @@ pub enum ScalarValue {
Date32(Option<i32>),
/// Date stored as a signed 64bit int
Date64(Option<i64>),
+ /// Timestamp Second
+ TimestampSecond(Option<i64>),
Review comment:
👍 Great to be consistent.
Not related to this PR necessarily but I'm going to work on support for
converting date strings to timestamps of different resolutions. Right now only
Nanos are supported, ideally we'd have `to_timestamp(...)` that allows choosing
nanos, micros, millis, or seconds under the hood.
##########
File path: rust/datafusion/src/execution/context.rs
##########
@@ -1403,6 +1403,121 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn aggregate_timestamps_sum() -> Result<()> {
+ let tmp_dir = TempDir::new()?;
+ let mut ctx = create_ctx(&tmp_dir, 1)?;
+ ctx.register_table("t", test::table_with_timestamps())
+ .unwrap();
+
+ let results = plan_and_collect(
+ &mut ctx,
+ "SELECT sum(nanos), sum(micros), sum(millis), sum(secs) FROM t",
+ )
+ .await
+ .unwrap_err();
+
+ assert_eq!(results.to_string(), "Error during planning: Coercion from
[Timestamp(Nanosecond, None)] to the signature Uniform(1, [Int8, Int16, Int32,
Int64, UInt8, UInt16, UInt32, UInt64, Float32, Float64]) failed.");
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn aggregate_timestamps_count() -> Result<()> {
+ let tmp_dir = TempDir::new()?;
+ let mut ctx = create_ctx(&tmp_dir, 1)?;
+ ctx.register_table("t", test::table_with_timestamps())
+ .unwrap();
+
+ let results = plan_and_collect(
+ &mut ctx,
+ "SELECT count(nanos), count(micros), count(millis), count(secs)
FROM t",
+ )
+ .await
+ .unwrap();
+
+ let expected = vec![
+ "+--------------+---------------+---------------+-------------+",
+ "| COUNT(nanos) | COUNT(micros) | COUNT(millis) | COUNT(secs) |",
+ "+--------------+---------------+---------------+-------------+",
+ "| 3 | 3 | 3 | 3 |",
+ "+--------------+---------------+---------------+-------------+",
+ ];
+ assert_batches_sorted_eq!(expected, &results);
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn aggregate_timestamps_min() -> Result<()> {
Review comment:
Min and max definitely makes sense.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
