changefeedccl: use new parquet library

pkg/ccl/changefeedccl/BUILD.bazel

@@ -96,7 +96,6 @@ go_library(
         "//pkg/sql/execinfrapb",
         "//pkg/sql/exprutil",
         "//pkg/sql/flowinfra",
-        "//pkg/sql/importer",
         "//pkg/sql/isql",
         "//pkg/sql/parser",
         "//pkg/sql/pgwire/pgcode",
@@ -154,9 +153,6 @@ go_library(
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_logtags//:logtags",
         "@com_github_cockroachdb_redact//:redact",
-        "@com_github_fraugster_parquet_go//:parquet-go",
-        "@com_github_fraugster_parquet_go//parquet",
-        "@com_github_fraugster_parquet_go//parquetschema",
         "@com_github_gogo_protobuf//jsonpb",
         "@com_github_gogo_protobuf//types",
         "@com_github_google_btree//:btree",
@@ -317,7 +313,6 @@ go_test(
         "@com_github_cockroachdb_cockroach_go_v2//crdb",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_dustin_go_humanize//:go-humanize",
-        "@com_github_fraugster_parquet_go//:parquet-go",
         "@com_github_gogo_protobuf//types",
         "@com_github_jackc_pgx_v4//:pgx",
         "@com_github_lib_pq//:pq",

pkg/ccl/changefeedccl/changefeed_test.go

@@ -7668,7 +7668,6 @@ func TestChangefeedPredicateWithSchemaChange(t *testing.T) {
 	defer log.Scope(t).Close(t)
 
 	skip.UnderRace(t, "takes too long under race")
-	defer TestingSetIncludeParquetMetadata()()
 
 	setupSQL := []string{
 		`CREATE TYPE status AS ENUM ('open', 'closed', 'inactive')`,

pkg/ccl/changefeedccl/changefeedbase/options.go

@@ -423,18 +423,31 @@ var NoLongerExperimental = map[string]string{
 	DeprecatedSinkSchemeCloudStorageS3:        SinkSchemeCloudStorageS3,
 }
 
+// OptionsSet is a test of changefeed option strings.
+type OptionsSet map[string]struct{}
+
 // InitialScanOnlyUnsupportedOptions is options that are not supported with the
 // initial scan only option
-var InitialScanOnlyUnsupportedOptions = makeStringSet(OptEndTime, OptResolvedTimestamps, OptDiff,
+var InitialScanOnlyUnsupportedOptions OptionsSet = makeStringSet(OptEndTime, OptResolvedTimestamps, OptDiff,
 	OptMVCCTimestamps, OptUpdatedTimestamps)
 
+// ParquetFormatUnsupportedOptions is options that are not supported with the
+// parquet format.
+//
+// OptKeyInValue is disallowed because parquet files have no concept of key
+// columns, so there is no reason to emit duplicate key datums.
+//
+// TODO(#103129): add support for some of these
+var ParquetFormatUnsupportedOptions OptionsSet = makeStringSet(OptEndTime, OptResolvedTimestamps, OptDiff,
+	OptMVCCTimestamps, OptUpdatedTimestamps, OptKeyInValue)
+
 // AlterChangefeedUnsupportedOptions are changefeed options that we do not allow
 // users to alter.
 // TODO(sherman): At the moment we disallow altering both the initial_scan_only
 // and the end_time option. However, there are instances in which it should be
 // allowed to alter either of these options. We need to support the alteration
 // of these fields.
-var AlterChangefeedUnsupportedOptions = makeStringSet(OptCursor, OptInitialScan,
+var AlterChangefeedUnsupportedOptions OptionsSet = makeStringSet(OptCursor, OptInitialScan,
 	OptNoInitialScan, OptInitialScanOnly, OptEndTime)
 
 // AlterChangefeedOptionExpectValues is used to parse alter changefeed options
@@ -1039,16 +1052,21 @@ func (s StatementOptions) ValidateForCreateChangefeed(isPredicateChangefeed bool
 	if err != nil {
 		return err
 	}
-	validateInitialScanUnsupportedOptions := func(errMsg string) error {
-		for o := range InitialScanOnlyUnsupportedOptions {
+
+	// validateUnsupportedOptions returns an error if any of the supplied are
+	// in the statement options. The error string should be the string
+	// representation of the option (ex. "key_in_value", or "initial_scan='only'").
+	validateUnsupportedOptions := func(unsupportedOptions OptionsSet, errorStr string) error {
+		for o := range unsupportedOptions {
 			if _, ok := s.m[o]; ok {
-				return errors.Newf(`cannot specify both %s='only' and %s`, OptInitialScan, o)
+				return errors.Newf(`cannot specify both %s and %s`, errorStr, o)
 			}
 		}
 		return nil
 	}
 	if scanType == OnlyInitialScan {
-		if err := validateInitialScanUnsupportedOptions(fmt.Sprintf("%s='only'", OptInitialScan)); err != nil {
+		if err := validateUnsupportedOptions(InitialScanOnlyUnsupportedOptions,
+			fmt.Sprintf("%s='only'", OptInitialScan)); err != nil {
 			return err
 		}
 	} else {
@@ -1058,17 +1076,8 @@ func (s StatementOptions) ValidateForCreateChangefeed(isPredicateChangefeed bool
 	}
 	// Right now parquet does not support any of these options
 	if s.m[OptFormat] == string(OptFormatParquet) {
-		if isPredicateChangefeed {
-			// Diff option is allowed when using predicate changefeeds with parquet format.
-			for o := range InitialScanOnlyUnsupportedOptions {
-				if _, ok := s.m[o]; ok && o != OptDiff {
-					return errors.Newf(`cannot specify both format='%s' and %s`, OptFormatParquet, o)
-				}
-			}
-		} else {
-			if err := validateInitialScanUnsupportedOptions(string(OptFormatParquet)); err != nil {
-				return err
-			}
+		if err := validateUnsupportedOptions(ParquetFormatUnsupportedOptions, fmt.Sprintf("format=%s", OptFormatParquet)); err != nil {
+			return err
 		}
 	}
 	for o := range s.m {

pkg/ccl/changefeedccl/changefeedbase/options_test.go

@@ -40,7 +40,6 @@ func TestOptionsValidations(t *testing.T) {
 		{map[string]string{"initial_scan_only": "", "resolved": ""}, true, "cannot specify both initial_scan='only'"},
 		{map[string]string{"initial_scan_only": "", "resolved": ""}, true, "cannot specify both initial_scan='only'"},
 		{map[string]string{"key_column": "b"}, false, "requires the unordered option"},
-		{map[string]string{"diff": "", "format": "parquet"}, true, ""},
 	}
 
 	for _, test := range tests {

pkg/ccl/changefeedccl/encoder_test.go

@@ -1068,8 +1068,6 @@ func TestParquetEncoder(t *testing.T) {
 	defer log.Scope(t).Close(t)
 
 	testFn := func(t *testing.T, s TestServer, f cdctest.TestFeedFactory) {
-		defer TestingSetIncludeParquetMetadata()()
-
 		tests := []struct {
 			name           string
 			changefeedStmt string
@@ -1089,22 +1087,22 @@ func TestParquetEncoder(t *testing.T) {
 		for _, test := range tests {
 			t.Run(test.name, func(t *testing.T) {
 				sqlDB := sqlutils.MakeSQLRunner(s.DB)
-				sqlDB.Exec(t, `CREATE TABLE foo (i INT PRIMARY KEY, x STRING, y INT, z FLOAT NOT NULL, a BOOL)`)
+				sqlDB.Exec(t, `CREATE TABLE foo (i INT PRIMARY KEY, x STRING, y INT, z FLOAT NOT NULL, a BOOL, c INT[])`)
 				defer sqlDB.Exec(t, `DROP TABLE FOO`)
-				sqlDB.Exec(t, `INSERT INTO foo VALUES (1, 'Alice', 3, 0.5032135844230652, true), (2, 'Bob',
-	2, CAST('nan' AS FLOAT),false),(3, NULL, NULL, 4.5, NULL)`)
+				sqlDB.Exec(t, `INSERT INTO foo VALUES (1, 'Alice', 3, 0.5032135844230652, true,  ARRAY[]), (2, 'Bob',
+	2, CAST('nan' AS FLOAT),false, NULL),(3, NULL, NULL, 4.5, NULL,  ARRAY[1,NULL,3])`)
 				foo := feed(t, f, test.changefeedStmt)
 				defer closeFeed(t, foo)
 
 				assertPayloads(t, foo, []string{
-					`foo: [1]->{"after": {"a": true, "i": 1, "x": "Alice", "y": 3, "z": 0.5032135844230652}}`,
-					`foo: [2]->{"after": {"a": false, "i": 2, "x": "Bob", "y": 2, "z": "NaN"}}`,
-					`foo: [3]->{"after": {"a": null, "i": 3, "x": null, "y": null, "z": 4.5}}`,
+					`foo: [1]->{"after": {"a": true, "c": [], "i": 1, "x": "Alice", "y": 3, "z": 0.5032135844230652}}`,
+					`foo: [2]->{"after": {"a": false, "c": null, "i": 2, "x": "Bob", "y": 2, "z": "NaN"}}`,
+					`foo: [3]->{"after": {"a": null, "c": [1, null, 3], "i": 3, "x": null, "y": null, "z": 4.5}}`,
 				})
 
 				sqlDB.Exec(t, `UPDATE foo SET x='wonderland' where i=1`)
 				assertPayloads(t, foo, []string{
-					`foo: [1]->{"after": {"a": true, "i": 1, "x": "wonderland", "y": 3, "z": 0.5032135844230652}}`,
+					`foo: [1]->{"after": {"a": true, "c": [], "i": 1, "x": "wonderland", "y": 3, "z": 0.5032135844230652}}`,
 				})
 
 				sqlDB.Exec(t, `DELETE from foo where i=1`)

pkg/ccl/changefeedccl/helpers_test.go

@@ -936,7 +936,6 @@ func makeFeedFactoryWithOptions(
 		userDB, cleanup := getInitialDBForEnterpriseFactory(t, s, db, options)
 		f.(*cloudFeedFactory).configureUserDB(userDB)
 		return f, func() {
-			TestingSetIncludeParquetMetadata()()
 			cleanup()
 		}
 	case "enterprise":
@@ -1081,12 +1080,20 @@ func cdcTestNamedWithSystem(
 		testLabel = fmt.Sprintf("%s/%s", sinkType, name)
 	}
 	t.Run(testLabel, func(t *testing.T) {
+		// Even if the parquet format is not being used, enable metadata
+		// in all tests for simplicity.
 		testServer, cleanupServer := makeServerWithOptions(t, options)
+		knobs := testServer.TestingKnobs.
+			DistSQL.(*execinfra.TestingKnobs).
+			Changefeed.(*TestingKnobs)
+		knobs.EnableParquetMetadata = true
+
 		feedFactory, cleanupSink := makeFeedFactoryWithOptions(t, sinkType, testServer.Server, testServer.DB, options)
 		feedFactory = maybeUseExternalConnection(feedFactory, testServer.DB, sinkType, options, t)
 		defer cleanupServer()
 		defer cleanupSink()
 		defer cleanupCloudStorage()
+
 		testFn(t, testServer, feedFactory)
 	})
 }
@@ -1277,15 +1284,6 @@ func waitForJobStatus(
 	})
 }
 
-// TestingSetIncludeParquetMetadata adds the option to turn on adding metadata
-// to the parquet file which is used in testing.
-func TestingSetIncludeParquetMetadata() func() {
-	includeParquetTestMetadata = true
-	return func() {
-		includeParquetTestMetadata = false
-	}
-}
-
 // ChangefeedJobPermissionsTestSetup creates entities and users with various permissions
 // for tests which test access control for changefeed jobs.
 //

pkg/ccl/changefeedccl/parquet.go

@@ -10,6 +10,7 @@ package changefeedccl
 
 import (
 	"io"
+	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/cdcevent"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -22,13 +23,13 @@ type parquetWriter struct {
 	datumAlloc []tree.Datum
 }
 
-// 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, opts ...parquet.Option,
-) (*parquetWriter, error) {
-	columnNames := make([]string, len(row.ResultColumns())+1)
-	columnTypes := make([]*types.T, len(row.ResultColumns())+1)
+// newParquetSchemaDefintion returns a parquet schema definition based on the
+// cdcevent.Row and the number of cols in the schema.
+func newParquetSchemaDefintion(row cdcevent.Row) (*parquet.SchemaDefinition, int, error) {
+	numCols := len(row.ResultColumns()) + 1
+
+	columnNames := make([]string, numCols)
+	columnTypes := make([]*types.T, numCols)
 
 	idx := 0
 	if err := row.ForEachColumn().Col(func(col cdcevent.ResultColumn) error {
@@ -37,27 +38,49 @@ func newParquetWriterFromRow(
 		idx += 1
 		return nil
 	}); err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 
 	columnNames[idx] = parquetCrdbEventTypeColName
 	columnTypes[idx] = types.String
 
 	schemaDef, err := parquet.NewSchema(columnNames, columnTypes)
+	if err != nil {
+		return nil, 0, err
+	}
+	return schemaDef, numCols, nil
+}
+
+// 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,
+	knobs *TestingKnobs, /* may be nil */
+	opts ...parquet.Option,
+) (*parquetWriter, error) {
+	schemaDef, numCols, err := newParquetSchemaDefintion(row)
 	if err != nil {
 		return nil, err
 	}
 
 	writerConstructor := parquet.NewWriter
-	if includeParquetTestMetadata {
+
+	if knobs.EnableParquetMetadata {
+		if opts, err = addParquetTestMetadata(row, opts); err != nil {
+			return nil, err
+		}
+
+		// To use parquet test utils for reading datums, the writer needs to be
+		// configured with additional metadata.
 		writerConstructor = parquet.NewWriterWithReaderMeta
 	}
 
 	writer, err := writerConstructor(schemaDef, sink, opts...)
 	if err != nil {
 		return nil, err
 	}
-	return &parquetWriter{inner: writer, datumAlloc: make([]tree.Datum, len(columnNames))}, nil
+	return &parquetWriter{inner: writer, datumAlloc: make([]tree.Datum, numCols)}, nil
 }
 
 // addData writes the updatedRow, adding the row's event type. There is no guarantee
@@ -70,7 +93,7 @@ func (w *parquetWriter) addData(updatedRow cdcevent.Row, prevRow cdcevent.Row) e
 	return w.inner.AddRow(w.datumAlloc)
 }
 
-// Close closes the writer and flushes any buffered data to the sink.
+// close closes the writer and flushes any buffered data to the sink.
 func (w *parquetWriter) close() error {
 	return w.inner.Close()
 }
@@ -88,3 +111,27 @@ func populateDatums(updatedRow cdcevent.Row, prevRow cdcevent.Row, datumAlloc []
 	datums = append(datums, getEventTypeDatum(updatedRow, prevRow).DString())
 	return nil
 }
+
+// addParquetTestMetadata appends options to the provided options to configure the
+// parquet writer to write metadata required by cdc test feed factories.
+func addParquetTestMetadata(row cdcevent.Row, opts []parquet.Option) ([]parquet.Option, error) {
+	keyCols := make([]string, 0)
+	if err := row.ForEachKeyColumn().Col(func(col cdcevent.ResultColumn) error {
+		keyCols = append(keyCols, col.Name)
+		return nil
+	}); err != nil {
+		return opts, err
+	}
+	opts = append(opts, parquet.WithMetadata(map[string]string{"keyCols": strings.Join(keyCols, ",")}))
+
+	allCols := make([]string, 0)
+	if err := row.ForEachColumn().Col(func(col cdcevent.ResultColumn) error {
+		allCols = append(allCols, col.Name)
+		return nil
+	}); err != nil {
+		return opts, err
+	}
+	allCols = append(allCols, parquetCrdbEventTypeColName)
+	opts = append(opts, parquet.WithMetadata(map[string]string{"allCols": strings.Join(allCols, ",")}))
+	return opts, nil
+}