<commit_before>package parquet.schema;

import java.util.concurrent.Callable;
import org.junit.Assert;
import org.junit.Test;
import parquet.schema.PrimitiveType.PrimitiveTypeName;

import static parquet.schema.PrimitiveType.PrimitiveTypeName.*;
import static parquet.schema.Type.Repetition.*;

public class TestTypeBuilders {
  @Test
  public void testPaperExample() {
    MessageType expected =
        new MessageType("Document",
            new PrimitiveType(REQUIRED, INT64, "DocId"),
            new GroupType(OPTIONAL, "Links",
                new PrimitiveType(REPEATED, INT64, "Backward"),
                new PrimitiveType(REPEATED, INT64, "Forward")),
            new GroupType(REPEATED, "Name",
                new GroupType(REPEATED, "Language",
                    new PrimitiveType(REQUIRED, BINARY, "Code"),
                    new PrimitiveType(REQUIRED, BINARY, "Country")),
                new PrimitiveType(OPTIONAL, BINARY, "Url")));
    MessageType builderType = Types.buildMessage()
        .required(INT64).named("DocId")
        .optionalGroup()
            .repeated(INT64).named("Backward")
            .repeated(INT64).named("Forward")
            .named("Links")
        .repeatedGroup()
            .repeatedGroup()
                .required(BINARY).named("Code")
                .required(BINARY).named("Country")
            .named("Language")
            .optional(BINARY).named("Url")
            .named("Name")
        .named("Document");
    Assert.assertEquals(expected, builderType);
  }

  @Test
  public void testFixedLengthEquals() {
    Type f4 = Types.required(FIXED_LEN_BYTE_ARRAY).length(4).named("f4");
    Type f8 = Types.required(FIXED_LEN_BYTE_ARRAY).length(8).named("f8");
    Assert.assertFalse("Types with different lengths should not be equal",
        f4.equals(f8));
  }

  @Test
  public void testDecimalAnnotationBinary() {
    MessageType expected = new MessageType("DecimalMessage",
        new PrimitiveType(REQUIRED, BINARY, 0, "aDecimal",
            OriginalType.DECIMAL, new OriginalTypeMeta(9, 2)));
    MessageType builderType = Types.buildMessage()
        .required(BINARY).as(OriginalType.DECIMAL).precision(9).scale(2)
            .named("aDecimal")
        .named("DecimalMessage");
    Assert.assertEquals(expected, builderType);
  }

  @Test
  public void testDecimalAnnotationFixed() {
    MessageType expected = new MessageType("DecimalMessage",
        new PrimitiveType(REQUIRED, FIXED_LEN_BYTE_ARRAY, 4, "aDecimal",
            OriginalType.DECIMAL, new OriginalTypeMeta(9, 2)));
    MessageType builderType = Types.buildMessage()
        .required(FIXED_LEN_BYTE_ARRAY).length(4)
            .as(OriginalType.DECIMAL).precision(9).scale(2)
            .named("aDecimal")
        .named("DecimalMessage");
    Assert.assertEquals(expected, builderType);
  }

