Functional relational joins [Java]

TableTest.java

import static org.assertj.core.api.Assertions.assertThat;

import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.function.Supplier;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;

import org.apache.commons.lang3.tuple.Pair;
import org.junit.Test;


/**
 * An implementation of functional relational joins in Java.
 *
 * Implementation inspired by: http://www.haskellforall.com/2014/12/a-very-general-api-for-relational-joins.html
 * Original author: Gabriel Gonzalez
 * Collected on: 2018-09-26
 * Implemented by: Michael-Keith Bernard
 * Implemented on: 2018-09-26
  */
public class TableTest {
    public abstract static class TableKeys<K> {
        public abstract Set<K> keys();
        public abstract <V> V match(Supplier<V> ifAll, Function<Set<K>, V> ifSome);

        public abstract TableKeys<K> union(TableKeys<K> other);
        public abstract TableKeys<K> intersect(TableKeys<K> other);

        public static <K> TableKeys empty() {
            return new Some<>(ImmutableSet.of());
        }

        public static <K> TableKeys<K> of(Set<K> keys) {
            return new Some<>(keys);
        }

        public static <K> TableKeys<K> all() {
            return new All<>();
        }

        public static final class All<K> extends TableKeys<K> {
            private All() { }

            @Override
            public Set<K> keys() {
                throw new IllegalStateException("cannot enumerate infinite set of keys");
            }

            @Override
            public <V> V match(Supplier<V> ifAll, Function<Set<K>, V> ifSome) {
                return ifAll.get();
            }

            @Override
            public TableKeys<K> union(TableKeys<K> other) {
                return TableKeys.all();
            }

            @Override
            public TableKeys<K> intersect(TableKeys<K> other) {
                return other.match(TableKeys::all, _k -> other);
            }
        }

        public static final class Some<K> extends TableKeys<K> {
            private final Set<K> keys;

            private Some(Set<K> keys) {
                this.keys = keys;
            }

            @Override
            public Set<K> keys() {
                return keys;
            }

            @Override
            public <V> V match(Supplier<V> ifAll, Function<Set<K>, V> ifSome) {
                return ifSome.apply(keys);
            }

            @Override
            public TableKeys<K> union(TableKeys<K> other) {
                return other.match(TableKeys::all, s -> TableKeys.of(Sets.union(keys, s)));
            }

            @Override
            public TableKeys<K> intersect(TableKeys<K> other) {
                return other.match(() -> this, s -> TableKeys.of(Sets.intersection(keys, s)));
            }
        }
    }

    public interface Table<K, V> {
        Optional<V> lookup(K key);
        TableKeys<K> keys();

        // Functor Interface (type class)
        static <K, V1, V2> Table<K, V2> map(Table<K, V1> t, Function<V1, V2> f) {
            return new Table<K, V2>() {
                @Override
                public Optional<V2> lookup(K key) {
                    return t.lookup(key).map(f);
                }

                @Override
                public TableKeys<K> keys() {
                    return t.keys();
                }
            };
        }

        // Apply Interface (type class)
        static <K, V1, V2> Table<K, V2> apply(Table<K, Function<V1, V2>> f, Table<K, V1> t) {
            return new Table<K, V2>() {
                @Override
                public Optional<V2> lookup(K key) {
                    Optional<Function<V1, V2>> fn = f.lookup(key);
                    return fn.isPresent() ? t.lookup(key).map(fn.get()) : Optional.empty();
                }

                @Override
                public TableKeys<K> keys() {
                    return f.keys().intersect(t.keys());
                }
            };
        }

        // Applicative Interface (type class)
        static <K, V> Table<K, V> pure(V value) {
            return new Table<K, V>() {
                @Override
                public Optional<V> lookup(K key) {
                    return Optional.of(value);
                }

                @Override
                public TableKeys<K> keys() {
                    return TableKeys.all();
                }
            };
        }

        // Alt Interface (type class)
        static <K, V> Table<K, V> alt(Table<K, V> t1, Table<K, V> t2) {
            return new Table<K, V>() {
                @Override
                public Optional<V> lookup(K key) {
                    Optional<V> v = t1.lookup(key);
                    return v.isPresent() ? v : t2.lookup(key);
                }

                @Override
                public TableKeys<K> keys() {
                    return t1.keys().union(t2.keys());
                }
            };
        }

