sql/sem/tree: fix Next() and Prev() for DTimeTZ

Prior to this commit, the DTimeTZ functions Next() and Prev()
could skip over valid values according to the ordering of DTimeTZ
values in an index (which matches the ordering defined by the
TimeTZ functions After() and Before()).

This commit fixes these functions so that Next() now returns smallest
valid DTimeTZ that is greater than the receiver, and Prev() returns
the largest valid DTimeTZ that is less than the receiver. This is
an important invariant that the optimizer relies on when building index
constraints.

Fixes cockroachdb#74912

Release note (bug fix): Fixed a bug that could occur when a TIMETZ
column was indexed, and a query predicate constrained that column using
a < or > operator with a timetz constant. If the column contained values
with time zones that did not match the time zone of the timetz constant,
it was possible that not all matching values could be returned by the
query. Specifically, the results may not have included values within one
microsecond of the predicate's absolute time. This bug was introduced
when the timetz datatype was first added in 20.1. It exists on all
versions of 20.1, 20.2, 21.1, and 21.2 prior to this patch.

pkg/sql/logictest/testdata/logic_test/timetz

@@ -670,3 +670,38 @@ SELECT a FROM regression_44774 ORDER BY a
 
 statement ok
 DROP TABLE regression_44774
+
+# Check that index scans of timetz columns do not miss values.
+subtest regression_74912
+
+# Create a table with two identical columns. One is indexed, and one is not.
+statement ok
+CREATE TABLE regression_74912 (t1 TIMETZ PRIMARY KEY, t2 TIMETZ);
+INSERT INTO regression_74912 VALUES
+  ('05:00:00.000001', '05:00:00.000001'),
+  ('07:00:00.000001+02:00:00', '07:00:00.000001+02:00:00'),
+  ('09:00:00.000001+04:00:00', '09:00:00.000001+04:00:00'),
+  ('20:59:00.000001+15:59:00', '20:59:00.000001+15:59:00');
+
+query I
+SELECT count(*) FROM regression_74912@{NO_FULL_SCAN} WHERE t1 > '05:00:00';
+----
+4
+
+query I
+SELECT count(*) FROM regression_74912@{NO_FULL_SCAN} WHERE t1 < '05:00:00.000001';
+----
+3
+
+query I
+SELECT count(*) FROM regression_74912 WHERE t2 > '05:00:00';
+----
+4
+
+query I
+SELECT count(*) FROM regression_74912 WHERE t2 < '05:00:00.000001';
+----
+3
+
+statement ok
+DROP TABLE regression_74912

pkg/sql/sem/tree/datum.go

@@ -2089,15 +2089,57 @@ func (d *DTimeTZ) Prev(ctx *EvalContext) (Datum, bool) {
 	if d.IsMin(ctx) {
 		return nil, false
 	}
-	return NewDTimeTZFromOffset(d.TimeOfDay-1, d.OffsetSecs), true
+	// In the common case, the absolute time doesn't change, we simply decrement
+	// the offset by one second and increment the time of day by one second. Once
+	// we hit the minimum offset for the current absolute time, then we decrement
+	// the absolute time by one microsecond and wrap around to the highest offset
+	// for the new absolute time. This aligns with how Before and After are
+	// defined for TimeTZ.
+	var newTimeOfDay timeofday.TimeOfDay
+	var newOffsetSecs int32
+	if d.OffsetSecs == timetz.MinTimeTZOffsetSecs ||
+		d.TimeOfDay+duration.MicrosPerSec > timeofday.Max {
+		newTimeOfDay = d.TimeOfDay - 1
+		shiftSeconds := int32((newTimeOfDay - timeofday.Min) / duration.MicrosPerSec)
+		if d.OffsetSecs+shiftSeconds > timetz.MaxTimeTZOffsetSecs {
+			shiftSeconds = timetz.MaxTimeTZOffsetSecs - d.OffsetSecs
+		}
+		newOffsetSecs = d.OffsetSecs + shiftSeconds
+		newTimeOfDay -= timeofday.TimeOfDay(shiftSeconds) * duration.MicrosPerSec
+	} else {
+		newTimeOfDay = d.TimeOfDay + duration.MicrosPerSec
+		newOffsetSecs = d.OffsetSecs - 1
+	}
+	return NewDTimeTZFromOffset(newTimeOfDay, newOffsetSecs), true
 }
 
 // Next implements the Datum interface.
 func (d *DTimeTZ) Next(ctx *EvalContext) (Datum, bool) {
 	if d.IsMax(ctx) {
 		return nil, false
 	}
-	return NewDTimeTZFromOffset(d.TimeOfDay+1, d.OffsetSecs), true
+	// In the common case, the absolute time doesn't change, we simply increment
+	// the offset by one second and decrement the time of day by one second. Once
+	// we hit the maximum offset for the current absolute time, then we increment
+	// the absolute time by one microsecond and wrap around to the lowest offset
+	// for the new absolute time. This aligns with how Before and After are
+	// defined for TimeTZ.
+	var newTimeOfDay timeofday.TimeOfDay
+	var newOffsetSecs int32
+	if d.OffsetSecs == timetz.MaxTimeTZOffsetSecs ||
+		d.TimeOfDay-duration.MicrosPerSec < timeofday.Min {
+		newTimeOfDay = d.TimeOfDay + 1
+		shiftSeconds := int32((timeofday.Max - newTimeOfDay) / duration.MicrosPerSec)
+		if d.OffsetSecs-shiftSeconds < timetz.MinTimeTZOffsetSecs {
+			shiftSeconds = d.OffsetSecs - timetz.MinTimeTZOffsetSecs
+		}
+		newOffsetSecs = d.OffsetSecs - shiftSeconds
+		newTimeOfDay += timeofday.TimeOfDay(shiftSeconds) * duration.MicrosPerSec
+	} else {
+		newTimeOfDay = d.TimeOfDay - duration.MicrosPerSec
+		newOffsetSecs = d.OffsetSecs + 1
+	}
+	return NewDTimeTZFromOffset(newTimeOfDay, newOffsetSecs), true
 }
 
 // IsMax implements the Datum interface.

