Skip to content

Commit babb4ac

Browse files
Praful Makanisduskis
Praful Makani
authored and
sduskis
committed
BigQuery: Added listPartitions. (#4923)
* added listPartitions * added unit test * modified code * modified unit test for listPartitions * added integration test * Fix table name * Fix integration test
1 parent 308bc15 commit babb4ac

File tree

4 files changed

+194
-0
lines changed

4 files changed

+194
-0
lines changed

‎google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -848,6 +848,12 @@ public int hashCode() {
848848
*/
849849
Page<Table> listTables(DatasetId datasetId, TableListOption... options);
850850

851+
/**
852+
* @param tableId
853+
* @return A list of the partition ids present in the partitioned table
854+
*/
855+
List<String> listPartitions(TableId tableId);
856+
851857
/**
852858
* Sends an insert all request.
853859
*

‎google-cloud-clients/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import com.google.common.collect.ImmutableList;
4747
import com.google.common.collect.Iterables;
4848
import com.google.common.collect.Maps;
49+
import java.util.ArrayList;
4950
import java.util.List;
5051
import java.util.Map;
5152
import java.util.concurrent.Callable;
@@ -507,6 +508,26 @@ public Page<Table> listTables(DatasetId datasetId, TableListOption... options) {
507508
return listTables(completeDatasetId, getOptions(), optionMap(options));
508509
}
509510

511+
@Override
512+
public List<String> listPartitions(TableId tableId) {
513+
List<String> partitions = new ArrayList<String>();
514+
Table metaTable =
515+
getTable(TableId.of(tableId.getDataset(), tableId.getTable() + "$__PARTITIONS_SUMMARY__"));
516+
Schema metaSchema = metaTable.getDefinition().getSchema();
517+
String partition_id = null;
518+
for (Field field : metaSchema.getFields()) {
519+
if (field.getName().equals("partition_id")) {
520+
partition_id = field.getName();
521+
break;
522+
}
523+
}
524+
TableResult result = metaTable.list(metaSchema);
525+
for (FieldValueList list : result.iterateAll()) {
526+
partitions.add(list.get(partition_id).getStringValue());
527+
}
528+
return partitions;
529+
}
530+
510531
private static Page<Table> listTables(
511532
final DatasetId datasetId,
512533
final BigQueryOptions serviceOptions,

‎google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,128 @@ public class BigQueryImplTest {
228228
private static final BigQuery.TableOption TABLE_OPTION_FIELDS =
229229
BigQuery.TableOption.fields(BigQuery.TableField.SCHEMA, BigQuery.TableField.ETAG);
230230

231+
// Table list partitions
232+
private static final Field PROJECT_ID_FIELD =
233+
Field.newBuilder("project_id", LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build();
234+
private static final Field DATASET_ID_FIELD =
235+
Field.newBuilder("dataset_id", LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build();
236+
private static final Field TABLE_ID_FIELD =
237+
Field.newBuilder("table_id", LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build();
238+
private static final Field PARTITION_ID_FIELD =
239+
Field.newBuilder("partition_id", LegacySQLTypeName.STRING)
240+
.setMode(Field.Mode.NULLABLE)
241+
.build();
242+
private static final Field CREATION_TIME_FIELD =
243+
Field.newBuilder("creation_time", LegacySQLTypeName.INTEGER)
244+
.setMode(Field.Mode.NULLABLE)
245+
.build();
246+
private static final Field CREATION_TIMESTAMP_FIELD =
247+
Field.newBuilder("creation_timestamp", LegacySQLTypeName.TIMESTAMP)
248+
.setMode(Field.Mode.NULLABLE)
249+
.build();
250+
private static final Field LAST_MODIFIED_FIELD =
251+
Field.newBuilder("last_modified_time", LegacySQLTypeName.INTEGER)
252+
.setMode(Field.Mode.NULLABLE)
253+
.build();
254+
private static final Field LAST_MODIFIED_TIMESTAMP_FIELD =
255+
Field.newBuilder("last_modified_timestamp", LegacySQLTypeName.TIMESTAMP)
256+
.setMode(Field.Mode.NULLABLE)
257+
.build();
258+
private static final Schema SCHEMA_PARTITIONS =
259+
Schema.of(
260+
PROJECT_ID_FIELD,
261+
DATASET_ID_FIELD,
262+
TABLE_ID_FIELD,
263+
PARTITION_ID_FIELD,
264+
CREATION_TIME_FIELD,
265+
CREATION_TIMESTAMP_FIELD,
266+
LAST_MODIFIED_FIELD,
267+
LAST_MODIFIED_TIMESTAMP_FIELD);
268+
private static final TableDefinition TABLE_DEFINITION_PARTITIONS =
269+
StandardTableDefinition.newBuilder()
270+
.setSchema(SCHEMA_PARTITIONS)
271+
.setNumBytes(0L)
272+
.setNumLongTermBytes(0L)
273+
.setNumRows(3L)
274+
.setLocation("unknown")
275+
.build();
276+
private static final TableInfo TABLE_INFO_PARTITIONS =
277+
TableInfo.newBuilder(TABLE_ID, TABLE_DEFINITION_PARTITIONS)
278+
.setEtag("ETAG")
279+
.setCreationTime(1553689573240L)
280+
.setLastModifiedTime(1553841163438L)
281+
.setNumBytes(0L)
282+
.setNumLongTermBytes(0L)
283+
.setNumRows(BigInteger.valueOf(3L))
284+
.build();
285+
private static final TableCell TABLE_CELL1_PROJECT_ID = new TableCell().setV(PROJECT);
286+
private static final TableCell TABLE_CELL1_DATASET_ID = new TableCell().setV(DATASET);
287+
private static final TableCell TABLE_CELL1_TABLE_ID = new TableCell().setV(TABLE);
288+
private static final TableCell TABLE_CELL1_PARTITION_ID = new TableCell().setV("20190327");
289+
private static final TableCell TABLE_CELL1_CREATION_TIME = new TableCell().setV("1553694932498");
290+
private static final TableCell TABLE_CELL1_CREATION_TIMESTAMP =
291+
new TableCell().setV("1553694932.498");
292+
private static final TableCell TABLE_CELL1_LAST_MODIFIED_TIME =
293+
new TableCell().setV("1553694932989");
294+
private static final TableCell TABLE_CELL1_LAST_MODIFIED_TIMESTAMP =
295+
new TableCell().setV("1553694932.989");
296+
297+
private static final TableCell TABLE_CELL2_PARTITION_ID = new TableCell().setV("20190328");
298+
private static final TableCell TABLE_CELL2_CREATION_TIME = new TableCell().setV("1553754224760");
299+
private static final TableCell TABLE_CELL2_CREATION_TIMESTAMP =
300+
new TableCell().setV("1553754224.76");
301+
private static final TableCell TABLE_CELL2_LAST_MODIFIED_TIME =
302+
new TableCell().setV("1553754225587");
303+
private static final TableCell TABLE_CELL2_LAST_MODIFIED_TIMESTAMP =
304+
new TableCell().setV("1553754225.587");
305+
306+
private static final TableCell TABLE_CELL3_PARTITION_ID = new TableCell().setV("20190329");
307+
private static final TableCell TABLE_CELL3_CREATION_TIME = new TableCell().setV("1553841162879");
308+
private static final TableCell TABLE_CELL3_CREATION_TIMESTAMP =
309+
new TableCell().setV("1553841162.879");
310+
private static final TableCell TABLE_CELL3_LAST_MODIFIED_TIME =
311+
new TableCell().setV("1553841163438");
312+
private static final TableCell TABLE_CELL3_LAST_MODIFIED_TIMESTAMP =
313+
new TableCell().setV("1553841163.438");
314+
315+
private static final TableDataList TABLE_DATA_WITH_PARTITIONS =
316+
new TableDataList()
317+
.setTotalRows(3L)
318+
.setRows(
319+
ImmutableList.of(
320+
new TableRow()
321+
.setF(
322+
ImmutableList.of(
323+
TABLE_CELL1_PROJECT_ID,
324+
TABLE_CELL1_DATASET_ID,
325+
TABLE_CELL1_TABLE_ID,
326+
TABLE_CELL1_PARTITION_ID,
327+
TABLE_CELL1_CREATION_TIME,
328+
TABLE_CELL1_CREATION_TIMESTAMP,
329+
TABLE_CELL1_LAST_MODIFIED_TIME,
330+
TABLE_CELL1_LAST_MODIFIED_TIMESTAMP)),
331+
new TableRow()
332+
.setF(
333+
ImmutableList.of(
334+
TABLE_CELL1_PROJECT_ID,
335+
TABLE_CELL1_DATASET_ID,
336+
TABLE_CELL1_TABLE_ID,
337+
TABLE_CELL2_PARTITION_ID,
338+
TABLE_CELL2_CREATION_TIME,
339+
TABLE_CELL2_CREATION_TIMESTAMP,
340+
TABLE_CELL2_LAST_MODIFIED_TIME,
341+
TABLE_CELL2_LAST_MODIFIED_TIMESTAMP)),
342+
new TableRow()
343+
.setF(
344+
ImmutableList.of(
345+
TABLE_CELL1_PROJECT_ID,
346+
TABLE_CELL1_DATASET_ID,
347+
TABLE_CELL1_TABLE_ID,
348+
TABLE_CELL3_PARTITION_ID,
349+
TABLE_CELL3_CREATION_TIME,
350+
TABLE_CELL3_CREATION_TIMESTAMP,
351+
TABLE_CELL3_LAST_MODIFIED_TIME,
352+
TABLE_CELL3_LAST_MODIFIED_TIMESTAMP))));
231353
// Table list options
232354
private static final BigQuery.TableListOption TABLE_LIST_PAGE_SIZE =
233355
BigQuery.TableListOption.pageSize(42L);
@@ -639,6 +761,20 @@ public void testGetTable() {
639761
assertEquals(new Table(bigquery, new TableInfo.BuilderImpl(TABLE_INFO_WITH_PROJECT)), table);
640762
}
641763

764+
@Test
765+
public void testListPartition() {
766+
EasyMock.expect(
767+
bigqueryRpcMock.getTable(
768+
PROJECT, DATASET, "table$__PARTITIONS_SUMMARY__", EMPTY_RPC_OPTIONS))
769+
.andReturn(TABLE_INFO_PARTITIONS.toPb());
770+
EasyMock.expect(bigqueryRpcMock.listTableData(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))
771+
.andReturn(TABLE_DATA_WITH_PARTITIONS);
772+
EasyMock.replay(bigqueryRpcMock);
773+
bigquery = options.getService();
774+
List<String> partition = bigquery.listPartitions(TABLE_ID_WITH_PROJECT);
775+
assertEquals(3, partition.size());
776+
}
777+
642778
@Test
643779
public void testGetTableNotFoundWhenThrowIsDisabled() {
644780
EasyMock.expect(bigqueryRpcMock.getTable(PROJECT, DATASET, TABLE, EMPTY_RPC_OPTIONS))

‎google-cloud-clients/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import static org.junit.Assert.fail;
2828

2929
import com.google.api.gax.paging.Page;
30+
import com.google.cloud.Date;
3031
import com.google.cloud.RetryOption;
3132
import com.google.cloud.bigquery.BigQuery;
3233
import com.google.cloud.bigquery.BigQuery.DatasetDeleteOption;
@@ -92,6 +93,7 @@
9293
import java.util.HashMap;
9394
import java.util.HashSet;
9495
import java.util.Iterator;
96+
import java.util.List;
9597
import java.util.Map;
9698
import java.util.Set;
9799
import java.util.UUID;
@@ -641,6 +643,35 @@ public void testListTablesWithPartitioning() {
641643
}
642644
}
643645

646+
@Test
647+
public void testListPartitions() throws InterruptedException {
648+
String tableName = "test_table_partitions";
649+
Date date = Date.fromJavaUtilDate(new java.util.Date());
650+
String partitionDate = date.toString().replaceAll("-", "");
651+
TableId tableId = TableId.of(DATASET, tableName + "$" + partitionDate);
652+
String query =
653+
String.format(
654+
"CREATE OR REPLACE TABLE %s.%s ( StringField STRING )"
655+
+ " PARTITION BY DATE(_PARTITIONTIME) "
656+
+ "OPTIONS( partition_expiration_days=1)",
657+
DATASET, tableName);
658+
Job job = bigquery.create(JobInfo.of(QueryJobConfiguration.newBuilder(query).build()));
659+
job.waitFor();
660+
assertTrue(job.isDone());
661+
try {
662+
Map<String, Object> row = new HashMap<String, Object>();
663+
row.put("StringField", "StringValue");
664+
InsertAllRequest request = InsertAllRequest.newBuilder(tableId).addRow(row).build();
665+
InsertAllResponse response = bigquery.insertAll(request);
666+
assertFalse(response.hasErrors());
667+
assertEquals(0, response.getInsertErrors().size());
668+
List<String> partitions = bigquery.listPartitions(TableId.of(DATASET, tableName));
669+
assertEquals(1, partitions.size());
670+
} finally {
671+
bigquery.delete(DATASET, tableName);
672+
}
673+
}
674+
644675
@Test
645676
public void testUpdateTable() {
646677
String tableName = "test_update_table";

0 commit comments

Comments
 (0)