diff --git a/bigquery/external.go b/bigquery/external.go index 3f8647e568a..36a1ab64ac3 100644 --- a/bigquery/external.go +++ b/bigquery/external.go @@ -100,19 +100,17 @@ type ExternalDataConfig struct { func (e *ExternalDataConfig) toBQ() bq.ExternalDataConfiguration { q := bq.ExternalDataConfiguration{ - SourceFormat: string(e.SourceFormat), - SourceUris: e.SourceURIs, - Autodetect: e.AutoDetect, - Compression: string(e.Compression), - IgnoreUnknownValues: e.IgnoreUnknownValues, - MaxBadRecords: e.MaxBadRecords, + SourceFormat: string(e.SourceFormat), + SourceUris: e.SourceURIs, + Autodetect: e.AutoDetect, + Compression: string(e.Compression), + IgnoreUnknownValues: e.IgnoreUnknownValues, + MaxBadRecords: e.MaxBadRecords, + HivePartitioningOptions: e.HivePartitioningOptions.toBQ(), } if e.Schema != nil { q.Schema = e.Schema.toBQ() } - if e.HivePartitioningOptions != nil { - q.HivePartitioningOptions = e.HivePartitioningOptions.toBQ() - } if e.Options != nil { e.Options.populateExternalDataConfig(&q) } @@ -121,13 +119,14 @@ func (e *ExternalDataConfig) toBQ() bq.ExternalDataConfiguration { func bqToExternalDataConfig(q *bq.ExternalDataConfiguration) (*ExternalDataConfig, error) { e := &ExternalDataConfig{ - SourceFormat: DataFormat(q.SourceFormat), - SourceURIs: q.SourceUris, - AutoDetect: q.Autodetect, - Compression: Compression(q.Compression), - IgnoreUnknownValues: q.IgnoreUnknownValues, - MaxBadRecords: q.MaxBadRecords, - Schema: bqToSchema(q.Schema), + SourceFormat: DataFormat(q.SourceFormat), + SourceURIs: q.SourceUris, + AutoDetect: q.Autodetect, + Compression: Compression(q.Compression), + IgnoreUnknownValues: q.IgnoreUnknownValues, + MaxBadRecords: q.MaxBadRecords, + Schema: bqToSchema(q.Schema), + HivePartitioningOptions: bqToHivePartitioningOptions(q.HivePartitioningOptions), } switch { case q.CsvOptions != nil: @@ -141,9 +140,6 @@ func bqToExternalDataConfig(q *bq.ExternalDataConfiguration) (*ExternalDataConfi return nil, err } } - if q.HivePartitioningOptions != nil { - e.HivePartitioningOptions = bqToHivePartitioningOptions(q.HivePartitioningOptions) - } return e, nil } diff --git a/bigquery/load.go b/bigquery/load.go index 9fa0190685b..1b2d3c6bfa7 100644 --- a/bigquery/load.go +++ b/bigquery/load.go @@ -65,6 +65,10 @@ type LoadConfig struct { // For ingestion from datastore backups, ProjectionFields governs which fields // are projected from the backup. The default behavior projects all fields. ProjectionFields []string + + // HivePartitioningOptions allows use of Hive partitioning based on the + // layout of objects in Cloud Storage. + HivePartitioningOptions *HivePartitioningOptions } func (l *LoadConfig) toBQ() (*bq.JobConfiguration, io.Reader) { @@ -81,6 +85,7 @@ func (l *LoadConfig) toBQ() (*bq.JobConfiguration, io.Reader) { SchemaUpdateOptions: l.SchemaUpdateOptions, UseAvroLogicalTypes: l.UseAvroLogicalTypes, ProjectionFields: l.ProjectionFields, + HivePartitioningOptions: l.HivePartitioningOptions.toBQ(), }, } media := l.Src.populateLoadConfig(config.Load) @@ -100,6 +105,7 @@ func bqToLoadConfig(q *bq.JobConfiguration, c *Client) *LoadConfig { SchemaUpdateOptions: q.Load.SchemaUpdateOptions, UseAvroLogicalTypes: q.Load.UseAvroLogicalTypes, ProjectionFields: q.Load.ProjectionFields, + HivePartitioningOptions: bqToHivePartitioningOptions(q.Load.HivePartitioningOptions), } var fc *FileConfig if len(q.Load.SourceUris) == 0 { diff --git a/bigquery/load_test.go b/bigquery/load_test.go index eb9ec38804d..ca3fe9da75f 100644 --- a/bigquery/load_test.go +++ b/bigquery/load_test.go @@ -342,6 +342,31 @@ func TestLoad(t *testing.T) { return j }(), }, + { + dst: c.Dataset("dataset-id").Table("table-id"), + src: func() *GCSReference { + g := NewGCSReference("uri") + g.SourceFormat = Parquet + return g + }(), + config: LoadConfig{ + HivePartitioningOptions: &HivePartitioningOptions{ + Mode: CustomHivePartitioningMode, + SourceURIPrefix: "source_uri", + RequirePartitionFilter: true, + }, + }, + want: func() *bq.Job { + j := defaultLoadJob() + j.Configuration.Load.SourceFormat = "PARQUET" + j.Configuration.Load.HivePartitioningOptions = &bq.HivePartitioningOptions{ + Mode: "CUSTOM", + RequirePartitionFilter: true, + SourceUriPrefix: "source_uri", + } + return j + }(), + }, } for i, tc := range testCases {