        static <K, V> Table<K, V> fromMap(Map<K, V> map) {
            return new Table<K, V>() {
                @Override
                public Optional<V> lookup(K key) {
                    return Optional.ofNullable(map.get(key));
                }

                @Override
                public TableKeys<K> keys() {
                    return TableKeys.of(map.keySet());
                }
            };
        }

        static <K, V> Map<K, V> toMap(Table<K, V> table) {
            ImmutableMap.Builder<K, V> builder = ImmutableMap.builder();
            table.keys().keys().forEach(k -> builder.put(k, table.lookup(k).get()));
            return builder.build();
        }

        static <K, A, B> Table<K, Pair<A, B>> join(Table<K, A> t1, Table<K, B> t2) {
            return apply(map(t1, a -> b -> Pair.of(a, b)), t2);
        }

        static <K, A, B> Table<K, Pair<A, Optional<B>>> leftJoin(Table<K, A> t1, Table<K, B> t2) {
            return join(t1, optional(t2));
        }

        static <K, A, B> Table<K, Pair<Optional<A>, B>> rightJoin(Table<K, A> t1, Table<K, B> t2) {
            return join(optional(t1), t2);
        }

        static <K, A, B> Table<K, Pair<Optional<A>, Optional<B>>> fullJoin(Table<K, A> t1, Table<K, B> t2) {
            Table<K, Pair<Optional<A>, Optional<B>>> table = join(optional(t1), optional(t2));
            return new Table<K, Pair<Optional<A>, Optional<B>>>() {
                @Override
                public Optional<Pair<Optional<A>, Optional<B>>> lookup(K key) {
                    return table.lookup(key);
                }

                @Override
                public TableKeys<K> keys() {
                    return t1.keys().union(t2.keys());
                }
            };
        }

        static <K, V> Table<K, Optional<V>> optional(Table<K, V> t) {
            return alt(Table.map(t, Optional::of), pure(Optional.empty()));
        }
    }

    private final Map<String, String> usernames = ImmutableMap
            .of("MK", "segfaultax", "Scott", "sjohnson");
    private final Map<String, String> lastnames = ImmutableMap
            .of("MK", "Bernard", "Scott", "Johnson");
    private final Map<String, String> nicknames = ImmutableMap
            .of("MK", "FunctionalWeenie", "Steven", "Anti-Functional");

    @Test
    public void testJoin() {
        Table<String, String> t1 = Table.fromMap(usernames);
        Table<String, String> t2 = Table.fromMap(lastnames);
        Table<String, Pair<String, String>> joined = Table.join(t1, t2);

        assertThat(Table.toMap(joined))
                .containsEntry("MK", Pair.of("segfaultax", "Bernard"))
                .containsEntry("Scott", Pair.of("sjohnson", "Johnson"));
    }

    @Test
    public void testLeftJoin() {
        Table<String, String> t1 = Table.fromMap(usernames);
        Table<String, String> t2 = Table.fromMap(nicknames);
        Table<String, Pair<String, Optional<String>>> joined = Table.leftJoin(t1, t2);

        assertThat(Table.toMap(joined))
                .doesNotContainKeys("Steven")
                .containsEntry("MK", Pair.of("segfaultax", Optional.of("FunctionalWeenie")))
                .containsEntry("Scott", Pair.of("sjohnson", Optional.empty()));
    }

    @Test
    public void testRightJoin() {
        Table<String, String> t1 = Table.fromMap(usernames);
        Table<String, String> t2 = Table.fromMap(nicknames);
        Table<String, Pair<Optional<String>, String>> joined = Table.rightJoin(t1, t2);

        assertThat(Table.toMap(joined))
                .doesNotContainKeys("Scott")
                .containsEntry("MK", Pair.of(Optional.of("segfaultax"), "FunctionalWeenie"))
                .containsEntry("Steven", Pair.of(Optional.empty(), "Anti-Functional"));
    }

    @Test
    public void testFullOuterJoin() {
        Table<String, String> t1 = Table.fromMap(usernames);
        Table<String, String> t2 = Table.fromMap(nicknames);
        Table<String, Pair<Optional<String>, Optional<String>>> joined = Table.fullJoin(t1, t2);

        assertThat(Table.toMap(joined))
                .containsEntry("MK", Pair.of(Optional.of("segfaultax"), Optional.of("FunctionalWeenie")))
                .containsEntry("Scott", Pair.of(Optional.of("sjohnson"), Optional.empty()))
                .containsEntry("Steven", Pair.of(Optional.empty(), Optional.of("Anti-Functional")));
    }

}