  @Test
  public void testDecimalAnnotationLengthCheck() {
    // maximum precision for 4 bytes is 9
    assertThrows("should reject precision 10 with length 4",
        IllegalStateException.class, new Callable<Type>() {<commit_msg>Add more tests for type builders.<commit_after>package parquet.schema;

import java.util.concurrent.Callable;
import org.junit.Assert;
import org.junit.Ignore;
import org.junit.Test;
import parquet.schema.PrimitiveType.PrimitiveTypeName;

import static parquet.schema.PrimitiveType.PrimitiveTypeName.*;
import static parquet.schema.Type.Repetition.*;

public class TestTypeBuilders {
  @Test
  public void testPaperExample() {
    MessageType expected =
        new MessageType("Document",
            new PrimitiveType(REQUIRED, INT64, "DocId"),
            new GroupType(OPTIONAL, "Links",
                new PrimitiveType(REPEATED, INT64, "Backward"),
                new PrimitiveType(REPEATED, INT64, "Forward")),
            new GroupType(REPEATED, "Name",
                new GroupType(REPEATED, "Language",
                    new PrimitiveType(REQUIRED, BINARY, "Code"),
                    new PrimitiveType(REQUIRED, BINARY, "Country")),
                new PrimitiveType(OPTIONAL, BINARY, "Url")));
    MessageType builderType = Types.buildMessage()
        .required(INT64).named("DocId")
        .optionalGroup()
            .repeated(INT64).named("Backward")
            .repeated(INT64).named("Forward")
            .named("Links")
        .repeatedGroup()
            .repeatedGroup()
                .required(BINARY).named("Code")
                .required(BINARY).named("Country")
            .named("Language")
            .optional(BINARY).named("Url")
            .named("Name")
        .named("Document");
    Assert.assertEquals(expected, builderType);
  }

  @Test
  public void testGroupTypeConstruction() {
    PrimitiveType f1 = Types.required(BINARY).as(OriginalType.UTF8).named("f1");
    PrimitiveType f2 = Types.required(INT32).named("f2");
    PrimitiveType f3 = Types.optional(INT32).named("f3");
    String name = "group";
    for (Type.Repetition repetition : Type.Repetition.values()) {
      GroupType expected = new GroupType(repetition, name,
          f1,
          new GroupType(repetition, "g1", f2, f3));
      GroupType built = Types.buildGroup()
          .repetition(repetition)
          .addField(f1)
          .group().repetition(repetition).addFields(f2, f3).named("g1")
          .named(name);
      Assert.assertEquals(expected, built);

      switch (repetition) {
        case REQUIRED:
          built = Types.requiredGroup()
              .addField(f1)
              .requiredGroup().addFields(f2, f3).named("g1")
              .named(name);
          break;
        case OPTIONAL:
          built = Types.optionalGroup()
              .addField(f1)
              .optionalGroup().addFields(f2, f3).named("g1")
              .named(name);
          break;
        case REPEATED:
          built = Types.repeatedGroup()
              .addField(f1)
              .repeatedGroup().addFields(f2, f3).named("g1")
              .named(name);
          break;
      }
      Assert.assertEquals(expected, built);
    }
  }

  @Test
  public void testPrimitiveTypeConstruction() {
    PrimitiveTypeName[] types = new PrimitiveTypeName[] {
        BOOLEAN, INT32, INT64, INT96, FLOAT, DOUBLE, BINARY
    };
    for (PrimitiveTypeName type : types) {
      String name = type.toString() + "_";
      for (Type.Repetition repetition : Type.Repetition.values()) {
        PrimitiveType expected = new PrimitiveType(repetition, type, name);
        PrimitiveType built = Types.primitive(type)
            .repetition(repetition).named(name);
        Assert.assertEquals(expected, built);
        switch (repetition) {
          case REQUIRED:
            built = Types.required(type).named(name);
            break;
          case OPTIONAL:
            built = Types.optional(type).named(name);
            break;
          case REPEATED:
            built = Types.repeated(type).named(name);
            break;
        }
        Assert.assertEquals(expected, built);
      }
    }
  }

  @Test
  public void testFixedTypeConstruction() {
    String name = "fixed_";
    int len = 5;
    for (Type.Repetition repetition : Type.Repetition.values()) {
      PrimitiveType expected = new PrimitiveType(
          repetition, FIXED_LEN_BYTE_ARRAY, len, name);
      PrimitiveType built = Types.primitive(FIXED_LEN_BYTE_ARRAY)
          .repetition(repetition).length(len).named(name);
      Assert.assertEquals(expected, built);
      switch (repetition) {
        case REQUIRED:
          built = Types.required(FIXED_LEN_BYTE_ARRAY).length(len).named(name);
          break;
        case OPTIONAL:
          built = Types.optional(FIXED_LEN_BYTE_ARRAY).length(len).named(name);
          break;
        case REPEATED:
          built = Types.repeated(FIXED_LEN_BYTE_ARRAY).length(len).named(name);
          break;
      }
      Assert.assertEquals(expected, built);
    }
  }

  @Test
  public void testEmptyGroup() {
    assertThrows("Should complain that required group is empty",
        IllegalStateException.class, new Callable<Type>() {
          @Override
          public Type call() throws Exception {
            return Types.requiredGroup().named("g");
          }
        });
    assertThrows("Should complain that optional group is empty",
        IllegalStateException.class, new Callable<Type>() {
          @Override
          public Type call() throws Exception {
            return Types.optionalGroup().named("g");
          }
        });
    assertThrows("Should complain that repeated group is empty",
        IllegalStateException.class, new Callable<Type>() {
          @Override
          public Type call() throws Exception {
            return Types.repeatedGroup().named("g");
          }
        });
  }

  @Test
  @Ignore(value="Enforcing this breaks tests in parquet-thrift")
  public void testEmptyMessage() {
    assertThrows("Should complain that message is empty",
        IllegalStateException.class, new Callable<Type>() {
          @Override
          public Type call() throws Exception {
            return Types.buildMessage().named("m");
          }
        });
  }

  @Test
  public void testMissingRepetition() {
    assertThrows("Should complain that repetition is missing",
        NullPointerException.class, new Callable<Type>() {
          @Override
          public Type call() throws Exception {
            return Types.primitive(INT32).named("num");
          }
        });
    assertThrows("Should complain that group repetition is missing",
        NullPointerException.class, new Callable<Type>() {
          @Override
          public Type call() throws Exception {
            return Types.buildGroup()
                .optional(INT32).named("num")
                .named("num");
          }
        });
  }

  @Test
  public void testFixedLengthEquals() {
    Type f4 = Types.required(FIXED_LEN_BYTE_ARRAY).length(4).named("f4");
    Type f8 = Types.required(FIXED_LEN_BYTE_ARRAY).length(8).named("f8");
    Assert.assertFalse("Types with different lengths should not be equal",
        f4.equals(f8));
  }

  @Test
  public void testDecimalAnnotationBinary() {
    MessageType expected = new MessageType("DecimalMessage",
        new PrimitiveType(REQUIRED, BINARY, 0, "aDecimal",
            OriginalType.DECIMAL, new OriginalTypeMeta(9, 2)));
    MessageType builderType = Types.buildMessage()
        .required(BINARY).as(OriginalType.DECIMAL).precision(9).scale(2)
            .named("aDecimal")
        .named("DecimalMessage");
    Assert.assertEquals(expected, builderType);
  }

  @Test
  public void testDecimalAnnotationMissingScale() {
    MessageType expected = new MessageType("DecimalMessage",
        new PrimitiveType(REQUIRED, BINARY, 0, "aDecimal",
            OriginalType.DECIMAL, new OriginalTypeMeta(9, 0)));
    MessageType builderType = Types.buildMessage()
        .required(BINARY).as(OriginalType.DECIMAL).precision(9)
            .named("aDecimal")
        .named("DecimalMessage");
    Assert.assertEquals(expected, builderType);

    expected = new MessageType("DecimalMessage",
        new PrimitiveType(REQUIRED, FIXED_LEN_BYTE_ARRAY, 7, "aDecimal",
            OriginalType.DECIMAL, new OriginalTypeMeta(9, 0)));
    builderType = Types.buildMessage()
        .required(FIXED_LEN_BYTE_ARRAY).length(7)
            .as(OriginalType.DECIMAL).precision(9)
            .named("aDecimal")
        .named("DecimalMessage");
    Assert.assertEquals(expected, builderType);
  }

  @Test
  public void testDecimalAnnotationMissingPrecision() {
    assertThrows("Should reject decimal annotation without precision",
        IllegalArgumentException.class, new Callable<Type>() {
          @Override
          public Type call() throws Exception {
            return Types.buildMessage()
                .required(BINARY).as(OriginalType.DECIMAL).scale(2)
                    .named("aDecimal")
                .named("DecimalMessage");
          }
        });
    assertThrows("Should reject decimal annotation without precision",
        IllegalArgumentException.class, new Callable<Type>() {
          @Override
          public Type call() throws Exception {
            return Types.buildMessage()
                .required(FIXED_LEN_BYTE_ARRAY).length(7)
                    .as(OriginalType.DECIMAL).scale(2)
                    .named("aDecimal")
                .named("DecimalMessage");
          }
        });
  }

  @Test
  public void testDecimalAnnotationFixed() {
    MessageType expected = new MessageType("DecimalMessage",
        new PrimitiveType(REQUIRED, FIXED_LEN_BYTE_ARRAY, 4, "aDecimal",
            OriginalType.DECIMAL, new OriginalTypeMeta(9, 2)));
    MessageType builderType = Types.buildMessage()
        .required(FIXED_LEN_BYTE_ARRAY).length(4)
            .as(OriginalType.DECIMAL).precision(9).scale(2)
            .named("aDecimal")
        .named("DecimalMessage");
    Assert.assertEquals(expected, builderType);
  }

  @Test
  public void testDecimalAnnotationLengthCheck() {
    // maximum precision for 4 bytes is 9
    assertThrows("should reject precision 10 with length 4",
        IllegalStateException.class, new Callable<Type>() {
{'fraction_non_alphanumeric': 0.11547221378682053, 'fraction_numerical': 0.009565745520801701, 'mean_word_length': 2.555465587044534, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6614479', 'n_tokens_mistral': 3695, 'n_tokens_neox': 3600, 'n_words': 737}