Skip to content

Commit 3d26091

Browse files
alvarowolfxalvarowolfx
authored
feat(bigquery): add reference file schema option for federated formats (#6693)
Resolves internal b/246809553
1 parent 17cceeb commit 3d26091

File tree

4 files changed

+123
-0
lines changed

4 files changed

+123
-0
lines changed

‎bigquery/external.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,10 @@ type ExternalDataConfig struct {
109109
// Connections are managed through the BigQuery Connection API:
110110
// https://pkg.go.dev/cloud.google.com/go/bigquery/connection/apiv1
111111
ConnectionID string
112+
113+
// When creating an external table, the user can provide a reference file with the table schema.
114+
// This is enabled for the following formats: AVRO, PARQUET, ORC.
115+
ReferenceFileSchemaURI string
112116
}
113117

114118
func (e *ExternalDataConfig) toBQ() bq.ExternalDataConfiguration {
@@ -121,6 +125,7 @@ func (e *ExternalDataConfig) toBQ() bq.ExternalDataConfiguration {
121125
MaxBadRecords: e.MaxBadRecords,
122126
HivePartitioningOptions: e.HivePartitioningOptions.toBQ(),
123127
ConnectionId: e.ConnectionID,
128+
ReferenceFileSchemaUri: e.ReferenceFileSchemaURI,
124129
}
125130
if e.Schema != nil {
126131
q.Schema = e.Schema.toBQ()
@@ -145,6 +150,7 @@ func bqToExternalDataConfig(q *bq.ExternalDataConfiguration) (*ExternalDataConfi
145150
Schema: bqToSchema(q.Schema),
146151
HivePartitioningOptions: bqToHivePartitioningOptions(q.HivePartitioningOptions),
147152
ConnectionID: q.ConnectionId,
153+
ReferenceFileSchemaURI: q.ReferenceFileSchemaUri,
148154
}
149155
for _, v := range q.DecimalTargetTypes {
150156
e.DecimalTargetTypes = append(e.DecimalTargetTypes, DecimalTargetType(v))

‎bigquery/integration_test.go‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1393,6 +1393,100 @@ func TestIntegration_Load(t *testing.T) {
13931393

13941394
}
13951395

1396+
func TestIntegration_LoadWithReferenceSchemaFile(t *testing.T) {
1397+
if client == nil {
1398+
t.Skip("Integration tests skipped")
1399+
}
1400+
1401+
formats := []DataFormat{Avro, Parquet}
1402+
for _, format := range formats {
1403+
ctx := context.Background()
1404+
table := dataset.Table(tableIDs.New())
1405+
defer table.Delete(ctx)
1406+
1407+
expectedSchema := Schema{
1408+
{Name: "username", Type: StringFieldType, Required: false},
1409+
{Name: "tweet", Type: StringFieldType, Required: false},
1410+
{Name: "timestamp", Type: StringFieldType, Required: false},
1411+
{Name: "likes", Type: IntegerFieldType, Required: false},
1412+
}
1413+
ext := strings.ToLower(string(format))
1414+
sourceURIs := []string{
1415+
"gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/a-twitter." + ext,
1416+
"gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/b-twitter." + ext,
1417+
"gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/c-twitter." + ext,
1418+
}
1419+
referenceURI := "gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/a-twitter." + ext
1420+
source := NewGCSReference(sourceURIs...)
1421+
source.SourceFormat = format
1422+
loader := table.LoaderFrom(source)
1423+
loader.ReferenceFileSchemaURI = referenceURI
1424+
job, err := loader.Run(ctx)
1425+
if err != nil {
1426+
t.Fatalf("loader.Run: %v", err)
1427+
}
1428+
err = wait(ctx, job)
1429+
if err != nil {
1430+
t.Fatalf("wait: %v", err)
1431+
}
1432+
metadata, err := table.Metadata(ctx)
1433+
if err != nil {
1434+
t.Fatalf("table.Metadata: %v", err)
1435+
}
1436+
diff := testutil.Diff(expectedSchema, metadata.Schema)
1437+
if diff != "" {
1438+
t.Errorf("got=-, want=+:\n%s", diff)
1439+
}
1440+
}
1441+
}
1442+
1443+
func TestIntegration_ExternalTableWithReferenceSchemaFile(t *testing.T) {
1444+
if client == nil {
1445+
t.Skip("Integration tests skipped")
1446+
}
1447+
1448+
formats := []DataFormat{Avro, Parquet}
1449+
for _, format := range formats {
1450+
ctx := context.Background()
1451+
externalTable := dataset.Table(tableIDs.New())
1452+
defer externalTable.Delete(ctx)
1453+
1454+
expectedSchema := Schema{
1455+
{Name: "username", Type: StringFieldType, Required: false},
1456+
{Name: "tweet", Type: StringFieldType, Required: false},
1457+
{Name: "timestamp", Type: StringFieldType, Required: false},
1458+
{Name: "likes", Type: IntegerFieldType, Required: false},
1459+
}
1460+
ext := strings.ToLower(string(format))
1461+
sourceURIs := []string{
1462+
"gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/a-twitter." + ext,
1463+
"gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/b-twitter." + ext,
1464+
"gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/c-twitter." + ext,
1465+
}
1466+
referenceURI := "gs://cloud-samples-data/bigquery/federated-formats-reference-file-schema/a-twitter." + ext
1467+
1468+
err := externalTable.Create(ctx, &TableMetadata{
1469+
ExternalDataConfig: &ExternalDataConfig{
1470+
SourceFormat: format,
1471+
SourceURIs: sourceURIs,
1472+
ReferenceFileSchemaURI: referenceURI,
1473+
},
1474+
})
1475+
if err != nil {
1476+
t.Fatalf("table.Create: %v", err)
1477+
}
1478+
1479+
metadata, err := externalTable.Metadata(ctx)
1480+
if err != nil {
1481+
t.Fatalf("table.Metadata: %v", err)
1482+
}
1483+
diff := testutil.Diff(expectedSchema, metadata.Schema)
1484+
if diff != "" {
1485+
t.Errorf("got=-, want=+:\n%s", diff)
1486+
}
1487+
}
1488+
}
1489+
13961490
func TestIntegration_DML(t *testing.T) {
13971491
if client == nil {
13981492
t.Skip("Integration tests skipped")

‎bigquery/load.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,10 @@ type LoadConfig struct {
8888
// Experimental: this option is experimental and may be modified or removed in future versions,
8989
// regardless of any other documented package stability guarantees.
9090
JobTimeout time.Duration
91+
92+
// When loading a table with external data, the user can provide a reference file with the table schema.
93+
// This is enabled for the following formats: AVRO, PARQUET, ORC.
94+
ReferenceFileSchemaURI string
9195
}
9296

9397
func (l *LoadConfig) toBQ() (*bq.JobConfiguration, io.Reader) {
@@ -105,6 +109,7 @@ func (l *LoadConfig) toBQ() (*bq.JobConfiguration, io.Reader) {
105109
UseAvroLogicalTypes: l.UseAvroLogicalTypes,
106110
ProjectionFields: l.ProjectionFields,
107111
HivePartitioningOptions: l.HivePartitioningOptions.toBQ(),
112+
ReferenceFileSchemaUri: l.ReferenceFileSchemaURI,
108113
},
109114
JobTimeoutMs: l.JobTimeout.Milliseconds(),
110115
}
@@ -129,6 +134,7 @@ func bqToLoadConfig(q *bq.JobConfiguration, c *Client) *LoadConfig {
129134
UseAvroLogicalTypes: q.Load.UseAvroLogicalTypes,
130135
ProjectionFields: q.Load.ProjectionFields,
131136
HivePartitioningOptions: bqToHivePartitioningOptions(q.Load.HivePartitioningOptions),
137+
ReferenceFileSchemaURI: q.Load.ReferenceFileSchemaUri,
132138
}
133139
if q.JobTimeoutMs > 0 {
134140
lc.JobTimeout = time.Duration(q.JobTimeoutMs) * time.Millisecond

‎bigquery/load_test.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,23 @@ func TestLoad(t *testing.T) {
388388
return j
389389
}(),
390390
},
391+
{
392+
dst: c.Dataset("dataset-id").Table("table-id"),
393+
src: func() *GCSReference {
394+
g := NewGCSReference("uri")
395+
g.SourceFormat = Parquet
396+
return g
397+
}(),
398+
config: LoadConfig{
399+
ReferenceFileSchemaURI: "schema.parquet",
400+
},
401+
want: func() *bq.Job {
402+
j := defaultLoadJob()
403+
j.Configuration.Load.SourceFormat = "PARQUET"
404+
j.Configuration.Load.ReferenceFileSchemaUri = "schema.parquet"
405+
return j
406+
}(),
407+
},
391408
}
392409

393410
for i, tc := range testCases {

0 commit comments

Comments
 (0)