@Override
 public void close(boolean abort) throws IOException {
   writer.close(abort);
 }
Exemplo n.º 2
0
 @Override
 public void close() throws IOException {
   if (writer != null) {
     writer.close(false);
   }
 }
  public static FileSplit createTestFile(
      String filePath,
      HiveOutputFormat<?, ?> outputFormat,
      @SuppressWarnings("deprecation") SerDe serDe,
      String compressionCodec,
      List<TestColumn> testColumns,
      int numRows)
      throws Exception {
    // filter out partition keys, which are not written to the file
    testColumns = ImmutableList.copyOf(filter(testColumns, not(TestColumn::isPartitionKey)));

    JobConf jobConf = new JobConf();
    ReaderWriterProfiler.setProfilerOptions(jobConf);

    Properties tableProperties = new Properties();
    tableProperties.setProperty(
        "columns", Joiner.on(',').join(transform(testColumns, TestColumn::getName)));
    tableProperties.setProperty(
        "columns.types", Joiner.on(',').join(transform(testColumns, TestColumn::getType)));
    serDe.initialize(new Configuration(), tableProperties);

    if (compressionCodec != null) {
      CompressionCodec codec =
          new CompressionCodecFactory(new Configuration()).getCodecByName(compressionCodec);
      jobConf.set(COMPRESS_CODEC, codec.getClass().getName());
      jobConf.set(COMPRESS_TYPE, SequenceFile.CompressionType.BLOCK.toString());
      jobConf.set("parquet.compression", compressionCodec);
      jobConf.set("parquet.enable.dictionary", "true");
    }

    RecordWriter recordWriter =
        outputFormat.getHiveRecordWriter(
            jobConf,
            new Path(filePath),
            Text.class,
            compressionCodec != null,
            tableProperties,
            new Progressable() {
              @Override
              public void progress() {}
            });

    try {
      serDe.initialize(new Configuration(), tableProperties);

      SettableStructObjectInspector objectInspector =
          getStandardStructObjectInspector(
              ImmutableList.copyOf(transform(testColumns, TestColumn::getName)),
              ImmutableList.copyOf(transform(testColumns, TestColumn::getObjectInspector)));

      Object row = objectInspector.create();

      List<StructField> fields = ImmutableList.copyOf(objectInspector.getAllStructFieldRefs());

      for (int rowNumber = 0; rowNumber < numRows; rowNumber++) {
        for (int i = 0; i < testColumns.size(); i++) {
          Object writeValue = testColumns.get(i).getWriteValue();
          if (writeValue instanceof Slice) {
            writeValue = ((Slice) writeValue).getBytes();
          }
          objectInspector.setStructFieldData(row, fields.get(i), writeValue);
        }

        Writable record = serDe.serialize(row, objectInspector);
        recordWriter.write(record);
      }
    } finally {
      recordWriter.close(false);
    }

    Path path = new Path(filePath);
    path.getFileSystem(new Configuration()).setVerifyChecksum(true);
    File file = new File(filePath);
    return new FileSplit(path, 0, file.length(), new String[0]);
  }