test_class.declarator: bqHelper
test_class.modifier: @Mock private
test_class.original_string: @Mock private BigQueryHelper bqHelper;
test_class.type: BigQueryHelper
test_class.var_name: bqHelper
test_class.declarator: bqClient
test_class.modifier: @Mock private
test_class.original_string: @Mock private BigQueryStorageClient bqClient;
test_class.type: BigQueryStorageClient
test_class.var_name: bqClient
test_class.declarator: taskContext
test_class.modifier: @Mock private
test_class.original_string: @Mock private TaskAttemptContext taskContext;
test_class.type: TaskAttemptContext
test_class.var_name: taskContext
test_class.declarator: config
test_class.modifier: private
test_class.original_string: private JobConf config;
test_class.type: JobConf
test_class.var_name: config
test_class.declarator: input
test_class.modifier: private
test_class.original_string: private DirectBigQueryInputFormat input;
test_class.type: DirectBigQueryInputFormat
test_class.var_name: input
test_class.declarator: tableRef
test_class.modifier: private
test_class.original_string: private TableReference tableRef;
test_class.type: TableReference
test_class.var_name: tableRef
test_class.declarator: jobProjectId = "foo-project"
test_class.modifier: private
test_class.original_string: private String jobProjectId = "foo-project";
test_class.type: String
test_class.var_name: jobProjectId
test_class.declarator: dataProjectId = "publicdata"
test_class.modifier: private
test_class.original_string: private String dataProjectId = "publicdata";
test_class.type: String
test_class.var_name: dataProjectId
test_class.declarator: datasetId = "test_dataset"
test_class.modifier: private
test_class.original_string: private String datasetId = "test_dataset";
test_class.type: String
test_class.var_name: datasetId
test_class.declarator: tableId = "test_table"
test_class.modifier: private
test_class.original_string: private String tableId = "test_table";
test_class.type: String
test_class.var_name: tableId
test_class.file: bigquery/src/test/java/com/google/cloud/hadoop/io/bigquery/DirectBigQueryInputFormatTest.java
test_class.identifier: DirectBigQueryInputFormatTest
test_case.body: @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) // request 3, but only get 2 back
              .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);
    }
test_case.class_method_signature: DirectBigQueryInputFormatTest.getSplits()
test_case.full_signature: @Test public void getSplits()
test_case.identifier: getSplits
test_case.invocations: build setFormat setReadOptions setParent setRequestedStreams setTableReference newBuilder setTableId setDatasetId setProjectId newBuilder build setRowRestriction addAllSelectedFields newBuilder of build addAllStreams setAvroSchema newBuilder build setSchema newBuilder of build setName newBuilder build setName newBuilder of thenReturn when createReadSession any getSplits containsExactlyElementsIn assertThat printStackTrace exit getTable verify createReadSession verify
test_case.modifiers: @Test public
test_case.parameters: ()
test_case.return: void
test_case.signature: void getSplits()
focal_class.declarator: DIRECT_PARALLELISM =
        new HadoopConfigurationProperty<>(MRJobConfig.NUM_MAPS, 10)
focal_class.modifier: private static final
focal_class.original_string: private static final HadoopConfigurationProperty<Integer> DIRECT_PARALLELISM =
        new HadoopConfigurationProperty<>(MRJobConfig.NUM_MAPS, 10);
focal_class.type: HadoopConfigurationProperty<Integer>
focal_class.var_name: DIRECT_PARALLELISM
focal_class.file: bigquery/src/main/java/com/google/cloud/hadoop/io/bigquery/DirectBigQueryInputFormat.java
focal_class.identifier: DirectBigQueryInputFormat
focal_class.class_method_signature: DirectBigQueryInputFormat.getSplits(JobContext context)
focal_class.full_signature: @Override public List<InputSplit> getSplits(JobContext context)
focal_class.identifier: getSplits
focal_class.modifiers: @Override public
focal_class.parameters: (JobContext context)
focal_class.return: List<InputSplit>
focal_class.signature: List<InputSplit> getSplits(JobContext context)
focal_class.class_method_signature: DirectBigQueryInputFormat.getTable(Configuration configuration, BigQueryHelper bigQueryHelper)
focal_class.full_signature: private static Table getTable(Configuration configuration, BigQueryHelper bigQueryHelper)
focal_class.identifier: getTable
focal_class.modifiers: private static
focal_class.parameters: (Configuration configuration, BigQueryHelper bigQueryHelper)
focal_class.return: Table
focal_class.signature: Table getTable(Configuration configuration, BigQueryHelper bigQueryHelper)
focal_class.class_method_signature: DirectBigQueryInputFormat.startSession(
        Configuration configuration, Table table, BigQueryStorageClient client)
focal_class.full_signature: private static ReadSession startSession(
        Configuration configuration, Table table, BigQueryStorageClient client)
focal_class.identifier: startSession
focal_class.modifiers: private static
focal_class.parameters: (
        Configuration configuration, Table table, BigQueryStorageClient client)
focal_class.return: ReadSession
focal_class.signature: ReadSession startSession(
        Configuration configuration, Table table, BigQueryStorageClient client)
focal_class.class_method_signature: DirectBigQueryInputFormat.createRecordReader(
        InputSplit inputSplit, TaskAttemptContext context)
focal_class.full_signature: @Override public RecordReader<NullWritable, GenericRecord> createRecordReader(
        InputSplit inputSplit, TaskAttemptContext context)
focal_class.identifier: createRecordReader
focal_class.modifiers: @Override public
focal_class.parameters: (
        InputSplit inputSplit, TaskAttemptContext context)
focal_class.return: RecordReader<NullWritable, GenericRecord>
focal_class.signature: RecordReader<NullWritable, GenericRecord> createRecordReader(
        InputSplit inputSplit, TaskAttemptContext context)
focal_class.class_method_signature: DirectBigQueryInputFormat.getClient(Configuration config)
focal_class.full_signature: protected BigQueryStorageClient getClient(Configuration config)
focal_class.identifier: getClient
focal_class.modifiers: protected
focal_class.parameters: (Configuration config)
focal_class.return: BigQueryStorageClient
focal_class.signature: BigQueryStorageClient getClient(Configuration config)
focal_class.class_method_signature: DirectBigQueryInputFormat.getBigQueryHelper(Configuration config)
focal_class.full_signature: protected BigQueryHelper getBigQueryHelper(Configuration config)
focal_class.identifier: getBigQueryHelper
focal_class.modifiers: protected
focal_class.parameters: (Configuration config)
focal_class.return: BigQueryHelper
focal_class.signature: BigQueryHelper getBigQueryHelper(Configuration config)
focal_class.superclass: extends InputFormat<NullWritable, GenericRecord>
focal_method.body: @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());
    }
focal_method.class_method_signature: DirectBigQueryInputFormat.getSplits(JobContext context)
focal_method.full_signature: @Override public List<InputSplit> getSplits(JobContext context)
focal_method.identifier: getSplits
focal_method.invocations: getConfiguration getClient getBigQueryHelper get checkArgument getKey getTable startSession longValue getNumRows round getStreamsCount collect map stream getStreamsList getName getSchema getAvroSchema toList
focal_method.modifiers: @Override public
focal_method.parameters: (JobContext context)
focal_method.return: List<InputSplit>
focal_method.signature: List<InputSplit> getSplits(JobContext context)