DirectBigQueryInputFormat extends InputFormat<NullWritable, GenericRecord> { @Override public List<InputSplit> getSplits(JobContext context) throws IOException { final Configuration configuration = context.getConfiguration(); BigQueryStorageClient client = getClient(configuration); BigQueryHelper bigQueryHelper; try { bigQueryHelper = getBigQueryHelper(configuration); } catch (GeneralSecurityException gse) { throw new IOException("Failed to create BigQuery client", gse); } double skewLimit = SKEW_LIMIT.get(configuration, configuration::getDouble); Preconditions.checkArgument( skewLimit >= 1.0, "%s is less than 1; not all records would be read. Exiting", SKEW_LIMIT.getKey()); Table table = getTable(configuration, bigQueryHelper); ReadSession session = startSession(configuration, table, client); long numRows = table.getNumRows().longValue(); long limit = Math.round(skewLimit * numRows / session.getStreamsCount()); return session.getStreamsList().stream() .map( stream -> new DirectBigQueryInputSplit( stream.getName(), session.getAvroSchema().getSchema(), limit)) .collect(Collectors.toList()); }  @Override List<InputSplit> getSplits(JobContext context); @Override RecordReader<NullWritable, GenericRecord> createRecordReader(
      InputSplit inputSplit, TaskAttemptContext context);  }
@Test public void getSplits() throws IOException { JobContext jobContext = new JobContextImpl(config, new JobID()); CreateReadSessionRequest request = CreateReadSessionRequest.newBuilder() .setTableReference( TableReferenceProto.TableReference.newBuilder() .setProjectId("publicdata") .setDatasetId("test_dataset") .setTableId("test_table")) .setRequestedStreams(3) .setParent("projects/foo-project") .setReadOptions( TableReadOptions.newBuilder() .addAllSelectedFields(ImmutableList.of("foo", "bar")) .setRowRestriction("foo == 0") .build()) .setFormat(DataFormat.AVRO) .build(); ReadSession session = ReadSession.newBuilder() .setAvroSchema(AvroSchema.newBuilder().setSchema("schema").build()) .addAllStreams( ImmutableList.of( Stream.newBuilder().setName("stream1").build(), Stream.newBuilder().setName("stream2").build())) .build(); ImmutableList<DirectBigQueryInputSplit> expected = ImmutableList.of( new DirectBigQueryInputSplit("stream1", "schema", 14), new DirectBigQueryInputSplit("stream2", "schema", 14)); when(bqClient.createReadSession(any(CreateReadSessionRequest.class))).thenReturn(session); try { List<InputSplit> splits = input.getSplits(jobContext); assertThat(splits).containsExactlyElementsIn(expected); } catch (Exception e) { e.printStackTrace(); System.exit(1); } verify(bqHelper).getTable(tableRef); verify(bqClient).createReadSession(request); }