Merge #37744 #37745

37744: sql: improve interval math for div and mul r=mjibson a=mjibson

In #37582 we added correct `interval * float` semantics. Here we extend that to
float division.

In addition, change int division to be a wrapper around float division so that
remainders are correctly handled. Int multiply has no need to change since it's
lossless.

Release note (sql change): Correct interval math when multiplying or dividing
by floats or ints.

37745: col{data,serde}: various fixes related to TestArrowBatchConverterRandom r=solongordon a=asubiotto

See individual commits for details

Fixes #37458
Fixes #37692

Co-authored-by: Matt Jibson <matt.jibson@gmail.com>
Co-authored-by: Alfonso Subiotto Marqués <alfonso@cockroachlabs.com>

pkg/sql/exec/coldata/nulls.go

@@ -241,12 +241,31 @@ func (n *Nulls) NullBitmap() []byte {
 	return n.nulls
 }
 
-// SetNullBitmap sets the null bitmap.
-func (n *Nulls) SetNullBitmap(bm []byte) {
+// SetNullBitmap sets the null bitmap. size corresponds to how many elements
+// this bitmap represents. The bits past the end of this size will be set to
+// valid.
+func (n *Nulls) SetNullBitmap(bm []byte, size int) {
 	n.nulls = bm
 	n.hasNulls = false
-	for _, i := range bm {
-		if i != 0 {
+	// Set all indices as valid past the last element.
+	if len(bm) > 0 && size != 0 {
+		// Set the last bits in the last element in which we want to preserve null
+		// information. mod, if non-zero, is the number of bits we don't want to
+		// overwrite (otherwise all bits are important). Note that we cast size to a
+		// uint64 to avoid extra instructions when modding.
+		mod := uint64(size) % 8
+		endIdx := size - 1
+		if mod != 0 {
+			bm[endIdx/8] |= onesMask << mod
+		}
+		// Fill the rest of the bitmap.
+		for i := (endIdx / 8) + 1; i < len(bm); {
+			i += copy(bm[i:], filledNulls[:])
+		}
+	}
+
+	for i := 0; i < len(bm); i++ {
+		if bm[i] != onesMask {
 			n.hasNulls = true
 			return
 		}

pkg/sql/exec/coldata/vec_tmpl.go

@@ -134,6 +134,16 @@ func (m *memColumn) CopyAt(src Vec, destStartIdx, srcStartIdx, srcEndIdx uint64,
 	// {{range .}}
 	case _TYPES_T:
 		copy(m._TemplateType()[destStartIdx:], src._TemplateType()[srcStartIdx:srcEndIdx])
+		// TODO(asubiotto): Improve this, there are cases where we don't need to
+		// allocate a new bitmap.
+		srcBitmap := src.Nulls().NullBitmap()
+		m.nulls.nulls = make([]byte, len(srcBitmap))
+		if src.HasNulls() {
+			m.nulls.hasNulls = true
+			copy(m.nulls.nulls, srcBitmap)
+		} else {
+			m.nulls.UnsetNulls()
+		}
 	// {{end}}
 	default:
 		panic(fmt.Sprintf("unhandled type %d", typ))

pkg/sql/exec/colserde/arrowbatchconverter.go

@@ -262,7 +262,7 @@ func (c *ArrowBatchConverter) ArrowToBatch(data []*array.Data) (coldata.Batch, e
 		}
 		arrowBitmap := arr.NullBitmapBytes()
 		if len(arrowBitmap) != 0 {
-			vec.Nulls().SetNullBitmap(arrowBitmap)
+			vec.Nulls().SetNullBitmap(arrowBitmap, n)
 		}
 	}
 	c.scratch.batch.SetLength(uint16(n))

pkg/sql/exec/colserde/arrowbatchconverter_test.go

@@ -50,15 +50,25 @@ func TestArrowBatchConverterRandom(t *testing.T) {
 	b := exec.RandomBatch(rng, typs, rng.Intn(coldata.BatchSize)+1, rng.Float64())
 	c := NewArrowBatchConverter(typs)
 
+	// Make a copy of the original batch because the converter modifies and casts
+	// data without copying for performance reasons.
+	expectedLength := b.Length()
+	expectedWidth := b.Width()
+	expectedColVecs := make([]coldata.Vec, len(b.ColVecs()))
+	for i := range typs {
+		expectedColVecs[i] = coldata.NewMemColumn(typs[i], int(b.Length()))
+		expectedColVecs[i].Copy(b.ColVec(i), 0, uint64(b.Length()), typs[i])
+	}
+
 	arrowData, err := c.BatchToArrow(b)
 	require.NoError(t, err)
 	result, err := c.ArrowToBatch(arrowData)
 	require.NoError(t, err)
 	if result.Selection() != nil {
 		t.Fatal("violated invariant that batches have no selection vectors")
 	}
-	require.Equal(t, b.Length(), result.Length())
-	require.Equal(t, b.Width(), result.Width())
+	require.Equal(t, expectedLength, result.Length())
+	require.Equal(t, expectedWidth, result.Width())
 	for i, typ := range typs {
 		// Verify equality of ColVecs (this includes nulls). Since the coldata.Vec
 		// backing array is always of coldata.BatchSize due to the scratch batch
@@ -67,7 +77,7 @@ func TestArrowBatchConverterRandom(t *testing.T) {
 		// fail.
 		require.Equal(
 			t,
-			b.ColVec(i).Slice(typ, 0, uint64(b.Length())),
+			expectedColVecs[i].Slice(typ, 0, uint64(b.Length())),
 			result.ColVec(i).Slice(typ, 0, uint64(result.Length())),
 		)
 	}

pkg/sql/logictest/testdata/logic_test/datetime

@@ -1350,3 +1350,31 @@ SELECT '-239852040018-04-28':::DATE
 
 statement error out of range
 SELECT(7133080445639580613::INT8 + '1977-11-03'::DATE) = '-239852040018-04-28':::DATE
+
+subtest interval_math
+
+query TTTTTTT
+SELECT
+    i,
+    i / 2::INT8,
+    i * 2::INT8,
+    i / 2::FLOAT8,
+    i * 2::FLOAT8,
+    i / .2362::FLOAT8,
+    i * .2362::FLOAT8
+FROM
+    (
+        VALUES
+            ('1 day'::INTERVAL),
+            ('1 month'::INTERVAL),
+            ('1 hour'::INTERVAL),
+            ('1 month 2 days 4 hours'::INTERVAL)
+    )
+        AS v (i)
+ORDER BY
+    i
+----
+01:00:00               00:30:00          02:00:00                00:30:00          02:00:00                04:14:01.320914                 00:14:10.32
+1 day                  12:00:00          2 days                  12:00:00          2 days                  4 days 05:36:31.701948          05:40:07.68
+1 mon                  15 days           2 mons                  15 days           2 mons                  4 mons 7 days 00:15:51.058425   7 days 02:03:50.4
+1 mon 2 days 04:00:00  16 days 02:00:00  2 mons 4 days 08:00:00  16 days 02:00:00  2 mons 4 days 08:00:00  4 mons 15 days 28:24:59.745978  7 days 14:20:47.04

pkg/util/duration/duration.go

@@ -480,7 +480,7 @@ func (d Duration) Mul(x int64) Duration {
 
 // Div returns a Duration representing a time length of d/x.
 func (d Duration) Div(x int64) Duration {
-	return MakeDuration(d.nanos/x, d.Days/x, d.Months/x)
+	return d.DivFloat(float64(x))
 }
 
 // MulFloat returns a Duration representing a time length of d*x.
@@ -497,10 +497,13 @@ func (d Duration) MulFloat(x float64) Duration {
 
 // DivFloat returns a Duration representing a time length of d/x.
 func (d Duration) DivFloat(x float64) Duration {
+	monthInt, monthFrac := math.Modf(float64(d.Months) / x)
+	dayInt, dayFrac := math.Modf((float64(d.Days) / x) + (monthFrac * daysInMonth))
+
 	return MakeDuration(
-		int64(float64(d.nanos)/x),
-		int64(float64(d.Days)/x),
-		int64(float64(d.Months)/x),
+		int64((float64(d.nanos)/x)+(dayFrac*float64(nanosInDay))),
+		int64(dayInt),
+		int64(monthInt),
 	)
 }
 

pkg/util/duration/duration_test.go

@@ -344,50 +344,65 @@ func TestAddMicros(t *testing.T) {
 	}
 }
 
-func TestMulFloat(t *testing.T) {
+func TestFloatMath(t *testing.T) {
 	const nanosInMinute = nanosInSecond * 60
 	const nanosInHour = nanosInMinute * 60
 
 	tests := []struct {
 		d   Duration
 		f   float64
-		res Duration
+		mul Duration
+		div Duration
 	}{
 		{
 			Duration{Months: 1, Days: 2, nanos: nanosInHour * 2},
 			0.15,
 			Duration{Days: 4, nanos: nanosInHour*19 + nanosInMinute*30},
+			Duration{Months: 6, Days: 33, nanos: nanosInHour*21 + nanosInMinute*20},
 		},
 		{
 			Duration{Months: 1, Days: 2, nanos: nanosInHour * 2},
 			0.3,
 			Duration{Days: 9, nanos: nanosInHour * 15},
+			Duration{Months: 3, Days: 16, nanos: nanosInHour*22 + nanosInMinute*40},
 		},
 		{
 			Duration{Months: 1, Days: 2, nanos: nanosInHour * 2},
 			0.5,
 			Duration{Days: 16, nanos: nanosInHour * 1},
+			Duration{Months: 2, Days: 4, nanos: nanosInHour * 4},
 		},
 		{
 			Duration{Months: 1, Days: 2, nanos: nanosInHour * 2},
 			0.8,
 			Duration{Days: 25, nanos: nanosInHour * 16},
+			Duration{Months: 1, Days: 10, nanos: nanosInHour*2 + nanosInMinute*30},
 		},
 		{
 			Duration{Months: 1, Days: 17, nanos: nanosInHour * 2},
 			2.0,
 			Duration{Months: 2, Days: 34, nanos: nanosInHour * 4},
+			Duration{Days: 23, nanos: nanosInHour * 13},
 		},
 	}
 
 	for i, test := range tests {
-		if res := test.d.MulFloat(test.f); test.res != res {
+		if res := test.d.MulFloat(test.f); test.mul != res {
 			t.Errorf(
 				"%d: expected %v.MulFloat(%f) = %v, found %v",
 				i,
 				test.d,
 				test.f,
-				test.res,
+				test.mul,
+				res)
+		}
+		if res := test.d.DivFloat(test.f); test.div != res {
+			t.Errorf(
+				"%d: expected %v.DivFloat(%f) = %v, found %v",
+				i,
+				test.d,
+				test.f,
+				test.div,
 				res)
 		}
 	}