dataframe.ts

class DataRow<T> {
  private constructor(public readonly value: T) {}

  static of<T>(value: DataRow<T>): DataRow<T>;
  static of<T>(value: T): DataRow<T>;
  static of<T>(value: T | DataRow<T>): DataRow<T> {
    return value instanceof DataRow ? value : new DataRow(value);
  }
}

class DataFrame<T> implements Iterable<T> {
  private constructor(private readonly rows: Iterable<DataRow<T>>) {}

  select<K extends keyof T>(...columns: K[]): DataFrame<Pick<T, K>> {
    return this.map(value => pick(value, columns));
  }

  map<U>(mapper: (value: T) => U): DataFrame<U> {
    return DataFrame.from(map(this, value => mapper(value)));
  }

  join<U>(
    other: DataFrame<U>,
    joiner: (left: T, right: U) => boolean,
  ): DataFrame<T & U> {
    return DataFrame.from(
      flatMap(this, left =>
        map(
          filter(other, right => joiner(left, right)),
          right => DataRow.of({ ...left, ...right }),
        ),
      ),
    );
  }

  [Symbol.iterator](): Iterator<T> {
    return map(this.rows, row => row.value)[Symbol.iterator]();
  }

  static of<T>(...rows: ReadonlyArray<DataRow<T>>): DataFrame<T>;
  static of<T>(...rows: ReadonlyArray<T>): DataFrame<T>;
  static of<T>(...rows: ReadonlyArray<T | DataRow<T>>): DataFrame<T> {
    return DataFrame.from(rows) as DataFrame<T>;
  }

  static from<T>(rows: Iterable<DataRow<T>>): DataFrame<T>;
  static from<T>(rows: Iterable<T>): DataFrame<T>;
  static from<T>(rows: Iterable<T | DataRow<T>>): DataFrame<T> {
    return new DataFrame(map<T | DataRow<T>, DataRow<T>>(rows, DataRow.of));
  }
}

const iterable = <T>(generator: () => Iterator<T>): Iterable<T> => ({
  [Symbol.iterator]: generator,
});

function iterableIterator<T>(source: Iterator<T>): IterableIterator<T> {
  return {
    ...source,
    [Symbol.iterator]() {
      return source;
    },
  } as IterableIterator<T>;
}

function iterator<T>(source: Iterable<T>): IterableIterator<T> {
  return iterableIterator(source[Symbol.iterator]());
}

const map = <T, U>(source: Iterable<T>, mapper: (value: T) => U): Iterable<U> =>
  iterable(function*() {
    for (const value of source) {
      yield mapper(value);
    }
  });

const flatMap = <T, U>(
  source: Iterable<T>,
  mapper: (value: T) => Iterable<U>,
): Iterable<U> =>
  iterable(function*() {
    for (const value of source) {
      yield* mapper(value);
    }
  });

const filter = <T>(
  source: Iterable<T>,
  predicate: (value: T) => boolean,
): Iterable<T> =>
  iterable(function*() {
    for (const value of source) {
      if (predicate(value)) {
        yield value;
      }
    }
  });

function first<T>(
  source: Iterable<T>,
): [IteratorResult<T>, IterableIterator<T>] {
  const iterator = source[Symbol.iterator]();

  return [iterator.next(), iterableIterator(iterator)];
}

function asResult<T>(value: T): IteratorResult<T> {
  return {
    value,
  };
}

function reduce<T>(
  source: Iterable<T>,
  reducer: (prev: T, current: T) => T,
): T | undefined;
function reduce<T>(
  source: Iterable<T>,
  reducer: (prev: T, current: T) => T,
  initialValue: T,
): T;
function reduce<T, U>(
  source: Iterable<T>,
  reducer: (prev: U, current: T) => U,
  initialValue: U,
): U;
function reduce<T, U>(
  source: Iterable<T>,
  reducer: (prev: U | T, current: T) => U | T,
  initialValue?: U | T,
): U | T | undefined {
  const [result, iterable] =
    arguments.length === 3
      ? [asResult(initialValue), iterator(source)]
      : first(source);
  if (result.done) {
    return;
  }
  let { value: previous } = result;
  for (const value of iterable) {
    previous = reducer(previous!, value);
  }
  return previous;
}

function entries<T>(obj: { [s: string]: T }): [string, T][] {
  return Object.keys(obj).map(key => [key, obj[key]]);
}

function fromEntries<T>(
  entries: Iterable<readonly [PropertyKey, T]>,
): { [k in PropertyKey]: T } {
  return reduce(
    entries,
    (obj, [key, value]) => {
      obj[key] = value;
      return obj;
    },
    {} as any,
  );
}

const pick = <T, K extends keyof T>(value: T, keys: K[]): Pick<T, K> =>
  fromEntries(
    entries(value as any).filter(([key]) => keys.includes(key)),
  ) as Pick<T, K>;

const people = DataFrame.of({ id: 1, first: 'Jack', last: 'Moore' });
const addresses = DataFrame.of(
  { personId: 1, x: 34.7, y: 47.9, state: 'CA' },
  { personId: 1, x: 39.2, y: 17.3, state: 'KS' },
  { personId: 2, x: 12.2, y: 15.3, state: 'CA' },
);

const joined = people
  .join(addresses, (person, address) => person.id === address.personId)
  .select('first', 'last', 'state');

console.log([...joined]);