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.apache.beam.sdk.util.Preconditions.checkStateNotNull;
  
  import java.io.Serializable;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.stream.Collectors;
  import org.apache.beam.sdk.schemas.Schema;
  import org.apache.beam.sdk.values.Row;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
  import org.checkerframework.checker.nullness.qual.Nullable;
  
  /**
   * A utility that filters fields from Beam {@link Row}s. This filter can be configured to indicate
   * what fields you would like to either <strong>keep</strong> or <strong>drop</strong>. You may also
   * specify a singular {@link Row} field to extract with <strong>only</strong>. Afterward, call
   * {@link #filter(Row)} on a Schema-compatible Row to filter it. An un-configured filter will simply
   * return the input row untouched.
   *
   * <p>A configured {@link RowFilter} will naturally produce {@link Row}s with a new Beam {@link
   * Schema}. You can access this new Schema via the filter's {@link #outputSchema()}.
   *
   * <p>Configure a {@link RowFilter} as follows:
   *
   * <pre>{@code
   * // this is an un-configured filter
   * RowFilter unconfigured = new RowFilter(beamSchema);
   *
   * // this filter will exclusively keep these fields and drop everything else
   * List<String> fields = Arrays.asList("foo", "bar", "baz");
   * RowFilter keepFilter = new RowFilter(beamSchema).keep(fields);
   *
   * // this filter will drop these fields
   * RowFilter dropFilter = new RowFilter(beamSchema).drop(fields);
   *
   * // this filter will only output the contents of row field "my_record"
   * String field = "my_record";
   * RowFilter onlyFilter = new RowFilter(beamSchema).only(field);
   *
   * // produces a filtered row
   * Row outputRow = keepFilter.filter(row);
   * }</pre>
   *
   * Check the documentation for {@link #keep(List)}, {@link #drop(List)}, and {@link #only(String)}
   * for further details on what an output Row can look like.
   */
  public class RowFilter implements Serializable {
    private final Schema rowSchema;
    private @Nullable Schema transformedSchema;
    // for 'only' case
    private @Nullable String onlyField;
  
    public RowFilter(Schema rowSchema) {
      this.rowSchema = rowSchema;
    }
  
    /**
     * Configures this {@link RowFilter} to filter {@link Row}s by keeping only the specified fields.
     * Nested fields can be specified using dot-notation.
     *
     * <p>For example, if we want to keep the list of fields {@code ["foo", "baz"]}, for the input
     * {@link Row}:
     *
     * <pre>{@code
     * foo: 123
     * bar: 456
     * baz:
     *   nested_1: abc
     *   nested_2: xyz
     * }</pre>
     *
     * we will get the following output {@link Row}:
     *
     * <pre>{@code
     * foo: 123
     * baz
     *   nested_1: abc
     *   nested_2: xyz
     * }</pre>
     */
    public RowFilter keep(List<String> fields) {
      checkUnconfigured();
      verifyNoNestedFields(fields, "keep");
      validateSchemaContainsFields(rowSchema, fields, "keep");
      transformedSchema = keepFields(rowSchema, fields);
      return this;
    }
  
    /**
     * Configures this {@link RowFilter} to filter {@link Row} by removing the specified fields.
     * Nested fields can be specified using dot-notation.
     *
     * <p>For example, if we want to drop the list of fields {@code ["foo", "baz"]}, for this input
     * {@link Row}:
     *
     * <pre>{@code
     * foo: 123
     * bar: 456
     * baz:
     *   nested_1: abc
     *   nested_2: xyz
     * }</pre>
     *
     * we will get the following output {@link Row}:
     *
     * <pre>{@code
     * bar: 456
     * }</pre>
     */
    public RowFilter drop(List<String> fields) {
      checkUnconfigured();
      verifyNoNestedFields(fields, "drop");
      validateSchemaContainsFields(rowSchema, fields, "drop");
      transformedSchema = dropFields(rowSchema, fields);
      return this;
    }
  
    /**
     * Configures this {@link RowFilter} to only output the contents of a single row field.
     *
     * <p>For example, if we want to only extract the contents of field "foo" for this input {@link
     * Row}:
     *
     * <pre>{@code
     * abc: 123
     * bar: my_str
     * foo:
     *   xyz:
     *     baz: 456
     *     qwe: 789
     * }</pre>
     *
     * we will get the following output {@link Row}:
     *
     * <pre>{@code
     * xyz:
     *   baz: 456
     *   qwe: 789
     * }</pre>
     *
     * <p>Note that this will fail if the field is not of type {@link Row}, e.g. if {@code "abc"} is
     * specified for the example above.
     */
    public RowFilter only(String field) {
      checkUnconfigured();
      verifyNoNestedFields(Collections.singletonList(field), "only");
      validateSchemaContainsFields(rowSchema, Collections.singletonList(field), "only");
      Schema.Field rowField = rowSchema.getField(field);
      Preconditions.checkArgument(
          rowField.getType().getTypeName().equals(Schema.TypeName.ROW),
          "Expected type '%s' for field '%s', but instead got type '%s'.",
          Schema.TypeName.ROW,
          rowField.getName(),
          rowField.getType().getTypeName());
  
      transformedSchema = rowField.getType().getRowSchema();
      onlyField = field;
      return this;
    }
  
    /**
     * Performs a filter operation (keep or drop) on the input {@link Row}. Must have already
     * configured a filter operation with {@link #drop(List)} or {@link #keep(List)} for this {@link
     * RowFilter}.
     *
     * <p>If not yet configured, will simply return the same {@link Row}.
     */
    public Row filter(Row row) {
      if (transformedSchema == null) {
        return row;
      }
  
      Preconditions.checkState(
          row.getSchema().assignableTo(rowSchema),
          "Encountered Row with schema that is incompatible with this filter's schema."
              + "\nRow schema: %s"
              + "\nSchema used to initialize this filter: %s",
          row.getSchema(),
          rowSchema);
  
      // 'only' case
      if (onlyField != null) {
        return checkStateNotNull(row.getRow(onlyField));
      }
  
      // 'keep' and 'drop'
      return Preconditions.checkNotNull(copyWithNewSchema(row, outputSchema()));
    }
  
    /** Applies {@link #filter(Row)} to each Row in a list. */
    public List<Row> filter(List<Row> rows) {
      return rows.stream().map(this::filter).collect(Collectors.toList());
    }
  
    /** Returns the output {@link Row}'s {@link Schema}. */
    public Schema outputSchema() {
      return transformedSchema != null ? transformedSchema : rowSchema;
    }
  
    private void checkUnconfigured() {
      Preconditions.checkState(
          transformedSchema == null,
          "Invalid filter configuration: Please set only one of 'keep', 'drop', or 'only'.");
    }
  
    /** Verifies that this selection contains no nested fields. */
    private void verifyNoNestedFields(List<String> fields, String operation) {
      List<String> nestedFields = new ArrayList<>();
      for (String field : fields) {
        if (field.contains(".")) {
          nestedFields.add(field);
        }
      }
      if (!nestedFields.isEmpty()) {
        throw new IllegalArgumentException(
            String.format("'%s' does not support nested fields: %s", operation, nestedFields));
      }
    }
  
    /**
     * Checks whether a {@link Schema} contains a list of field names. Nested fields can be expressed
     * with dot-notation. Throws a helpful error in the case where a field doesn't exist, or if a
     * nested field could not be reached.
     */
    @VisibleForTesting
    static void validateSchemaContainsFields(
        Schema schema, List<String> specifiedFields, String operation) {
      Set<String> notFound = new HashSet<>();
      Set<String> notRowField = new HashSet<>();
  
      for (String field : specifiedFields) {
        List<String> levels = Splitter.on(".").splitToList(field);
  
        Schema currentSchema = schema;
  
        for (int i = 0; i < levels.size(); i++) {
          String currentFieldName = String.join(".", levels.subList(0, i + 1));
  
          if (!currentSchema.hasField(levels.get(i))) {
            notFound.add(currentFieldName);
            break;
          }
  
          if (i + 1 < levels.size()) {
            Schema.Field nextField = currentSchema.getField(levels.get(i));
            if (!nextField.getType().getTypeName().equals(Schema.TypeName.ROW)) {
              notRowField.add(currentFieldName);
              break;
            }
            currentSchema = Preconditions.checkNotNull(nextField.getType().getRowSchema());
          }
        }
      }
  
      if (!notFound.isEmpty() || !notRowField.isEmpty()) {
        String message = "Validation failed for '" + operation + "'.";
        if (!notFound.isEmpty()) {
          message += "\nRow Schema does not contain the following specified fields: " + notFound;
        }
        if (!notRowField.isEmpty()) {
          message +=
              "\nThe following specified fields are not of type Row. Their nested fields could not be reached: "
                  + notRowField;
        }
        throw new IllegalArgumentException(message);
      }
    }
  
    /**
     * Creates a field tree, separating each top-level field from its (potential) nested fields. E.g.
     * ["foo.bar.baz", "foo.abc", "xyz"] --> {"foo": ["bar.baz", "abc"], "xyz": []}
     */
    @VisibleForTesting
    static Map<String, List<String>> getFieldTree(List<String> fields) {
      Map<String, List<String>> fieldTree = Maps.newHashMap();
  
      for (String field : fields) {
        List<String> components = Splitter.on(".").splitToList(field);
        String root = components.get(0);
        fieldTree.computeIfAbsent(root, r -> new ArrayList<>());
  
        if (components.size() > 1) {
          String nestedFields = String.join(".", components.subList(1, components.size()));
          Preconditions.checkNotNull(fieldTree.get(root)).add(nestedFields);
        }
      }
      return fieldTree;
    }
  
    /**
     * Returns a new {@link Row} containing only the fields that intersect with the new {@link Schema}
     * Relies on a previous step to have validated the compatibility of the new {@link Schema}.
     */
    @VisibleForTesting
    @Nullable
    static Row copyWithNewSchema(@Nullable Row row, Schema newSchema) {
      if (row == null) {
        return null;
      }
      Map<String, Object> values = new HashMap<>(newSchema.getFieldCount());
  
      for (Schema.Field field : newSchema.getFields()) {
        String name = field.getName();
        Object value = row.getValue(name);
        if (field.getType().getTypeName().equals(Schema.TypeName.ROW)) {
          Schema nestedRowSchema = Preconditions.checkNotNull(field.getType().getRowSchema());
          value = copyWithNewSchema(row.getRow(name), nestedRowSchema);
        }
        if (value != null) {
          values.put(name, value);
        }
      }
      return Row.withSchema(newSchema).withFieldValues(values).build();
    }
  
    /**
     * Returns a new {@link Schema} with the specified fields removed.
     *
     * <p>No guarantee that field ordering will remain the same.
     */
    @VisibleForTesting
    static Schema dropFields(Schema schema, List<String> fieldsToDrop) {
      if (fieldsToDrop.isEmpty()) {
        return schema;
      }
      List<Schema.Field> newFieldsList = new ArrayList<>(schema.getFields());
      Map<String, List<String>> fieldTree = getFieldTree(fieldsToDrop);
  
      for (Map.Entry<String, List<String>> fieldAndDescendents : fieldTree.entrySet()) {
        String root = fieldAndDescendents.getKey();
        List<String> nestedFields = fieldAndDescendents.getValue();
        Schema.Field fieldToRemove = schema.getField(root);
        Schema.FieldType typeToRemove = fieldToRemove.getType();
  
        // Base case: we're at the specified field to remove.
        if (nestedFields.isEmpty()) {
          newFieldsList.remove(fieldToRemove);
        } else {
          // Otherwise, we're asked to remove a nested field. Verify current field is ROW type
          Preconditions.checkArgument(
              typeToRemove.getTypeName().equals(Schema.TypeName.ROW),
              "Expected type %s for specified nested field '%s', but instead got type %s.",
              Schema.TypeName.ROW,
              root,
              typeToRemove.getTypeName());
  
          Schema nestedSchema = Preconditions.checkNotNull(typeToRemove.getRowSchema());
          Schema newNestedSchema = dropFields(nestedSchema, nestedFields);
          Schema.Field modifiedField =
              Schema.Field.of(root, Schema.FieldType.row(newNestedSchema))
                  .withNullable(typeToRemove.getNullable());
  
          // Replace with modified field
          newFieldsList.set(newFieldsList.indexOf(fieldToRemove), modifiedField);
        }
      }
  
      // re-order based on original schema's ordering
      Map<String, Integer> indexMap = new HashMap<>();
      for (int i = 0; i < schema.getFieldCount(); i++) {
        indexMap.put(schema.getField(i).getName(), i);
      }
      newFieldsList.sort(Comparator.comparingInt(f -> checkStateNotNull(indexMap.get(f.getName()))));
  
      return new Schema(newFieldsList);
    }
  
    /**
     * Returns a new {@link Schema} with only the specified fields kept.
     *
     * <p>No guarantee that field ordering will remain the same.
     */
    @VisibleForTesting
    static Schema keepFields(Schema schema, List<String> fieldsToKeep) {
      if (fieldsToKeep.isEmpty()) {
        return schema;
      }
      List<Schema.Field> newFieldsList = new ArrayList<>(fieldsToKeep.size());
      Map<String, List<String>> fieldTree = getFieldTree(fieldsToKeep);
  
      for (Map.Entry<String, List<String>> fieldAndDescendents : fieldTree.entrySet()) {
        String root = fieldAndDescendents.getKey();
        List<String> nestedFields = fieldAndDescendents.getValue();
        Schema.Field fieldToKeep = schema.getField(root);
        Schema.FieldType typeToKeep = fieldToKeep.getType();
  
        // Base case: we're at the specified field to keep, and we can skip this conditional.
        // Otherwise: we're asked to keep a nested field, so we dig deeper to determine which nested
        // fields to keep
        if (!nestedFields.isEmpty()) {
          Preconditions.checkArgument(
              typeToKeep.getTypeName().equals(Schema.TypeName.ROW),
              "Expected type %s for specified nested field '%s', but instead got type %s.",
              Schema.TypeName.ROW,
              root,
              typeToKeep.getTypeName());
  
          Schema nestedSchema = Preconditions.checkNotNull(typeToKeep.getRowSchema());
          Schema newNestedSchema = keepFields(nestedSchema, nestedFields);
          fieldToKeep =
              Schema.Field.of(root, Schema.FieldType.row(newNestedSchema))
                  .withNullable(typeToKeep.getNullable());
        }
        newFieldsList.add(fieldToKeep);
      }
  
      // re-order based on original schema's ordering
      Map<String, Integer> indexMap = new HashMap<>();
      for (int i = 0; i < schema.getFieldCount(); i++) {
        indexMap.put(schema.getField(i).getName(), i);
      }
      newFieldsList.sort(Comparator.comparingInt(f -> checkStateNotNull(indexMap.get(f.getName()))));
  
      return new Schema(newFieldsList);
    }
  }
  