Search in sources :

Example 16 with TableReference

use of com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference in project java-bigquerystorage by googleapis.

the class ITBigQueryStorageTest method testBasicSqlTypes.

@Test
public void testBasicSqlTypes() throws InterruptedException, IOException {
    String table_name = "test_basic_sql_types";
    String createTableStatement = String.format(" CREATE TABLE %s.%s " + " (int_field INT64 NOT NULL," + "  num_field NUMERIC NOT NULL," + "  float_field FLOAT64 NOT NULL," + "  bool_field BOOL NOT NULL," + "  str_field STRING NOT NULL," + "  bytes_field BYTES NOT NULL) " + " OPTIONS( " + "   description=\"a table with basic column types\" " + " ) " + "AS " + "   SELECT " + "     17," + "     CAST(1234.56 AS NUMERIC)," + "     6.547678," + "     TRUE," + "     \"String field value\"," + "     b\"абвгд\"", DATASET, table_name);
    RunQueryJobAndExpectSuccess(QueryJobConfiguration.newBuilder(createTableStatement).build());
    TableReference tableReference = TableReference.newBuilder().setTableId(table_name).setDatasetId(DATASET).setProjectId(ServiceOptions.getDefaultProjectId()).build();
    List<GenericData.Record> rows = ReadAllRows(/* tableReference = */
    tableReference, /* filter = */
    null);
    assertEquals("Actual rows read: " + rows.toString(), 1, rows.size());
    GenericData.Record record = rows.get(0);
    Schema avroSchema = record.getSchema();
    String actualSchemaMessage = String.format("Unexpected schema. Actual schema:%n%s", avroSchema.toString(/* pretty = */
    true));
    String rowAssertMessage = String.format("Row not matching expectations: %s", record.toString());
    assertEquals(actualSchemaMessage, Schema.Type.RECORD, avroSchema.getType());
    assertEquals(actualSchemaMessage, "__root__", avroSchema.getName());
    assertEquals(actualSchemaMessage, 6, avroSchema.getFields().size());
    assertEquals(actualSchemaMessage, Schema.Type.LONG, avroSchema.getField("int_field").schema().getType());
    assertEquals(rowAssertMessage, 17L, (long) record.get("int_field"));
    assertEquals(actualSchemaMessage, Schema.Type.BYTES, avroSchema.getField("num_field").schema().getType());
    assertEquals(actualSchemaMessage, LogicalTypes.decimal(/* precision = */
    38, /* scale = */
    9), avroSchema.getField("num_field").schema().getLogicalType());
    BigDecimal actual_num_field = new Conversions.DecimalConversion().fromBytes((ByteBuffer) record.get("num_field"), avroSchema, avroSchema.getField("num_field").schema().getLogicalType());
    assertEquals(rowAssertMessage, BigDecimal.valueOf(/* unscaledVal = */
    1_234_560_000_000L, /* scale = */
    9), actual_num_field);
    assertEquals(actualSchemaMessage, Schema.Type.DOUBLE, avroSchema.getField("float_field").schema().getType());
    assertEquals(rowAssertMessage, /* expected = */
    6.547678d, /* actual = */
    (double) record.get("float_field"), /* delta = */
    0.0001);
    assertEquals(actualSchemaMessage, Schema.Type.BOOLEAN, avroSchema.getField("bool_field").schema().getType());
    assertEquals(rowAssertMessage, true, record.get("bool_field"));
    assertEquals(actualSchemaMessage, Schema.Type.STRING, avroSchema.getField("str_field").schema().getType());
    assertEquals(rowAssertMessage, new Utf8("String field value"), record.get("str_field"));
    assertEquals(actualSchemaMessage, Schema.Type.BYTES, avroSchema.getField("bytes_field").schema().getType());
    assertArrayEquals(rowAssertMessage, Utf8.getBytesFor("абвгд"), ((ByteBuffer) (record.get("bytes_field"))).array());
}
Also used : TableReference(com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference) Conversions(org.apache.avro.Conversions) Schema(org.apache.avro.Schema) Utf8(org.apache.avro.util.Utf8) GenericData(org.apache.avro.generic.GenericData) BigDecimal(java.math.BigDecimal) Test(org.junit.Test)

Example 17 with TableReference

use of com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference in project java-bigquerystorage by googleapis.

the class ITBigQueryStorageTest method testSimpleReadAndResume.

@Test
public void testSimpleReadAndResume() {
    TableReference tableReference = TableReference.newBuilder().setProjectId("bigquery-public-data").setDatasetId("samples").setTableId("shakespeare").build();
    ReadSession session = client.createReadSession(/* tableReference = */
    tableReference, /* parent = */
    parentProjectId, /* requestedStreams = */
    1);
    assertEquals(String.format("Did not receive expected number of streams for table reference '%s' CreateReadSession response:%n%s", TextFormat.shortDebugString(tableReference), session.toString()), 1, session.getStreamsCount());
    // We have to read some number of rows in order to be able to resume. More details:
    // https://cloud.google.com/bigquery/docs/reference/storage/rpc/google.cloud.bigquery.storage.v1beta1#google.cloud.bigquery.storage.v1beta1.ReadRowsRequest
    long rowCount = ReadStreamToOffset(session.getStreams(0), /* rowOffset = */
    34_846);
    StreamPosition readPosition = StreamPosition.newBuilder().setStream(session.getStreams(0)).setOffset(rowCount).build();
    ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder().setReadPosition(readPosition).build();
    ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
    for (ReadRowsResponse response : stream) {
        rowCount += response.getRowCount();
    }
    // Verifies that the number of rows skipped and read equals to the total number of rows in the
    // table.
    assertEquals(164_656, rowCount);
}
Also used : TableReference(com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference) ReadRowsResponse(com.google.cloud.bigquery.storage.v1beta1.Storage.ReadRowsResponse) ReadSession(com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession) StreamPosition(com.google.cloud.bigquery.storage.v1beta1.Storage.StreamPosition) ReadRowsRequest(com.google.cloud.bigquery.storage.v1beta1.Storage.ReadRowsRequest) Test(org.junit.Test)

