content: /*
   * 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.beam.sdk.util;
  
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.containsString;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertThrows;
  import static org.junit.Assert.assertTrue;
  
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import org.apache.beam.sdk.schemas.Schema;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.Row;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.rules.ExpectedException;
  
  /** Tests for {@link RowFilter}. */
  public class RowFilterTest {
    @Rule public ExpectedException thrown = ExpectedException.none();
  
    private static final Schema DOUBLY_NESTED_ROW_SCHEMA =
        Schema.builder()
            .addStringField("doubly_nested_str")
            .addInt32Field("doubly_nested_int")
            .build();
  
    private static final Schema NESTED_ROW_SCHEMA =
        Schema.builder()
            .addStringField("nested_str")
            .addInt32Field("nested_int")
            .addFloatField("nested_float")
            .addRowField("nested_row", DOUBLY_NESTED_ROW_SCHEMA)
            .build();
    private static final Schema ROW_SCHEMA =
        Schema.builder()
            .addStringField("str")
            .addBooleanField("bool")
            .addNullableInt32Field("nullable_int")
            .addArrayField("arr_int", Schema.FieldType.INT32)
            .addRowField("row", NESTED_ROW_SCHEMA)
            .addNullableRowField("nullable_row", NESTED_ROW_SCHEMA)
            .build();
  
    @Test
    public void testSchemaValidation() {
      List<List<String>> goodFields =
          Arrays.asList(
              Arrays.asList("str", "bool", "nullable_row"),
              Arrays.asList("nullable_int", "arr_int"),
              Arrays.asList("row.nested_str", "row.nested_row.doubly_nested_str"),
              Arrays.asList("nullable_row.nested_row.doubly_nested_int"));
  
      for (List<String> fields : goodFields) {
        RowFilter.validateSchemaContainsFields(ROW_SCHEMA, fields, "test-operation");
      }
    }
  
    @Test
    public void testSchemaValidationFailsWithHelpfulErrorForMissingFields() {
      List<KV<List<String>, List<String>>> nonExistentFields =
          Arrays.asList(
              KV.of(
                  Arrays.asList("nonexistent_1", "nonexistent_2", "nonexistent_3"),
                  Arrays.asList("nonexistent_1", "nonexistent_2", "nonexistent_3")),
              KV.of(
                  Arrays.asList("nullable_int", "arr_int", "nonexistent"),
                  Collections.singletonList("nonexistent")),
              KV.of(
                  Arrays.asList(
                      "nullable_row.nested_row.nonexistent", "row.nonexistent", "row.nested_float"),
                  Arrays.asList("nullable_row.nested_row.nonexistent", "row.nonexistent")));
  
      for (KV<List<String>, List<String>> fields : nonExistentFields) {
        List<String> allFields = fields.getKey();
        List<String> badFields = fields.getValue();
  
        IllegalArgumentException e =
            assertThrows(
                IllegalArgumentException.class,
                () ->
                    RowFilter.validateSchemaContainsFields(ROW_SCHEMA, allFields, "test-operation"));
  
        assertThat(e.getMessage(), containsString("Validation failed for 'test-operation'"));
        assertThat(
            e.getMessage(),
            containsString("Row Schema does not contain the following specified fields"));
        for (String badField : badFields) {
          assertThat(e.getMessage(), containsString(badField));
        }
      }
    }
  
    @Test
    public void testSchemaValidationFailsWithHelpfulErrorForInvalidNestedFields() {
      List<KV<List<String>, List<String>>> nonNestedFields =
          Arrays.asList(
              KV.of(
                  Arrays.asList(
                      "row.nested_row", "row.nested_int", "row.nested_str.unexpected_nested"),
                  Collections.singletonList("row.nested_str")),
              KV.of(
                  Arrays.asList(
                      "nullable_row.nested_str",
                      "nullable_row.nested_str.unexpected",
                      "row.nested_int.unexpected_2"),
                  Arrays.asList("nullable_row.nested_str", "row.nested_int")));
  
      for (KV<List<String>, List<String>> fields : nonNestedFields) {
        List<String> allFields = fields.getKey();
        List<String> badFields = fields.getValue();
  
        IllegalArgumentException e =
            assertThrows(
                IllegalArgumentException.class,
                () ->
                    RowFilter.validateSchemaContainsFields(ROW_SCHEMA, allFields, "test-operation"));
  
        assertThat(e.getMessage(), containsString("Validation failed for 'test-operation'"));
        assertThat(
            e.getMessage(),
            containsString(
                "The following specified fields are not of type Row. Their nested fields could not be reached"));
        for (String badField : badFields) {
          assertThat(e.getMessage(), containsString(badField));
        }
      }
    }
  
    @Test
    public void testGetFieldTree() {
      List<String> fields =
          Arrays.asList(
              "top-level",
              "top-level-2",
              "top-level.nested-level",
              "top-level.nested-level-2",
              "top-level.nested-level.doubly-nested-level",
              "top-level.nested-level.doubly-nested-level-2");
      List<String> nestedLayer =
          Arrays.asList(
              "nested-level",
              "nested-level-2",
              "nested-level.doubly-nested-level",
              "nested-level.doubly-nested-level-2");
  
      Map<String, List<String>> expectedTree =
          ImmutableMap.<String, List<String>>builder()
              .put("top-level-2", Collections.emptyList())
              .put("top-level", nestedLayer)
              .build();
  
      assertEquals(expectedTree, RowFilter.getFieldTree(fields));
  
      List<String> doublyNestedLayer = Arrays.asList("doubly-nested-level", "doubly-nested-level-2");
  
      Map<String, List<String>> expectedNestedTree =
          ImmutableMap.<String, List<String>>builder()
              .put("nested-level-2", Collections.emptyList())
              .put("nested-level", doublyNestedLayer)
              .build();
  
      assertEquals(expectedNestedTree, RowFilter.getFieldTree(nestedLayer));
    }
  
    @Test
    public void testDropSchemaFields() {
      List<String> fieldsToDrop =
          Arrays.asList(
              "str",
              "arr_int",
              "nullable_int",
              "row.nested_int",
              "row.nested_float",
              "row.nested_row.doubly_nested_int",
              "nullable_row.nested_str",
              "nullable_row.nested_row");
  
      Schema expectedDroppedSchema =
          Schema.builder()
              .addBooleanField("bool")
              .addRowField(
                  "row",
                  Schema.builder()
                      .addStringField("nested_str")
                      .addRowField(
                          "nested_row", Schema.builder().addStringField("doubly_nested_str").build())
                      .build())
              .addNullableRowField(
                  "nullable_row",
                  Schema.builder().addInt32Field("nested_int").addFloatField("nested_float").build())
              .build();
  
      assertTrue(expectedDroppedSchema.equivalent(RowFilter.dropFields(ROW_SCHEMA, fieldsToDrop)));
    }
  
    @Test
    public void testKeepSchemaFields() {
      List<String> fieldsToKeep =
          Arrays.asList(
              "str",
              "arr_int",
              "nullable_int",
              "row.nested_int",
              "row.nested_float",
              "row.nested_row.doubly_nested_int",
              "nullable_row.nested_str",
              "nullable_row.nested_row");
  
      Schema expectedKeptSchema =
          Schema.builder()
              .addStringField("str")
              .addArrayField("arr_int", Schema.FieldType.INT32)
              .addNullableInt32Field("nullable_int")
              .addRowField(
                  "row",
                  Schema.builder()
                      .addInt32Field("nested_int")
                      .addFloatField("nested_float")
                      .addRowField(
                          "nested_row", Schema.builder().addInt32Field("doubly_nested_int").build())
                      .build())
              .addNullableRowField(
                  "nullable_row",
                  Schema.builder()
                      .addStringField("nested_str")
                      .addRowField("nested_row", DOUBLY_NESTED_ROW_SCHEMA)
                      .build())
              .build();
  
      assertTrue(expectedKeptSchema.equivalent(RowFilter.keepFields(ROW_SCHEMA, fieldsToKeep)));
    }
  
    @Test
    public void testDropNestedFieldsFails() {
      thrown.expect(IllegalArgumentException.class);
      thrown.expectMessage("'drop' does not support nested fields");
  
      new RowFilter(ROW_SCHEMA)
          .drop(
              Arrays.asList(
                  "bool",
                  "nullable_int",
                  "row.nested_int",
                  "row.nested_float",
                  "row.nested_row.doubly_nested_int",
                  "nullable_row"));
    }
  
    @Test
    public void testKeepNestedFieldsFails() {
      thrown.expect(IllegalArgumentException.class);
      thrown.expectMessage("'keep' does not support nested fields");
  
      new RowFilter(ROW_SCHEMA)
          .keep(
              Arrays.asList("str", "arr_int", "row.nested_str", "row.nested_row.doubly_nested_str"));
    }
  
    @Test
    public void testOnlyFailsWhenSpecifyingNonRowField() {
      thrown.expect(IllegalArgumentException.class);
      thrown.expectMessage(
          "Expected type 'ROW' for field 'nullable_int', but instead got type 'INT32'");
  
      new RowFilter(ROW_SCHEMA).only("nullable_int");
    }
  
    private static final Row ORIGINAL_ROW =
        Row.withSchema(ROW_SCHEMA)
            .addValue("str_value")
            .addValue(true)
            .addValue(123)
            .addValue(Arrays.asList(1, 2, 3, 4, 5))
            .addValue(
                Row.withSchema(NESTED_ROW_SCHEMA)
                    .addValue("nested_str_value")
                    .addValue(456)
                    .addValue(1.234f)
                    .addValue(
                        Row.withSchema(DOUBLY_NESTED_ROW_SCHEMA)
                            .addValue("doubly_nested_str_value")
                            .addValue(789)
                            .build())
                    .build())
            .addValue(null)
            .build();
  
    private static final Schema FILTERED_DOUBLY_NESTED_SCHEMA =
        Schema.builder().addStringField("doubly_nested_str").build();
    private static final Schema FILTERED_NESTED_SCHEMA =
        Schema.builder()
            .addStringField("nested_str")
            .addRowField("nested_row", FILTERED_DOUBLY_NESTED_SCHEMA)
            .build();
    private static final Schema FILTERED_SCHEMA =
        Schema.builder()
            .addStringField("str")
            .addArrayField("arr_int", Schema.FieldType.INT32)
            .addRowField("row", FILTERED_NESTED_SCHEMA)
            .build();
  
    private static final Row FILTERED_ROW =
        Row.withSchema(FILTERED_SCHEMA)
            .addValue("str_value")
            .addValue(Arrays.asList(1, 2, 3, 4, 5))
            .addValue(
                Row.withSchema(FILTERED_NESTED_SCHEMA)
                    .addValue("nested_str_value")
                    .addValue(
                        Row.withSchema(FILTERED_DOUBLY_NESTED_SCHEMA)
                            .addValue("doubly_nested_str_value")
                            .build())
                    .build())
            .build();
  
    @Test
    public void testCopyRowWithNewSchema() {
      assertEquals(FILTERED_ROW, RowFilter.copyWithNewSchema(ORIGINAL_ROW, FILTERED_SCHEMA));
    }
  
    @Test
    public void testOnlyRowField() {
      RowFilter rowFilter = new RowFilter(ROW_SCHEMA).only("row");
  
      Row expecedRow =
          Row.withSchema(rowFilter.outputSchema())
              .addValues(ORIGINAL_ROW.getRow("row").getValues())
              .build();
  
      assertEquals(expecedRow, rowFilter.filter(ORIGINAL_ROW));
    }
  }
  