Omega359 commented on code in PR #12400:
URL: https://github.com/apache/datafusion/pull/12400#discussion_r1754801798
##########
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_bound: bool) ->
Result<ArrayRef> {
+ if args.len() != 3 {
+ return exec_err!(
+ "Arguments length must be 3 for {}",
+ if include_upper_bound {
+ "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_bound)?;
+ let (stop_arr, stop_tz_opt) = cast_timestamp_arg(&args[1],
include_upper_bound)?;
+ 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 = start_tz_opt
+ .clone()
+ .map_or_else(TimestampNanosecondBuilder::new, |start_tz_str| {
+ TimestampNanosecondBuilder::new().with_timezone(start_tz_str)
+ });
+ 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_bound {
+ "GENERATE_SERIES"
+ } else {
+ "RANGE"
+ }
+ );
+ }
+
+ let neg = TSNT::add_month_day_nano(start, step, start_tz)
+ .ok_or(exec_datafusion_err!(
+ "Cannot generate timestamp range where start + step overflows"
+ ))?
+ .cmp(&start)
+ == Ordering::Less;
+
+ let stop_dt = as_datetime_with_timezone::<TSNT>(stop, stop_tz).ok_or(
+ exec_datafusion_err!(
+ "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(
+ exec_datafusion_err!(
+ "Cannot generate timestamp for start: {}: {:?}",
+ current,
+ start_tz
+ ),
+ )?;
+
+ let values = from_fn(|| {
+ if (include_upper_bound
+ && ((neg && current_dt < stop_dt) || (!neg && current_dt >
stop_dt)))
+ || (!include_upper_bound
+ && ((neg && current_dt <= stop_dt)
+ || (!neg && current_dt >= stop_dt)))
+ {
+ return None;
+ }
+
+ let prev_current = current;
+
+ if let Some(ts) = TSNT::add_month_day_nano(current, step,
start_tz) {
+ current = ts;
+ current_dt = as_datetime_with_timezone::<TSNT>(current,
start_tz)?;
+
+ Some(Some(prev_current))
+ } else {
+ // we failed to parse the timestamp here so terminate the
series
+ None
+ }
+ });
+
+ list_builder.append_value(values);
+ }
+
+ let arr = Arc::new(list_builder.finish());
+
+ Ok(arr)
+}
+
+fn cast_timestamp_arg(
Review Comment:
My programming roots are showing through :) I can take a look at that
--
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]
