Introduce AttributeMap interface

polaris-core/src/main/java/org/apache/polaris/core/collection/AbstractAttributeMap.java

@@ -0,0 +1,185 @@
+/*
+ * 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.polaris.core.collection;
+
+import com.google.common.collect.Iterators;
+import java.util.AbstractSet;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import org.jspecify.annotations.NonNull;
+import org.jspecify.annotations.Nullable;
+
+/** Base implementation of {@link AttributeMap} backed by a {@link HashMap}. */
+public abstract sealed class AbstractAttributeMap implements AttributeMap
+    permits MutableAttributeMap, ImmutableAttributeMap {
+
+  protected abstract static class Builder<M extends AttributeMap, B extends Builder<M, B>> {
+
+    protected final MutableAttributeMap map = new MutableAttributeMap();
+
+    /** Inserts the given {@link Attribute} in the map. */
+    public <T> B put(Attribute<T> entry) {
+      return put(entry.key(), entry.value());
+    }
+
+    /** Associates {@code key} with {@code value}, replacing any existing mapping. */
+    public <T> B put(AttributeKey<T> key, @Nullable T value) {
+      map.put(key, value);
+      return self();
+    }
+
+    /** Puts all attributes from the given map. */
+    public B putAll(AttributeMap other) {
+      map.putAll(other);
+      return self();
+    }
+
+    protected abstract B self();
+
+    /** Builds and returns the attribute map. */
+    public abstract M build();
+  }
+
+  private final Map<AttributeKey<?>, Object> delegate = new HashMap<>();
+
+  protected AbstractAttributeMap() {}
+
+  protected AbstractAttributeMap(AttributeMap toCopy) {
+    toCopy.entrySet().forEach(entry -> delegate.put(entry.key(), entry.value()));
+  }
+
+  protected final Map<AttributeKey<?>, Object> delegate() {
+    return delegate;
+  }
+
+  @Override
+  public int size() {
+    return delegate.size();
+  }
+
+  @Override
+  public boolean isEmpty() {
+    return delegate.isEmpty();
+  }
+
+  @Override
+  @Nullable
+  public <T> T get(AttributeKey<T> key) {
+    return cast(delegate.get(key));
+  }
+
+  @Override
+  public @Nullable <V> V getOrDefault(AttributeKey<V> key, @Nullable V defaultValue) {
+    return cast(delegate.getOrDefault(key, defaultValue));
+  }
+
+  @Override
+  public <T> T put(AttributeKey<T> key, @Nullable T value) {
+    return cast(delegate.put(key, value));
+  }
+
+  @Nullable
+  @Override
+  public <T> T remove(AttributeKey<T> key) {
+    return cast(delegate.remove(key));
+  }
+
+  @Override
+  public void putAll(AttributeMap map) {
+    map.entrySet().forEach(entry -> delegate.put(entry.key(), entry.value()));
+  }
+
+  @Override
+  public boolean containsKey(AttributeKey<?> key) {
+    return delegate.containsKey(key);
+  }
+
+  @Override
+  public boolean containsValue(@Nullable Object value) {
+    return delegate.containsValue(value);
+  }
+
+  @Override
+  public void clear() {
+    delegate.clear();
+  }
+
+  @Override
+  @NonNull
+  public Set<AttributeKey<?>> keySet() {
+    return delegate.keySet();
+  }
+
+  @Override
+  @NonNull
+  public Collection<Object> values() {
+    return delegate.values();
+  }
+
+  @Override
+  @NonNull
+  public Set<Attribute<?>> entrySet() {
+    return new AttributeSet();
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (!(o instanceof AbstractAttributeMap that)) {
+      return false;
+    }
+    return Objects.equals(this.delegate, that.delegate);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(delegate);
+  }
+
+  @Override
+  public String toString() {
+    // Do not print attribute values!
+    return this.getClass().getSimpleName() + "{keys=" + delegate.keySet() + '}';
+  }
+
+  private class AttributeSet extends AbstractSet<Attribute<?>> {
+
+    @Override
+    public int size() {
+      return delegate.size();
+    }
+
+    @Override
+    @NonNull
+    public Iterator<Attribute<?>> iterator() {
+      return Iterators.transform(
+          delegate.entrySet().iterator(),
+          key -> new Attribute<>(cast(key.getKey()), key.getValue()));
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  private static <T> T cast(Object o) {
+    return (T) o;
+  }
+}

polaris-core/src/main/java/org/apache/polaris/core/collection/AttributeMap.java

@@ -0,0 +1,207 @@
+/*
+ * 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.polaris.core.collection;
+
+import java.util.Collection;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.function.Consumer;
+import org.jspecify.annotations.NonNull;
+import org.jspecify.annotations.Nullable;
+
+/**
+ * A map-like structure containing typed attributes, keyed by {@link AttributeKey}.
+ *
+ * <p>Unlike a raw {@code Map<String, Object>}, attribute maps guarantee that, for any attribute
+ * key, the corresponding value has the right type, eliminating unchecked casts in callers.
+ *
+ * <p>Note: this guarantee is only valid as long as callers never create two instances of {@link
+ * AttributeKey} having the same string key, but different types.
+ *
+ * <p>Two implementations are available: {@link MutableAttributeMap} and {@link
+ * ImmutableAttributeMap}.
+ */
+public sealed interface AttributeMap
+    permits AbstractAttributeMap, ImmutableAttributeMap, MutableAttributeMap {
+
+  /** A shared empty, immutable instance. */
+  ImmutableAttributeMap EMPTY = new ImmutableAttributeMap();
+
+  /** Returns an immutable copy of {@code map}. */
+  static ImmutableAttributeMap copyOf(AttributeMap map) {
+    return map instanceof ImmutableAttributeMap immutable
+        ? immutable
+        : new ImmutableAttributeMap(map);
+  }
+
+  /** Returns an immutable map containing the given attributes. */
+  static ImmutableAttributeMap of(Attribute<?> attribute, Attribute<?>... attributes) {
+    ImmutableAttributeMap.Builder builder = ImmutableAttributeMap.builder().put(attribute);
+    for (Attribute<?> attr : attributes) {
+      builder.put(attr);
+    }
+    return builder.build();
+  }
+
+  /**
+   * A typed attribute key.
+   *
+   * @param <T> the type of the value associated with this key
+   * @param key the string identifier for this attribute key.
+   */
+  @SuppressWarnings({"UnusedTypeParameter", "unused"})
+  record AttributeKey<T>(String key) {
+
+    @Override
+    @NonNull
+    public String toString() {
+      return key;
+    }
+  }
+
+  /**
+   * An entry associating an {@link AttributeKey} with a value of its type. Used to iterate over
+   * attributes in the map.
+   *
+   * @see #entrySet()
+   * @see #forEach(Consumer)
+   * @see #put(Attribute)
+   */
+  record Attribute<T>(AttributeKey<T> key, T value) {
+
+    @Override
+    @NonNull
+    public String toString() {
+      // Do not print attribute values!
+      return "Attribute{key=" + key + "}";
+    }
+  }
+
+  /** Returns the number of key-value mappings in this map. */
+  int size();
+
+  /** Returns {@code true} if this map contains no key-value mappings. */
+  boolean isEmpty();
+
+  /** Returns the value associated with {@code key}, or {@code null} if absent. */
+  @Nullable <T> T get(AttributeKey<T> key);
+
+  /**
+   * Returns the value for {@code key} if present, otherwise {@code defaultValue}.
+   *
+   * <p>Unlike {@link #get}, this distinguishes between a missing key and a key explicitly mapped to
+   * {@code null}.
+   */
+  @Nullable <V> V getOrDefault(AttributeKey<V> key, @Nullable V defaultValue);
+
+  /**
+   * Returns the value associated with {@code key} wrapped in an {@link Optional}, or an empty
+   * {@link Optional} if the key is absent.
+   */
+  default <T> Optional<T> getOptional(AttributeKey<T> key) {
+    return Optional.ofNullable(get(key));
+  }
+
+  /**
+   * Returns the value associated with {@code key}, throwing {@link IllegalStateException} if the
+   * key is absent.
+   */
+  default <T> T getRequired(AttributeKey<T> key) {
+    if (!containsKey(key)) {
+      throw new IllegalStateException("Required attribute " + key.key() + " not found");
+    }
+    return get(key);
+  }
+
+  /**
+   * Associates the specified value with the specified key in this map. Returns the previous value
+   * associated with {@code key}, or {@code null} if there was no mapping for {@code key}.
+   */
+  @Nullable <T> T put(AttributeKey<T> key, @Nullable T value);
+
+  /**
+   * Adds the given attribute to this map. Returns the previous value associated with the
+   * attribute's key, or {@code null} if there was no mapping for that key. This is just a shortcut
+   * for {@code putAttribute(attribute.getKey(), attribute.getValue())}.
+   */
+  @Nullable
+  default <T> T put(Attribute<T> attribute) {
+    return put(attribute.key(), attribute.value());
+  }
+
+  /** Copies all the attributes from the specified map to this map. */
+  void putAll(AttributeMap map);
+
+  /**
+   * Removes the mapping for the given typed {@code key}, if present. Returns the previous value
+   * associated with {@code key}, or {@code null} if there was no mapping for {@code key}.
+   */
+  @Nullable <T> T remove(AttributeKey<T> key);
+
+  /** Returns {@code true} if this map contains an attribute with the specified {@code key}. */
+  boolean containsKey(AttributeKey<?> key);
+
+  /** Returns {@code true} if this map contains an attribute with the specified {@code value} */
+  boolean containsValue(@Nullable Object value);
+
+  /**
+   * Returns a {@link Set} view of the keys contained in this map. The set is backed by the map, so
+   * changes to the map are reflected in the set, and vice versa. If the map is modified while an
+   * iteration over the set is in progress (except through the iterator's own {@code remove}
+   * operation), the results of the iteration are undefined. The set supports element removal, which
+   * removes the corresponding mapping from the map, via the {@code Iterator.remove}, {@code
+   * Set.remove}, {@code removeAll}, {@code retainAll}, and {@code clear} operations. It does not
+   * support the {@code add} or {@code addAll} operations.
+   */
+  Set<AttributeKey<?>> keySet();
+
+  /**
+   * Returns a {@link Collection} view of the values contained in this map. The collection is backed
+   * by the map, so changes to the map are reflected in the collection, and vice versa. If the map
+   * is modified while an iteration over the collection is in progress (except through the
+   * iterator's own {@code remove} operation), the results of the iteration are undefined. The
+   * collection supports element removal, which removes the corresponding mapping from the map, via
+   * the {@code Iterator.remove}, {@code Collection.remove}, {@code removeAll}, {@code retainAll}
+   * and {@code clear} operations. It does not support the {@code add} or {@code addAll} operations.
+   */
+  Collection<Object> values();
+
+  /**
+   * Returns a {@link Set} view of the mappings contained in this map. The set is backed by the map,
+   * so changes to the map are reflected in the set, and vice versa. If the map is modified while an
+   * iteration over the set is in progress (except through the iterator's own {@code remove}
+   * operation, or through the {@code setValue} operation on a map entry returned by the iterator)
+   * the results of the iteration are undefined. The set supports element removal, which removes the
+   * corresponding mapping from the map, via the {@code Iterator.remove}, {@code Set.remove}, {@code
+   * removeAll}, {@code retainAll} and {@code clear} operations. It does not support the {@code add}
+   * or {@code addAll} operations.
+   */
+  Set<Attribute<?>> entrySet();
+
+  /** Invokes {@code action} for every attribute in this map. */
+  default void forEach(Consumer<Attribute<?>> action) {
+    Objects.requireNonNull(action);
+    entrySet().forEach(action);
+  }
+
+  /** Removes all mappings from this map. The map will be empty after this call returns. */
+  void clear();
+}