avro操作工具类


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;
    }
}

posted @ 2026-07-14 15:07  骑白马走三关  阅读(5)  评论(0)    收藏  举报