@Override
public void close(boolean abort) throws IOException {
writer.close(abort);
}
@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]);
}
