Merge cockroachdb#103131 cockroachdb#103181

103131: util/parquet: add crdb metadata to parquet files r=miretskiy a=jayshrivastava

util/parquet: add option to write kv metadata to files

This change adds an option to the writer which allows the caller
to write arbitrary kv metadata to parquet files. This is useful
for testing purposes.

Informs: cockroachdb#99028
Epic: https://cockroachlabs.atlassian.net/browse/CRDB-15071
Release note: None

---

util/parquet: remove dependency on writer to read parquet files

Previously, the test utils used to read parquet files
would require the writer as an argument. The main reason
the writer was required is that the writer contained
crdb-specific type information which could be used to
decode raw data until crdb datums.

With this change, the writer is updated to write this
crdb-specific type information to the parquet file in
its metadata. The reader is updated to the read type
information from the file metadata. There is a new
test utility function `ReadFile(parquetFile string)`
which can be used to read all datums from a parquet
file without providing any additional type information.
The function also returns metadata since it is possible
for users of the `Writer` to write arbitrary metadata
and such users may need this metadata in testing.

Informs: cockroachdb#99028
Epic: https://cockroachlabs.atlassian.net/browse/CRDB-15071
Release note: None

103181: server: Default job wait time to 10s r=miretskiy a=miretskiy

This is a revert of cockroachdb#103134 to reset job wait period to 10s Epic: None

Release note: None

Co-authored-by: Jayant Shrivastava <[email protected]>
Co-authored-by: Yevgeniy Miretskiy <[email protected]>

docs/generated/settings/settings-for-tenants.txt

@@ -74,7 +74,7 @@ server.oidc_authentication.scopes	string	openid	sets OIDC scopes to include with
 server.rangelog.ttl	duration	720h0m0s	if nonzero, entries in system.rangelog older than this duration are periodically purged	tenant-rw
 server.shutdown.connection_wait	duration	0s	the maximum amount of time a server waits for all SQL connections to be closed before proceeding with a drain. (note that the --drain-wait parameter for cockroach node drain may need adjustment after changing this setting)	tenant-rw
 server.shutdown.drain_wait	duration	0s	the amount of time a server waits in an unready state before proceeding with a drain (note that the --drain-wait parameter for cockroach node drain may need adjustment after changing this setting. --drain-wait is to specify the duration of the whole draining process, while server.shutdown.drain_wait is to set the wait time for health probes to notice that the node is not ready.)	tenant-rw
-server.shutdown.jobs_wait	duration	0s	the maximum amount of time a server waits for all currently executing jobs to notice drain request and to perform orderly shutdown	tenant-rw
+server.shutdown.jobs_wait	duration	10s	the maximum amount of time a server waits for all currently executing jobs to notice drain request and to perform orderly shutdown	tenant-rw
 server.shutdown.query_wait	duration	10s	the timeout for waiting for active queries to finish during a drain (note that the --drain-wait parameter for cockroach node drain may need adjustment after changing this setting)	tenant-rw
 server.time_until_store_dead	duration	5m0s	the time after which if there is no new gossiped information about a store, it is considered dead	tenant-rw
 server.user_login.cert_password_method.auto_scram_promotion.enabled	boolean	true	whether to automatically promote cert-password authentication to use SCRAM	tenant-rw

docs/generated/settings/settings.html

@@ -104,7 +104,7 @@
 <tr><td><div id="setting-server-secondary-tenants-redact-trace-enabled" class="anchored"><code>server.secondary_tenants.redact_trace.enabled</code></div></td><td>boolean</td><td><code>true</code></td><td>controls if server side traces are redacted for tenant operations</td><td>Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-server-shutdown-connection-wait" class="anchored"><code>server.shutdown.connection_wait</code></div></td><td>duration</td><td><code>0s</code></td><td>the maximum amount of time a server waits for all SQL connections to be closed before proceeding with a drain. (note that the --drain-wait parameter for cockroach node drain may need adjustment after changing this setting)</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-server-shutdown-drain-wait" class="anchored"><code>server.shutdown.drain_wait</code></div></td><td>duration</td><td><code>0s</code></td><td>the amount of time a server waits in an unready state before proceeding with a drain (note that the --drain-wait parameter for cockroach node drain may need adjustment after changing this setting. --drain-wait is to specify the duration of the whole draining process, while server.shutdown.drain_wait is to set the wait time for health probes to notice that the node is not ready.)</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
