comphead commented on code in PR #12400:
URL: https://github.com/apache/datafusion/pull/12400#discussion_r1750835005
##########
datafusion/functions-nested/src/range.rs:
##########
@@ -394,3 +412,136 @@ fn gen_range_date(args: &[ArrayRef], include_upper: bool)
-> Result<ArrayRef> {
Ok(arr)
}
+
+fn gen_range_timestamp(args: &[ArrayRef], include_upper: bool) ->
Result<ArrayRef> {
+ if args.len() != 3 {
+ return exec_err!(
+ "arguments length must be 3 for {}",
+ if include_upper {
+ "generate_series"
+ } else {
+ "range"
+ }
+ );
+ }
+
+ // coerce_types fn should coerce all types to Timestamp(Nanosecond, tz)
+ let (start_arr, start_tz_opt) = cast_timestamp_arg(&args[0],
include_upper)?;
+ let (stop_arr, stop_tz_opt) = cast_timestamp_arg(&args[1], include_upper)?;
+ let step_arr = as_interval_mdn_array(&args[2])?;
+ let start_tz = parse_tz(start_tz_opt)?;
+ let stop_tz = parse_tz(stop_tz_opt)?;
+
+ // values are timestamps
+ let values_builder = if let Some(start_tz_str) = start_tz_opt {
+ TimestampNanosecondBuilder::new().with_timezone(&**start_tz_str)
+ } else {
+ TimestampNanosecondBuilder::new()
+ };
+ let mut list_builder = ListBuilder::new(values_builder);
+
+ for idx in 0..start_arr.len() {
+ if start_arr.is_null(idx) || stop_arr.is_null(idx) ||
step_arr.is_null(idx) {
+ list_builder.append_null();
+ continue;
+ }
+
+ let start = start_arr.value(idx);
+ let stop = stop_arr.value(idx);
+ let step = step_arr.value(idx);
+
+ let (months, days, ns) = IntervalMonthDayNanoType::to_parts(step);
+ if months == 0 && days == 0 && ns == 0 {
+ return exec_err!(
+ "Interval argument to {} must not be 0",
+ if include_upper {
+ "generate_series"
+ } else {
+ "range"
+ }
+ );
+ }
+
+ let neg = TSNT::add_month_day_nano(start, step, start_tz)
+ .ok_or(DataFusionError::Execution(
+ "Cannot generate timestamp range where start + step overflows"
+ .to_string(),
+ ))?
+ .cmp(&start)
+ == Ordering::Less;
+
+ let stop_dt = as_datetime_with_timezone::<TSNT>(stop, stop_tz).ok_or(
+ DataFusionError::Execution(format!(
Review Comment:
exec_datafusion_error!(
##########
datafusion/functions-nested/src/range.rs:
##########
@@ -394,3 +412,136 @@ fn gen_range_date(args: &[ArrayRef], include_upper: bool)
-> Result<ArrayRef> {
Ok(arr)
}
+
+fn gen_range_timestamp(args: &[ArrayRef], include_upper: bool) ->
Result<ArrayRef> {
+ if args.len() != 3 {
+ return exec_err!(
+ "arguments length must be 3 for {}",
+ if include_upper {
+ "generate_series"
+ } else {
+ "range"
+ }
+ );
+ }
+
+ // coerce_types fn should coerce all types to Timestamp(Nanosecond, tz)
+ let (start_arr, start_tz_opt) = cast_timestamp_arg(&args[0],
include_upper)?;
+ let (stop_arr, stop_tz_opt) = cast_timestamp_arg(&args[1], include_upper)?;
+ let step_arr = as_interval_mdn_array(&args[2])?;
+ let start_tz = parse_tz(start_tz_opt)?;
+ let stop_tz = parse_tz(stop_tz_opt)?;
+
+ // values are timestamps
+ let values_builder = if let Some(start_tz_str) = start_tz_opt {
+ TimestampNanosecondBuilder::new().with_timezone(&**start_tz_str)
+ } else {
+ TimestampNanosecondBuilder::new()
+ };
+ let mut list_builder = ListBuilder::new(values_builder);
+
+ for idx in 0..start_arr.len() {
+ if start_arr.is_null(idx) || stop_arr.is_null(idx) ||
step_arr.is_null(idx) {
+ list_builder.append_null();
+ continue;
+ }
+
+ let start = start_arr.value(idx);
+ let stop = stop_arr.value(idx);
+ let step = step_arr.value(idx);
+
+ let (months, days, ns) = IntervalMonthDayNanoType::to_parts(step);
+ if months == 0 && days == 0 && ns == 0 {
+ return exec_err!(
+ "Interval argument to {} must not be 0",
+ if include_upper {
+ "generate_series"
+ } else {
+ "range"
+ }
+ );
+ }
+
+ let neg = TSNT::add_month_day_nano(start, step, start_tz)
+ .ok_or(DataFusionError::Execution(
+ "Cannot generate timestamp range where start + step overflows"
+ .to_string(),
+ ))?
+ .cmp(&start)
+ == Ordering::Less;
+
+ let stop_dt = as_datetime_with_timezone::<TSNT>(stop, stop_tz).ok_or(
+ DataFusionError::Execution(format!(
+ "Cannot generate timestamp for stop: {}: {:?}",
+ stop, stop_tz
+ )),
+ )?;
+
+ let mut current = start;
+ let mut current_dt = as_datetime_with_timezone::<TSNT>(current,
start_tz).ok_or(
+ DataFusionError::Execution(format!(
+ "Cannot generate timestamp for start: {}: {:?}",
Review Comment:
exec_datafusion_error!(
--
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]
