sjwiesman / flink-parquet-writer Goto Github PK

testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
testImplementation group: 'junit', name: 'junit', version: '4.12'
testImplementation group: 'org.junit.jupiter', name: 'junit-jupiter-engine', version: '5.6.2'
testImplementation group: 'org.junit.jupiter', name: 'junit-jupiter-api', version: '5.6.2'
testImplementation group: 'org.slf4j', name: 'slf4j-simple', version: '1.7.36'

private static Schema formSubmitSchema = new Schema.Parser().parse(formSubmitSchemaString);

private static String formSubmitSchemaString = "{\n" +
            "  \"type\": \"record\",\n" +
            "  \"name\": \"FormSubmitEvent\",\n" +
            "  \"fields\": [\n" +
            "    { \"name\": \"event\", \"type\": [\"null\", \"string\"] },\n" +
            "    { \"name\": \"tenantId\", \"type\": [\"null\", \"string\"] },\n" +
            "    { \"name\": \"formData\", \"type\": [\"null\", \"string\"] },\n" +
            "    { \"name\": \"formID\", \"type\": [\"null\", \"string\"] },\n" +
            "    { \"name\": \"startTime\", \"type\": [\"null\", \"string\"] },\n" +
            "    { \"name\": \"year\", \"type\": [\"null\", \"int\"] },\n" +
            "    { \"name\": \"month\", \"type\": [\"null\", \"int\"] },\n" +
            "    { \"name\": \"day\", \"type\": [\"null\", \"int\"] }\n" +
            "  ]\n" +
            "}\n";

private static Logger logger = LoggerFactory.getLogger(App.class);


@SuppressWarnings("deprecation")
public static void main(String[] args) {
    try {
        final StreamExecutionEnvironment see = StreamExecutionEnvironment.getExecutionEnvironment();
        Properties properties = new Properties();
        properties.setProperty("bootstrap.servers", "localhost:9092");
        properties.setProperty("auto.offset.reset", "earliest");
        
        Map<String, Object> configs = new HashMap<>();
        configs.put(AWSSchemaRegistryConstants.AWS_REGION, "ap-south-1");
        configs.put(AWSSchemaRegistryConstants.SCHEMA_AUTO_REGISTRATION_SETTING, true);
        configs.put(AWSSchemaRegistryConstants.REGISTRY_NAME, "KafkaSchemaRegistry");
        configs.put(AWSSchemaRegistryConstants.SCHEMA_NAME, "testSchema");
        configs.put(AWSSchemaRegistryConstants.AVRO_RECORD_TYPE, AvroRecordType.GENERIC_RECORD.getName());
        
        FlinkKafkaConsumer<GenericRecord> consumer = new FlinkKafkaConsumer<>(
            "testTopic1",
            GlueSchemaRegistryAvroDeserializationSchema.forGeneric(formSubmitSchema, configs),
            properties);
        
        DataStream<GenericRecord> stream = see.addSource(consumer);

        Path outputBasePath = new Path("s3://qa-nsl-ge-data/topics/");

        final StreamingFileSink<GenericRecord> sink = StreamingFileSink
            .forBulkFormat(outputBasePath, ParquetAvroWriters.forGenericRecord(formSubmitSchema))
            .withBucketAssigner(new BucketAssigner<GenericRecord, String>() {
                
                @Override
                public String getBucketId(GenericRecord record, Context context) {
                    int year = (int) record.get("year");
                    int month = (int) record.get("month");
                    int day = (int) record.get("day");
                    logger.info("year: {}, month: {}, day: {}", year, month, day);
                    return String.format("year=%d/month=%d/day=%d", year, month, day);
                }

                @Override
                public SimpleVersionedSerializer<String> getSerializer() {
                    return SimpleVersionedStringSerializer.INSTANCE;
                }
            })
            .build();
        
        stream.print();
        logger.info("flink job for tagmanager is started");

        stream.addSink(sink);
        
        see.execute();
    } catch (Exception e) {
        System.out.println(e.getMessage());
        e.printStackTrace();
    }
}