Example 18 with TableReference

use of com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference in project java-bigquerystorage by googleapis.

the class ITBigQueryStorageTest method testStructAndArraySqlTypes.

@Test
public void testStructAndArraySqlTypes() throws InterruptedException, IOException {
    String table_name = "test_struct_and_array_sql_types";
    String createTableStatement = String.format(" CREATE TABLE %s.%s " + " (array_field ARRAY<INT64>," + "  struct_field STRUCT<int_field INT64 NOT NULL, str_field STRING NOT NULL> NOT NULL)" + " OPTIONS( " + "   description=\"a table with array and time column types\" " + " ) " + "AS " + "   SELECT " + "     [1, 2, 3]," + "     (10, 'abc')", DATASET, table_name);
    RunQueryJobAndExpectSuccess(QueryJobConfiguration.newBuilder(createTableStatement).build());
    TableReference tableReference = TableReference.newBuilder().setTableId(table_name).setDatasetId(DATASET).setProjectId(ServiceOptions.getDefaultProjectId()).build();
    List<GenericData.Record> rows = ReadAllRows(/* tableReference = */
    tableReference, /* filter = */
    null);
    assertEquals("Actual rows read: " + rows.toString(), 1, rows.size());
    GenericData.Record record = rows.get(0);
    Schema avroSchema = record.getSchema();
    String actualSchemaMessage = String.format("Unexpected schema. Actual schema:%n%s", avroSchema.toString(/* pretty = */
    true));
    String rowAssertMessage = String.format("Row not matching expectations: %s", record.toString());
    assertEquals(actualSchemaMessage, Schema.Type.RECORD, avroSchema.getType());
    assertEquals(actualSchemaMessage, "__root__", avroSchema.getName());
    assertEquals(actualSchemaMessage, 2, avroSchema.getFields().size());
    assertEquals(actualSchemaMessage, Schema.Type.ARRAY, avroSchema.getField("array_field").schema().getType());
    assertEquals(actualSchemaMessage, Schema.Type.LONG, avroSchema.getField("array_field").schema().getElementType().getType());
    assertArrayEquals(rowAssertMessage, new Long[] { 1L, 2L, 3L }, ((GenericData.Array<Long>) record.get("array_field")).toArray(new Long[0]));
    // Validate the STRUCT field and its members.
    Schema structSchema = avroSchema.getField("struct_field").schema();
    assertEquals(actualSchemaMessage, Schema.Type.RECORD, structSchema.getType());
    GenericData.Record structRecord = (GenericData.Record) record.get("struct_field");
    assertEquals(actualSchemaMessage, Schema.Type.LONG, structSchema.getField("int_field").schema().getType());
    assertEquals(rowAssertMessage, 10L, (long) structRecord.get("int_field"));
    assertEquals(actualSchemaMessage, Schema.Type.STRING, structSchema.getField("str_field").schema().getType());
    assertEquals(rowAssertMessage, new Utf8("abc"), structRecord.get("str_field"));
}
Also used : TableReference(com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference) Schema(org.apache.avro.Schema) Utf8(org.apache.avro.util.Utf8) GenericData(org.apache.avro.generic.GenericData) Test(org.junit.Test)

Aggregations

TableReference (com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference)14 Test (org.junit.Test)14 GenericData (org.apache.avro.generic.GenericData)8 ReadSession (com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession)7 ReadRowsRequest (com.google.cloud.bigquery.storage.v1beta1.Storage.ReadRowsRequest)5 ReadRowsResponse (com.google.cloud.bigquery.storage.v1beta1.Storage.ReadRowsResponse)5 StreamPosition (com.google.cloud.bigquery.storage.v1beta1.Storage.StreamPosition)5 Schema (org.apache.avro.Schema)5 Utf8 (org.apache.avro.util.Utf8)5 CreateReadSessionRequest (com.google.cloud.bigquery.storage.v1beta1.Storage.CreateReadSessionRequest)4 AvroRowConsumer (com.google.cloud.bigquery.storage.v1beta1.it.SimpleRowReader.AvroRowConsumer)4 ArrayList (java.util.ArrayList)3 Field (com.google.cloud.bigquery.Field)2 TableId (com.google.cloud.bigquery.TableId)2 TableInfo (com.google.cloud.bigquery.TableInfo)2 BigQueryStorageClient (com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageClient)2 TableReadOptions (com.google.cloud.bigquery.storage.v1beta1.ReadOptions.TableReadOptions)2 Storage (com.google.cloud.bigquery.storage.v1beta1.Storage)2 TableReferenceProto (com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto)2 InvalidArgumentException (com.google.api.gax.rpc.InvalidArgumentException)1