import org.apache.avro.Schema;
import org.apache.avro.file.DataFileReader;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.file.SeekableByteArrayInput;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericDatumReader;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.io.Decoder;
import org.apache.avro.io.DecoderFactory;
import org.apache.avro.util.Utf8;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
public class AvroUtils {
// 写多条记录的完整Avro文件
public static byte[] writeMultipleWithSchema(Iterable<GenericRecord> records, Schema schema)
throws IOException {
ByteArrayOutputStream baos = new ByteArrayOutputStream();
try (DataFileWriter<GenericRecord> writer = new DataFileWriter<>(
new GenericDatumWriter<>(schema))) {
writer.create(schema, baos);
for (GenericRecord record : records) {
writer.append(record);
}
}
return baos.toByteArray();
}
/**
* 读取第一条记录
*
* @param data 带schema数据
* @return GenericRecord
* @throws IOException IOException
*/
public static GenericRecord readFirstRecord(byte[] data) throws IOException {
try (SeekableByteArrayInput bais = new SeekableByteArrayInput(data);
DataFileReader<GenericRecord> reader = new DataFileReader<>(bais,
new GenericDatumReader<>()
)) {
if (reader.hasNext()) {
return reader.next();
}
return null;
}
}
/**
* 读取所有记录返回List
*
* @param data 带schema数据
* @return List<GenericRecord>
* @throws IOException IOException
*/
public static List<GenericRecord> readAllRecords(byte[] data) throws IOException {
List<GenericRecord> records = new ArrayList<>();
try (SeekableByteArrayInput bais = new SeekableByteArrayInput(data);
DataFileReader<GenericRecord> reader = new DataFileReader<>(bais,
new GenericDatumReader<>()
)) {
while (reader.hasNext()) {
records.add(reader.next());
}
}
return records;
}
public static GenericRecord deserializeAvro(byte[] data, Schema schema) throws IOException {
GenericDatumReader<GenericRecord> reader = new GenericDatumReader<>(schema);
Decoder decoder = DecoderFactory.get().binaryDecoder(data, null);
return reader.read(null, decoder);
}
public static Map<String, Object> avro2map(GenericRecord record) {
Map<String, Object> res = new LinkedHashMap<>();
if (record == null) {
return res;
}
for (Schema.Field field : record.getSchema().getFields()) {
Object value = record.get(field.pos());
res.put(field.name(), unwrapAvro(value));
}
return res;
}
public static Object unwrapAvro(Object value) {
if (value == null) {
return null;
}
// Record
if (value instanceof GenericRecord) {
return avro2map((GenericRecord) value);
}
// Array
if (value instanceof GenericData.Array) {
GenericData.Array<?> array = (GenericData.Array<?>) value;
List<Object> list = new ArrayList<>(array.size());
for (Object o : array) {
list.add(unwrapAvro(o));
}
return list;
}
// String
if (value instanceof Utf8) {
return value.toString();
}
// Bytes
if (value instanceof ByteBuffer) {
ByteBuffer bb = (ByteBuffer) value;
byte[] bytes = new byte[bb.remaining()];
bb.get(bytes);
return bytes;
}
// primitive / Long / Integer / Boolean / Double ...
return value;
}
}