alamb commented on code in PR #22114:
URL: https://github.com/apache/datafusion/pull/22114#discussion_r4286757605


##########
datafusion/common/src/scalar/mod.rs:
##########
@@ -1433,6 +1434,143 @@ impl ScalarValue {
         }
     }
 
+    /// Returns a [`ScalarValue`] of type `target_type` representing `date`.
+    ///
+    /// For `Timestamp` types the result is midnight on `date` in the type's
+    /// timezone (or UTC when the type has no timezone).
+    ///
+    /// Returns `None` if `target_type` is not a date or timestamp type, or if
+    /// midnight on `date` cannot be represented such as when the timezone 
skips
+    /// midnight that day (for example a daylight saving transition at 00:00).
+    ///
+    /// # Example
+    /// ```
+    /// # use std::sync::Arc;
+    /// # use arrow::datatypes::{DataType, TimeUnit};
+    /// # use chrono::NaiveDate;
+    /// # use datafusion_common::ScalarValue;
+    /// let date = NaiveDate::from_ymd_opt(2024, 5, 1).unwrap();
+    ///
+    /// // Dates
+    /// assert_eq!(
+    ///     ScalarValue::try_from_date(date, &DataType::Date32),
+    ///     Some(ScalarValue::Date32(Some(19844)))
+    /// );
+    ///
+    /// // Timestamps without a timezone are midnight UTC
+    /// assert_eq!(
+    ///     ScalarValue::try_from_date(date, 
&DataType::Timestamp(TimeUnit::Second, None)),
+    ///     Some(ScalarValue::TimestampSecond(Some(1714521600), None))
+    /// );
+    ///
+    /// // Timestamps with a timezone are midnight in that timezone
+    /// let tz: Arc<str> = Arc::from("+02:00");
+    /// assert_eq!(
+    ///     ScalarValue::try_from_date(
+    ///         date,
+    ///         &DataType::Timestamp(TimeUnit::Second, Some(Arc::clone(&tz)))
+    ///     ),
+    ///     Some(ScalarValue::TimestampSecond(Some(1714514400), Some(tz)))
+    /// );
+    ///
+    /// // Other types are not supported
+    /// assert_eq!(ScalarValue::try_from_date(date, &DataType::Int32), None);
+    /// ```
+    pub fn try_from_date(date: NaiveDate, target_type: &DataType) -> 
Option<Self> {
+        Some(match target_type {
+            DataType::Date32 => {
+                ScalarValue::Date32(Some(Date32Type::from_naive_date(date)))
+            }
+            DataType::Date64 => {
+                ScalarValue::Date64(Some(Date64Type::from_naive_date(date)))
+            }
+            DataType::Timestamp(unit, tz_opt) => {
+                let naive_midnight = date.and_hms_opt(0, 0, 0)?;
+                let tz: Option<Tz> = tz_opt.clone().and_then(|s| 
s.parse().ok());
+
+                // Midnight has no representation when the date falls outside
+                // the unit's range, or when the zone skips it.
+                let value = match unit {
+                    TimeUnit::Second => 
TimestampSecondType::from_naive_datetime(
+                        naive_midnight,
+                        tz.as_ref(),
+                    )?,
+                    TimeUnit::Millisecond => {
+                        TimestampMillisecondType::from_naive_datetime(
+                            naive_midnight,
+                            tz.as_ref(),
+                        )?
+                    }
+                    TimeUnit::Microsecond => {
+                        TimestampMicrosecondType::from_naive_datetime(
+                            naive_midnight,
+                            tz.as_ref(),
+                        )?
+                    }
+                    TimeUnit::Nanosecond => 
TimestampNanosecondType::from_naive_datetime(
+                        naive_midnight,
+                        tz.as_ref(),
+                    )?,
+                };
+                match unit {
+                    TimeUnit::Second => {
+                        ScalarValue::TimestampSecond(Some(value), 
tz_opt.clone())
+                    }
+                    TimeUnit::Millisecond => {
+                        ScalarValue::TimestampMillisecond(Some(value), 
tz_opt.clone())
+                    }
+                    TimeUnit::Microsecond => {
+                        ScalarValue::TimestampMicrosecond(Some(value), 
tz_opt.clone())
+                    }
+                    TimeUnit::Nanosecond => {
+                        ScalarValue::TimestampNanosecond(Some(value), 
tz_opt.clone())
+                    }
+                }
+            }
+            _ => return None,
+        })
+    }
+
+    /// Returns the raw `i64` value of a non-null timestamp with exactly the
+    /// given `unit` and `tz`.
+    ///
+    /// Returns `None` if this value is not a timestamp, is `NULL`, or has a
+    /// different unit or timezone.
+    ///
+    /// # Example
+    /// ```
+    /// # use std::sync::Arc;
+    /// # use arrow::datatypes::TimeUnit;
+    /// # use datafusion_common::ScalarValue;
+    /// let utc: Arc<str> = Arc::from("UTC");
+    /// let ts = ScalarValue::TimestampSecond(Some(1714521600), 
Some(Arc::clone(&utc)));
+    ///
+    /// // Same unit and timezone: the raw value
+    /// assert_eq!(ts.try_as_timestamp(TimeUnit::Second, Some(&utc)), 
Some(1714521600));
+    ///
+    /// // Different unit
+    /// assert_eq!(ts.try_as_timestamp(TimeUnit::Millisecond, Some(&utc)), 
None);
+    ///
+    /// // Different timezone
+    /// assert_eq!(ts.try_as_timestamp(TimeUnit::Second, None), None);
+    /// ```
+    pub fn try_as_timestamp(&self, unit: TimeUnit, tz: Option<&Arc<str>>) -> 
Option<i64> {

Review Comment:
   agreed -- this was just a refactor from an internal function. I have 
reworked it



-- 
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.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to