/*
 * 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.parquet.thrift;

import org.apache.parquet.thrift.test.compat.ListOfUnions;
import org.apache.parquet.thrift.test.compat.MapWithUnionKey;
import org.apache.parquet.thrift.test.compat.MapWithUnionValue;
import org.apache.parquet.thrift.test.compat.NestedNestedUnion;
import org.apache.parquet.thrift.test.compat.NestedUnion;
import org.apache.parquet.thrift.test.compat.OptionalInsideRequired;
import org.apache.parquet.thrift.test.compat.RequiredInsideOptional;
import org.apache.parquet.thrift.test.compat.StructWithNestedUnion;
import org.apache.parquet.thrift.test.compat.StructWithOptionalUnionOfStructs;
import org.apache.parquet.thrift.test.compat.UnionOfStructs;
import org.apache.parquet.thrift.test.compat.UnionV2;
import org.junit.Test;

import static org.apache.parquet.thrift.TestThriftSchemaConverter.shouldGetProjectedSchema;

public class TestThriftSchemaConverterProjectUnion {

  /**
   * Test projecting into a top level union
   */
  @Test
  public void testTopLevelUnions() {

    // very simple case, a union of structs where each struct has only 1 field
    shouldGetProjectedSchema("aLong/**", "aLong",
        "message ParquetSchema {\n" +
        "  optional group aString = 1 {\n" +
        "    required binary s (UTF8) = 1;\n" +
        "  }\n" +
        "  optional group aLong = 2 {\n" +
        "    required int64 l = 1;\n" +
        "  }\n" +
        "  optional group aNewBool = 3 {\n" +
        "    required boolean b = 1;\n" +
        "  }\n" +
        "}", UnionV2.class);

    // a union of structs, where each struct has more than one field
    // we should still only get one field per child here though
    shouldGetProjectedSchema("aNewBool/**", "aNewBool",
        "message ParquetSchema {\n" +
        "  optional group structV3 = 1 {\n" +
        "    required binary name (UTF8) = 1;\n" +
        "  }\n" +
        "  optional group structV4 = 2 {\n" +
        "    required binary name (UTF8) = 1;\n" +
        "  }\n" +
        "  optional group aNewBool = 3 {\n" +
        "    required boolean b = 1;\n" +
        "  }\n" +
        "}", UnionOfStructs.class);
  }

  /**
   * An optional union should be dropped if not selected.
   */
  @Test
  public void optionalUnionShouldBeDropped() {
    shouldGetProjectedSchema("name", "name",
        "message ParquetSchema {\n" +
        "  required binary name (UTF8) = 1;\n" +
        "}", StructWithOptionalUnionOfStructs.class);
  }

  /**
   * An optional union inside a required struct should still be
   * dropped if not selected.
   */
  @Test
  public void optionalUnionInRequiredStructShouldBeDropped() {
    shouldGetProjectedSchema("name", "name",
        "message ParquetSchema {\n" +
        "  required binary name (UTF8) = 1;\n" +
        "}", OptionalInsideRequired.class);
  }

  /**
   * A required union inside of an un-selected optional struct
   * should get dropped along with the optional struct.
   */
  @Test
  public void requiredUnionInsideOptionalStructShouldBeDropped() {
    // test a required union inside an optional struct
    shouldGetProjectedSchema("name", "name",
        "message ParquetSchema {\n" +
        "  required binary name (UTF8) = 1;\n" +
        "}", RequiredInsideOptional.class);
  }

  /**
   * A required union inside a selected optional struct should keep
   * one sentinel column for each "kind" of union member.
   */
  @Test
  public void requiredUnionInsideOptionalStructShouldBeKeptIfParentSelected() {
    shouldGetProjectedSchema("aStruct/name", "aStruct.name",
        "message ParquetSchema {\n" +
        "  optional group aStruct = 2 {\n" +
        "    required binary name (UTF8) = 1;\n" +
        "    required group aUnion = 2 {\n" +
        "      optional group structV3 = 1 {\n" +
        "        required binary name (UTF8) = 1;\n" +
        "      }\n" +
        "      optional group structV4 = 2 {\n" +
        "        required binary name (UTF8) = 1;\n" +
        "      }\n" +
        "      optional group aNewBool = 3 {\n" +
        "        required boolean b = 1;\n" +
        "      }\n" +
        "    }\n" +
        "  }\n" +
        "}", RequiredInsideOptional.class);
  }

  /**
   * Selecting only one "kind" of a union should trigger keeping one sentinel
   * column from the rest of the "kinds" of the union.
   */
  @Test
  public void selectingOneUnionMemberKeepsSentinels() {
    shouldGetProjectedSchema(
        "aUnion/structV4/addedStruct/gender",
        "aUnion.structV4.addedStruct.gender",
        "message ParquetSchema {\n" +
        "  optional group aUnion = 2 {\n" +
        "    optional group structV3 = 1 {\n" +
        "      required binary name (UTF8) = 1;\n" +
        "    }\n" +
        "    optional group structV4 = 2 {\n" +
        "      optional group addedStruct = 4 {\n" +
        "        optional binary gender (UTF8) = 3;\n" +
        "      }\n" +
        "    }\n" +
        "    optional group aNewBool = 3 {\n" +
        "      required boolean b = 1;\n" +
        "    }\n" +
        "  }\n" +
        "}",
        StructWithOptionalUnionOfStructs.class);
  }

  /**
   * In the case of a union inside a union (and union inside union inside union),
   * even though the fields will be "optional" because of how unions
   * are stored, we still need to grab sentinel columns from
   * the members of the union.
   */
  @Test
  public void testUnionInsideUnion() {

    shouldGetProjectedSchema("structV3/age", "structV3.age",
        "message ParquetSchema {\n" +
        "  optional group structV3 = 1 {\n" +
        "    optional binary age (UTF8) = 2;\n" +
        "  }\n" +
        "  optional group unionOfStructs = 2 {\n" +
        "    optional group structV3 = 1 {\n" +
        "      required binary name (UTF8) = 1;\n" +
        "    }\n" +
        "    optional group structV4 = 2 {\n" +
        "      required binary name (UTF8) = 1;\n" +
        "    }\n" +
        "    optional group aNewBool = 3 {\n" +
        "      required boolean b = 1;\n" +
        "    }\n" +
        "  }\n" +
        "  optional group aLong = 3 {\n" +
        "    required int64 l = 1;\n" +
        "  }\n" +
        "}", NestedUnion.class);

    shouldGetProjectedSchema(
        "unionOfStructs/structV4/addedStruct/gender",
        "unionOfStructs.structV4.addedStruct.gender",
        "message ParquetSchema {\n" +
        "  optional group structV3 = 1 {\n" +
        "    required binary name (UTF8) = 1;\n" +
        "  }\n" +
        "  optional group unionOfStructs = 2 {\n" +
        "    optional group structV3 = 1 {\n" +
        "      required binary name (UTF8) = 1;\n" +
        "    }\n" +
        "    optional group structV4 = 2 {\n" +
        "      optional group addedStruct = 4 {\n" +
        "        optional binary gender (UTF8) = 3;\n" +
        "      }\n" +
        "    }\n" +
        "    optional group aNewBool = 3 {\n" +
        "      required boolean b = 1;\n" +
        "    }\n" +
        "  }\n" +
        "  optional group aLong = 3 {\n" +
        "    required int64 l = 1;\n" +
        "  }\n" +
        "}\n", NestedUnion.class);

    shouldGetProjectedSchema(
        "unionV2/aLong/**",
        "unionV2.aLong",
        "message ParquetSchema {\n" +
        "  optional group nestedUnion = 1 {\n" +
        "    optional group structV3 = 1 {\n" +
        "      required binary name (UTF8) = 1;\n" +
        "    }\n" +
        "    optional group unionOfStructs = 2 {\n" +
        "      optional group structV3 = 1 {\n" +
        "        required binary name (UTF8) = 1;\n" +
        "      }\n" +
        "      optional group structV4 = 2 {\n" +
        "        required binary name (UTF8) = 1;\n" +
        "      }\n" +
        "      optional group aNewBool = 3 {\n" +
        "        required boolean b = 1;\n" +
        "      }\n" +
        "    }\n" +
        "    optional group aLong = 3 {\n" +
        "      required int64 l = 1;\n" +
        "    }\n" +
        "  }\n" +
        "  optional group unionV2 = 2 {\n" +
        "    optional group aString = 1 {\n" +
        "      required binary s (UTF8) = 1;\n" +
        "    }\n" +
        "    optional group aLong = 2 {\n" +
        "      required int64 l = 1;\n" +
        "    }\n" +
        "    optional group aNewBool = 3 {\n" +
        "      required boolean b = 1;\n" +
        "    }\n" +
        "  }\n" +
        "}", NestedNestedUnion.class);

  }

  @Test
  public void testListOfUnions() {

    // selecting a field from an optional list of unions should also choose a sentinel field
    // for the rest of the union members
    // at the same time, this should also drop the required list of unions, because it is safe to
    // drop even a required but repeated group
    shouldGetProjectedSchema(
      "optListUnion/structV3/age",
      "optListUnion.structV3.age",
      "message ParquetSchema {\n" +
      "  optional group optListUnion (LIST) = 1 {\n" +
      "    repeated group optListUnion_tuple {\n" +
      "      optional group structV3 = 1 {\n" +
      "        optional binary age (UTF8) = 2;\n" +
      "      }\n" +
      "      optional group structV4 = 2 {\n" +
      "        required binary name (UTF8) = 1;\n" +
      "      }\n" +
      "      optional group aNewBool = 3 {\n" +
      "        required boolean b = 1;\n" +
      "      }\n" +
      "    }\n" +
      "  }\n" +
      "}", ListOfUnions.class);

    // same goes for selecting a field from a required list of unions
    // and at the same time, the optional list of unions should be dropped too
    shouldGetProjectedSchema(
        "reqListUnion/structV3/age",
        "reqListUnion.structV3.age",
        "message ParquetSchema {\n" +
        "  required group reqListUnion (LIST) = 2 {\n" +
        "    repeated group reqListUnion_tuple {\n" +
        "      optional group structV3 = 1 {\n" +
        "        optional binary age (UTF8) = 2;\n" +
        "      }\n" +
        "      optional group structV4 = 2 {\n" +
        "        required binary name (UTF8) = 1;\n" +
        "      }\n" +
        "      optional group aNewBool = 3 {\n" +
        "        required boolean b = 1;\n" +
        "      }\n" +
        "    }\n" +
        "  }\n" +
        "}", ListOfUnions.class);

  }

  @Test
  public void testMapWithUnionKey() {
    shouldGetProjectedSchema(
        "optMapWithUnionKey/key/**",
        "optMapWithUnionKey.key",
        "message ParquetSchema {\n" +
        "  optional group optMapWithUnionKey (MAP) = 1 {\n" +
        "    repeated group key_value (MAP_KEY_VALUE) {\n" +
        "      required group key {\n" +
        "        optional group structV3 = 1 {\n" +
        "          required binary name (UTF8) = 1;\n" +
        "          optional binary age (UTF8) = 2;\n" +
        "          optional binary gender (UTF8) = 3;\n" +
        "        }\n" +
        "        optional group structV4 = 2 {\n" +
        "          required binary name (UTF8) = 1;\n" +
        "          optional binary age (UTF8) = 2;\n" +
        "          optional binary gender (UTF8) = 3;\n" +
        "          optional group addedStruct = 4 {\n" +
        "            required binary name (UTF8) = 1;\n" +
        "            optional binary age (UTF8) = 2;\n" +
        "            optional binary gender (UTF8) = 3;\n" +
        "          }\n" +
        "        }\n" +
        "        optional group aNewBool = 3 {\n" +
        "          required boolean b = 1;\n" +
        "        }\n" +
        "      }\n" +
        "      optional group value {\n" +
        "        required binary name (UTF8) = 1;\n" +
        "      } " +
        "    }\n" +
        "  }\n" +
        "}", MapWithUnionKey.class);

    shouldGetProjectedSchema(
        "optMapWithUnionKey/key/**;optMapWithUnionKey/value/gender",
        "optMapWithUnionKey.{key,value.gender}",
        "message ParquetSchema {\n" +
        "  optional group optMapWithUnionKey (MAP) = 1 {\n" +
        "    repeated group key_value (MAP_KEY_VALUE) {\n" +
        "      required group key {\n" +
        "        optional group structV3 = 1 {\n" +
        "          required binary name (UTF8) = 1;\n" +
        "          optional binary age (UTF8) = 2;\n" +
        "          optional binary gender (UTF8) = 3;\n" +
        "        }\n" +
        "        optional group structV4 = 2 {\n" +
        "          required binary name (UTF8) = 1;\n" +
        "          optional binary age (UTF8) = 2;\n" +
        "          optional binary gender (UTF8) = 3;\n" +
        "          optional group addedStruct = 4 {\n" +
        "            required binary name (UTF8) = 1;\n" +
        "            optional binary age (UTF8) = 2;\n" +
        "            optional binary gender (UTF8) = 3;\n" +
        "          }\n" +
        "        }\n" +
        "        optional group aNewBool = 3 {\n" +
        "          required boolean b = 1;\n" +
        "        }\n" +
        "      }\n" +
        "      optional group value {\n" +
        "        optional binary gender (UTF8) = 3;\n" +
        "      }\n" +
        "    }\n" +
        "  }\n" +
        "}", MapWithUnionKey.class);
  }

  @Test
  public void testMapWithUnionValue() {
    shouldGetProjectedSchema(
        "optMapWithUnionValue/key/**;optMapWithUnionValue/value/structV4/addedStruct/gender",
        "optMapWithUnionValue.{key,value.structV4.addedStruct.gender}",
        "message ParquetSchema {\n" +
        "  optional group optMapWithUnionValue (MAP) = 1 {\n" +
        "    repeated group key_value (MAP_KEY_VALUE) {\n" +
        "      required group key {\n" +
        "        required binary name (UTF8) = 1;\n" +
        "        optional binary age (UTF8) = 2;\n" +
        "        optional binary gender (UTF8) = 3;\n" +
        "      }\n" +
        "      optional group value {\n" +
        "        optional group structV3 = 1 {\n" +
        "          required binary name (UTF8) = 1;\n" +
        "        }\n" +
        "        optional group structV4 = 2 {\n" +
        "          optional group addedStruct = 4 {\n" +
        "            optional binary gender (UTF8) = 3;\n" +
        "          }\n" +
        "        }\n" +
        "        optional group aNewBool = 3 {\n" +
        "          required boolean b = 1;\n" +
        "        }\n" +
        "      }\n" +
        "    }\n" +
        "  }\n" +
        "}", MapWithUnionValue.class);
  }

  /**
   * Tests a complicated struct that contains required, optional, and unspecified
   * members of a few different corner cases.
   */
  @Test
  public void testMessyNestedUnions() {
    shouldGetProjectedSchema("reqStructWithUnionV2/name", "reqStructWithUnionV2.name",
        "message ParquetSchema {\n" +
            "  required group reqUnionOfStructs = 2 {\n" +
            "    optional group structV3 = 1 {\n" +
            "      required binary name (UTF8) = 1;\n" +
            "    }\n" +
            "    optional group structV4 = 2 {\n" +
            "      required binary name (UTF8) = 1;\n" +
            "    }\n" +
            "    optional group aNewBool = 3 {\n" +
            "      required boolean b = 1;\n" +
            "    }\n" +
            "  }\n" +
            "  required group reqNestedUnion = 5 {\n" +
            "    optional group structV3 = 1 {\n" +
            "      required binary name (UTF8) = 1;\n" +
            "    }\n" +
            "    optional group unionOfStructs = 2 {\n" +
            "      optional group structV3 = 1 {\n" +
            "        required binary name (UTF8) = 1;\n" +
            "      }\n" +
            "      optional group structV4 = 2 {\n" +
            "        required binary name (UTF8) = 1;\n" +
            "      }\n" +
            "      optional group aNewBool = 3 {\n" +
            "        required boolean b = 1;\n" +
            "      }\n" +
            "    }\n" +
            "    optional group aLong = 3 {\n" +
            "      required int64 l = 1;\n" +
            "    }\n" +
            "  }\n" +
            "  required group reqStructWithUnionV2 = 8 {\n" +
            "    required binary name (UTF8) = 1;\n" +
            "    required group aUnion = 2 {\n" +
            "      optional group aString = 1 {\n" +
            "        required binary s (UTF8) = 1;\n" +
            "      }\n" +
            "      optional group aLong = 2 {\n" +
            "        required int64 l = 1;\n" +
            "      }\n" +
            "      optional group aNewBool = 3 {\n" +
            "        required boolean b = 1;\n" +
            "      }\n" +
            "    }\n" +
            "  }\n" +
            "  required group reqUnionStructUnion = 11 {\n" +
            "    optional group structV3 = 1 {\n" +
            "      required binary name (UTF8) = 1;\n" +
            "    }\n" +
            "    optional group structWithUnionOfStructs = 2 {\n" +
            "      required binary name (UTF8) = 1;\n" +
            "      required group aUnion = 2 {\n" +
            "        optional group structV3 = 1 {\n" +
            "          required binary name (UTF8) = 1;\n" +
            "        }\n" +
            "        optional group structV4 = 2 {\n" +
            "          required binary name (UTF8) = 1;\n" +
            "        }\n" +
            "        optional group aNewBool = 3 {\n" +
            "          required boolean b = 1;\n" +
            "        }\n" +
            "      }\n" +
            "    }\n" +
            "    optional group aLong = 3 {\n" +
            "      required int64 l = 1;\n" +
            "    }\n" +
            "  }\n" +
            "}", StructWithNestedUnion.class);
  }
}
