gobblin-modules/gobblin-parquet/src/test/java/org/apache/gobblin/writer/TestConstants.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.gobblin.writer;

import parquet.example.data.Group;
import parquet.example.data.simple.SimpleGroup;
import parquet.schema.MessageType;
import parquet.schema.OriginalType;
import parquet.schema.PrimitiveType;
import parquet.schema.Types;

import org.apache.gobblin.parquet.writer.test.TestConstantsBase;
import org.apache.gobblin.test.TestRecord;


public class TestConstants extends TestConstantsBase<Group> {

  public static final MessageType PARQUET_SCHEMA = Types.buildMessage()
      .addFields(
          Types.required(PrimitiveType.PrimitiveTypeName.BINARY).as(OriginalType.UTF8)
              .named(TestConstants.PAYLOAD_FIELD_NAME),
          Types.required(PrimitiveType.PrimitiveTypeName.INT32).named(TestConstants.PARTITION_FIELD_NAME),
          // Sequence field is INT32 instead of INT64, because this version of parquet only supports INT32
          Types.required(PrimitiveType.PrimitiveTypeName.INT32).named(TestConstants.SEQUENCE_FIELD_NAME))
      .named("Data");

  @Override
  public Group convertToParquetGroup(TestRecord record) {
    Group group = new SimpleGroup(PARQUET_SCHEMA);
    group.add(PAYLOAD_FIELD_NAME, record.getPayload());
    group.add(SEQUENCE_FIELD_NAME, Long.valueOf(record.getSequence()).intValue());
    group.add(PARTITION_FIELD_NAME, record.getPartition());
    return group;
  }

}