class 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());
  }

}