pkg/sql/sem/tree/datum_integration_test.go

@@ -20,11 +20,13 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
+	"github.com/cockroachdb/cockroach/pkg/sql/randgen"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeofday"
 	"github.com/cockroachdb/cockroach/pkg/util/timetz"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -948,6 +950,166 @@ func TestDTimeTZ(t *testing.T) {
 	}
 }
 
+func checkTimeTZ(t *testing.T, d *tree.DTimeTZ) {
+	t.Helper()
+	if d.OffsetSecs < timetz.MinTimeTZOffsetSecs || d.OffsetSecs > timetz.MaxTimeTZOffsetSecs {
+		t.Fatalf("d.OffsetSecs out of range: %d", d.OffsetSecs)
+	}
+	if d.TimeOfDay < timeofday.Min || d.TimeOfDay > timeofday.Max {
+		t.Fatalf("d.TimeOfDay out of range: %d", d.TimeOfDay)
+	}
+}
+
+func TestDTimeTZPrev(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	rng, _ := randutil.NewTestRand()
+	evalCtx := &tree.EvalContext{
+		SessionDataStack: sessiondata.NewStack(&sessiondata.SessionData{
+			Location: time.UTC,
+		}),
+	}
+
+	// Check a few specific values.
+	closeToMidnight, depOnCtx, err := tree.ParseDTimeTZ(evalCtx, "23:59:59.865326-03:15:29", time.Microsecond)
+	require.NoError(t, err)
+	require.False(t, depOnCtx)
+	prev, ok := closeToMidnight.Prev(evalCtx)
+	require.True(t, ok)
+	require.Equal(t, "'11:16:28.865325-15:59:00'", prev.String())
+	prevPrev, ok := prev.Prev(evalCtx)
+	require.True(t, ok)
+	assert.Equal(t, "'11:16:29.865325-15:58:59'", prevPrev.String())
+
+	maxTime, depOnCtx, err := tree.ParseDTimeTZ(evalCtx, "24:00:00-1559", time.Microsecond)
+	require.NoError(t, err)
+	require.False(t, depOnCtx)
+	prev, ok = maxTime.Prev(evalCtx)
+	require.True(t, ok)
+	assert.Equal(t, "'23:59:59.999999-15:59:00'", prev.String())
+
+	minTime, depOnCtx, err := tree.ParseDTimeTZ(evalCtx, "00:00:00+1559", time.Microsecond)
+	require.NoError(t, err)
+	require.False(t, depOnCtx)
+	_, ok = minTime.Prev(evalCtx)
+	assert.False(t, ok)
+
+	minTimePlusOne, depOnCtx, err := tree.ParseDTimeTZ(evalCtx, "00:00:00.000001+1559", time.Microsecond)
+	require.NoError(t, err)
+	require.False(t, depOnCtx)
+	prev, ok = minTimePlusOne.Prev(evalCtx)
+	require.True(t, ok)
+	assert.Equal(t, minTime, prev)
+
+	// Choose a random start time, and run Prev for 10000 iterations.
+	startTime := randgen.RandDatum(rng, types.TimeTZ, false /* nullOk */)
+	var total int
+	for datum, ok := startTime, true; ok && total < 10000; datum, ok = datum.Prev(evalCtx) {
+		total++
+
+		// Check that the result of calling Prev is valid.
+		timeTZ, ok := datum.(*tree.DTimeTZ)
+		require.True(t, ok)
+		checkTimeTZ(t, timeTZ)
+
+		// Check that the result of calling Next on this new value is valid.
+		nextDatum, nextOk := timeTZ.Next(evalCtx)
+		if !nextOk {
+			assert.Equal(t, timeTZ, tree.DMaxTimeTZ)
+			continue
+		}
+		nextTimeTZ, ok := nextDatum.(*tree.DTimeTZ)
+		require.True(t, ok)
+		checkTimeTZ(t, nextTimeTZ)
+
+		// Check that the two datums have the expected relationship to one another.
+		assert.True(t, nextTimeTZ.After(timeTZ.TimeTZ))
+		assert.True(t, timeTZ.Before(nextTimeTZ.TimeTZ))
+		if nextTimeTZ.OffsetSecs == timetz.MinTimeTZOffsetSecs ||
+			nextTimeTZ.TimeOfDay+duration.MicrosPerSec > timeofday.Max {
+			assert.True(t, nextTimeTZ.ToTime().Sub(timeTZ.ToTime()) == time.Microsecond)
+		} else {
+			assert.True(t, nextTimeTZ.ToTime().Equal(timeTZ.ToTime()))
+		}
+	}
+}
+
+func TestDTimeTZNext(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	rng, _ := randutil.NewTestRand()
+	evalCtx := &tree.EvalContext{
+		SessionDataStack: sessiondata.NewStack(&sessiondata.SessionData{
+			Location: time.UTC,
+		}),
+	}
+
+	// Check a few specific values.
+	closeToMidnight, depOnCtx, err := tree.ParseDTimeTZ(evalCtx, "00:00:00.865326+03:15:29", time.Microsecond)
+	require.NoError(t, err)
+	require.False(t, depOnCtx)
+	next, ok := closeToMidnight.Next(evalCtx)
+	require.True(t, ok)
+	require.Equal(t, "'12:43:31.865327+15:59:00'", next.String())
+	nextNext, ok := next.Next(evalCtx)
+	require.True(t, ok)
+	assert.Equal(t, "'12:43:30.865327+15:58:59'", nextNext.String())
+
+	minTime, depOnCtx, err := tree.ParseDTimeTZ(evalCtx, "00:00:00+1559", time.Microsecond)
+	require.NoError(t, err)
+	require.False(t, depOnCtx)
+	next, ok = minTime.Next(evalCtx)
+	require.True(t, ok)
+	assert.Equal(t, "'00:00:00.000001+15:59:00'", next.String())
+
+	maxTime, depOnCtx, err := tree.ParseDTimeTZ(evalCtx, "24:00:00-1559", time.Microsecond)
+	require.NoError(t, err)
+	require.False(t, depOnCtx)
+	_, ok = maxTime.Next(evalCtx)
+	assert.False(t, ok)
+
+	maxTimeMinusOne, depOnCtx, err := tree.ParseDTimeTZ(evalCtx, "23:59:59.999999-1559", time.Microsecond)
+	require.NoError(t, err)
+	require.False(t, depOnCtx)
+	next, ok = maxTimeMinusOne.Next(evalCtx)
+	require.True(t, ok)
+	assert.Equal(t, maxTime, next)
+
+	// Choose a random start time, and run Next for 10000 iterations.
+	startTime := randgen.RandDatum(rng, types.TimeTZ, false /* nullOk */)
+	var total int
+	for datum, ok := startTime, true; ok && total < 10000; datum, ok = datum.Next(evalCtx) {
+		total++
+
+		// Check that the result of calling Next is valid.
+		timeTZ, ok := datum.(*tree.DTimeTZ)
+		require.True(t, ok)
+		checkTimeTZ(t, timeTZ)
+
+		// Check that the result of calling Prev on this new value is valid.
+		prevDatum, prevOk := timeTZ.Prev(evalCtx)
+		if !prevOk {
+			assert.Equal(t, timeTZ, tree.DMinTimeTZ)
+			continue
+		}
+		prevTimeTZ, ok := prevDatum.(*tree.DTimeTZ)
+		require.True(t, ok)
+		checkTimeTZ(t, prevTimeTZ)
+
+		// Check that the two datums have the expected relationship to one another.
+		assert.True(t, prevTimeTZ.Before(timeTZ.TimeTZ))
+		assert.True(t, timeTZ.After(prevTimeTZ.TimeTZ))
+		if prevTimeTZ.OffsetSecs == timetz.MaxTimeTZOffsetSecs ||
+			prevTimeTZ.TimeOfDay-duration.MicrosPerSec < timeofday.Min {
+			assert.True(t, timeTZ.ToTime().Sub(prevTimeTZ.ToTime()) == time.Microsecond)
+		} else {
+			assert.True(t, timeTZ.ToTime().Equal(prevTimeTZ.ToTime()))
+		}
+	}
+}
+
 func TestIsDistinctFrom(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)

pkg/util/duration/duration.go

@@ -27,8 +27,10 @@ import (
 const (
 	// MicrosPerMilli is the amount of microseconds in a millisecond.
 	MicrosPerMilli = 1000
-	// MillisPerSec is the amount of seconds in a millisecond.
+	// MillisPerSec is the amount of milliseconds in a second.
 	MillisPerSec = 1000
+	// MicrosPerSec is the amount of microseconds in a second.
+	MicrosPerSec = MicrosPerMilli * MillisPerSec
 	// SecsPerMinute is the amount of seconds in a minute.
 	SecsPerMinute = 60
 	// SecsPerHour is the amount of seconds in an hour.