// The bigquery_storage_quickstart application demonstrates usage of the // BigQuery Storage read API. It demonstrates API features such as column // projection (limiting the output to a subset of a table's columns), // column filtering (using simple predicates to filter records on the server // side), establishing the snapshot time (reading data from the table at a // specific point in time), decoding Avro row blocks using the third party // "github.com/linkedin/goavro" library, and decoding Arrow row blocks using // the third party "github.com/apache/arrow/go" library. package main import ( "bytes" "context" "encoding/json" "flag" "fmt" "io" "log" "sort" "strings" "sync" "time" bqStorage "cloud.google.com/go/bigquery/storage/apiv1" "cloud.google.com/go/bigquery/storage/apiv1/storagepb" "github.com/apache/arrow/go/v10/arrow" "github.com/apache/arrow/go/v10/arrow/ipc" "github.com/apache/arrow/go/v10/arrow/memory" gax "github.com/googleapis/gax-go/v2" goavro "github.com/linkedin/goavro/v2" "google.golang.org/genproto/googleapis/rpc/errdetails" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "google.golang.org/protobuf/types/known/timestamppb" ) // rpcOpts is used to configure the underlying gRPC client to accept large // messages. The BigQuery Storage API may send message blocks up to 128MB // in size. var rpcOpts = gax.WithGRPCOptions( grpc.MaxCallRecvMsgSize(1024 * 1024 * 129), ) // Available formats const ( AVRO_FORMAT = "avro" ARROW_FORMAT = "arrow" ) // Command-line flags. var ( projectID = flag.String("project_id", "", "Cloud Project ID, used for session creation.") snapshotMillis = flag.Int64("snapshot_millis", 0, "Snapshot time to use for reads, represented in epoch milliseconds format. Default behavior reads current data.") format = flag.String("format", AVRO_FORMAT, "format to read data from storage API. Default is avro.") ) func main() { flag.Parse() ctx := context.Background() bqReadClient, err := bqStorage.NewBigQueryReadClient(ctx) if err != nil { log.Fatalf("NewBigQueryStorageClient: %v", err) } defer bqReadClient.Close() // Verify we've been provided a parent project which will contain the read session. The // session may exist in a different project than the table being read. if *projectID == "" { log.Fatalf("No parent project ID specified, please supply using the --project_id flag.") } // This example uses baby name data from the public datasets. srcProjectID := "bigquery-public-data" srcDatasetID := "usa_names" srcTableID := "usa_1910_current" readTable := fmt.Sprintf("projects/%s/datasets/%s/tables/%s", srcProjectID, srcDatasetID, srcTableID, ) // We limit the output columns to a subset of those allowed in the table, // and set a simple filter to only report names from the state of // Washington (WA). tableReadOptions := &storagepb.ReadSession_TableReadOptions{ SelectedFields: []string{"name", "number", "state"}, RowRestriction: `state = "WA"`, } dataFormat := storagepb.DataFormat_AVRO if *format == ARROW_FORMAT { dataFormat = storagepb.DataFormat_ARROW } createReadSessionRequest := &storagepb.CreateReadSessionRequest{ Parent: fmt.Sprintf("projects/%s", *projectID), ReadSession: &storagepb.ReadSession{ Table: readTable, DataFormat: dataFormat, ReadOptions: tableReadOptions, }, MaxStreamCount: 1, } // Set a snapshot time if it's been specified. if *snapshotMillis > 0 { ts := timestamppb.New(time.Unix(0, *snapshotMillis*1000)) if !ts.IsValid() { log.Fatalf("Invalid snapshot millis (%d): %v", *snapshotMillis, err) } createReadSessionRequest.ReadSession.TableModifiers = &storagepb.ReadSession_TableModifiers{ SnapshotTime: ts, } } // Create the session from the request. session, err := bqReadClient.CreateReadSession(ctx, createReadSessionRequest, rpcOpts) if err != nil { log.Fatalf("CreateReadSession: %v", err) } fmt.Printf("Read session: %s\n", session.GetName()) if len(session.GetStreams()) == 0 { log.Fatalf("no streams in session. if this was a small query result, consider writing to output to a named table.") } // We'll use only a single stream for reading data from the table. Because // of dynamic sharding, this will yield all the rows in the table. However, // if you wanted to fan out multiple readers you could do so by having a // increasing the MaxStreamCount. readStream := session.GetStreams()[0].Name ch := make(chan *storagepb.ReadRowsResponse) // Use a waitgroup to coordinate the reading and decoding goroutines. var wg sync.WaitGroup // Start the reading in one goroutine. wg.Add(1) go func() { defer wg.Done() if err := processStream(ctx, bqReadClient, readStream, ch); err != nil { log.Fatalf("processStream failure: %v", err) } close(ch) }() // Start Avro processing and decoding in another goroutine. wg.Add(1) go func() { defer wg.Done() var err error switch *format { case ARROW_FORMAT: err = processArrow(ctx, session.GetArrowSchema().GetSerializedSchema(), ch) case AVRO_FORMAT: err = processAvro(ctx, session.GetAvroSchema().GetSchema(), ch) } if err != nil { log.Fatalf("error processing %s: %v", *format, err) } }() // Wait until both the reading and decoding goroutines complete. wg.Wait() } // printDatum prints the decoded row datum. func printDatum(d interface{}) { m, ok := d.(map[string]interface{}) if !ok { log.Printf("failed type assertion: %v", d) } // Go's map implementation returns keys in a random ordering, so we sort // the keys before accessing. keys := make([]string, len(m)) i := 0 for k := range m { keys[i] = k i++ } sort.Strings(keys) for _, key := range keys { fmt.Printf("%s: %-20v ", key, valueFromTypeMap(m[key])) } fmt.Println() } // printRecordBatch prints the arrow record batch func printRecordBatch(record arrow.Record) error { out, err := record.MarshalJSON() if err != nil { return err } list := []map[string]interface{}{} err = json.Unmarshal(out, &list) if err != nil { return err } if len(list) == 0 { return nil } first := list[0] keys := make([]string, len(first)) i := 0 for k := range first { keys[i] = k i++ } sort.Strings(keys) builder := strings.Builder{} for _, m := range list { for _, key := range keys { builder.WriteString(fmt.Sprintf("%s: %-20v ", key, m[key])) } builder.WriteString("\n") } fmt.Print(builder.String()) return nil } // valueFromTypeMap returns the first value/key in the type map. This function // is only suitable for simple schemas, as complex typing such as arrays and // records necessitate a more robust implementation. See the goavro library // and the Avro specification for more information. func valueFromTypeMap(field interface{}) interface{} { m, ok := field.(map[string]interface{}) if !ok { return nil } for _, v := range m { // Return the first key encountered. return v } return nil } // processStream reads rows from a single storage Stream, and sends the Storage Response // data blocks to a channel. This function will retry on transient stream // failures and bookmark progress to avoid re-reading data that's already been // successfully transmitted. func processStream(ctx context.Context, client *bqStorage.BigQueryReadClient, st string, ch chan<- *storagepb.ReadRowsResponse) error { var offset int64 // Streams may be long-running. Rather than using a global retry for the // stream, implement a retry that resets once progress is made. retryLimit := 3 retries := 0 for { // Send the initiating request to start streaming row blocks. rowStream, err := client.ReadRows(ctx, &storagepb.ReadRowsRequest{ ReadStream: st, Offset: offset, }, rpcOpts) if err != nil { return fmt.Errorf("couldn't invoke ReadRows: %w", err) } // Process the streamed responses. for { r, err := rowStream.Recv() if err == io.EOF { return nil } if err != nil { // If there is an error, check whether it is a retryable // error with a retry delay and sleep instead of increasing // retries count. var retryDelayDuration time.Duration if errorStatus, ok := status.FromError(err); ok && errorStatus.Code() == codes.ResourceExhausted { for _, detail := range errorStatus.Details() { retryInfo, ok := detail.(*errdetails.RetryInfo) if !ok { continue } retryDelay := retryInfo.GetRetryDelay() retryDelayDuration = time.Duration(retryDelay.Seconds)*time.Second + time.Duration(retryDelay.Nanos)*time.Nanosecond break } } if retryDelayDuration != 0 { log.Printf("processStream failed with a retryable error, retrying in %v", retryDelayDuration) time.Sleep(retryDelayDuration) } else { retries++ if retries >= retryLimit { return fmt.Errorf("processStream retries exhausted: %w", err) } } // break the inner loop, and try to recover by starting a new streaming // ReadRows call at the last known good offset. break } else { // Reset retries after a successful response. retries = 0 } rc := r.GetRowCount() if rc > 0 { // Bookmark our progress in case of retries and send the rowblock on the channel. offset = offset + rc // We're making progress, reset retries. retries = 0 ch <- r } } } } // processArrow receives row blocks from a channel, and uses the provided Arrow // schema to decode the blocks into individual row messages for printing. Will // continue to run until the channel is closed or the provided context is // cancelled. func processArrow(ctx context.Context, schema []byte, ch <-chan *storagepb.ReadRowsResponse) error { mem := memory.NewGoAllocator() buf := bytes.NewBuffer(schema) r, err := ipc.NewReader(buf, ipc.WithAllocator(mem)) if err != nil { return err } aschema := r.Schema() for { select { case <-ctx.Done(): // Context was cancelled. Stop. return ctx.Err() case rows, ok := <-ch: if !ok { // Channel closed, no further arrow messages. Stop. return nil } undecoded := rows.GetArrowRecordBatch().GetSerializedRecordBatch() if len(undecoded) > 0 { buf = bytes.NewBuffer(schema) buf.Write(undecoded) r, err = ipc.NewReader(buf, ipc.WithAllocator(mem), ipc.WithSchema(aschema)) if err != nil { return err } for r.Next() { rec := r.Record() err = printRecordBatch(rec) if err != nil { return err } } } } } } // processAvro receives row blocks from a channel, and uses the provided Avro // schema to decode the blocks into individual row messages for printing. Will // continue to run until the channel is closed or the provided context is // cancelled. func processAvro(ctx context.Context, schema string, ch <-chan *storagepb.ReadRowsResponse) error { // Establish a decoder that can process blocks of messages using the // reference schema. All blocks share the same schema, so the decoder // can be long-lived. codec, err := goavro.NewCodec(schema) if err != nil { return fmt.Errorf("couldn't create codec: %w", err) } for { select { case <-ctx.Done(): // Context was cancelled. Stop. return ctx.Err() case rows, ok := <-ch: if !ok { // Channel closed, no further avro messages. Stop. return nil } undecoded := rows.GetAvroRows().GetSerializedBinaryRows() for len(undecoded) > 0 { datum, remainingBytes, err := codec.NativeFromBinary(undecoded) if err != nil { if err == io.EOF { break } return fmt.Errorf("decoding error with %d bytes remaining: %v", len(undecoded), err) } printDatum(datum) undecoded = remainingBytes } } } }