-<tr><td><div id="setting-server-shutdown-jobs-wait" class="anchored"><code>server.shutdown.jobs_wait</code></div></td><td>duration</td><td><code>0s</code></td><td>the maximum amount of time a server waits for all currently executing jobs to notice drain request and to perform orderly shutdown</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
+<tr><td><div id="setting-server-shutdown-jobs-wait" class="anchored"><code>server.shutdown.jobs_wait</code></div></td><td>duration</td><td><code>10s</code></td><td>the maximum amount of time a server waits for all currently executing jobs to notice drain request and to perform orderly shutdown</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-server-shutdown-lease-transfer-wait" class="anchored"><code>server.shutdown.lease_transfer_wait</code></div></td><td>duration</td><td><code>5s</code></td><td>the timeout for a single iteration of the range lease transfer phase of draining (note that the --drain-wait parameter for cockroach node drain may need adjustment after changing this setting)</td><td>Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-server-shutdown-query-wait" class="anchored"><code>server.shutdown.query_wait</code></div></td><td>duration</td><td><code>10s</code></td><td>the timeout for waiting for active queries to finish during a drain (note that the --drain-wait parameter for cockroach node drain may need adjustment after changing this setting)</td><td>Serverless/Dedicated/Self-Hosted</td></tr>
 <tr><td><div id="setting-server-time-until-store-dead" class="anchored"><code>server.time_until_store_dead</code></div></td><td>duration</td><td><code>5m0s</code></td><td>the time after which if there is no new gossiped information about a store, it is considered dead</td><td>Serverless/Dedicated/Self-Hosted</td></tr>

pkg/ccl/changefeedccl/helpers_test.go

@@ -1278,8 +1278,7 @@ func waitForJobStatus(
 }
 
 // TestingSetIncludeParquetMetadata adds the option to turn on adding metadata
-// (primary key column names) to the parquet file which is used to convert parquet
-// data to JSON format
+// to the parquet file which is used in testing.
 func TestingSetIncludeParquetMetadata() func() {
 	includeParquetTestMetadata = true
 	return func() {

pkg/ccl/changefeedccl/parquet.go

@@ -25,7 +25,7 @@ type parquetWriter struct {
 // newParquetWriterFromRow constructs a new parquet writer which outputs to
 // the given sink. This function interprets the schema from the supplied row.
 func newParquetWriterFromRow(
-	row cdcevent.Row, sink io.Writer, maxRowGroupSize int64,
+	row cdcevent.Row, sink io.Writer, opts ...parquet.Option,
 ) (*parquetWriter, error) {
 	columnNames := make([]string, len(row.ResultColumns())+1)
 	columnTypes := make([]*types.T, len(row.ResultColumns())+1)
@@ -48,7 +48,12 @@ func newParquetWriterFromRow(
 		return nil, err
 	}
 
-	writer, err := parquet.NewWriter(schemaDef, sink, parquet.WithMaxRowGroupLength(maxRowGroupSize))
+	writerConstructor := parquet.NewWriter
+	if includeParquetTestMetadata {
+		writerConstructor = parquet.NewWriterWithReaderMeta
+	}
+
+	writer, err := writerConstructor(schemaDef, sink, opts...)
 	if err != nil {
 		return nil, err
 	}

pkg/ccl/changefeedccl/parquet_test.go

@@ -38,6 +38,8 @@ func TestParquetRows(t *testing.T) {
 	// Rangefeed reader can time out under stress.
 	skip.UnderStress(t)
 
+	defer TestingSetIncludeParquetMetadata()()
+
 	ctx := context.Background()
 	s, db, _ := serverutils.StartServer(t, base.TestServerArgs{
 		// TODO(#98816): cdctest.GetHydratedTableDescriptor does not work with tenant dbs.
@@ -110,7 +112,7 @@ func TestParquetRows(t *testing.T) {
 				require.NoError(t, err)
 
 				if writer == nil {
-					writer, err = newParquetWriterFromRow(updatedRow, f, maxRowGroupSize)
+					writer, err = newParquetWriterFromRow(updatedRow, f, parquet.WithMaxRowGroupLength(maxRowGroupSize))
 					if err != nil {
 						t.Fatalf(err.Error())
 					}

pkg/ccl/serverccl/server_controller_test.go

@@ -499,6 +499,12 @@ func TestServerStartStop(t *testing.T) {
 		if err := db2.Ping(); err != nil {
 			return err
 		}
+
+		// Don't wait for graceful jobs shutdown in this test since
+		// we want to make sure test completes reasonably quickly.
+		_, err = db2.Exec("SET CLUSTER SETTING server.shutdown.jobs_wait='0s'")
+		require.NoError(t, err)
+
 		return nil
 	})
 

pkg/server/drain.go

@@ -72,7 +72,7 @@ var (
 		"server.shutdown.jobs_wait",
 		"the maximum amount of time a server waits for all currently executing jobs "+
 			"to notice drain request and to perform orderly shutdown",
-		0*time.Second,
+		10*time.Second,
 		settings.NonNegativeDurationWithMaximum(10*time.Hour),
 	).WithPublic()
 )

pkg/util/parquet/BUILD.bazel

@@ -29,10 +29,10 @@ go_library(
         "@com_github_apache_arrow_go_v11//parquet",
         "@com_github_apache_arrow_go_v11//parquet/compress",
         "@com_github_apache_arrow_go_v11//parquet/file",
+        "@com_github_apache_arrow_go_v11//parquet/metadata",
         "@com_github_apache_arrow_go_v11//parquet/schema",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_lib_pq//oid",
-        "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/util/parquet/decoders.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/geo"
 	"github.com/cockroachdb/cockroach/pkg/geo/geopb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/bitarray"
 	"github.com/cockroachdb/cockroach/pkg/util/duration"
 	"github.com/cockroachdb/cockroach/pkg/util/timeofday"
@@ -236,6 +237,76 @@ func (collatedStringDecoder) decode(v parquet.ByteArray) (tree.Datum, error) {
 	return &tree.DCollatedString{Contents: string(v)}, nil
 }
 
+// decoderFromFamilyAndType returns the decoder to use based on the type oid and
+// family. Note the logical similarity to makeColumn in schema.go. This is
+// intentional as each decoder returned by this function corresponds to a
+// particular colWriter determined by makeColumn.
+// TODO: refactor to remove the code duplication with makeColumn
+func decoderFromFamilyAndType(typOid oid.Oid, family types.Family) (decoder, error) {
+	switch family {
+	case types.BoolFamily:
+		return boolDecoder{}, nil
+	case types.StringFamily:
+		return stringDecoder{}, nil
+	case types.IntFamily:
+		typ, ok := types.OidToType[typOid]
+		if !ok {
+			return nil, errors.AssertionFailedf("could not determine type from oid %d", typOid)
+		}
+		if typ.Oid() == oid.T_int8 {
+			return int64Decoder{}, nil
+		}
+		return int32Decoder{}, nil
+	case types.DecimalFamily:
+		return decimalDecoder{}, nil
+	case types.TimestampFamily:
+		return timestampDecoder{}, nil
+	case types.TimestampTZFamily:
+		return timestampTZDecoder{}, nil
+	case types.UuidFamily:
+		return uUIDDecoder{}, nil
+	case types.INetFamily:
+		return iNetDecoder{}, nil
+	case types.JsonFamily:
+		return jsonDecoder{}, nil
+	case types.BitFamily:
+		return bitDecoder{}, nil
+	case types.BytesFamily:
+		return bytesDecoder{}, nil
+	case types.EnumFamily:
+		return enumDecoder{}, nil
+	case types.DateFamily:
+		return dateDecoder{}, nil
+	case types.Box2DFamily:
+		return box2DDecoder{}, nil
+	case types.GeographyFamily:
+		return geographyDecoder{}, nil
+	case types.GeometryFamily:
+		return geometryDecoder{}, nil
+	case types.IntervalFamily:
+		return intervalDecoder{}, nil
+	case types.TimeFamily:
+		return timeDecoder{}, nil
+	case types.TimeTZFamily:
+		return timeTZDecoder{}, nil
+	case types.FloatFamily:
+		typ, ok := types.OidToType[typOid]
+		if !ok {
+			return nil, errors.AssertionFailedf("could not determine type from oid %d", typOid)
+		}
+		if typ.Oid() == oid.T_float4 {
+			return float32Decoder{}, nil
+		}
+		return float64Decoder{}, nil
+	case types.OidFamily:
+		return oidDecoder{}, nil
+	case types.CollatedStringFamily:
+		return collatedStringDecoder{}, nil
+	default:
+		return nil, errors.AssertionFailedf("could not find decoder for type oid %d and family %d", typOid, family)
+	}
+}
+
 // Defeat the linter's unused lint errors.
 func init() {
 	var _, _ = boolDecoder{}.decode(false)
@@ -253,7 +324,6 @@ func init() {
 	var _, _ = enumDecoder{}.decode(parquet.ByteArray{})
 	var _, _ = dateDecoder{}.decode(parquet.ByteArray{})
 	var _, _ = box2DDecoder{}.decode(parquet.ByteArray{})
-	var _, _ = box2DDecoder{}.decode(parquet.ByteArray{})
 	var _, _ = geographyDecoder{}.decode(parquet.ByteArray{})
 	var _, _ = geometryDecoder{}.decode(parquet.ByteArray{})
 	var _, _ = intervalDecoder{}.decode(parquet.ByteArray{})

pkg/util/parquet/schema.go

@@ -42,7 +42,6 @@ const defaultTypeLength = -1
 type column struct {
 	node      schema.Node
 	colWriter colWriter
-	decoder   decoder
 	typ       *types.T
 }
 
@@ -99,7 +98,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 	case types.BoolFamily:
 		result.node = schema.NewBooleanNode(colName, repetitions, defaultSchemaFieldID)
 		result.colWriter = scalarWriter(writeBool)
-		result.decoder = boolDecoder{}
 		result.typ = types.Bool
 		return result, nil
 	case types.StringFamily:
@@ -111,7 +109,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeString)
-		result.decoder = stringDecoder{}
 		return result, nil
 	case types.IntFamily:
 		// Note: integer datums are always signed: https://www.cockroachlabs.com/docs/stable/int.html
@@ -124,13 +121,11 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 				return result, err
 			}
 			result.colWriter = scalarWriter(writeInt64)
-			result.decoder = int64Decoder{}
 			return result, nil
 		}
 
 		result.node = schema.NewInt32Node(colName, repetitions, defaultSchemaFieldID)
 		result.colWriter = scalarWriter(writeInt32)
-		result.decoder = int32Decoder{}
 		return result, nil
 	case types.DecimalFamily:
 		// According to PostgresSQL docs, scale or precision of 0 implies max
@@ -155,7 +150,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeDecimal)
-		result.decoder = decimalDecoder{}
 		return result, nil
 	case types.UuidFamily:
 		result.node, err = schema.NewPrimitiveNodeLogical(colName,
@@ -165,7 +159,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeUUID)
-		result.decoder = uUIDDecoder{}
 		return result, nil
 	case types.TimestampFamily:
 		// We do not use schema.TimestampLogicalType because the library will enforce
@@ -177,7 +170,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeTimestamp)
-		result.decoder = timestampDecoder{}
 		return result, nil
 	case types.TimestampTZFamily:
 		// We do not use schema.TimestampLogicalType because the library will enforce
@@ -189,7 +181,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeTimestampTZ)
-		result.decoder = timestampTZDecoder{}
 		return result, nil
 	case types.INetFamily:
 		result.node, err = schema.NewPrimitiveNodeLogical(colName,
@@ -199,7 +190,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeINet)
-		result.decoder = iNetDecoder{}
 		return result, nil
 	case types.JsonFamily:
 		result.node, err = schema.NewPrimitiveNodeLogical(colName,
@@ -209,7 +199,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeJSON)
-		result.decoder = jsonDecoder{}
 		return result, nil
 	case types.BitFamily:
 		result.node, err = schema.NewPrimitiveNode(colName,
@@ -219,7 +208,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeBit)
-		result.decoder = bitDecoder{}
 		return result, nil
 	case types.BytesFamily:
 		result.node, err = schema.NewPrimitiveNode(colName,
@@ -229,7 +217,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeBytes)
-		result.decoder = bytesDecoder{}
 		return result, nil
 	case types.EnumFamily:
 		result.node, err = schema.NewPrimitiveNodeLogical(colName,
@@ -239,7 +226,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeEnum)
-		result.decoder = enumDecoder{}
 		return result, nil
 	case types.DateFamily:
 		// We do not use schema.DateLogicalType because the library will enforce
@@ -251,7 +237,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeDate)
-		result.decoder = dateDecoder{}
 		return result, nil
 	case types.Box2DFamily:
 		result.node, err = schema.NewPrimitiveNodeLogical(colName,
@@ -261,7 +246,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeBox2D)
-		result.decoder = box2DDecoder{}
 		return result, nil
 	case types.GeographyFamily:
 		result.node, err = schema.NewPrimitiveNode(colName,
@@ -271,7 +255,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeGeography)
-		result.decoder = geographyDecoder{}
 		return result, nil
 	case types.GeometryFamily:
 		result.node, err = schema.NewPrimitiveNode(colName,
@@ -281,7 +264,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeGeometry)
-		result.decoder = geometryDecoder{}
 		return result, nil
 	case types.IntervalFamily:
 		result.node, err = schema.NewPrimitiveNodeLogical(colName,
@@ -291,7 +273,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeInterval)
-		result.decoder = intervalDecoder{}
 		return result, nil
 	case types.TimeFamily:
 		// CRDB stores time datums in microseconds, adjusted to UTC.
@@ -303,7 +284,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeTime)
-		result.decoder = timeDecoder{}
 		return result, nil
 	case types.TimeTZFamily:
 		// We cannot use the schema.NewTimeLogicalType because it does not support
@@ -315,7 +295,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeTimeTZ)
-		result.decoder = timeTZDecoder{}
 		return result, nil
 	case types.FloatFamily:
 		if typ.Oid() == oid.T_float4 {
@@ -326,7 +305,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 				return result, err
 			}
 			result.colWriter = scalarWriter(writeFloat32)
-			result.decoder = float32Decoder{}
 			return result, nil
 		}
 		result.node, err = schema.NewPrimitiveNode(colName,
@@ -336,12 +314,10 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeFloat64)
-		result.decoder = float64Decoder{}
 		return result, nil
 	case types.OidFamily:
 		result.node = schema.NewInt32Node(colName, repetitions, defaultSchemaFieldID)
 		result.colWriter = scalarWriter(writeOid)
-		result.decoder = oidDecoder{}
 		return result, nil
 	case types.CollatedStringFamily:
 		result.node, err = schema.NewPrimitiveNodeLogical(colName,
@@ -351,7 +327,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 			return result, err
 		}
 		result.colWriter = scalarWriter(writeCollatedString)
-		result.decoder = collatedStringDecoder{}
 		return result, nil
 	case types.ArrayFamily:
 		// Arrays for type T are represented by the following:
@@ -383,7 +358,6 @@ func makeColumn(colName string, typ *types.T, repetitions parquet.Repetition) (c
 		if err != nil {
 			return result, err
 		}
-		result.decoder = elementCol.decoder
 		scalarColWriter, ok := elementCol.colWriter.(scalarWriter)
 		if !ok {
 			return result, errors.AssertionFailedf("expected scalar column writer")