An SQL Database Built on GraphBLAS
#!/usr/bin/env python3

import base64
import collections
import contextlib
import inspect
import os
import random
import sys
import time
import urllib
import uuid
import pathlib

import scipy as sp
import numpy as np
import scipy.sparse

import sqlalchemy
from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy import Column, Integer, Float, create_engine
from sqlalchemy import insert, select, and_
from sqlalchemy.orm import Session, aliased, as_declarative
from sqlalchemy.sql import func
from sqlalchemy.sql.expression import alias

from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.expression import ClauseElement, Executable


class CreateTableAs(Executable, ClauseElement):
    def __init__(self, name, query):
        self.name = name
        self.query = query


@compiles(CreateTableAs)
def _create_table_as(element, compiler, **kw):
    return "CREATE TABLE %s AS %s" % (
        element.name,
        compiler.process(element.query)
    )

class Timer:
    def __init__(self, name):
        self.start_time = None
        self.name = name

    def __enter__(self):
        self.start_time = time.time()

    def __exit__(self, *args):
        end_time = time.time()
        print(f"{self.name},{end_time - self.start_time},seconds")


class Counter:
    def __init__(self):
        self.count = 0

    def __call__(self):
        self.count = (ret := self.count) + 1
        return ret


class Condenser(collections.defaultdict):
    def __init__(self):
        super().__init__(Counter())


class UniqueChecker:
    def __init__(self):
        self.data = {}

    def __call__(self, *args):
        val = random.randint(0, 2**63 - 1)
        return self.data.setdefault(args, val) == val


@as_declarative()
class Base(object):
    @declared_attr
    def __tablename__(cls):
        return cls.__name__.lower()

class Node(Base):
    idnode = Column(Integer(), primary_key=True, nullable=False)
    guid = Column(String(36), nullable=False)

class Edge(Base):
    first = Column(Integer(), primary_key=True, nullable=False)
    second = Column(Integer(), primary_key=True, nullable=False)
    value = Column(Float(), nullable=False)


def sqlite_orm(data, uuids):
    (ridx, cidx), vals = data
    N = len(uuids)
    out = pathlib.Path(os.getenv('PROJECT_DIR')) / "db.sqlite"
    out.unlink(missing_ok=True)
    engine = create_engine("sqlite:///" + str(out), echo=True)

    Base.metadata.create_all(engine)

    with Timer("SQLite Operations"):
        with Session(engine) as session:
            with open(sys.argv[1], "r") as fin:
                with Timer("Reading Data and Building ORM objects"):
                    objects = []
                    for r, c, v in zip(ridx, cidx, vals):
                        objects.append(Edge(first=r, second=c, value=v))
                    for k,v in uuids.items():
                        objects.append(Node(idnode=k, guid=v))
                with Timer("Saving Table,sqlite orm"):
                    session.bulk_save_objects(objects)
                with Timer("Commit"):
                    session.commit()

        with Session(engine) as session:
            A = alias(Edge, name="A")
            x = alias(Node, name="x")
            y = alias(Node, name="y")

            stmt = (
                select(x.c.guid.label('first'), y.c.guid.label('second')).select_from(
                    A.join(x, A.c.first == x.c.idnode).join(y, A.c.second == y.c.idnode)
                )
            )
            print(stmt, file=sys.stderr)
            with Timer("NamedEdges,sqlite orm"):
                result = session.execute(stmt)
                with Timer("RetrieveResults,sqlite orm"):
                    result = list(result)
            with Timer("Count"):
                print("Result NNZ", len(result))

    print("")

@contextlib.contextmanager
def postgres_blank_database(unix_socket, tables=None):
    qs = urllib.parse.urlencode(dict(host=unix_socket))
    default_engine = create_engine("postgresql://@/postgres?" + qs, echo=True)
    engine = None
    with default_engine.connect() as conn:
        conn.execute("commit")

        dbname = "tmpdb" + base64.b64encode(os.urandom(12), altchars=b"__").decode(
            "utf-8"
        )

        try:
            conn.execute(f'create database "{dbname}"')

            engine = create_engine(f"postgresql://@/{dbname}?" + qs, echo=True)
            Base.metadata.create_all(engine, tables=tables)

            yield engine

        finally:
            try:
                engine.dispose()
            except:
                pass
            del engine
            conn.execute("commit")
            conn.execute(f'drop database if exists "{dbname}"')

def postgres_orm(data, uuids):
    (ridx, cidx), vals = data
    N = len(uuids)
    with postgres_blank_database(os.getenv("POSTGRES_SOCKETDIR")) as engine:
        with Timer("Postgres Operations"):
            with Session(engine) as session:
                with open(sys.argv[1], "r") as fin:
                    with Timer("Reading Data and Building ORM objects"):
                        objects = []
                        for r, c, v in zip(ridx, cidx, vals):
                            objects.append(Edge(first=r, second=c, value=v))
                        for k,v in uuids.items():
                            objects.append(Node(idnode=k, guid=v))
                    with Timer("Saving Table,postgres orm"):
                        session.bulk_save_objects(objects)
                    with Timer("Commit"):
                        session.commit()

            with Session(engine) as session:
                A = alias(Edge, name="A")
                x = alias(Node, name="x")
                y = alias(Node, name="y")

                stmt = (
                    select(x.c.guid.label('first'), y.c.guid.label('second')).select_from(
                        A.join(x, A.c.first == x.c.idnode).join(y, A.c.second == y.c.idnode)
                    )
                )
                print(stmt, file=sys.stderr)
                with Timer("NamedEdges,postgres orm"):
                    result = session.execute(stmt)
                    with Timer("RetrieveResults,postgres orm"):
                        result = list(result)
                with Timer("Count"):
                    print("Result NNZ", len(result))

        print("")

def tensql(data, uuids):
    (ridx, cidx), vals = data
    N = len(uuids)

    import pygraphblas.types
    import tensql
    import tensql.QIR

    tString = tensql.Types.Text()

    with Timer("TenSQL Operations"):
        with Timer("Building pygraphblas data"):
            A = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )

            uuid_objs = tuple({'idx': key, 'value': value} for key, value in uuids.items())

        with Timer("Saving Table,tensql"):
            db = tensql.Database()
            qdb = tensql.QIR.QirDatabase(db)
            t_edge = db.create_table_from_tensor("edge", A)
            t_node = db.create_table_from_definition(
                "node",
                {
                    'idx': tensql.Types.BigInt(),
                    'value': tensql.Types.Text()
                },
                ['idx'],
                shape = [N]
            )
            qdb.insert(qdb.node).values(uuid_objs).run()

            print(f"{len(t_edge)=}, {len(t_node)=}")
            db.save(f"{os.getenv('PROJECT_DIR')}/db.h5")

        with Timer("NamedEdges,tensql"):
            A = qdb.edge.aliased("A")
            x = qdb.node.aliased("x")
            y = qdb.node.aliased("y")

            result = (
                qdb.query(A)
                .join(x, x.idx == A.ridx)
                .join(y, y.idx == A.cidx)
                .select(A.ridx, A.cidx, x.value.alias('first'), y.value.alias('second'))
                .run()
            )

            with Timer("RetrieveResults,tensql"):
                results = [x for x in result]

    print("Result NNZ", result.count())

def tensql_records(data, uuids):
    (ridx, cidx), vals = data
    N = len(uuids)

    import pygraphblas.types
    import tensql
    import tensql.QIR

    tString = tensql.Types.Text()

    with Timer("TenSQL Operations"):
        with Timer("Building pygraphblas data"):
            A = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )

            uuid_objs = tuple({'idx': key, 'value': value} for key, value in uuids.items())

        with Timer("Saving Table,tensql_records"):
            db = tensql.Database()
            qdb = tensql.QIR.QirDatabase(db)
            t_edge = db.create_table_from_tensor("edge", A)
            t_node = db.create_table_from_definition(
                "node",
                {
                    'idx': tensql.Types.BigInt(),
                    'value': tensql.Types.Text()
                },
                ['idx'],
                shape = [N]
            )
            qdb.insert(qdb.node).values(uuid_objs).run()

            print(f"{len(t_edge)=}, {len(t_node)=}")
            db.save(f"{os.getenv('PROJECT_DIR')}/db.h5")

        with Timer("NamedEdges,tensql_records"):
            A = qdb.edge.aliased("A")
            x = qdb.node.aliased("x")
            y = qdb.node.aliased("y")

            result = (
                qdb.query(A)
                .join(x, x.idx == A.ridx)
                .join(y, y.idx == A.cidx)
                .select(A.ridx, A.cidx, x.value.alias('first'), y.value.alias('second'))
                .run()
            )

            with Timer("RetrieveResults,tensql_records"):
                results = result.to_records()

    print("Result NNZ", results.size)


def pandasql(data, uuids):
    (ridx, cidx), vals = data
    N = len(uuids)
    import pandasql
    import pandas as pd

    with Timer("PandaSQL Operations"):
        with Timer("Building dataframe"):
            A = pd.DataFrame(
                zip(ridx, cidx, vals),
                columns=["ridx", "cidx", "value"],
            )
            x = pd.DataFrame(
                uuids.items(),
                columns=["idx", "name"]
            )
            y = pd.DataFrame(
                uuids.items(),
                columns=["idx", "name"]
            )

        with Timer("NamedEdges,pandasql"):
            q = """
      SELECT x.name, y.name
      FROM A
        INNER JOIN x ON A.ridx == x.idx
        INNER JOIN y ON A.cidx == y.idx
      """

            result: pd.DataFrame = pandasql.sqldf(q)

    print("Result NNZ", len(result))


def main():
    with Timer("Reading data"):
        maxidx = 0
        ridx, cidx, vals = [], [], []
        condenser = Condenser()
        edge_unique = UniqueChecker()
        node_unique = UniqueChecker()
        uuids = {}
        with open(sys.argv[1], "r") as fin:
            for line in fin:
                first, second = tuple(condenser[int(x)] for x in line.strip().split())
                if node_unique(first):
                    uuids[first] = str(uuid.uuid4())
                if node_unique(second):
                    uuids[second] = str(uuid.uuid4())
                if edge_unique(first, second):
                    ridx.append(first)
                    cidx.append(second)
                    vals.append(random.random())
                    maxidx = max(maxidx, first, second)
        N = len(uuids)
    print("")

    all_benchmarks = [
        tensql,
        postgres_orm,
        sqlite_orm,
        pandasql,
    ]

    benchmark_lookup = {bm.__name__: bm for bm in all_benchmarks}
    run_benchmarks = set(sys.argv[2:])
    if 'all' in run_benchmarks:
        run_benchmarks = set(all_benchmarks)
    else:
        run_benchmarks = set(benchmark_lookup[name] for name in run_benchmarks)

    for benchmark in run_benchmarks:
        benchmark(((ridx, cidx), vals), uuids)
        print("")
        sys.stdout.flush()


if __name__ == "__main__":
    main()

benchmark/extract_named_edges.py
sandialabs/TenSQL
#!/usr/bin/env python3

import os
import random
import sys
import time
import collections

import sqlalchemy
from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy import Column, Integer, Float, create_engine
from sqlalchemy import insert, select
from sqlalchemy.orm import Session, aliased, as_declarative
from sqlalchemy.sql import func
from sqlalchemy.sql.expression import alias, and_

import numpy as np
import scipy as sp
import scipy.sparse


class Timer:
    def __init__(self, name):
        self.start_time = None
        self.name = name

    def __enter__(self):
        self.start_time = time.time()

    def __exit__(self, *args):
        end_time = time.time()
        print(f"{self.name} ran for {end_time - self.start_time} seconds")


class Counter:
    def __init__(self):
        self.count = 0

    def __call__(self):
        self.count = (ret := self.count) + 1
        return ret


class Condenser(collections.defaultdict):
    def __init__(self):
        super().__init__(Counter())


class UniqueChecker:
    def __init__(self):
        self.data = {}

    def __call__(self, *args):
        val = random.randint(0, 2**63 - 1)
        return self.data.setdefault(args, val) == val


@as_declarative()
class Base(object):
    @declared_attr
    def __tablename__(cls):
        return cls.__name__.lower()


class Edge(Base):
    first = Column(Integer(), primary_key=True, nullable=False)
    second = Column(Integer(), primary_key=True, nullable=False)
    value = Column(Float(), nullable=False)


class Result(Base):
    first = Column(Integer(), primary_key=True, nullable=False)
    second = Column(Integer(), primary_key=True, nullable=False)
    value = Column(Float(), nullable=False)


# if os.path.exists(sys.argv[2]):
#  os.remove(sys.argv[2])


def sqlite():
    engine = create_engine("sqlite://")

    Base.metadata.create_all(engine)

    with Timer("SQLite Operations"):
        with Session(engine) as session:
            with open(sys.argv[1], "r") as fin:
                with Timer("Reading Data and Building ORM objects"):
                    condenser = Condenser()
                    is_unique = UniqueChecker()
                    objects = []
                    for line in fin:
                        first, second = [
                            condenser[int(x)] for x in line.strip().split()
                        ]
                        if is_unique(first, second):
                            objects.append(
                                Edge(first=first, second=second, value=random.random())
                            )
                with Timer("Bulk Save Objects"):
                    session.bulk_save_objects(objects)
                with Timer("Commit"):
                    session.commit()

        with Session(engine) as session:
            A = alias(Edge, name="A")
            B = alias(Edge, name="B")
            stmt = insert(Result).from_select(
                [Result.first, Result.second, Result.value],
                select(A.c.first, A.c.second, A.c.value + B.c.value)
                .select_from(
                    A.join(B, and_(A.c.first == B.c.second, A.c.second == B.c.first))
                )
                .group_by(A.c.first, A.c.second),
            )
            print(stmt, file=sys.stderr)
            with Timer("Matrix eWiseAdd"):
                result = session.execute(stmt)
            session.commit()

    with Session(engine) as session:
        print(
            "Result NNZ",
            session.execute(select(func.count()).select_from(Result)).scalar(),
        )


def scipy_csr():
    with Timer("Scipy CSR Operations"):
        with Timer("Reading data"):
            ridx, cidx, vals = [], [], []
            maxidx = 0
            condenser = Condenser()
            is_unique = UniqueChecker()
            with open(sys.argv[1], "r") as fin:
                for line in fin:
                    first, second = tuple(
                        condenser[int(x)] for x in line.strip().split()
                    )
                    if is_unique(first, second):
                        ridx.append(first)
                        cidx.append(second)
                        vals.append(random.random())
                        maxidx = max(maxidx, first, second)

        N = maxidx + 1

        with Timer("Building sparse matrix"):
            X = sp.sparse.csr_matrix(
                (vals, (ridx, cidx)), dtype=np.float64, shape=(N, N)
            )
        print(len(ridx), len(cidx), len(vals), N)

        with Timer("Matrix eWiseAdd"):
            ret = X + X.T
    print("Result NNZ", ret.nnz)


def scipy_coo():
    with Timer("Scipy COO Operations"):
        with Timer("Reading data"):
            ridx, cidx, vals = [], [], []
            maxidx = 0
            condenser = Condenser()
            is_unique = UniqueChecker()
            with open(sys.argv[1], "r") as fin:
                for line in fin:
                    first, second = tuple(
                        condenser[int(x)] for x in line.strip().split()
                    )
                    if is_unique(first, second):
                        ridx.append(first)
                        cidx.append(second)
                        vals.append(random.random())
                        maxidx = max(maxidx, first, second)

        N = maxidx + 1

        with Timer("Building sparse matrix"):
            X = sp.sparse.coo_matrix(
                (vals, (ridx, cidx)), dtype=np.float64, shape=(N, N)
            )
        print(len(ridx), len(cidx), len(vals), N)

        with Timer("Matrix eWiseAdd"):
            ret = X + X.T
    print("Result NNZ", ret.nnz)


def pygraphblas():
    import pygraphblas

    with Timer("PyGraphBLAS Operations"):
        with Timer("Reading data"):
            maxidx = 0
            ridx, cidx, vals = [], [], []
            condenser = Condenser()
            is_unique = UniqueChecker()
            with open(sys.argv[1], "r") as fin:
                for line in fin:
                    first, second = tuple(
                        condenser[int(x)] for x in line.strip().split()
                    )
                    if is_unique(first, second):
                        ridx.append(first)
                        cidx.append(second)
                        vals.append(random.random())
                        maxidx = max(maxidx, first, second)

        N = maxidx + 1

        with Timer("Building sparse matrix"):
            X = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )

        with Timer("Matrix eWiseAdd"):
            ret = X.eadd(X.T)
    print("Result NNZ", ret.nvals)


# TODO
# def tensql():
# import pygraphblas.types
# import tensql
# import tensql.QIR
# with Timer("TenSQL Operations"):
#   with Timer("Reading data"):
#     maxidx = 0
#     ridx, cidx, vals = [], [], []
#     condenser = Condenser()
#     is_unique = UniqueChecker()
#     with open(sys.argv[1], 'r') as fin:
#       for line in fin:
#         first, second = tuple(condenser[int(x)] for x in line.strip().split())
#         if is_unique(first, second):
#           ridx.append(first)
#           cidx.append(second)
#           vals.append(random.random())
#           maxidx = max(maxidx, first, second)

#   with Timer("Building sparse matrix"):
#     N = maxidx + 1
#     X = pygraphblas.Matrix.from_lists(ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64)

#   with Timer("Building Table"):
#     db = tensql.Database()
#     print(f"{X.nvals=}")
#     t_matrix = db.create_table_from_tensor("matrix", X)
#     print(f"{len(t_matrix)=}")

#   with Timer("Matrix eWiseAdd"):
#     qdb = tensql.QIR.QirDatabase(db)
#     A = qdb.matrix.aliased("A")
#     B = qdb.matrix.aliased("B")

#     result = (
#       qdb.query(A)
#         .join(B, tensql.QIR.LogicalAnd(A.ridx == B.cidx, A.cidx == B.ridx, name='value'))
#         .select(A.ridx, A.cidx, A.value + B.value)
#         .run()
#     )
# print("Result NNZ", result.count())


def main():
    for run_benchmark in [pygraphblas, scipy_coo, scipy_csr, sqlite]:
        run_benchmark()
        print("")
        sys.stdout.flush()


if __name__ == "__main__":
    main()

benchmark/matrix_add.py
sandialabs/TenSQL
#!/usr/bin/env python3

import collections
import contextlib
import os
import random
import sys
import time
import urllib
import base64
import inspect

import scipy as sp
import numpy as np
import scipy.sparse

import sqlalchemy
from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy import Column, Integer, Float, create_engine
from sqlalchemy import insert, select, and_
from sqlalchemy.orm import Session, aliased, as_declarative
from sqlalchemy.sql import func
from sqlalchemy.sql.expression import alias

from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.expression import ClauseElement, Executable


class CreateTableAs(Executable, ClauseElement):
    def __init__(self, name, query):
        self.name = name
        self.query = query


@compiles(CreateTableAs)
def _create_table_as(element, compiler, **kw):
    return "CREATE TABLE %s AS %s" % (
        element.name,
        compiler.process(element.query)
    )

class Timer:
    def __init__(self, name):
        self.start_time = None
        self.name = name

    def __enter__(self):
        self.start_time = time.time()

    def __exit__(self, *args):
        end_time = time.time()
        print(f"{self.name},{end_time - self.start_time},seconds")


class Counter:
    def __init__(self):
        self.count = 0

    def __call__(self):
        self.count = (ret := self.count) + 1
        return ret


class Condenser(collections.defaultdict):
    def __init__(self):
        super().__init__(Counter())


class UniqueChecker:
    def __init__(self):
        self.data = {}

    def __call__(self, *args):
        val = random.randint(0, 2**63 - 1)
        return self.data.setdefault(args, val) == val


@as_declarative()
class Base(object):
    @declared_attr
    def __tablename__(cls):
        return cls.__name__.lower()

class Edge(Base):
    first = Column(Integer(), primary_key=True, nullable=False, index=True)
    second = Column(Integer(), primary_key=True, nullable=False, index=True)
    value = Column(Float(), nullable=False)

class Result(Base):
    first = Column(Integer(), primary_key=True, nullable=False)
    second = Column(Integer(), primary_key=True, nullable=False)
    value = Column(Float(), nullable=False)


# if os.path.exists(sys.argv[2]):
#  os.remove(sys.argv[2])


def sqlite_insert_into(data, N):
    (ridx, cidx), vals = data
    engine = create_engine("sqlite:///:memory:", echo=True)

    Base.metadata.create_all(engine)

    with Timer("SQLite Operations"):
        with Session(engine) as session:
            with open(sys.argv[1], "r") as fin:
                with Timer("Reading Data and Building ORM objects"):
                    objects = []
                    for r, c, v in zip(ridx, cidx, vals):
                        objects.append(Edge(first=r, second=c, value=v))
                with Timer("Saving Table,sqlite insert into"):
                    session.bulk_save_objects(objects)
                with Timer("Commit"):
                    session.commit()

        with Session(engine) as session:
            A = alias(Edge, name="A")
            B = alias(Edge, name="B")
            q_stmt = (
                select(A.c.first, A.c.second, A.c.value * B.c.value).select_from(
                    A.join(B, and_(A.c.first == B.c.second, A.c.second == B.c.first))
                )
                #          .group_by(A.c.first, A.c.second)
            )
            stmt = insert(Result).from_select(
                [Result.first, Result.second, Result.value], q_stmt
            )
            print(stmt, file=sys.stderr)
            with Timer("eWiseMult,sqlite insert into"):
                session.execute(stmt)
            with Timer("Commit"):
                session.commit()
            with Timer("Count"):
                print("Result NNZ", session.query(Result).count())

    print("")


def sqlite_orm(data, N):
    (ridx, cidx), vals = data
    engine = create_engine("sqlite:///:memory:")

    Base.metadata.create_all(engine)

    with Timer("SQLite Operations"):
        with Session(engine) as session:
            with open(sys.argv[1], "r") as fin:
                with Timer("Reading Data and Building ORM objects"):
                    objects = []
                    for r, c, v in zip(ridx, cidx, vals):
                        objects.append(Edge(first=r, second=c, value=v))
                with Timer("Saving Table,sqlite orm"):
                    session.bulk_save_objects(objects)
                with Timer("Commit"):
                    session.commit()

        with Session(engine) as session:
            A = alias(Edge, name="A")
            B = alias(Edge, name="B")
            stmt = (
                select(A.c.first, A.c.second, A.c.value * B.c.value).select_from(
                    A.join(B, and_(A.c.first == B.c.second, A.c.second == B.c.first))
                )
                #          .group_by(A.c.first, A.c.second)
            )
            print(stmt, file=sys.stderr)
            with Timer("eWiseMult,sqlite orm"):
                result = list(session.execute(stmt))
            with Timer("Count"):
                print("Result NNZ", len(result))

    print("")

def sqlite_create_table_as(data, N):
    (ridx, cidx), vals = data
    engine = create_engine("sqlite:///:memory:", echo=True)

    Base.metadata.create_all(engine, tables=[Edge.__table__])

    with Timer("SQLite Operations"):
        with Session(engine) as session:
            with open(sys.argv[1], "r") as fin:
                with Timer("Reading Data and Building ORM objects"):
                    objects = []
                    for r, c, v in zip(ridx, cidx, vals):
                        objects.append(Edge(first=r, second=c, value=v))
                with Timer("Saving Table,sqlite create table as"):
                    session.bulk_save_objects(objects)
                with Timer("Commit"):
                    session.commit()


        with Session(engine) as session:
            A = alias(Edge, name="A")
            B = alias(Edge, name="B")
            stmt = CreateTableAs(
                "Result",
                select(A.c.first, A.c.second, (A.c.value * B.c.value).label("value")).select_from(
                    A.join(B, and_(A.c.first == B.c.second, A.c.second == B.c.first))
                )
            )
            print(stmt, file=sys.stderr)
            with Timer("eWiseMult,sqlite create table as"):
                session.execute(stmt)
            with Timer("Count"):
                print("Result NNZ", session.query(Result).count())

    print("")


@contextlib.contextmanager
def postgres_blank_database(unix_socket, tables=None):
    qs = urllib.parse.urlencode(dict(host=unix_socket))
    default_engine = create_engine("postgresql://@/postgres?" + qs, echo=True)
    engine = None
    with default_engine.connect() as conn:
        conn.execute("commit")

        dbname = "tmpdb" + base64.b64encode(os.urandom(12), altchars=b"__").decode(
            "utf-8"
        )

        try:
            conn.execute(f'create database "{dbname}"')

            engine = create_engine(f"postgresql://@/{dbname}?" + qs)
            Base.metadata.create_all(engine, tables=tables)

            yield engine

        finally:
            try:
                engine.dispose()
            except:
                pass
            del engine
            conn.execute("commit")
            conn.execute(f'drop database if exists "{dbname}"')


def postgres_insert_into(data, N):
    (ridx, cidx), vals = data
    with postgres_blank_database(os.getenv("POSTGRES_SOCKETDIR")) as engine:
        with Timer("Postgres Operations"):
            with Session(engine) as session:
                with open(sys.argv[1], "r") as fin:
                    with Timer("Reading Data and Building ORM objects"):
                        objects = []
                        for r, c, v in zip(ridx, cidx, vals):
                            objects.append(Edge(first=r, second=c, value=v))
                    with Timer("Saving Table,postgres insert into"):
                        session.bulk_save_objects(objects)
                    with Timer("Commit"):
                        session.commit()

            with Session(engine) as session:
                A = alias(Edge, name="A")
                B = alias(Edge, name="B")
                q_stmt = (
                    select(A.c.first, A.c.second, A.c.value * B.c.value).select_from(
                        A.join(
                            B, and_(A.c.first == B.c.second, A.c.second == B.c.first)
                        )
                    )
                    #            .group_by(A.c.first, A.c.second)
                )
                stmt = insert(Result).from_select(
                    [Result.first, Result.second, Result.value], q_stmt
                )
                print(stmt, file=sys.stderr)
                with Timer("eWiseMult,postgres insert into"):
                    session.execute(stmt)
                with Timer("Commit"):
                    session.commit()
                with Timer("Count"):
                    print("Result NNZ", session.query(Result).count())

        print("")


def postgres_orm(data, N):
    (ridx, cidx), vals = data
    with postgres_blank_database(os.getenv("POSTGRES_SOCKETDIR")) as engine:
        with Timer("Postgres Operations"):
            with Session(engine) as session:
                with open(sys.argv[1], "r") as fin:
                    with Timer("Reading Data and Building ORM objects"):
                        objects = []
                        for r, c, v in zip(ridx, cidx, vals):
                            objects.append(Edge(first=r, second=c, value=v))
                    with Timer("Saving Table,postgres orm"):
                        session.bulk_save_objects(objects)
                    with Timer("Commit"):
                        session.commit()

            with Session(engine) as session:
                A = alias(Edge, name="A")
                B = alias(Edge, name="B")
                stmt = (
                    select(A.c.first, A.c.second, (A.c.value * B.c.value).label("value")).select_from(
                        A.join(
                            B, and_(A.c.first == B.c.second, A.c.second == B.c.first)
                        )
                    )
                    #            .group_by(A.c.first, A.c.second)
                )
                print(stmt, file=sys.stderr)
                with Timer("eWiseMult,postgres orm"):
                    result = list(session.execute(stmt))
                with Timer("Count"):
                    print("Result NNZ", len(result))

        print("")

def postgres_create_table_as(data, N):
    (ridx, cidx), vals = data
    with postgres_blank_database(os.getenv("POSTGRES_SOCKETDIR"), tables=[Edge.__table__]) as engine:
        with Timer("Postgres Operations"):
            with Session(engine) as session:
                with open(sys.argv[1], "r") as fin:
                    with Timer("Reading Data and Building ORM objects"):
                        objects = []
                        for r, c, v in zip(ridx, cidx, vals):
                            objects.append(Edge(first=r, second=c, value=v))
                    with Timer("Saving Table,postgres create table as"):
                        session.bulk_save_objects(objects)
                    with Timer("Commit"):
                        session.commit()

            with Session(engine) as session:
                A = alias(Edge, name="A")
                B = alias(Edge, name="B")
                stmt = CreateTableAs(
                    "Result",
                    select(A.c.first, A.c.second, (A.c.value * B.c.value).label("value")).select_from(
                        A.join(
                            B, and_(A.c.first == B.c.second, A.c.second == B.c.first)
                        )
                    )
                )
                print(stmt, file=sys.stderr)
                with Timer("eWiseMult,postgres create table as"):
                    session.execute(stmt)
                with Timer("Count"):
                    print("Result NNZ", session.query(Result).count())

        print("")


def scipy_csr(data, N):
    (ridx, cidx), vals = data
    with Timer("Scipy CSR Operations"):
        with Timer("Building sparse matrix"):
            X = sp.sparse.csr_matrix(
                (vals, (ridx, cidx)), dtype=np.float64, shape=(N, N)
            )
        print(len(ridx), len(cidx), len(vals), N)

        with Timer("eWiseMult,scipy csr"):
            ret = X.multiply(X.T)
    print("Result NNZ", ret.nnz)


def scipy_coo(data, N):
    (ridx, cidx), vals = data
    with Timer("Scipy COO Operations"):
        with Timer("Building sparse matrix"):
            X = sp.sparse.coo_matrix(
                (vals, (ridx, cidx)), dtype=np.float64, shape=(N, N)
            )
        print(len(ridx), len(cidx), len(vals), N)

        with Timer("eWiseMult,scipy coo"):
            ret = X.multiply(X.T)
    print("Result NNZ", ret.nnz)


def pygraphblas(data, N):
    (ridx, cidx), vals = data
    import pygraphblas

    with Timer("PyGraphBLAS Operations"):
        with Timer("Building sparse matrix"):
            X = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )
            X = X.dup().cast(pygraphblas.types.FP64)

        with Timer("eWiseMult,pygraphblas"):
            ret = X.emult(X.T)
    print("Result NNZ", ret.nvals)

def tensql_create_table_as(data, N):
    (ridx, cidx), vals = data

    import pygraphblas.types
    import tensql
    import tensql.QIR

    with Timer("TenSQL Operations"):
        with Timer("Building sparse matrix"):
            X = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )

        with Timer("Saving Table,tensql_create_table_as"):
            db = tensql.Database()
            print(f"{X.nvals=}")
            t_matrix = db.create_table_from_tensor("matrix", X)
            print(f"{len(t_matrix)=}")
            db.save(f"{os.getenv('PROJECT_DIR')}/db.h5")

        with Timer("eWiseMult,tensql_create_table_as"):
            qdb = tensql.QIR.QirDatabase(db)
            A = qdb.matrix.aliased("A")
            B = qdb.matrix.aliased("B")

            result = (
                qdb.query(A)
                .join(
                    B,
                    tensql.QIR.LogicalAnd(
                        A.ridx == B.cidx, A.cidx == B.ridx, name="value"
                    ),
                )
                .select(A.ridx, A.cidx, A.value * B.value)
                .run()
            )

            db.save(f"{os.getenv('PROJECT_DIR')}/db.h5")
    print("Result NNZ", result.count())

def tensql(data, N):
    (ridx, cidx), vals = data

    import pygraphblas.types
    import tensql
    import tensql.QIR

    with Timer("TenSQL Operations"):
        with Timer("Building sparse matrix"):
            X = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )

        with Timer("Saving Table,tensql"):
            db = tensql.Database()
            print(f"{X.nvals=}")
            t_matrix = db.create_table_from_tensor("matrix", X)
            print(f"{len(t_matrix)=}")
            db.save(f"{os.getenv('PROJECT_DIR')}/db.h5")

        with Timer("eWiseMult,tensql"):
            qdb = tensql.QIR.QirDatabase(db)
            A = qdb.matrix.aliased("A")
            B = qdb.matrix.aliased("B")

            result = (
                qdb.query(A)
                .join(
                    B,
                    tensql.QIR.LogicalAnd(
                        A.ridx == B.cidx, A.cidx == B.ridx, name="value"
                    ),
                )
                .select(A.ridx, A.cidx, A.value * B.value)
                .run()
            )

    print("Result NNZ", result.count())


def pandasql(data, N):
    (ridx, cidx), vals = data
    import pandasql
    import pandas as pd

    with Timer("PandaSQL Operations"):
        with Timer("Building dataframe"):
            matrix = pd.DataFrame(
                zip(ridx, cidx, vals),
                columns=["ridx", "cidx", "value"],
            )

        with Timer("eWiseMult,pandasql"):
            q = """
      SELECT a.ridx, a.cidx, a.value * b.value AS value
      FROM matrix AS a
        INNER JOIN matrix AS b ON a.ridx == b.cidx AND a.cidx == b.ridx
      """

            result: pd.DataFrame = pandasql.sqldf(q)

    print("Result NNZ", len(result))


def main():
    with Timer("Reading data"):
        maxidx = 0
        ridx, cidx, vals = [], [], []
        condenser = Condenser()
        is_unique = UniqueChecker()
        with open(sys.argv[1], "r") as fin:
            for line in fin:
                first, second = tuple(condenser[int(x)] for x in line.strip().split())
                if is_unique(first, second):
                    ridx.append(first)
                    cidx.append(second)
                    vals.append(random.random())
                    maxidx = max(maxidx, first, second)
        N = maxidx + 1
    print("")

    all_benchmarks = [
        pygraphblas,
        scipy_coo,
        scipy_csr,
        tensql,
        postgres_orm,
        postgres_insert_into,
        postgres_create_table_as,
        sqlite_orm,
        sqlite_insert_into,
        sqlite_create_table_as,
        pandasql,
    ]

    benchmark_lookup = {bm.__name__: bm for bm in all_benchmarks}
    run_benchmarks = set(sys.argv[2:])
    if 'all' in run_benchmarks:
        run_benchmarks = set(all_benchmarks)
    else:
        run_benchmarks = set(benchmark_lookup[name] for name in run_benchmarks)

    for benchmark in run_benchmarks:
        benchmark(((ridx, cidx), vals), N)
        print("")
        sys.stdout.flush()


if __name__ == "__main__":
    main()

benchmark/matrix_hprod.py
sandialabs/TenSQL
#!/usr/bin/env python3

import collections
import contextlib
import os
import random
import sys
import time
import urllib
import base64
import pathlib

import scipy as sp
import numpy as np
import scipy.sparse

import sqlalchemy
from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy import Column, Integer, Float, create_engine
from sqlalchemy import insert, select
from sqlalchemy.orm import Session, aliased, as_declarative
from sqlalchemy.sql import func
from sqlalchemy.sql.expression import alias

from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.expression import ClauseElement, Executable
from sqlalchemy import text as sql_text


class CreateTableAs(Executable, ClauseElement):
    def __init__(self, name, query):
        self.name = name
        self.query = query


@compiles(CreateTableAs)
def _create_table_as(element, compiler, **kw):
    return "CREATE TABLE %s AS %s" % (
        element.name,
        compiler.process(element.query)
    )

class Timer:
    def __init__(self, name):
        self.start_time = None
        self.name = name

    def __enter__(self):
        self.start_time = time.time()

    def __exit__(self, *args):
        end_time = time.time()
        print(f"{self.name},{end_time - self.start_time},seconds")


class Counter:
    def __init__(self):
        self.count = 0

    def __call__(self):
        self.count = (ret := self.count) + 1
        return ret


class Condenser(collections.defaultdict):
    def __init__(self):
        super().__init__(Counter())


class UniqueChecker:
    def __init__(self):
        self.data = {}

    def __call__(self, *args):
        val = random.randint(0, 2**63 - 1)
        return self.data.setdefault(args, val) == val


@as_declarative()
class Base(object):
    @declared_attr
    def __tablename__(cls):
        return cls.__name__.lower()

class Node(Base):
    idnode = Column(Integer(), primary_key=True, nullable=False)
    name = Column(String(36), nullable=False)

class Edge(Base):
    first = Column(Integer(), primary_key=True, nullable=False)
    second = Column(Integer(), primary_key=True, nullable=False)
    value = Column(Float(), nullable=False)

class Result(Base):
    first = Column(Integer(), primary_key=True, nullable=False)
    second = Column(Integer(), primary_key=True, nullable=False)

def sqlite_orm(data, N):
    (ridx, cidx), vals = data
    out = pathlib.Path(os.getenv('PROJECT_DIR')) / "db.sqlite"
    out.unlink(missing_ok=True)
    engine = create_engine("sqlite:///" + str(out), echo=True)

    Base.metadata.create_all(engine)

    with Timer("SQLite Operations"):
        with Session(engine) as session:
            with open(sys.argv[1], "r") as fin:
                with Timer("Reading Data and Building ORM objects"):
                    objects = []
                    for r, c, v in zip(ridx, cidx, vals):
                        objects.append(Edge(first=r, second=c, value=v))
                with Timer("Saving Table,sqlite orm"):
                    session.bulk_save_objects(objects)
                with Timer("Commit"):
                    session.commit()

        with Session(engine) as session:
            A = alias(Edge, name="A")
            B = alias(Edge, name="B")
            stmt = (
                select(A.c.first, B.c.second)
                .select_from(A.join(B, A.c.second == B.c.first))
                .group_by(A.c.first, B.c.second)
            )
            print(stmt, file=sys.stderr)
            with Timer("Matrix Multiply,sqlite orm"):
                result = session.execute(stmt)
                with Timer("RetriveResults,sqlite orm"):
                    result = list(result)
            with Timer("Count"):
                print("Result NNZ", len(result))

    print("")

@contextlib.contextmanager
def postgres_blank_database(unix_socket, tables=None):
    qs = urllib.parse.urlencode(dict(host=unix_socket))
    default_engine = create_engine("postgresql://@/postgres?" + qs, echo=True)
    engine = None
    with default_engine.connect() as conn:
        conn.execute(sql_text("commit"))

        dbname = "tmpdb" + base64.b64encode(os.urandom(12), altchars=b"__").decode(
            "utf-8"
        )

        try:
            conn.execute(sql_text(f'create database "{dbname}"'))

            engine = create_engine(f"postgresql://@/{dbname}?" + qs, echo=True)
            Base.metadata.create_all(engine, tables=tables)

            yield engine

        finally:
            try:
                engine.dispose()
            except:
                pass
            del engine
            conn.execute(sql_text("commit"))
            conn.execute(sql_text(f'drop database if exists "{dbname}"'))


def postgres_orm(data, N):
    (ridx, cidx), vals = data
    with postgres_blank_database(os.getenv("POSTGRES_SOCKETDIR")) as engine:
        with Timer("Postgres Operations"):
            with Session(engine) as session:
                with open(sys.argv[1], "r") as fin:
                    with Timer("Reading Data and Building ORM objects"):
                        objects = []
                        for r, c, v in zip(ridx, cidx, vals):
                            objects.append(Edge(first=r, second=c, value=v))
                    with Timer("Saving Table,postgres orm"):
                        session.bulk_save_objects(objects)
                    with Timer("Commit"):
                        session.commit()

            with Session(engine) as session:
                A = alias(Edge, name="A")
                B = alias(Edge, name="B")
                stmt = (
                    select(A.c.first, B.c.second)
                    .select_from(A.join(B, A.c.second == B.c.first))
                    .group_by(A.c.first, B.c.second)
                )
                with Timer("Matrix Multiply,postgres orm"):
                    result = session.execute(stmt)
                    with Timer("RetrieveResults,postgres orm"):
                        result = list(result)
                with Timer("Count"):
                    print("Result NNZ", len(result))

        print("")

def tensql(data, N):
    (ridx, cidx), vals = data

    import pygraphblas.types
    import tensql
    import tensql.QIR

    with Timer("TenSQL Operations"):
        with Timer("Building sparse matrix"):
            X = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )

        with Timer("Saving Table,tensql"):
            db = tensql.Database()
            print(f"{X.nvals=}")
            t_matrix = db.create_table_from_tensor("matrix", X)
            print(f"{len(t_matrix)=}")
            print(f"{X.shape=}")
            db.save(f"{os.getenv('PROJECT_DIR')}/db.h5")

        with Timer("Matrix Multiply,tensql"):
            qdb = tensql.QIR.QirDatabase(db)
            A = qdb.matrix.aliased("A")
            B = qdb.matrix.aliased("B")

            result = (
                qdb.query(A)
                .join(B, A.cidx == B.ridx)
                .select(A.ridx, B.cidx)
                .group_by(A.ridx, B.cidx)
                .run()
            )
            with Timer("RetrieveResults,tensql"):
                results = list(result)
    print("Result NNZ", len(result))

def tensql_fullmxm(data, N):
    (ridx, cidx), vals = data

    import pygraphblas.types
    import tensql
    import tensql.QIR

    with Timer("TenSQL Operations"):
        with Timer("Building sparse matrix"):
            X = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )

        with Timer("Saving Table,tensql"):
            db = tensql.Database()
            print(f"{X.nvals=}")
            t_matrix = db.create_table_from_tensor("matrix", X)
            print(f"{len(t_matrix)=}")
            print(f"{X.shape=}")
            db.save(f"{os.getenv('PROJECT_DIR')}/db.h5")

        with Timer("Matrix Multiply,tensql"):
            qdb = tensql.QIR.QirDatabase(db)
            A = qdb.matrix.aliased("A")
            B = qdb.matrix.aliased("B")

            result = (
                qdb.query(A)
                .join(B, A.cidx == B.ridx)
                .select(A.ridx, B.cidx, tensql.QIR.Sum(A.value * B.value))
                .group_by(A.ridx, B.cidx)
                .run()
            )
            with Timer("RetrieveResults,tensql"):
                results = list(result)
    print("Result NNZ", len(result))

def tensql_fullmxm_records(data, N):
    (ridx, cidx), vals = data

    import pygraphblas.types
    import tensql
    import tensql.QIR

    with Timer("TenSQL Operations"):
        with Timer("Building sparse matrix"):
            X = pygraphblas.Matrix.from_lists(
                ridx, cidx, vals, nrows=N, ncols=N, typ=pygraphblas.types.FP64
            )

        with Timer("Saving Table,tensql_records"):
            db = tensql.Database()
            print(f"{X.nvals=}")
            t_matrix = db.create_table_from_tensor("matrix", X)
            print(f"{len(t_matrix)=}")
            print(f"{X.shape=}")
            db.save(f"{os.getenv('PROJECT_DIR')}/db.h5")

        with Timer("Matrix Multiply,tensql_records"):
            qdb = tensql.QIR.QirDatabase(db)
            A = qdb.matrix.aliased("A")
            B = qdb.matrix.aliased("B")

            result = (
                qdb.query(A)
                .join(B, A.cidx == B.ridx)
                .select(A.ridx, B.cidx, tensql.QIR.Sum(A.value * B.value))
                .group_by(A.ridx, B.cidx)
                .run()
            )
            with Timer("RetrieveResults,tensql_records"):
                results = result.to_records()
    print("Result NNZ", results.size)


def pandasql(data, N):
    (ridx, cidx), vals = data
    import pandasql
    import pandas as pd

    with Timer("PandaSQL Operations"):
        with Timer("Building dataframe"):
            matrix = pd.DataFrame(
                zip(ridx, cidx, vals),
                columns=["ridx", "cidx", "value"],
            )

        with Timer("Matrix Multiply,pandasql"):
            q = """
      SELECT a.ridx, b.cidx
      FROM matrix AS a
        INNER JOIN matrix AS b ON a.cidx == b.ridx
      GROUP BY a.ridx, b.cidx
      """

            result: pd.DataFrame = pandasql.sqldf(q)

    print("Result NNZ", len(result))


def main():
    with Timer("Reading data"):
        maxidx = 0
        ridx, cidx, vals = [], [], []
        condenser = Condenser()
        is_unique = UniqueChecker()
        with open(sys.argv[1], "r") as fin:
            for line in fin:
                first, second = tuple(condenser[int(x)] for x in line.strip().split())
                if is_unique(first, second):
                    ridx.append(first)
                    cidx.append(second)
                    vals.append(random.random())
                    maxidx = max(maxidx, first, second)
        N = maxidx + 1
    print("")

    all_benchmarks = [
        tensql,
        postgres_orm,
        sqlite_orm,
        pandasql,
    ]

    benchmark_lookup = {bm.__name__: bm for bm in all_benchmarks}
    run_benchmarks = sys.argv[2:]
    if 'all' in run_benchmarks:
        run_benchmarks = all_benchmarks
    else:
        run_benchmarks = [benchmark_lookup[name] for name in run_benchmarks]

    for benchmark in run_benchmarks:
        benchmark(((ridx, cidx), vals), N)
        print("")
        sys.stdout.flush()

if __name__ == "__main__":
    main()

benchmark/matrix_multiply.py
sandialabs/TenSQL
import cffi
import pathlib

src_path = pathlib.Path('src/PyObjectHelper.c')
with src_path.open('r') as fin:
  src = fin.read()

ffibuilder = cffi.FFI()
ffibuilder.set_source("tensql._PyObjectHelper", src)
ffibuilder.cdef("""
void PyObject_EQ(bool* z, int64_t* x, int64_t* y);
void PyObject_NE(bool* z, int64_t* x, int64_t* y);
void PyObject_LT(bool* z, int64_t* x, int64_t* y);
void PyObject_GT(bool* z, int64_t* x, int64_t* y);
void PyObject_LE(bool* z, int64_t* x, int64_t* y);
void PyObject_GE(bool* z, int64_t* x, int64_t* y);

void PyObject_ANY(int64_t* z, int64_t* x, int64_t* y);
void PyObject_FIRST(int64_t* z, int64_t* x, int64_t* y);
void PyObject_SECOND(int64_t* z, int64_t* x, int64_t* y);
""")

if __name__ == "__main__":
  ffibuilder.compile()

build_PyObjectHelper.py
sandialabs/TenSQL
#!/usr/bin/env python3
"""Custom test harness for tensql"""

import os
import unittest
import coverage
import pathlib
import warnings

def main():
    #warnings.simplefilter('always', DeprecationWarning)
    testsdir = (pathlib.Path(__file__).parent / 'tensql' / 'test').resolve()
    packagedir = testsdir.parent

    covdir = pathlib.Path("./coverage").resolve()
    covdir.mkdir(exist_ok=True)

    covhtml = covdir / "html"
    covhtml.mkdir(exist_ok=True)

    covdata = covdir / "datafile"

    cov = coverage.Coverage(
        data_file = str(covdata.resolve()),
        include = str((packagedir / "*").resolve()),
        omit = [
            str((testsdir / "*").resolve()),
            str((packagedir / "_version.py").resolve())
        ],
        config_file = str((testsdir / "coveragerc").resolve())
    )

    cov.start()
    suite = unittest.TestSuite()
    import tensql
    from tensql.test.test_util import test_registry

    for test in test_registry['all']:
        suite.addTest(test())

    runner = unittest.TextTestRunner(verbosity=3)
    runner.run(suite)
    cov.stop()
    cov.save()
    cov.html_report(directory=str(covhtml))

main()


run_tests.py
sandialabs/TenSQL
#!/usr/bin/env python
"""
Setup script for TenSQL
"""

from setuptools import setup, find_packages, Extension

extra_params = {}
setup_requires = ['cython', 'numpy<1.21', 'wheel']
try:
    import pip
    pip.main(['install', '-v'] + setup_requires)
    setup_requires = []
except Exception:
    # Going to use easy_install for
    traceback.print_exc()

import versioneer
import numpy as np

_DEBUG = False

extra_compile_args = [
  '-fopenmp', '-Wall'
]
if _DEBUG:
  extra_compile_args += [
    '-g'
  ]
else:
  extra_compile_args += [
    '-O3', '-DNDEBUG'
  ]

ext_modules = []
ext_modules.append(
  Extension(
    "tensql._PyHashOCNT",
    ["src/PyHashOCNT.pyx"],
    extra_compile_args=["-std=c++17", "-O3"],
    include_dirs=[np.get_include()],
    language="c++"
  )
)


cffi_modules = []

cffi_modules.append("build_PyObjectHelper.py:ffibuilder")

setup(
    version=versioneer.get_version(),
    cmdclass=versioneer.get_cmdclass(),
    packages=['tensql', 'tensql.QIR', 'tensql.LAIR', 'tensql.LAIR.visitors'],
    #packages=find_packages(include=['tensql'], exclude=['tensql.tests', 'tensql.tests.*']),
    ext_modules = ext_modules,
    cffi_modules = cffi_modules
)

setup.py
sandialabs/TenSQL
"""
The Column class represents the data in a column of a table.
"""

from __future__ import annotations

import copy
import contextlib
from typing import Optional, List, Any, Tuple

from pygraphblas import Matrix, Vector, Scalar
import numpy as np

from .util import GBTensor, GBType, ndim, initialize_tensor, grb_shape, tensor_to_numpy
from . import Types
from ._PyHashOCNT import PyHashOCNT

class Column:
    def __init__(
        self,
        name: str,
        data: Optional[GBTensor],
        *,
        type_: Types.Type,
        hidden: bool = False,
        refcounts: Optional[PyHashOCNT] = None
    ) -> None:
        assert isinstance(name, str)
        assert isinstance(data, (Matrix, Vector, Scalar)) or data is None
        assert isinstance(type_, Types.Type)

        if data is not None:
          assert isinstance(type_.as_pygrb, data.type) or issubclass(type_.as_pygrb, data.type)

        self._type = type_
        self._name = name
        self._data = data
        self._hidden = hidden
        if self._type.sideload and refcounts is None:
            self._refcounts = PyHashOCNT(1024)
        elif self._type.sideload and refcounts is not None:
            self._refcounts = refcounts.copy()
        else:
            self._refcounts = None

    def __copy__(self):
        """
        Creates a copy of the column that points to a copy of the original data
        """
        if self._data is None:
            data = None
        else:
            data = self._data.dup()

        return Column(self._name, data, type_=self.type_, hidden=self.hidden, refcounts=self._refcounts)

    def copy(self):
        """
        Creates a copy of the column that points to a copy of the original data
        """
        return copy.copy(self)

    @property
    def name(self) -> str:
        """
        Returns the name of the column
        """
        return self._name

    @property
    def data(self) -> Optional[GBTensor]:
        """
        Returns the underlying graphblas data of the column
        """
        return self._data

    @property
    def values(self) -> np.ndarray:
        ret = tensor_to_numpy(self._data)[-1]
        if self.is_sideloaded:
            ret = self._refcounts.decode_many(ret)

        return ret

    @property
    def data_template(self) -> Optional[GBTensor]:
        """
        Returns an empty copy of the data tensor
        """
        if self._data is None:
            return None
        elif self.is_scalar:
            return Scalar.from_type(self.type_.as_pygrb)
        elif self.is_vector:
            return Vector.sparse(self.type_.as_pygrb, size=self.data.size)
        elif self.is_matrix:
            return Matrix.sparse(self.type_.as_pygrb, nrows=self.data.nrows, ncols=self.data.ncols)
        else:
            raise NotImplementedError("Prototype is only implemented for scalars, vectors, and matrices")

    @property
    def is_scalar(self) -> bool:
        """
        Returns true if the column has zero dimensions, otherwise returns false
        """
        return isinstance(self._data, Scalar)

    @property
    def is_vector(self) -> bool:
        """
        Returns true if the column has one dimension, otherwise returns false
        """
        return isinstance(self._data, Vector)

    @property
    def is_matrix(self) -> bool:
        """
        Returns true if the column has two dimensions, otherwise returns false
        """
        return isinstance(self._data, Matrix)

    @property
    def is_boolean(self) -> bool:
        """
        Returns true if the column has a boolean type, otherwise returns false
        """
        return isinstance(self._type, Types.Boolean)

    @property
    def is_sideloaded(self) -> bool:
        """
        Returns true if the column has sideloaded data
        """
        return self._type.sideload

    @property
    def hidden(self) -> bool:
        """
        Returns true if the column's hidden flag is set, otherwise returns false
        """
        return self._hidden

    @property
    def type_(self) -> Types.Type:
        """
        Returns the column's type
        """
        return self._type

    @property
    def gbtype(self) -> GBType:
        """
        Returns the column's PyGraphBLAS type
        """
        return self._type.as_pygrb

    @property
    def nptype(self) -> np.dtype:
        """
        Returns the column's NumPy type
        """
        return self._type.as_numpy

    @property
    def ndim(self) -> Optional[int]:
        "Returns the number of primary keys in the table"
        if self._data is not None:
          return ndim(self._data)
        else:
          return None

    @property
    def schema(self) -> Dict[Any, Any]:
        return {
            "hidden": self.hidden,
            "type": self.type_.name,
            "empty": self._data is None
        }

    def save_hdf5(self, col_group, pk_cols, verbose=0) -> None:
        """
        Saves column data into h5py gruop
        """
        N = self.ndim
        if N is not None:
            pks = list(tensor_to_numpy(self._data))
            values = pks.pop()

            if self.is_sideloaded:
                if verbose > 0: print("Decoding sideloaded data")
                raw_values = self._refcounts.decode_many(values)
                if verbose > 0: print("Serializing sideloaded data")
                serialized_values = self.type_.serialize(raw_values)
                if verbose > 0: print("Creating array")
                values = np.array(serialized_values, dtype=self.type_.as_numpy, copy=False)

            if verbose > 0: print("Creating hdf5 dataset for value")
            col_group.create_dataset("value", data=values, dtype=self.type_.as_h5py)
            if verbose > 0: print("Creating hdf5 group for primary keys")
            pk_group = col_group.create_group("primary_keys")
            num_digits = len(str(N))
            for it_pk, pk in enumerate(pks):
                if verbose > 0: print("Creating hdf5 dataset for pk", it_pk)
                pk_group.create_dataset(
                  str(it_pk).zfill(num_digits),
                  data = np.array(pk),
                  dtype = pk_cols[it_pk].nptype
                )

    @classmethod
    def load_hdf5(cls, col_group, col_name, shape, **col_info) -> Column:
        """
        Loads column data from an h5py gruop
        """

        col_type = Types.from_string(col_info['type'])

        if not col_info['empty']:
            pk_arrays = list(sorted(col_group['primary_keys'].items(), key=lambda x: x[0]))
            N = len(pk_arrays)
            assert N == len(shape)

            data = []
            for it_pk, (id_pk, pk_arr) in enumerate(pk_arrays):
                data.append(pk_arr)

            values = col_group['value'][:]
            if col_type.sideload:
                refcounts = PyHashOCNT(max(values.size,1)*2)

                values = col_type.deserialize(values)
                values = refcounts.insert_many_numpy(values)
            else:
                refcounts = None
            data.append(values)

            tensor = initialize_tensor(col_type.as_pygrb, shape, data)
        else:
            tensor = None
            refcounts = None

        return Column(col_name, tensor, type_=col_type, hidden=col_info['hidden'], refcounts=refcounts)

    def _insert(self, pk: Optional[Tuple[int]], value: Any):
        if self._data is None:
            raise ValueError("Cannot directly insert primary keys")

        if pk is not None and len(pk) == 0:
            pk = None

        if self.is_scalar:
            assert pk is None
        elif self.is_vector:
            assert len(pk) == 1
        elif self.is_matrix:
            assert len(pk) == 2
        else:
            raise NotImplementedError("Prototype is only implemented for scalars, vectors, and matrices")

        with contextlib.ExitStack() as ocnt_stack:
          if self.is_sideloaded:
              ocnt_stack.enter_context(self._refcounts.checkpoint())
              try:
                current_value_ptr = self._data[pk]
              except KeyError:
                current_value_ptr = 0
              if current_value_ptr != 0:
                  current_value = self._refcounts.decode(current_value_ptr)
                  self._refcounts.remove(current_value)
              value, count = self._refcounts.insert(value)

          if self.is_scalar:
              self._data[None] = value
          elif self.is_vector:
              self._data[pk[0]] = value
          elif self.is_matrix:
              self._data[pk[0], pk[1]] = value

    def _insert_bulk(self, pks: List[np.ndarray], value: np.ndarray) -> None:
        assert isinstance(pks, list)
        assert all(pk.ndim == 1 for pk in pks)
        assert value.ndim == 1
        assert all(isinstance(pk, np.ndarray) for pk in pks)
        assert isinstance(value, np.ndarray)
        assert value.dtype == self.type_.as_numpy
        assert all(pk.size == value.size for pk in pks)

        if self.is_scalar:
            assert len(pks) == 0
        elif self.is_vector:
            assert len(pks) == 1
        elif self.is_matrix:
            assert len(pks) == 2
        else:
            raise NotImplementedError("Prototype is only implemented for scalars, vectors, and matrices")

        with contextlib.ExitStack() as ocnt_stack:
            if self.is_sideloaded:
                ocnt_stack.enter_context(self._refcounts.checkpoint())
                value = self._refcounts.insert_many_numpy(value)
            addtensor = initialize_tensor(self.gbtype, grb_shape(self._data), pks + [value])
            result = self._data.eadd(addtensor, add_op = self.gbtype.FIRST)
            if addtensor.nvals + self._data.nvals != result.nvals:
                raise ValueError("Primary key uniqueness constraint violated")
            else:
                self._data = result







            





















tensql/Column.py
sandialabs/TenSQL
import pathlib
from typing import Union

import h5py

from .Table import Table


class Database:
    def __init__(self):
        self.tables = {}

    def create_table(self, name: str, *args):
        assert name not in self.tables
        table = Table(*args)
        self.tables[name] = table
        return table

    def create_table_from_definition(self, name: str, *args, **kwargs) -> Table:
        assert name not in self.tables
        table = Table.from_definition(*args, **kwargs)
        self.tables[name] = table
        return table

    def create_table_from_tensor(self, name: str, *args, **kwargs) -> Table:
        assert name not in self.tables
        table = Table.from_tensor(*args, **kwargs)
        self.tables[name] = table
        return table

    def add_table(self, name: str, table: Table) -> Table:
        self.tables[name] = table
        return table

    def __getitem__(self, name: str) -> Table:
        return self.tables[name]

    def __setitem__(self, name: str, table: Table) -> Table:
        return self.add_table(name, table)

    def __delitem__(self, name: str) -> None:
        del self.tables[name]

    def __contains__(self, name: str) -> bool:
        return name in self.tables

    def save(self, path: Union[pathlib.Path, str], verbose=0):
        if verbose >= 1:
            print("Saving", "Database", str(path))
        with h5py.File(str(path), "w") as fout:
            for table_name, table in self.tables.items():
                if verbose >= 1:
                  print("Saving", "Table", table_name, "with", table.stencil.nvals, "rows")
                table.save_hdf5(fout.create_group(table_name), verbose=verbose-1)

    @classmethod
    def load(cls, path: Union[pathlib.Path, str]):
        db = cls()

        with h5py.File(str(path), "r") as fin:
            for table_name, table_group in fin.items():
                db.tables[table_name] = Table.load_hdf5(table_group)

        return db

tensql/Database.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class AntiMask(Node):
    def __init__(self, operand: Node, mask: Node, name: Optional[str] = None) -> None:
        self._operand = operand
        self._mask = mask
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def mask(self) -> Node:
        return self._mask

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)
        yield ("mask", self._mask)

    @property
    def op(self) -> str:
        return "FIRST"

    @property
    def left(self) -> Node:
        return self._operand

    @property
    def right(self) -> Node:
        return self._mask

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "mask": self._mask, "name": self._name}

    @property
    def dsl(self) -> str:
        return f"AntiMask({self.operand.dsl}, {self.mask.dsl})"

tensql/LAIR/AntiMask.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node

class Cast(Node):
    def __init__(self, operand: Node, type_: Any, name: Optional[str] = None) -> None:
        self._operand = operand
        self._type = type_
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def type_(self) -> Node:
        return self._type

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "type_": self._type, "name": self._name}

    @property
    def dsl(self) -> str:
        return f"Cast({self.operand.dsl}, {self.type_.__name__})"

tensql/LAIR/Cast.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class CastLike(Node):
    def __init__(self, operand: Node, like: Node, name: Optional[str] = None) -> None:
        self._operand = operand
        self._like = like
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def like(self) -> Node:
        return self._like

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)
        yield ("like", self._like)

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "like": self._like, "name": self._name}

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self._operand.dsl}, {self._like.dsl})"

tensql/LAIR/CastLike.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class DotMxM(Node):
    def __init__(
        self,
        left: Node,
        right: Node,
        add_op: str,
        mul_op: str,
        name: Optional[str] = None,
    ) -> None:
        self._left = left
        self._right = right
        self._add_op = add_op
        self._mul_op = mul_op
        if name is None:
            name = self.description
        super().__init__(name)

    @property
    def left(self) -> Node:
        return self._left

    @property
    def right(self) -> Node:
        return self._right

    @property
    def add_op(self) -> str:
        return self._add_op

    @property
    def mul_op(self) -> str:
        return self._mul_op

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self._left.dsl}, {self._right.dsl}, {self._add_op!r}, {self._mul_op!r})"

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._add_op!r}, {self._mul_op!r})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"left": self._left, "right": self._right, "name": self._name}

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("left", self._left)
        yield ("right", self._right)

tensql/LAIR/DotMxM.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class DotMxV(Node):
    def __init__(
        self,
        left: Node,
        right: Node,
        add_op: str,
        mul_op: str,
        name: Optional[str] = None,
    ) -> None:
        self._left = left
        self._right = right
        self._add_op = add_op
        self._mul_op = mul_op
        if name is None:
            name = self.description
        super().__init__(name)

    @property
    def left(self) -> Node:
        return self._left

    @property
    def right(self) -> Node:
        return self._right

    @property
    def add_op(self) -> str:
        return self._add_op

    @property
    def mul_op(self) -> str:
        return self._mul_op

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self._left.dsl}, {self._right.dsl}, {self._add_op!r}, {self._mul_op!r})"

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._add_op!r}, {self._mul_op!r})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"left": self._left, "right": self._right, "name": self._name}

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("left", self._left)
        yield ("right", self._right)

tensql/LAIR/DotMxV.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class DotVxM(Node):
    def __init__(
        self,
        left: Node,
        right: Node,
        add_op: str,
        mul_op: str,
        name: Optional[str] = None,
    ) -> None:
        self._left = left
        self._right = right
        self._add_op = add_op
        self._mul_op = mul_op
        if name is None:
            name = self.description
        super().__init__(name)

    @property
    def left(self) -> Node:
        return self._left

    @property
    def right(self) -> Node:
        return self._right

    @property
    def add_op(self) -> str:
        return self._add_op

    @property
    def mul_op(self) -> str:
        return self._mul_op

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self._left.dsl}, {self._right.dsl}, {self._add_op!r}, {self._mul_op!r})"

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._add_op!r}, {self._mul_op!r})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"left": self._left, "right": self._right, "name": self._name}

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("left", self._left)
        yield ("right", self._right)

tensql/LAIR/DotVxM.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class ElementwiseAdd(Node):
    def __init__(
        self, op: str, left: Node, right: Node, name: Optional[str] = None
    ) -> None:
        self._left = left
        self._right = right
        self._op = op
        if name is None:
            name = self.description
        super().__init__(name)

    @property
    def op(self) -> str:
        return self._op

    @property
    def left(self) -> Node:
        return self._left

    @property
    def right(self) -> Node:
        return self._right

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._op!r})"

    @property
    def dsl(self) -> str:
        return f"({self._left.dsl} {self._op} {self._right.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "left": self._left,
            "right": self._right,
            "op": self._op,
            "name": self._name,
        }

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("left", self._left)
        yield ("right", self._right)

tensql/LAIR/ElementwiseAdd.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class ElementwiseApply(Node):
    def __init__(self, op: str, operand: Node, name: Optional[str] = None) -> None:
        self._operand = operand
        self._op = op
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def op(self) -> str:
        return self._op

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._op!r})"

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self._op!r}, {self._operand.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "op": self._op, "name": self._name}


class ElementwiseApplyBindFirst(Node):
    def __init__(
        self, op: str, operand: Node, scalar: Node, name: Optional[str] = None
    ) -> None:
        self._operand = operand
        self._scalar = scalar
        self._op = op
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def scalar(self) -> Node:
        return self._scalar

    @property
    def op(self) -> str:
        return self._op

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)
        yield ("scalar", self._scalar)

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._op!r})"

    @property
    def dsl(self) -> str:
        return f"({self._operand.dsl} {self.op} {self._scalar.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "operand": self._operand,
            "scalar": self._scalar,
            "op": self._op,
            "name": self._name,
        }


class ElementwiseApplyBindSecond(Node):
    def __init__(
        self, op: str, operand: Node, scalar: Node, name: Optional[str] = None
    ) -> None:
        self._operand = operand
        self._scalar = scalar
        self._op = op
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def scalar(self) -> Node:
        return self._scalar

    @property
    def op(self):
        return self._op

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)
        yield ("scalar", self._scalar)

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._op!r})"

    @property
    def dsl(self) -> str:
        return f"({self._operand.dsl} {self.op} {self._scalar.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "operand": self._operand,
            "scalar": self._scalar,
            "op": self._op,
            "name": self._name,
        }

tensql/LAIR/ElementwiseApply.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class ElementwiseMultiply(Node):
    def __init__(
        self, op: str, left: Node, right: Node, name: Optional[str] = None
    ) -> None:
        self._left = left
        self._right = right
        self._op = op
        if name is None:
            name = self.description
        super().__init__(name)

    @property
    def op(self) -> str:
        return self._op

    @property
    def left(self) -> Node:
        return self._left

    @property
    def right(self) -> Node:
        return self._right

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._op!r})"

    @property
    def dsl(self) -> str:
        return f"({self._left.dsl} {self._op} {self._right.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "left": self._left,
            "right": self._right,
            "op": self._op,
            "name": self._name,
        }

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("left", self._left)
        yield ("right", self._right)

tensql/LAIR/ElementwiseMultiply.py
sandialabs/TenSQL
from typing import Generator, Optional, Tuple, Dict, Any, Sequence

from ..util import EncoderDecoder, SequentialIdGenerator

from .DotMxM import DotMxM

from .Node import Node
from .Mask import Mask
from .ElementwiseMultiply import ElementwiseMultiply
from .CastLike import CastLike
from .PermuteIndices import PermuteIndices
from .ElementwiseApply import ElementwiseApplyBindSecond
from .VectorToDiagMatrix import VectorToDiagMatrix

AxesSpecifier = Sequence[int]

class InnerBroadcast(Node):
    def __init__(
        self,
        op: str,
        left: Node,
        left_indices: AxesSpecifier,
        right: Node,
        right_indices: AxesSpecifier,
        out_indices: AxesSpecifier,
        name: Optional[str] = None,
    ) -> None:
        self._op = op
        self._left = left
        self._right = right

        idx_re_encoder = EncoderDecoder(SequentialIdGenerator())
        self._out_indices = tuple(idx_re_encoder[idx] for idx in out_indices)
        self._left_indices = tuple(idx_re_encoder[idx] for idx in left_indices)
        self._right_indices = tuple(idx_re_encoder[idx] for idx in right_indices)

        super().__init__(name)
    
    @property
    def op(self) -> str:
        return self._op

    @property
    def left(self) -> Node:
        return self._left

    @property
    def left_indices(self) -> Tuple[int, ...]:
        return self._left_indices

    @property
    def right(self) -> Node:
        return self._right

    @property
    def right_indices(self) -> Tuple[int, ...]:
        return self._right_indices

    @property
    def out_indices(self) -> Tuple[int, ...]:
        return self._out_indices

    @property
    def pattern(self) -> Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]]:
        return (self._left_indices, self._right_indices, self._out_indices)

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("left", self._left)
        yield ("right", self._right)

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self._op!r}, {self._left.dsl}, {self._left_indices}, {self._right.dsl}, {self._right_indices}, {self._out_indices})"

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._op!r}, {self.pattern})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "op": self._op,
            "left": self._left,
            "left_indices": self._left_indices,
            "right": self._right,
            "right_indices": self._right_indices,
            "out_indices": self._out_indices,
            "name": self._name,
        }



tensql/LAIR/InnerBroadcast.py
sandialabs/TenSQL
from typing import Generator, Optional, Tuple, Dict, Any, Sequence

from ..util import EncoderDecoder, SequentialIdGenerator

from .DotMxM import DotMxM

from .Node import Node
from .Mask import Mask
from .ElementwiseMultiply import ElementwiseMultiply
from .CastLike import CastLike
from .PermuteIndices import PermuteIndices
from .ElementwiseApply import ElementwiseApplyBindSecond
from .VectorToDiagMatrix import VectorToDiagMatrix

AxesSpecifier = Sequence[int]

class InnerBroadcastMask(Node):
    def __init__(
        self,
        left: Node,
        left_indices: AxesSpecifier,
        right: Node,
        right_indices: AxesSpecifier,
        out_indices: AxesSpecifier,
        name: Optional[str] = None,
    ) -> None:
        self._left = left
        self._right = right

        idx_re_encoder = EncoderDecoder(SequentialIdGenerator())
        self._out_indices = tuple(idx_re_encoder[idx] for idx in out_indices)
        self._left_indices = tuple(idx_re_encoder[idx] for idx in left_indices)
        self._right_indices = tuple(idx_re_encoder[idx] for idx in right_indices)

        super().__init__(name)

    @property
    def op(self) -> str:
        return "FIRST"

    @property
    def left(self) -> Node:
        return self._left

    @property
    def left_indices(self) -> Tuple[int, ...]:
        return self._left_indices

    @property
    def right(self) -> Node:
        return self._right

    @property
    def right_indices(self) -> Tuple[int, ...]:
        return self._right_indices

    @property
    def out_indices(self) -> Tuple[int, ...]:
        return self._out_indices

    @property
    def pattern(self) -> Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]]:
        return (self._left_indices, self._right_indices, self._out_indices)

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("left", self._left)
        yield ("right", self._right)

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self._left.dsl}, {self._left_indices}, {self._right.dsl}, {self._right_indices}, {self._out_indices})"

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self.pattern})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "left": self._left,
            "left_indices": self._left_indices,
            "right": self._right,
            "right_indices": self._right_indices,
            "out_indices": self._out_indices,
            "name": self._name,
        }



tensql/LAIR/InnerBroadcastMask.py
sandialabs/TenSQL
from typing import Generator, Optional, Tuple, Dict, Any
from ..util import EncoderDecoder, SequentialIdGenerator

from .DotMxM import DotMxM

from .Node import Node
from .Mask import Mask
from .Pattern import Pattern
from .Promote import Promote
from .CastLike import CastLike
from .PermuteIndices import PermuteIndices
from .ElementwiseApply import ElementwiseApplyBindSecond
from .VectorToDiagMatrix import VectorToDiagMatrix
from .ElementwiseMultiply import ElementwiseMultiply
from .VectorToColumnMatrix import VectorToColumnMatrix
from .VectorToRowMatrix import VectorToRowMatrix

AxesSpecifier = Tuple[int, ...]

# class InnerBroadcastStencil(Node):
#   def __init__(
#       self,
#       left: Node, left_indices: AxesSpecifier,
#       right: Node, right_indices: AxesSpecifier,
#       out_indices: AxesSpecifier,
#       name: Optional[str] = None
#   ) -> None:
#     self._left = left
#     self._left_indices = left_indices
#     self._right = right
#     self._right_indices = right_indices
#     self._out_indices = out_indices
#     super().__init__(name)

#   @property
#   def left(self) -> Node:
#     return self._left

#   @property
#   def left_indices(self) -> AxesSpecifier:
#     return self._left_indices

#   @property
#   def right(self) -> Node:
#     return self._right

#   @property
#   def right_indices(self) -> AxesSpecifier:
#     return self._right_indices

#   @property
#   def out_indices(self) -> AxesSpecifier:
#     return self._out_indices

#   def __iter__(self) -> Generator[Node, None, None]:
#     yield self._left
#     yield self._right

#   @property
#   def dsl(self) -> str:
#     return f"{type(self).__name__}({self.left.dsl}, {self._left_indices}, {self.right.dsl}, {self._right_indices}, {self._out_indices})"

#   @property
#   def description(self) -> str:
#     return f"{type(self).__name__}({self._left_indices}, {self._right_indices}, {self._out_indices})"

#   def get_attributes(self) -> Dict[str, Any]:
#     return {
#       'left': self._left,
#       'left_indices': self._left_indices,
#       'right': self._right,
#       'right_indices': self._right_indices,
#       'out_indices': self._out_indices,
#       'name': self._name
#     }


def OptimizedInnerBroadcastStencil(
    left: Node,
    left_indices: AxesSpecifier,
    right: Node,
    right_indices: AxesSpecifier,
    out_indices: AxesSpecifier,
) -> Node:
    assert set(left_indices) - set(left_indices) == set()

    idx_re_encoder = EncoderDecoder(SequentialIdGenerator())
    left_indices = tuple(idx_re_encoder[idx] for idx in left_indices)
    right_indices = tuple(idx_re_encoder[idx] for idx in right_indices)
    out_indices = tuple(idx_re_encoder[idx] for idx in out_indices)
    idx_desc = (left_indices, right_indices, out_indices)

    return Mask(
        Promote(left, left_indices, out_indices),
        Promote(right, right_indices, out_indices),
    )


#  return ElementwiseMultiply(
#    'LAND',
#    Promote(left, left_indices, out_indices),
#    Promote(right, right_indices, out_indices)
#  )
# return InnerBroadcastStencil(left, left_indices, right, right_indices, out_indices)

#  NODIM = tuple()
#
#  left = Pattern(left)
#  right = Pattern(right)
#
#  if idx_desc == (NODIM, NODIM, NODIM):
#    sO = ElementwiseMultiply("AND", left, right)
#
#  elif idx_desc == ((0,), NODIM, (0,)):
#    sO = ElementwiseApplyBindSecond("AND", left, right)
#  elif idx_desc == ((0,1), NODIM, (0,1)):
#    sO = ElementwiseApplyBindSecond("AND", left, right)
#  elif idx_desc == ((0,1), NODIM, (1,0)):
#    sO = ElementwiseApplyBindSecond("AND", PermuteIndices(left), right)
#
#  elif idx_desc == ((0,), (0,), (0,)):
#    sO = ElementwiseMultiply("AND", left, right)
#  elif idx_desc == ((0,), (1,), (0,1)):
#    mA = VectorToColumnMatrix(left)
#    mB = VectorToRowMatrix(right)
#    sO = DotMxM(mA, mB)
#  elif idx_desc == ((0,), (1,), (1,0)):
#    mA = VectorToRowMatrix(left)
#    mB = VectorToColumnMatrix(right)
#    sO = DotMxM(mB, mA)
#
#  elif idx_desc == ((0,1), (0,), (0,1)):
#    mB = VectorToDiagMatrix(right)
#    sO = DotMxM(mB, left)
#  elif idx_desc == ((0,1), (0,), (1,0)):
#    mB = VectorToDiagMatrix(right)
#    sO = DotMxM(PermuteIndices(left), mB)
#  elif idx_desc == ((0,1), (1,), (0,1)):
#    mB = VectorToDiagMatrix(right)
#    sO = DotMxM(left, mB)
#  elif idx_desc == ((0,1), (1,), (1,0)):
#    mB = VectorToDiagMatrix(right)
#    sO = DotMxM(mB, PermuteIndices(left))
#
#  elif idx_desc == ((0,1), (0,1), (0,1)):
#    sO = ElementwiseMultiply("AND", left, right)
#  elif idx_desc == ((0,1), (0,1), (1,0)):
#    sO = PermuteIndices(ElementwiseMultiply("AND", left, right))
#  elif idx_desc == ((0,1), (1,0), (0,1)):
#    sO = ElementwiseMultiply("AND", left, PermuteIndices(right))
#  elif idx_desc == ((0,1), (1,0), (1,0)):
#    sO = ElementwiseMultiply("AND", PermuteIndices(left), right)
#
#  else:
#    sO = InnerBroadcastStencil(left, left_indices, right, right_indices, out_indices)
#
#  return sO

tensql/LAIR/InnerBroadcastStencil.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class Mask(Node):
    def __init__(self, operand: Node, mask: Node, name: Optional[str] = None) -> None:
        self._operand = operand
        self._mask = mask
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def mask(self) -> Node:
        return self._mask

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)
        yield ("mask", self._mask)

    @property
    def op(self) -> str:
        return "FIRST"

    @property
    def left(self) -> Node:
        return self._operand

    @property
    def right(self) -> Node:
        return self._mask

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "mask": self._mask, "name": self._name}

    @property
    def dsl(self) -> str:
        return f"Mask({self.operand.dsl}, {self.mask.dsl})"

tensql/LAIR/Mask.py
sandialabs/TenSQL
from __future__ import annotations

import abc
import pathlib
import json
from typing import Optional, Generator, Tuple, TYPE_CHECKING, Dict, Any, List

if TYPE_CHECKING:
    pass


class Node(abc.ABC):
    def __init__(self, name: Optional[str]) -> None:
        if name is not None:
            self._name = name
        else:
            self._name = self.description

    @property
    def name(self) -> str:
        return self._name

    @property
    @abc.abstractmethod
    def dsl(self) -> str:
        pass

    @property
    def description(self) -> str:
        return type(self).__name__

    @abc.abstractmethod
    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        pass

    @abc.abstractmethod
    def get_attributes(self) -> Dict[str, Any]:
        pass

    def describe_as_edges(self) -> Generator[Tuple[Node, Node], None, None]:
        # print(type(self))
        for relationship, parent in self:
            # print(type(parent), type(self))
            yield (parent, self, relationship)
            yield from parent.describe_as_edges()

    @staticmethod
    def save_graph(path: str, *ir_nodes: Node) -> None:
        path = pathlib.Path(path)
        with path.open("w") as fout:
            all_edges: List[Tuple[Node, Node]] = []
            for node in ir_nodes:
                all_edges.extend(node.describe_as_edges())

            print("digraph G {", file=fout)
            nodes = set()
            for head, tail, edge_name in all_edges:
                nodes.add((head.description, id(head)))
                nodes.add((tail.description, id(tail)))
                print(
                    f"    n{id(head)} -> n{id(tail)} [label={json.dumps(edge_name)}]",
                    file=fout,
                )

            print("", file=fout)
            for node_type, node_id in nodes:
                print(f"    n{node_id} [label={json.dumps(node_type)}]", file=fout)

            print("}", file=fout)

tensql/LAIR/Node.py
sandialabs/TenSQL
from typing import Generator, Tuple, Callable, Any, Tuple, Iterable, Dict, Optional
import abc

from .ElementwiseMultiply import ElementwiseMultiply

from .Node import Node
from .Promote import Promote
from .DotMxV import DotMxV
from .DotVxM import DotVxM
from .DotMxM import DotMxM
from .visitors import IndexVisitor, NodeReplacer, ReverseIndexVisitor, ShapeVisitor
from .visitors.index_algebra import intersect_sop
from .. import LAIR
from .. import util

def replace_node(root: Node, old_node: Node, new_node: Node) -> Node:
    #print("Replacing", old_node.dsl)
    #print("with", new_node.dsl)
    visitor = NodeReplacer(old_node, new_node)
    return visitor.visit(root)

def optimize(root: Node, debug: bool = False):
    keep_optimizing = True

    valid_semirings = {
        ("+", "*"),
        ("LOR", "LAND"),
        ("LOR", "FIRST"),
        ("LOR", "SECOND"),
    }

    while keep_optimizing:
        keep_optimizing = False
        replacement = None

        #print("Re-running analysis")

        index_visitor = IndexVisitor()
        index_visitor.visit(root)
        reverse_index_visitor = ReverseIndexVisitor(index_visitor.registry)
        reverse_index_visitor.visit(root)
        shape_visitor = ShapeVisitor()
        # shape_visitor.visit(root)
        for node, node_indices in index_visitor.registry.items():
            if isinstance(node, (LAIR.PermuteIndices,)):
                if tuple(node.new_axes) == tuple(sorted(node.new_axes)):
                    root = replace_node(root, node, node.operand)
                    keep_optimizing = True
            elif isinstance(node, (LAIR.ElementwiseApply,)) and node.op in ('ONE', 'ZERO'):
                # Move elementwise ones and zeros so they are applied before InnerBroadcastMask operations
                broadcast = node.operand
                if isinstance(broadcast, (LAIR.InnerBroadcastMask,)):
                    operand = broadcast.left
                    replacement = LAIR.InnerBroadcastMask(
                        LAIR.ElementwiseApply(node.op, operand),
                        broadcast.left_indices,
                        broadcast.right,
                        broadcast.right_indices,
                        broadcast.out_indices
                    )
                    root = replace_node(root, node, replacement)
                    keep_optimizing = True
                    replacement = None
            elif isinstance(node, (LAIR.ElementwiseMultiply, LAIR.Mask, LAIR.AntiMask)) and node.op in ("FIRST", "SECOND"):
                #if debug:
                #    print("Optimizing", node.description)
                #    print(node.description, "Forward\n\t", index_visitor.registry[node], "\n\t", index_visitor.registry[node.left], "\n\t", index_visitor.registry[node.right], "\n")
                #    print(node.description, "Reverse\n\t", reverse_index_visitor.registry[node], "\n\t", reverse_index_visitor.registry[node.left], "\n\t", reverse_index_visitor.registry[node.right], "\n")

                if node.op == "FIRST":
                    combined_operand_indices = tuple(
                        intersect_sop(
                            forward_left_idx, reverse_out_idx
                        )
                        for forward_left_idx, reverse_out_idx in zip(
                            index_visitor.registry[node.left],
                            reverse_index_visitor.registry[node],
                        )
                    )

                    redundant = all(
                        all(
                            any(
                                right_index_term <= out_index_term
                                for out_index_term in out_idxset
                            )
                            for right_index_term in right_idxset
                        )
                        for right_idxset, out_idxset in zip(
                            index_visitor.registry[node.right], combined_operand_indices
                        )
                    )
                    unnecessary = node_indices == index_visitor.registry[node.left]
                    if unnecessary or redundant:
                        #print("Removing unnecessary", node.description, redundant, unnecessary)
                        root = replace_node(root, node, node.left)
                        keep_optimizing = True
                elif node.op == "SECOND":
                    combined_operand_indices = tuple(
                        intersect_sop(
                            forward_right_idx, reverse_out_idx
                        )
                        for forward_right_idx, reverse_out_idx in zip(
                            index_visitor.registry[node.right],
                            reverse_index_visitor.registry[node],
                        )
                    )

                    redundant = all(
                        all(
                            any(
                                left_index_term <= out_index_term
                                for out_index_term in out_idxset
                            )
                            for left_index_term in left_idxset
                        )
                        for left_idxset, out_idxset in zip(
                            index_visitor.registry[node.left], combined_operand_indices
                        )
                    )
                    unnecessary = node_indices == index_visitor.registry[node.right]
                    # print(redundant, unnecessary)
                    if unnecessary or redundant:
                        #print("Removing unnecessary SECOND", redundant, unnecessary)
                        root = replace_node(root, node, node.right)
                        keep_optimizing = True
                elif isinstance(node.left, (LAIR.InnerBroadcastMask, LAIR.InnerBroadcast)) and isinstance(node.right, (LAIR.InnerBroadcastMask, LAIR.InnerBroadcast)) and node.left.op == "FIRST" and node.right.op == "FIRST" and node.left.right is node.right.right and node.op not in util.BinaryBooleanGrbOps:
                    left_decoder = dict(zip(node.left.out_indices, range(len(node.left.out_indices))))
                    right_decoder = dict(zip(node.right.out_indices, range(len(node.right.out_indices))))

                    out_indices = tuple(range(len(left_decoder)))
                    left_indices = tuple(left_decoder[idx] for idx in node.left.left_indices)
                    right_indices = tuple(right_decoder[idx] for idx in node.right.left_indices)

                    #print("&"*80)
                    #print(node.op)
                    #print("&"*80)

                    later_mask = LAIR.Mask(
                        LAIR.InnerBroadcast(node.op, node.left.left, left_indices, node.right.left, right_indices, out_indices),
                        node.right.right
                    )
                    root = replace_node(root, node, later_mask)
                    keep_optimizing = True
            elif isinstance(node, LAIR.InnerBroadcastMask):
                if set(node.left_indices) == set(node.out_indices) and set(node.right_indices) == set(node.out_indices):
                    left, right = node.left, node.right
                    if node.left_indices != node.out_indices:
                        new_indices = [list(node.out_indices).index(i) for i in node.left_indices]
                        left = LAIR.PermuteIndices(left, new_indices)
                    if node.right_indices != node.out_indices:
                        new_indices = [list(node.out_indices).index(i) for i in node.right_indices]
                        right = LAIR.PermuteIndices(right, new_indices)
                    root = replace_node(root, node, LAIR.Mask(left, right))
                    keep_optimizing = True
            elif isinstance(node, LAIR.InnerBroadcast):
                #print("Checking", node.pattern, isinstance(node.left, (LAIR.InnerBroadcastMask, LAIR.InnerBroadcast)))
                if set(node.left_indices) == set(node.out_indices) and set(node.right_indices) == set(node.out_indices):
                    left, right = node.left, node.right
                    if node.left_indices != node.out_indices:
                        left = LAIR.PermuteIndices(left, node.left_indices)
                    if node.right_indices != node.out_indices:
                        right = LAIR.PermuteIndices(right, node.right_indices)
                    root = replace_node(root, node, LAIR.ElementwiseMultiply(node.op, left, right))
                    keep_optimizing = True
                elif isinstance(node.left, (LAIR.InnerBroadcastMask, LAIR.InnerBroadcast)) and node.left.op == "FIRST" and len(node.left.right_indices) == len(node.out_indices):
                    left_out_encoder = dict(zip(node.left.out_indices, range(len(node.left.out_indices))))
                    left_out_decoder = dict(zip(range(len(node.left.out_indices)), node.left.out_indices))
                    out_encoder = dict(zip(node.out_indices, range(len(node.out_indices))))
                    out_decoder = dict(zip(range(len(node.out_indices)), node.out_indices))

                    out_indices = tuple(range(len(node.out_indices)))
                    left_indices = tuple(out_encoder[left_out_encoder[idx]] for idx in node.left.left_indices)
                    right_indices = tuple(out_encoder[idx] for idx in node.right_indices)
                    mask_indices = tuple(out_encoder[left_out_encoder[idx]] for idx in node.left.right_indices)

                    later_mask = LAIR.Mask(
                        LAIR.InnerBroadcast(node.op, node.left.left, left_indices, node.right, right_indices, out_indices),
                        LAIR.PermuteIndices(node.left.right, mask_indices)
                    )
                    root = replace_node(root, node, later_mask)
                    keep_optimizing = True
                elif isinstance(node.right, (LAIR.InnerBroadcastMask, LAIR.InnerBroadcast)) and node.right.op == "FIRST" and len(node.right.right_indices) == len(node.out_indices):
                    right_out_encoder = dict(zip(node.right.out_indices, range(len(node.right.out_indices))))
                    right_out_decoder = dict(zip(range(len(node.right.out_indices)), node.right.out_indices))
                    out_encoder = dict(zip(node.out_indices, range(len(node.out_indices))))
                    out_decoder = dict(zip(range(len(node.out_indices)), node.out_indices))

                    out_indices = tuple(range(len(node.out_indices)))
                    right_indices = tuple(out_encoder[right_out_encoder[idx]] for idx in node.right.left_indices)
                    left_indices = tuple(out_encoder[idx] for idx in node.left_indices)
                    mask_indices = tuple(out_encoder[right_out_encoder[idx]] for idx in node.right.right_indices)

                    later_mask = LAIR.Mask(
                        LAIR.InnerBroadcast(node.op, node.left, left_indices, node.right.left, right_indices, out_indices),
                        LAIR.PermuteIndices(node.right.right, mask_indices)
                    )
                    root = replace_node(root, node, later_mask)
                    keep_optimizing = True
            elif isinstance(node, LAIR.ElementwiseMultiply):
                if isinstance(node.left, (LAIR.InnerBroadcastMask, LAIR.InnerBroadcast)) and node.left.op == "FIRST" and node.op not in util.BinaryBooleanGrbOps:
                    left_decoder = dict(zip(node.left.out_indices, range(len(node.left.out_indices))))
                    left_indices = tuple(left_decoder[idx] for idx in node.left.left_indices)
                    out_indices = tuple(range(len(left_decoder)))
                    right_indices = out_indices

                    later_mask = LAIR.Mask(
                        LAIR.InnerBroadcast(node.op, node.left.left, left_indices, node.right, right_indices, out_indices),
                        node.left.right
                    )
                    root = replace_node(root, node, later_mask)
                    keep_optimizing = True
                elif isinstance(node.right, (LAIR.InnerBroadcastMask, LAIR.InnerBroadcast)) and node.right.op == "FIRST" and node.op not in util.BinaryBooleanGrbOps:
                    right_decoder = dict(zip(node.right.out_indices, range(len(node.right.out_indices))))
                    right_indices = tuple(right_decoder[idx] for idx in node.right.left_indices)
                    out_indices = tuple(range(len(right_decoder)))
                    left_indices = out_indices

                    later_mask = LAIR.Mask(
                        LAIR.InnerBroadcast(node.op, node.left, left_indices, node.right.left, right_indices, out_indices),
                        node.right.right
                    )
                    root = replace_node(root, node, later_mask)
                    keep_optimizing = True
            elif isinstance(node, LAIR.Reduction):
                reduction = node
                broadcast = None
                operand = None

                if isinstance(node.operand, (LAIR.InnerBroadcast, LAIR.InnerBroadcastMask)):
                    broadcast = node.operand
                    operand = broadcast.left
#                elif isinstance(node.operand, (LAIR.ElementwiseApply,)):
#                    if isinstance(node.operand.operand, (LAIR.InnerBroadcast, LAIR.InnerBroadcastMask)) and node.operand.operand.op == 'FIRST':
#                        ewiseapply = node.operand
#                        broadcast = ewiseapply.operand
#                        operand = LAIR.ElementwiseApply(ewiseapply.op, broadcast.left)


                if isinstance(broadcast, (LAIR.InnerBroadcast, LAIR.InnerBroadcastMask)):
                    add_op, mul_op = node.op, broadcast.op
                    replacement = None
                    if debug:
                        print("")
                        print("=" * 120)
                        print(add_op, mul_op, broadcast.pattern, (add_op, mul_op) in valid_semirings, broadcast.pattern == ((0, 2), (2, 1), (0, 1, 2)))
                        print(reduction.dsl)
                        print("=" * 120)
                    if (add_op, mul_op) in valid_semirings:
                        remaining_reduction_axes = set(reduction.axes)

                        def is_matrix_multiply(pattern, reduction_axes):
                            if tuple(len(x) for x in pattern) != (2, 2, 3):
                                return False
                            elif set(pattern[0]) == set(pattern[1]):
                                return False

                            shared_items = set(pattern[0]) & set(pattern[1])
                            if len(shared_items) != 1:
                                return False

                            shared_item = tuple(shared_items)[0]
                            if shared_item not in reduction_axes:
                                return False
                            
                            return True

                        #print(f"{is_matrix_multiply(broadcast.pattern, reduction.axes)=}")
                        if is_matrix_multiply(broadcast.pattern, reduction.axes):
                            shared_item = tuple(set(broadcast.pattern[0]) & set(broadcast.pattern[1]))[0]
                            
                            left, right = operand, broadcast.right
                            if shared_item == broadcast.pattern[0][0]:
                                left = LAIR.PermuteIndices(left, (1, 0))
                            if shared_item == broadcast.pattern[1][1]:
                                right = LAIR.PermuteIndices(right, (1, 0))

                            remaining_reduction_axes -= {shared_item}
                            replacement = LAIR.DotMxM(left, right, add_op, mul_op)
#                        elif broadcast.pattern == ((0, 2), (2, 1), (0, 1, 2)) and 2 in reduction.axes:
#                            replacement = LAIR.DotMxM(
#                                operand, broadcast.right, add_op, mul_op
#                            )
#                            remaining_reduction_axes -= {2,}
#                        elif broadcast.pattern == ((0, 2), (1, 2), (0, 1, 2)) and 2 in reduction.axes:
#                            replacement = LAIR.DotMxM(
#                                operand, LAIR.PermuteIndices(broadcast.right, (1, 0)), add_op, mul_op
#                            )
#                            remaining_reduction_axes -= {2,}
#                        elif broadcast.pattern == ((2, 0), (1, 2), (0, 1, 2)) and 2 in reduction.axes:
#                            replacement = LAIR.DotMxM(
#                                LAIR.PermuteIndices(operand, (1, 0)),
#                                LAIR.PermuteIndices(broadcast.right, (1, 0)),
#                                add_op, mul_op
#                            )
#                            remaining_reduction_axes -= {2,}
#                        elif broadcast.pattern == ((2, 0), (2, 1), (0, 1, 2)) and 2 in reduction.axes:
#                            replacement = LAIR.DotMxM(
#                                LAIR.PermuteIndices(operand, (1, 0)),
#                                broadcast.right,
#                                add_op, mul_op
#                            )
#                            remaining_reduction_axes -= {2,}
                        elif (len(broadcast.pattern[0]), len(broadcast.pattern[1]), len(broadcast.pattern[2])) in ((2, 1, 2), (1, 2, 2)) and (0, 1) == reduction.axes:
                            if len(broadcast.pattern[0]) == 2:
                                matrix, vector = operand, broadcast.right
                                matrix_pattern, vector_pattern, out_pattern = broadcast.pattern
                            else:
                                vector, matrix = operand, broadcast.right
                                vector_pattern, matrix_pattern, out_pattern = broadcast.pattern

                            n_matrix = LAIR.Reduction(
                                reduction.op, matrix,
                                axes=tuple(set(matrix_pattern) - set(vector_pattern))
                            )
                            replacement = LAIR.InnerBroadcast(
                                mul_op, n_matrix, (0,), vector, (0,), (0,)
                            )
                            remaining_reduction_axes = {0}
                        elif (len(broadcast.pattern[0]), len(broadcast.pattern[1]), len(broadcast.pattern[2])) in ((2, 1, 2),) and broadcast.pattern[1] == reduction.axes:
                            if broadcast.pattern[0] == (0, 1):
                                matrix = operand
                            else:
                                matrix = LAIR.PermuteIndices(operand)
                            vector = broadcast.right

                            if broadcast.pattern[1] == (0,):
                                replacement = LAIR.DotVxM(
                                    vector, matrix, add_op, mul_op
                                )
                            elif broadcast.pattern[1] == (1,):
                                replacement = LAIR.DotMxV(
                                    matrix, vector, add_op, mul_op
                                )
                            else:
                                assert False

                            remaining_reduction_axes = set()
                        elif (len(broadcast.pattern[0]), len(broadcast.pattern[1]), len(broadcast.pattern[2])) in ((1, 2, 2),) and broadcast.pattern[0] == reduction.axes:
                            if broadcast.pattern[1] == (0, 1):
                                matrix = broadcast.right
                            else:
                                matrix = LAIR.PermuteIndices(broadcast.right)
                            vector = operand

                            if broadcast.pattern[0] == (0,):
                                replacement = LAIR.DotVxM(
                                    vector, matrix, add_op, mul_op
                                )
                            elif broadcast.pattern[0] == (1,):
                                replacement = LAIR.DotMxV(
                                    matrix, vector, add_op, mul_op
                                )
                            else:
                                assert False

                            remaining_reduction_axes = set()
                        elif broadcast.pattern == ((0, 1), (1, 2), (0, 1, 2)) and 1 in reduction.axes:
                            replacement = LAIR.DotMxM(
                                operand, broadcast.right, add_op, mul_op
                            )
                            remaining_reduction_axes -= {1,}
                        elif broadcast.pattern == ((1,), (1, 0), (0, 1)) and 1 in reduction.axes:
                            replacement = LAIR.DotVxM(
                                operand, broadcast.right, add_op, mul_op
                            )
                            remaining_reduction_axes -= {1,}
                        elif broadcast.pattern == ((1,), (0, 1), (0, 1)) and 1 in reduction.axes:
                            replacement = LAIR.DotVxM(
                                operand, LAIR.PermuteIndices(broadcast.right), add_op, mul_op
                            )
                            remaining_reduction_axes -= {1,}
                        elif broadcast.pattern == ((0,1), (1,), (0,1)) and 1 in reduction.axes:
                            replacement = LAIR.DotMxV(
                                operand, broadcast.right, add_op, mul_op
                            )
                            remaining_reduction_axes -= {1,}
                        elif broadcast.pattern == ((1,0), (1,), (0,1)) and 1 in reduction.axes:
                            replacement = LAIR.DotMxV(
                                LAIR.PermuteIndices(operand), broadcast.right, add_op, mul_op
                            )
                            remaining_reduction_axes -= {0,}

                        if replacement is not None:
                            #if ewiseapply is not None:
                            #    replacement = LAIR.ElementwiseApply(ewiseapply.op, replacement)

                            new_reduce_axes = set()
                            removed_count = 0
                            for ax in reduction.axes:
                                if ax in remaining_reduction_axes:
                                    new_reduce_axes.add(ax - removed_count)
                                else:
                                    removed_count += 1

                            if len(remaining_reduction_axes) > 0:
                                replacement = LAIR.Reduction(reduction.op, replacement, tuple(new_reduce_axes))

                if replacement is not None:
                    root = replace_node(root, node, replacement)
                    keep_optimizing = True
            if keep_optimizing:
                break

    return root

tensql/LAIR/Optimizer.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class Output(Node):
    def __init__(self, operand: Node, name: str) -> None:
        self._operand = operand
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def description(self):
        return f"{type(self).__name__}({self.name!r})"

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self.operand.dsl}, name={self.name!r})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "name": self._name}

tensql/LAIR/Output.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class Pattern(Node):
    def __init__(self, operand: Node, name: Optional[str] = None) -> None:
        self._operand = operand
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self.operand.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "name": self._name}

tensql/LAIR/Pattern.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node
from ..util import Axes


class PermuteIndices(Node):
    def __init__(
        self, operand: Node, new_axes: Optional[Axes], name: Optional[str] = None
    ) -> None:
        self._operand = operand
        if new_axes is None:
            self._new_axes = None
        else:
            self._new_axes = tuple(new_axes)
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def new_axes(self) -> Optional[Axes]:
        return self._new_axes

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._new_axes})"

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self.operand.dsl}, new_axes={self._new_axes})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "operand": self._operand,
            "new_axes": self._new_axes,
            "name": self._name,
        }

tensql/LAIR/PermuteIndices.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple
from ..util import GBTensor

from .Node import Node

from ..Types import Type

class PrimaryKey(Node):
    def __init__(
        self, stencil: GBTensor, colnum: int, type_: Type, *, name: str
    ) -> None:
        super().__init__(name)
        self._type = type_
        self._stencil = stencil
        self._colnum = colnum

    @property
    def stencil(self) -> GBTensor:
        return self._stencil

    @property
    def colnum(self) -> int:
        return self._colnum

    @property
    def type_(self) -> Type:
        return self._type

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("stencil", self._stencil)

    @property
    def dsl(self) -> str:
        return (
            f"{type(self).__name__}(name={self._name!r}, stencil={self._stencil.dsl})"
        )

    @property
    def description(self) -> str:
        return f"{type(self).__name__}(name={self._name!r})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "colnum": self._colnum,
            "stencil": self._stencil,
            "name": self._name,
            "type_": self._type,
        }

tensql/LAIR/PrimaryKey.py
sandialabs/TenSQL
from typing import Generator, Optional, Tuple, Dict, Any

from ..util import EncoderDecoder, SequentialIdGenerator

from .DotMxM import DotMxM

from .Node import Node
from .Mask import Mask
from .ElementwiseMultiply import ElementwiseMultiply
from .CastLike import CastLike
from .PermuteIndices import PermuteIndices
from .ElementwiseApply import ElementwiseApplyBindSecond
from .VectorToDiagMatrix import VectorToDiagMatrix

AxesSpecifier = Tuple[int, ...]

class Promote(Node):
    def __init__(
        self,
        operand: Node,
        operand_indices: AxesSpecifier,
        out_indices: AxesSpecifier,
        name: Optional[str] = None,
    ) -> None:
        self._operand = operand

        idx_re_encoder = EncoderDecoder(SequentialIdGenerator())
        self._operand_indices = tuple(idx_re_encoder[idx] for idx in operand_indices)
        self._out_indices = tuple(idx_re_encoder[idx] for idx in out_indices)

        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def signature(self) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
        return (self._operand_indices, self._out_indices)

    @property
    def operand_indices(self) -> Node:
        return self._operand_indices

    @property
    def out_indices(self) -> Node:
        return self._out_indices

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self.operand.dsl}, {self._operand_indices}, {self._out_indices})"

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._operand_indices}, {self._out_indices})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "operand": self._operand,
            "operand_indices": self._operand_indices,
            "out_indices": self._out_indices,
            "name": self._name,
        }


def OptimizedPromote(
    operand: Node,
    operand_indices: AxesSpecifier,
    mask: Node,
    mask_indices: AxesSpecifier,
) -> Node:
    assert set(operand_indices) - set(mask_indices) == set()

    return Mask(Promote(operand, operand_indices, mask_indices), mask)

tensql/LAIR/Promote.py
sandialabs/TenSQL
from typing import Generator, Optional, Sequence, Dict, Any, Tuple

from .Node import Node
from ..util import Axes


class Reduction(Node):
    def __init__(
        self, op: str, operand: Node, axes: Axes, name: Optional[str] = None
    ) -> None:
        self._operand = operand
        self._op = op
        self._axes = tuple(axes)
        super().__init__(name)

    @property
    def axes(self) -> Axes:
        return self._axes

    @property
    def operand(self) -> Node:
        return self._operand

    @property
    def op(self) -> str:
        return self._op

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def description(self) -> str:
        return f"{type(self).__name__}({self._op!r}, axes={self.axes})"

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self.operand.dsl}, axes={self.axes})"

    def get_attributes(self) -> Dict[str, Any]:
        return {
            "op": self._op,
            "operand": self._operand,
            "axes": self._axes,
            "name": self._name,
        }


class Sum(Reduction):
    def __init__(self, operand: Node, axes: Axes, name: Optional[str] = None):
        super().__init__("+", operand, axes, name)


class Min(Reduction):
    def __init__(self, operand: Node, axes: Axes, name: Optional[str] = None):
        super().__init__("MIN", operand, axes, name)


class Max(Reduction):
    def __init__(self, operand: Node, axes: Axes, name: Optional[str] = None):
        super().__init__("MAX", operand, axes, name)

tensql/LAIR/Reduction.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple
from ..util import GBTensor

from .Node import Node


class Tensor(Node):
    def __init__(
        self, tensor: GBTensor, stencil: Optional[GBTensor] = None, *, name: str
    ) -> None:
        self._tensor = tensor
        self._stencil = stencil
        super().__init__(name)

    @property
    def tensor(self) -> GBTensor:
        return self._tensor

    @property
    def stencil(self) -> Optional[GBTensor]:
        return self._stencil

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        if self._stencil is not None:
            yield ("stencil", self._stencil)
        else:
            yield from tuple()

    @property
    def dsl(self) -> str:
        if self._stencil is None:
            return f"{type(self).__name__}(name={self._name!r})"
        else:
            return f"{type(self).__name__}(name={self._name!r}, stencil={self._stencil.dsl})"

    @property
    def description(self) -> str:
        return f"{type(self).__name__}(name={self._name!r})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"tensor": self._tensor, "stencil": self._stencil, "name": self._name}

tensql/LAIR/Tensor.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class VectorToColumnMatrix(Node):
    def __init__(self, operand: Node, name: Optional[str] = None) -> None:
        self._operand = operand
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self.operand.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "name": self._name}

tensql/LAIR/VectorToColumnMatrix.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class VectorToDiagMatrix(Node):
    def __init__(self, operand: Node, name: Optional[str] = None) -> None:
        self._operand = operand
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self.operand.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "name": self._name}

tensql/LAIR/VectorToDiagMatrix.py
sandialabs/TenSQL
from typing import Generator, Optional, Dict, Any, Tuple

from .Node import Node


class VectorToRowMatrix(Node):
    def __init__(self, operand: Node, name: Optional[str] = None) -> None:
        self._operand = operand
        super().__init__(name)

    @property
    def operand(self) -> Node:
        return self._operand

    def __iter__(self) -> Generator[Tuple[str, Node], None, None]:
        yield ("operand", self._operand)

    @property
    def dsl(self) -> str:
        return f"{type(self).__name__}({self.operand.dsl})"

    def get_attributes(self) -> Dict[str, Any]:
        return {"operand": self._operand, "name": self._name}

tensql/LAIR/VectorToRowMatrix.py
sandialabs/TenSQL
from __future__ import annotations

from typing import Any, List, Dict, Sequence, Optional
import weakref
import sys
import time

from .. import LAIR
from .. import util
from ..util import GBTensor, grb_shape, ndim, extract_indices, BinaryBooleanGrbOps, build_tensor

import pygraphblas
import pygraphblas.types
from pygraphblas import Scalar, Vector, Matrix
import numpy as np

VisitorReturnType = Optional[GBTensor]

class TempResult:
    def __init__(
        self,
        node: LAIR.Node,
        result: VisitorReturnType,
        children: Sequence[TempResult],
        is_output: bool,
    ):
        self.node = node
        self.result = result
        self.children = tuple(children)
        self.is_output = is_output

class Visitor:
    def visit(self, node: LAIR.Node, *args: Any) -> Any:
        for cls in type(node).mro():
            if issubclass(cls, LAIR.Node):
                visit_name = f"visit_{cls.__name__}"
                if hasattr(self, visit_name):
                    return getattr(self, visit_name)(node, *args)
        return self.generic_visit(node)

    def generic_visit(self, node: LAIR.Node) -> None:
        classes = [cls for cls in type(node).mro() if issubclass(cls, LAIR.Node)]
        if len(classes) > 0:
            raise NotImplementedError(
                f"{type(self).__name__} does not implement a visit method for any of {classes}"
            )
        else:
            raise ValueError(f"Invalid argument: {node!r}")


class ExpressionVisitor(Visitor):
    UnaryPyOps = util.UnaryPyOps
    UnaryGrbOps = util.UnaryGrbOps
    BinaryPyOps = util.BinaryPyOps
    BinaryGrbOps = util.BinaryGrbOps

    def __init__(self, verbose=False) -> None:
        self._expression_cache: Optional[
            Dict[LAIR.Node, TempResult]
        ] = weakref.WeakValueDictionary()
        self.verbose = verbose
        self.recursion = 0

    def _dfs_flatten(self, node, *, is_top_level=True) -> Sequence[TempResult]:
        children = []
        for relationship, child in node:
            children.append(self._dfs_flatten(child, is_top_level=False))
        tmp_result = self._expression_cache.get(node)
        if tmp_result is None:
            tmp_result = TempResult(node, None, children, is_top_level)
            self._expression_cache[node] = tmp_result
        return tmp_result

    def run(self, *nodes: LAIR.Node) -> Sequence[VisitorReturnType]:
        # Perform some reference counting witchcraft with the expression cache to avoid
        # recomputing results, but also avoid wasting memory by storing all temporary results.
        # The expression cache is a weakref.WeakValueDictionary, so Python frees memory of the
        # temporary results when they're no longer in use.  The cache only does anything useful
        # if self._expression_cache has been prepopulated by self.run().

        old_expression_cache_size: int = len(self._expression_cache)

        outputs: List[TempResult] = []
        for node in nodes:
            outputs.append(self._dfs_flatten(node))

        while len(outputs) > 0:
            tmp_result = outputs.pop(0)
            yield self.visit(tmp_result.node)
            del tmp_result

        # Make sure we're not leaking memory
        assert len(self._expression_cache) == old_expression_cache_size

    def visit(self, node: LAIR.Node, *args: Any) -> VisitorReturnType:
        if self.verbose:
            print(" " * self.recursion + "Visiting", node.description)
            sys.stdout.flush()
        self.recursion += 1

        # Perform some reference counting witchcraft with the expression cache to avoid
        # recomputing results, but also avoid wasting memory by storing all temporary results.
        # The expression cache is a weakref.WeakValueDictionary, so Python frees memory of the
        # temporary results when they're no longer in use.  The cache only does anything useful
        # if self._expression_cache has been prepopulated by self.run().
        tmp_result = self._expression_cache.get(node)
        if tmp_result is not None:
            if tmp_result.result is None:
                if self.verbose:
                    print(
                        " " * self.recursion + "Computing result for", node.description
                    )
                start_time = time.time()
                tmp_result.result = super().visit(node, *args)
                stop_time = time.time()
                if self.verbose:
                    print(
                        " " * self.recursion
                        + f"Completed in {stop_time - start_time} seconds",
                        node.description,
                    )

                tmp_result.children = None
            else:
                if self.verbose:
                    print(
                        " " * self.recursion + "Found cached result for",
                        node.description,
                    )
            ret = tmp_result.result
        else:
            if self.verbose:
                print(
                    " " * self.recursion + "Expression cache was missing an entry for",
                    node.description,
                )
            ret = super().visit(node, *args)

        self.recursion -= 1

        if self.verbose:
            print(" " * self.recursion + "Finished", node.description)
            sys.stdout.flush()

        return ret

    def visit_CastLike(self, node: LAIR.CastLike) -> VisitorReturnType:
        x = self.visit(node.operand)
        l = self.visit(node.like)
        if isinstance(x, Scalar):
            ret = Scalar.from_type(l.type)
            if x.nvals > 0:
                ret[None] = x[0]
            return ret
        elif isinstance(x, (Vector, Matrix)):
            return x.cast(l.type)
        else:
            raise ValueError(
                "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

    def visit_Cast(self, node: LAIR.Cast) -> VisitorReturnType:
        x = self.visit(node.operand)
        if isinstance(x, Scalar):
            ret = Scalar.from_type(node.type_)
            if x.nvals > 0:
                ret[None] = x[0]
            return ret
        elif isinstance(x, Vector):
            ret = Vector.sparse(typ=node.type_, size=x.size)
            identity = getattr(x.type, 'IDENTITY')
            return x.apply(identity, out=ret)
        elif isinstance(x, Matrix):
            ret = Matrix.sparse(node.type_, nrows=x.nrows, ncols=x.ncols)
            identity = getattr(x.type, 'IDENTITY')
            return x.apply(identity, out=ret)
        else:
            raise ValueError(
                "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

    def visit_DotMxM(self, node: LAIR.DotMxM) -> VisitorReturnType:
        A = self.visit(node.left)
        B = self.visit(node.right)

        assert isinstance(A, Matrix)
        assert isinstance(B, Matrix)

        if (node.add_op, node.mul_op) != ("+", "*"):
            typ: pygraphblas.types.Type = A.type
            add_op = getattr(typ, self.BinaryGrbOps[node.add_op])
            mul_op = getattr(typ, self.BinaryGrbOps[node.mul_op])
            monoid = typ.new_monoid(add_op, typ.default_zero)
            semiring = typ.new_semiring(monoid, mul_op)
            return A.mxm(B, semiring=semiring)
        else:
            return A.mxm(B)

    def visit_DotMxV(self, node: LAIR.DotMxM) -> VisitorReturnType:
        A = self.visit(node.left)
        x = self.visit(node.right)

        assert isinstance(A, Matrix)
        assert isinstance(x, Vector)

        typ: pygraphblas.types.Type = A.type
        add_op = getattr(typ, self.BinaryGrbOps[node.add_op])
        mul_op = getattr(typ, self.BinaryGrbOps[node.mul_op])
        monoid = typ.new_monoid(add_op, typ.default_zero)
        semiring = typ.new_semiring(monoid, mul_op)

        return A.mxv(x, semiring=semiring)

    def visit_DotVxM(self, node: LAIR.DotMxM) -> VisitorReturnType:
        x = self.visit(node.left)
        A = self.visit(node.right)

        assert isinstance(x, Vector)
        assert isinstance(A, Matrix)

        typ: pygraphblas.types.Type = A.type
        add_op = getattr(typ, self.BinaryGrbOps[node.add_op])
        mul_op = getattr(typ, self.BinaryGrbOps[node.mul_op])
        monoid = typ.new_monoid(add_op, typ.default_zero)
        semiring = typ.new_semiring(monoid, mul_op)

        return x.vxm(A, semiring=semiring)

    def visit_ElementwiseAdd(self, node: LAIR.ElementwiseAdd) -> VisitorReturnType:
        x = self.visit(node.left)
        y = self.visit(node.right)

        assert x.type == y.type

        op_name = self.BinaryGrbOps[node.op]
        grb_op = getattr(x.type, op_name)

        if isinstance(x, Scalar) and isinstance(y, Scalar):
            ret = Scalar.from_value(self.BinaryPyOps[node.op](x[0], y[0]))
        elif (not isinstance(x, Scalar)) and isinstance(y, Scalar):
            ret = x.apply_first(y, grb_op)
        elif isinstance(x, Scalar) and (not isinstance(y, Scalar)):
            ret = y.apply_second(y, grb_op)
        elif (
            isinstance(x, (Vector, Matrix))
            and isinstance(y, (Vector, Matrix))
            and ndim(x) == ndim(y)
        ):
            ret = x.eadd(y, grb_op)
        else:
            raise ValueError(
                "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

        return ret

    def visit_ElementwiseApply(self, node: LAIR.ElementwiseApply) -> VisitorReturnType:
        x = self.visit(node.operand)
        if isinstance(x, Scalar):
            if x.nvals > 0:
                ret = Scalar.from_value(self.UnaryPyOps[node.op](x[0]))
            else:
                ret = Scalar.from_type(x.type)
        elif isinstance(x, (Vector, Matrix)):
            ret = x.apply(getattr(x.type, self.UnaryGrbOps[node.op]))
        else:
            raise ValueError(
                "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )
        return ret

    def visit_ElementwiseApplyBindFirst(
        self, node: LAIR.ElementwiseApplyBindFirst
    ) -> VisitorReturnType:
        x = self.visit(node.operand)
        s = self.visit(node.scalar)

        assert isinstance(s, Scalar)

        if isinstance(x, Scalar):
            if x.nvals > 0:
                ret = Scalar.from_value(self.BinaryPyOps[node.op](s[0], x[0]))
            else:
                ret = Scalar.from_type(x.type)
        elif isinstance(x, (Vector, Matrix)):
            op_name = self.BinaryGrbOps[node.op]
            grb_op = getattr(x.type, op_name)
            ret = x.apply_first(s, grb_op)
        else:
            raise ValueError(
                "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )
        return ret

    def visit_ElementwiseApplyBindSecond(
        self, node: LAIR.ElementwiseApplyBindSecond
    ) -> VisitorReturnType:
        x = self.visit(node.operand)
        s = self.visit(node.scalar)

        assert isinstance(s, Scalar)

        if isinstance(x, Scalar):
            if x.nvals > 0:
                ret = Scalar.from_value(self.BinaryPyOps[node.op](x[0], s[0]))
            else:
                ret = Scalar.from_type(x.type)
        elif isinstance(x, (Vector, Matrix)):
            op_name = self.BinaryGrbOps[node.op]
            grb_op = getattr(x.type, op_name)
            ret = x.apply_second(s, grb_op)
        else:
            raise ValueError(
                "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )
        return ret

    def visit_InnerBroadcastMask(self, node: LAIR.InnerBroadcastMask) -> VisitorReturnType:
        NDIM = tuple()

        x = self.visit(node.left)
        y = self.visit(node.right).nonzero()

        #print(f"{node.pattern=}, {x.type=}, {y.type=}")

        #BOOL = pygraphblas.types.BOOL
        if y.type != x.type:
            if isinstance(y, Scalar):
                ret = Scalar.from_type(x.type)
                if y.nvals > 0:
                    ret[None] = y[0]
            elif isinstance(y, Vector):
                ret = Vector.sparse(typ=x.type, size=y.size)
                identity = getattr(y.type, 'IDENTITY')
                y.apply(identity, out=ret)
            elif isinstance(y, Matrix):
                ret = Matrix.sparse(x.type, nrows=y.nrows, ncols=y.ncols)
                identity = getattr(y.type, 'IDENTITY')
                y.apply(identity, out=ret)
            else:
                raise ValueError(
                    "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
                )
            y = ret
            del ret
        #print("Executing InnerBroadcastMask", node.pattern)
        #print(f"{node.pattern=}, {x.type=}, {y.type=}")

        op_name = self.BinaryGrbOps[node.op]
        grb_op = getattr(x.type, op_name)

        typ: pygraphblas.types.Type = x.type
        add_op = getattr(typ, "ANY")
        mul_op = getattr(typ, op_name)
        monoid = typ.new_monoid(add_op, typ.default_zero)
        mxm_semiring = typ.new_semiring(monoid, mul_op)
#        monoid.print(level=5)
#        mxm_semiring.print(level=5)
#        print(f"{node.pattern=}, {typ=}, {x.type=}, {y.type=}, {typ.default_zero=}, {op_name=}")

        out_shape = [None for _ in range(len(node.pattern[2]))]
        for it_in, dim_in in enumerate(node.pattern[0]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    out_shape[it_out] = grb_shape(x)[it_in]
        for it_in, dim_in in enumerate(node.pattern[1]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    out_shape[it_out] = grb_shape(y)[it_in]

        assert None not in out_shape

        ret = build_tensor(x.type, out_shape)

        if x.nvals == 0 or y.nvals == 0:
            pass

        elif node.pattern == (NDIM, NDIM, NDIM):
            ret[None] = self.BinaryPyOps[op_name](x[0], y[0])

        elif node.pattern == (NDIM, (0,), (0,)):
            y.apply_first(x, grb_op, out=ret)
        elif node.pattern == (NDIM, (0,1), (0,1)):
            y.apply_first(x, grb_op, out=ret)
        elif node.pattern == (NDIM, (1,0), (0,1)):
            y.apply_first(x, grb_op, out=ret)
            ret = ret.T

        elif node.pattern == ((0,), NDIM, (0,)):
            x.apply_second(y, grb_op, out=ret)
        elif node.pattern == ((0,1), NDIM, (0,1)):
            x.apply_second(y, grb_op, out=ret)
        elif node.pattern == ((1,0), NDIM, (0,1)):
            x.apply_second(y, grb_op, out=ret)
            ret = ret.T

        elif node.pattern == ((0,), (0,), (0,)):
            x.emult(y, mult_op=grb_op, out=ret)
        elif node.pattern == ((0,), (1,), (0, 1)):
            mx = Matrix.sparse(x.type, x.size, 1)
            mx.assign_col(0, x)
            my = Matrix.sparse(y.type, 1, y.size)
            my.assign_row(0, y)
            mx.mxm(my, semiring=mxm_semiring, out=ret)
        elif node.pattern == ((1,), (0,), (0, 1)):
            mx = Matrix.sparse(x.type, x.size, 1)
            mx.assign_col(0, x)
            my = Matrix.sparse(y.type, 1, y.size)
            my.assign_row(0, y)
            mx.mxm(my, semiring=mxm_semiring, out=ret)
            ret = ret.T

        elif node.pattern == ((0,), (0, 1), (0, 1)):
            mx = Matrix.from_diag(x)
            mx.mxm(y, semiring=mxm_semiring, out=ret)
        elif node.pattern == ((1,), (0, 1), (0, 1)):
            mx = Matrix.from_diag(x)
            mx.mxm(y.T, semiring=mxm_semiring, out=ret)
            ret = ret.T
        elif node.pattern == ((0,), (1, 0), (0, 1)):
            mx = Matrix.from_diag(x)
            mx.mxm(y.T, semiring=mxm_semiring, out=ret)
        elif node.pattern == ((1,), (1, 0), (0, 1)):
            mx = Matrix.from_diag(x)
            mx.mxm(y, semiring=mxm_semiring, out=ret)
            ret = ret.T

        elif node.pattern == ((0, 1), (0,), (0, 1)):
            my = Matrix.from_diag(y)
            x.T.mxm(my, semiring=mxm_semiring, out=ret)
            ret = ret.T
        elif node.pattern == ((0, 1), (1,), (0, 1)):
            my = Matrix.from_diag(y)
            x.mxm(my, semiring=mxm_semiring, out=ret)
        elif node.pattern == ((1, 0), (0,), (0, 1)):
            my = Matrix.from_diag(y)
            x.mxm(my, semiring=mxm_semiring, out=ret)
            ret = ret.T
        elif node.pattern == ((1, 0), (1,), (0, 1)):
            my = Matrix.from_diag(y)
            x.T.mxm(my, semiring=mxm_semiring, out=ret)

        elif node.pattern == ((0,1), (0,1), (0,1)):
            x.emult(y, mult_op=grb_op, out=ret)
        elif node.pattern == ((0,1), (1,0), (0,1)):
            x.emult(y.T, mult_op=grb_op, out=ret)
        elif node.pattern == ((1,0), (0,1), (0,1)):
            x.T.emult(y, mult_op=grb_op, out=ret)
        elif node.pattern == ((1,0), (1,0), (0,1)):
            x.emult(y, mult_op=grb_op, out=ret)
            ret = ret.T

        else:
            print(node.pattern)
            raise ValueError(
                f"Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

        #print(f"{node.pattern=},{ret.type=},{x.nvals=},{y.nvals=},{ret.nvals=}")
        
        return ret

    def visit_InnerBroadcast(self, node: LAIR.InnerBroadcast) -> VisitorReturnType:
        NDIM = tuple()

        x = self.visit(node.left)
        y = self.visit(node.right)

#        if x.type != y.type:
#            if isinstance(y, (Vector, Matrix)):
#                y = y.cast(x.type)
#            elif isinstance(y, Scalar):
#                new_y = Scalar.from_type(x.type)
#                new_y[None] = y[None]
#                y = new_y
#            else:
#                raise ValueError(
#                    "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
#                )
        if x.type != y.type:
            if isinstance(y, Scalar):
                ret = Scalar.from_type(x.type)
                if y.nvals > 0:
                    ret[None] = y[0]
            elif isinstance(y, Vector):
                ret = Vector.sparse(typ=x.type, size=y.size)
                identity = getattr(y.type, 'IDENTITY')
                y.apply(identity, out=ret)
            elif isinstance(y, Matrix):
                ret = Matrix.sparse(x.type, nrows=y.nrows, ncols=y.ncols)
                identity = getattr(y.type, 'IDENTITY')
                y.apply(identity, out=ret)
            else:
                raise ValueError(
                    "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
                )
            y = ret

        op_name = self.BinaryGrbOps[node.op]

        typ: pygraphblas.types.Type = x.type
        grb_op = getattr(typ, op_name)
        is_boolean = node.op in util.BinaryBooleanGrbOps
        if is_boolean:
          add_op = getattr(pygraphblas.types.BOOL, "ANY")
          mul_op = grb_op
          zero = False
        else:
          add_op = getattr(typ, "ANY")
          mul_op = grb_op
          zero = typ.default_zero
        #print(f"{typ=},ANY,{add_op.get_op()=},{op_name},{mul_op.get_op()=},{zero=},{is_boolean=}")
        monoid = typ.new_monoid(add_op, zero)
        mxm_semiring = typ.new_semiring(monoid, mul_op)

        out_shape = [None for _ in range(len(node.pattern[2]))]
        for it_in, dim_in in enumerate(node.pattern[0]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    out_shape[it_out] = grb_shape(x)[it_in]
        for it_in, dim_in in enumerate(node.pattern[1]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    out_shape[it_out] = grb_shape(y)[it_in]

        assert None not in out_shape

        #print("Executing InnerBroadcast", node.pattern, mxm_semiring, f"{mxm_semiring.pls=}", f"{mxm_semiring.mul=}", f"{mxm_semiring.name=}")

        if x.nvals == 0 or y.nvals == 0:
            if len(node.pattern[2]) == 0:
                ret = Scalar.from_type(x.type)
            elif len(node.pattern[2]) == 1:
                ret = Vector.sparse(x.type, *out_shape)
            elif len(node.pattern[2]) == 2:
                ret = Matrix.sparse(x.type, *out_shape)
            else:
                raise ValueError(
                    f"Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
                )

        elif node.pattern == (NDIM, NDIM, NDIM):
            ret = Scalar.from_value(self.BinaryPyOps[op_name](x[0], y[0]))

        elif node.pattern == (NDIM, (0,), (0,)):
            ret = y.apply_first(first=x, op=grb_op)
        elif node.pattern == (NDIM, (0,1), (0,1)):
            ret = y.apply_first(first=x, op=grb_op)
        elif node.pattern == (NDIM, (1,0), (0,1)):
            ret = y.apply_first(first=x, op=grb_op).T

        elif node.pattern == ((0,), NDIM, (0,)):
            ret = x.apply_second(second=y, op=grb_op)
        elif node.pattern == ((0,1), NDIM, (0,1)):
            ret = x.apply_second(second=y, op=grb_op)
        elif node.pattern == ((1,0), NDIM, (0,1)):
            ret = x.apply_second(second=y, op=grb_op).T

        elif node.pattern == ((0,), (0,), (0,)):
            ret = x.emult(y, mult_op=grb_op)
        elif node.pattern == ((0,), (1,), (0, 1)):
            mx = Matrix.sparse(x.type, x.size, 1)
            mx.assign_col(0, x)
            my = Matrix.sparse(y.type, 1, y.size)
            my.assign_row(0, y)
            ret = mx.mxm(my, semiring=mxm_semiring)
        elif node.pattern == ((1,), (0,), (0, 1)):
            mx = Matrix.sparse(x.type, x.size, 1)
            mx.assign_col(0, x)
            my = Matrix.sparse(y.type, 1, y.size)
            my.assign_row(0, y)
            ret = mx.mxm(my, semiring=mxm_semiring).T

        elif node.pattern == ((0,), (0, 1), (0, 1)):
            mx = Matrix.from_diag(x)
            ret = mx.mxm(y, semiring=mxm_semiring)
        elif node.pattern == ((1,), (0, 1), (0, 1)):
            mx = Matrix.from_diag(x)
            ret = mx.mxm(y.T, semiring=mxm_semiring).T
        elif node.pattern == ((0,), (1, 0), (0, 1)):
            mx = Matrix.from_diag(x)
            ret = mx.mxm(y.T, semiring=mxm_semiring)
        elif node.pattern == ((1,), (1, 0), (0, 1)):
            mx = Matrix.from_diag(x)
            ret = mx.mxm(y, semiring=mxm_semiring).T

        elif node.pattern == ((0, 1), (0,), (0, 1)):
            my = Matrix.from_diag(y)
            ret = x.T.mxm(my, semiring=mxm_semiring).T
            #print("&" * 80)
            #print(f"{x.nvals=},{my.nvals=},{ret.nvals=},{x.nonzero().nvals=},{my.nonzero().nvals=}")
            #print("&" * 80)
        elif node.pattern == ((0, 1), (1,), (0, 1)):
            my = Matrix.from_diag(y)
            ret = x.mxm(my, semiring=mxm_semiring)
        elif node.pattern == ((1, 0), (0,), (0, 1)):
            my = Matrix.from_diag(y)
            ret = x.mxm(my, semiring=mxm_semiring).T
        elif node.pattern == ((1, 0), (1,), (0, 1)):
            my = Matrix.from_diag(y)
            ret = x.T.mxm(my, semiring=mxm_semiring)

        elif node.pattern == ((0,1), (0,1), (0,1)):
            ret = x.emult(y, mult_op=grb_op)
        elif node.pattern == ((0,1), (1,0), (0,1)):
            ret = x.emult(y.T, mult_op=grb_op)
        elif node.pattern == ((1,0), (0,1), (0,1)):
            ret = x.T.emult(y, mult_op=grb_op)
        elif node.pattern == ((1,0), (1,0), (0,1)):
            ret = x.emult(y, mult_op=grb_op).T

        else:
            print(node.pattern)
            raise ValueError(
                f"Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

        #print(f"InnerBroadcast,{node.pattern=},{x.nvals=},{y.nvals=},{ret.nvals=}")
        
        return ret


    def visit_ElementwiseMultiply(
        self, node: LAIR.ElementwiseMultiply
    ) -> VisitorReturnType:
        x = self.visit(node.left)
        y = self.visit(node.right)

        if x.type != y.type:
            if isinstance(y, (Vector, Matrix)):
                y = y.cast(x.type)
            elif isinstance(y, Scalar):
                new_y = Scalar.from_type(x.type)
                new_y[None] = y[None]
                y = new_y
            else:
                raise ValueError(
                    "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
                )

        op_name = self.BinaryGrbOps[node.op]

        grb_op = getattr(x.type, op_name)

        if isinstance(x, Scalar) and isinstance(y, Scalar):
            if x.nvals > 0 and y.nvals > 0:
                ret = Scalar.from_value(self.BinaryPyOps[node.op](x[0], y[0]))
            else:
                ret = Scalar.from_type(x.type)
        elif (not isinstance(x, Scalar)) and isinstance(y, Scalar):
            if x.nvals > 0 and y.nvals > 0:
                ret = x.apply_second(grb_op, y)
            elif isinstance(x, Vector):
                ret = Vector.sparse(x.type, x.size)
            elif isinstance(x, Matrix):
                ret = Matrix.sparse(x.type, x.nrows, x.ncols)
        elif isinstance(x, Scalar) and (not isinstance(y, Scalar)):
            if x.nvals > 0 and y.nvals > 0:
                ret = y.apply_first(x, grb_op)
            elif isinstance(y, Vector):
                ret = Vector.sparse(y.type, y.size)
            elif isinstance(y, Matrix):
                ret = Matrix.sparse(y.type, y.nrows, y.ncols)
        elif (
            isinstance(x, (Vector, Matrix))
            and isinstance(y, (Vector, Matrix))
            and ndim(x) == ndim(y)
        ):
            ret = x.emult(y, grb_op)
        elif (
            isinstance(x, (Scalar, Vector, Matrix))
            and isinstance(y, (Scalar, Vector, Matrix))
            and ndim(x) != ndim(y)
        ):
            raise ValueError(f"Ndim of x and y must be identical")
        else:
            raise ValueError(
                f"Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

        return ret

    def visit_InnerBroadcastStencil(
        self, node: LAIR.InnerBroadcastStencil
    ) -> VisitorReturnType:
        raise ValueError(
            f"The requested promotion operation is not supported by GraphBLAS {node.left_indices}, {node.right_indices} -> {node.out_indices}"
        )

    def visit_AntiMask(self, node: LAIR.AntiMask) -> VisitorReturnType:
        operand = self.visit(node.operand)
        mask = self.visit(node.mask)

        assert mask.type == pygraphblas.types.BOOL
        assert grb_shape(operand) == grb_shape(mask)

        if isinstance(operand, Scalar) and isinstance(mask, Scalar):
            if operand.nvals > 0 and mask.nvals == 0:
                ret = operand
            else:
                ret = Scalar.from_type(operand.type)
        elif isinstance(operand, Vector) and isinstance(mask, Vector):
            ret = operand.extract(mask, desc=pygraphblas.descriptor.C)
        elif isinstance(operand, Matrix) and isinstance(mask, Matrix):
            ret = operand.extract_matrix(mask=mask, desc=pygraphblas.descriptor.C)
        elif (
            isinstance(operand, (Scalar, Vector, Matrix))
            and isinstance(mask, (Scalar, Vector, Matrix))
            and ndim(operand) != ndim(mask)
        ):
            raise ValueError(f"Ndim of x and y must be identical")
        else:
            raise ValueError(
                f"Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

        return ret

    def visit_Mask(self, node: LAIR.Mask) -> VisitorReturnType:
        operand = self.visit(node.operand)
        mask = self.visit(node.mask)

        assert mask.type == pygraphblas.types.BOOL
        assert ndim(operand) == ndim(mask)

        if isinstance(operand, Scalar) and isinstance(mask, Scalar):
            if operand.nvals > 0 and mask.nvals > 0:
                ret = operand
            else:
                ret = Scalar.from_type(operand.type)
        elif isinstance(operand, Vector) and isinstance(mask, Vector):
            ret = operand.extract(mask)
        elif isinstance(operand, Matrix) and isinstance(mask, Matrix):
            ret = operand.extract_matrix(mask=mask)
        elif (
            isinstance(operand, (Scalar, Vector, Matrix))
            and isinstance(mask, (Scalar, Vector, Matrix))
            and ndim(operand) != ndim(mask)
        ):
            raise ValueError(f"Ndim of x and y must be identical")
        else:
            raise ValueError(
                f"Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

        return ret

    def visit_Output(self, node: LAIR.Output) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_Pattern(self, node: LAIR.Pattern) -> VisitorReturnType:
        x = self.visit(node.operand)

        if isinstance(x, Scalar):
            if x.nvals > 0 and float(x[0]) != 0.0:
                ret = Scalar.from_value(True)
            else:
                ret = Scalar.from_type(x.type)
        elif isinstance(x, Vector):
            ret = x.nonzero().pattern()
        elif isinstance(x, Matrix):
            ret = x.nonzero().pattern()
        else:
            raise ValueError(
                "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
            )

        # print(node.description,"result", ret.nvals)

        return ret

    def visit_PrimaryKey(self, node: LAIR.PrimaryKey) -> VisitorReturnType:
        operand = self.visit(node.stencil)
        return extract_indices(operand, axis=node.colnum, typ=node.type_)

    def visit_Promote(self, node: LAIR.Promote) -> VisitorReturnType:
#        if node._operand_indices == node._out_indices:
#            return node
#
#        assert all(idx in self._out_indices for idx in self._operand_indices)
#
#        sig = node.signature
#        NDIM = tuple()
#        if sig == (NDIM, (0,)):
#
#        elif sig == (NDIM, (0, 1)):
#        elif sig == ((0,), (0, 1)):
#        elif sig == ((0,), (1, 0)):
#        else:
#            raise ValueError(
#                "Prototype only supports PyGraphBLAS scalars, vectors, and matrices"
#            )
#
#        ndim_in = len(node._operand_indices)
#        ndim_out = len(node._out_indices)
#        if ndim_in > ndim_out:
#            raise
        raise ValueError(
            f"The requested promotion operation is not supported by GraphBLAS {node.operand_indices} -> {node.out_indices}"
        )

    def visit_Reduction(self, node: LAIR.Reduction) -> VisitorReturnType:
        operand = self.visit(node.operand)

        NODIM = tuple()

        orig_axes = tuple(range(ndim(operand)))
        pattern = (orig_axes, tuple(ax for ax in orig_axes if ax not in node.axes))
        if 0:
            #Doesn't work, see: https://github.com/Graphegon/pygraphblas/issues/60
            grb_op = getattr(operand.type, f"{self.BinaryGrbOps[node.op]}_MONOID")
        else:
            add_op = getattr(operand.type, self.BinaryGrbOps[node.op])
            grb_op = operand.type.new_monoid(add_op, operand.type.default_zero)

        # Scalar Input
        if isinstance(operand, Scalar) and pattern == (NODIM, NODIM):
            ret = operand

        # Vector Input
        elif isinstance(operand, Vector) and pattern == ((0,), NODIM):
            ret = Scalar.from_type(operand.type)
            if operand.nvals > 0:
                ret[None] = operand.reduce(grb_op)
        elif isinstance(operand, Vector) and pattern == ((0,), (0,)):
            ret = operand

        # Matrix Input
        elif isinstance(operand, Matrix) and pattern == ((0, 1), NODIM):
            ret = Scalar.from_type(operand.type)
            if operand.nvals > 0:
                retval = operand.T.reduce(grb_op)
                ret[None] = retval
        elif isinstance(operand, Matrix) and pattern == ((0, 1), (0,)):
            ret = operand.reduce_vector(grb_op)
        elif isinstance(operand, Matrix) and pattern == ((0, 1), (1,)):
            ret = operand.T.reduce_vector(grb_op)
        elif isinstance(operand, Matrix) and pattern == ((0, 1), (0, 1)):
            ret = operand

        else:
            raise ValueError(
                f"Invalid reduction pattern {pattern[0]}->{pattern[1]} for operand with type={type(operand)}"
            )

        return ret

    def visit_Tensor(self, node: LAIR.Tensor) -> VisitorReturnType:
        return node.tensor

    def visit_PermuteIndices(self, node: LAIR.PermuteIndices) -> VisitorReturnType:
        operand = self.visit(node.operand)
        axes = node.new_axes or tuple(reversed(range(ndim(operand))))

        if isinstance(operand, Scalar):
            ret = operand
        elif isinstance(operand, Vector):
            ret = operand
        elif isinstance(operand, Matrix) and axes == (0, 1):
            ret = operand
        elif isinstance(operand, Matrix) and axes == (1, 0):
            ret = operand.T
        else:
            raise ValueError(
                f"Unknown transpose type with new_axes={axes} and operand type {type(operand)}"
            )

        return ret

    def visit_VectorToColumnMatrix(
        self, node: LAIR.VectorToColumnMatrix
    ) -> VisitorReturnType:
        arg = self.visit(node.operand)
        assert isinstance(arg, Vector)
        ret: Matrix = Matrix.sparse(typ=arg.type, nrows=arg.size, ncols=1)
        ret.assign_col(0, arg)
        return ret

    def visit_VectorToRowMatrix(
        self, node: LAIR.VectorToRowMatrix
    ) -> VisitorReturnType:
        arg = self.visit(node.operand)
        assert isinstance(arg, Vector)
        ret: Matrix = Matrix.sparse(typ=arg.type, ncols=arg.size, nrows=1)
        ret.assign_row(0, arg)
        return ret

    def visit_VectorToDiagMatrix(
        self, node: LAIR.VectorToDiagMatrix
    ) -> VisitorReturnType:
        arg = self.visit(node.operand)
        assert isinstance(arg, Vector)
        ret: Matrix = Matrix.from_diag(arg)
        return ret

tensql/LAIR/Visitor.py
sandialabs/TenSQL
from .InnerBroadcast import InnerBroadcast
from .InnerBroadcastMask import InnerBroadcastMask
from .CastLike import CastLike
from .Cast import Cast
from .DotMxM import DotMxM
from .DotMxV import DotMxV
from .DotVxM import DotVxM
from .ElementwiseAdd import ElementwiseAdd
from .ElementwiseApply import (
    ElementwiseApply,
    ElementwiseApplyBindFirst,
    ElementwiseApplyBindSecond,
)
from .ElementwiseMultiply import ElementwiseMultiply
from .AntiMask import AntiMask
from .Mask import Mask
from .Node import Node
from .Output import Output
from .Pattern import Pattern
from .PrimaryKey import PrimaryKey
from .Reduction import Reduction
from .Tensor import Tensor
from .PermuteIndices import PermuteIndices
from .VectorToColumnMatrix import VectorToColumnMatrix
from .VectorToDiagMatrix import VectorToDiagMatrix
from .VectorToRowMatrix import VectorToRowMatrix

# Functions that provide shortcuts to generate compositions of the above
from .Promote import OptimizedPromote, Promote
from .InnerBroadcastStencil import (
    OptimizedInnerBroadcastStencil,
)  # , InnerBroadcastStencil

tensql/LAIR/__init__.py
sandialabs/TenSQL
from typing import Optional, Tuple, Set, FrozenSet, List, Dict

from ...util import GBTensor, ndim
from ..Visitor import Visitor
from ... import LAIR
from .index_algebra import intersect_sop, union_sop, complement_sop, SOPExpression

VisitorReturnType = List[SOPExpression]

class IndexVisitor(Visitor):
    def __init__(self):
        self.registry = {}

    def visit(self, node: LAIR.Node) -> VisitorReturnType:
        ret = super().visit(node)
        self.registry[node] = ret
        return ret

    def visit_InnerBroadcastMask(self, node: LAIR.InnerBroadcastMask) -> VisitorReturnType:
        left_indices = self.visit(node.left)
        right_indices = self.visit(node.right)

        ret = [[] for it_out in node.pattern[2]]
        for it_in, dim_in in enumerate(node.pattern[0]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    ret[it_out].append(left_indices[it_in])
        for it_in, dim_in in enumerate(node.pattern[1]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    ret[it_out].append(right_indices[it_in])
        return [intersect_sop(*idxset) for idxset in ret]

    def visit_InnerBroadcast(self, node: LAIR.InnerBroadcast) -> VisitorReturnType:
        left_indices = self.visit(node.left)
        right_indices = self.visit(node.right)

        ret = [[] for it_out in node.pattern[2]]
        for it_in, dim_in in enumerate(node.pattern[0]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    ret[it_out].append(left_indices[it_in])
        for it_in, dim_in in enumerate(node.pattern[1]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    ret[it_out].append(right_indices[it_in])
        return [intersect_sop(*idxset) for idxset in ret]

    def visit_CastLike(self, node: LAIR.CastLike) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_DotMxM(self, node: LAIR.DotMxM) -> VisitorReturnType:
        left_indices = self.visit(node.left)
        right_indices = self.visit(node.right)
        return [left_indices[0], right_indices[1]]

    def visit_DotMxV(self, node: LAIR.DotMxV) -> VisitorReturnType:
        left_indices = self.visit(node.left)
        right_indices = self.visit(node.right)
        return [
            left_indices[0],
            #intersect_sop(left_indices[1], right_indices[0]),
        ]

    def visit_DotVxM(self, node: LAIR.DotVxM) -> VisitorReturnType:
        left_indices = self.visit(node.left)
        right_indices = self.visit(node.right)
        return [
            #intersect_sop(left_indices[0], right_indices[0]),
            right_indices[1],
        ]

    def visit_ElementwiseAdd(self, node: LAIR.ElementwiseAdd) -> VisitorReturnType:
        left_indices = self.visit(node.left)
        right_indices = self.visit(node.right)
        assert len(left_indices) == len(right_indices)
        return [
            union_sop(lidx, ridx)
            for lidx, ridx in zip(left_indices, right_indices)
        ]

    def visit_ElementwiseApply(self, node: LAIR.ElementwiseApply) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_Cast(self, node: LAIR.Cast) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_ElementwiseApplyBindFirst(
        self, node: LAIR.ElementwiseApplyBindFirst
    ) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_ElementwiseApplyBindSecond(
        self, node: LAIR.ElementwiseApplyBindSecond
    ) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_ElementwiseMultiply(
        self, node: LAIR.ElementwiseMultiply
    ) -> VisitorReturnType:
        left_indices = self.visit(node.left)
        right_indices = self.visit(node.right)
        # print(len(left_indices), len(right_indices))
        assert (len(left_indices) == len(right_indices) or min(len(left_indices), len(right_indices)) == 0)
        if len(left_indices) == len(right_indices):
            out_indices = [
                intersect_sop(lidx, ridx)
                for lidx, ridx in zip(left_indices, right_indices)
            ]
        elif len(left_indices) == 0:
            out_indices = right_indices
        elif len(right_indices) == 0:
            out_indices = left_indices
        # print("IndexVisitor.visit_ElementwiseMultiply", out_indices, left_indices, right_indices)
        return out_indices

    def visit_AntiMask(self, node: LAIR.AntiMask) -> VisitorReturnType:
        operand_indices = self.visit(node.operand)
        mask_indices = self.visit(node.mask)
        assert len(operand_indices) == len(mask_indices)
        return [
            intersect_sop(oidx, complement_sop(midx))
            for oidx, midx in zip(operand_indices, mask_indices)
        ]

    def visit_Mask(self, node: LAIR.Mask) -> VisitorReturnType:
        operand_indices = self.visit(node.operand)
        mask_indices = self.visit(node.mask)
        assert len(operand_indices) == len(mask_indices)
        return [
            intersect_sop(oidx, midx)
            for oidx, midx in zip(operand_indices, mask_indices)
        ]

    def visit_Output(self, node: LAIR.Output) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_Pattern(self, node: LAIR.Pattern) -> VisitorReturnType:
        operand_indices = self.visit(node.operand)
        
        return [
            intersect_sop(oidx, frozenset({frozenset({(id(node), it, False)})}))
            for it, oidx in enumerate(operand_indices)
        ]

    def visit_PrimaryKey(self, node: LAIR.PrimaryKey) -> VisitorReturnType:
        return self.visit(node.stencil)

    def visit_Promote(self, node: LAIR.Promote) -> VisitorReturnType:
        ret = []
        operand_indices = self.visit(node.operand)
        for oidx in node.out_indices:
            if oidx in node.operand_indices:
                ret.append(operand_indices[oidx])
            else:
                ret.append(frozenset())
        return ret

    def visit_Reduction(self, node: LAIR.Reduction) -> VisitorReturnType:
        operand_indices = self.visit(node.operand)
        return [
            oidx
            for it_oidx, oidx in enumerate(operand_indices)
            if it_oidx not in node.axes
        ]

    def visit_Tensor(self, node: LAIR.Tensor) -> VisitorReturnType:
        if node.stencil is None:
            return [
                frozenset({frozenset({(id(node.tensor), it, False)})})
                for it in range(ndim(node.tensor))
            ]
        else:
            stencil_indices = self.visit(node.stencil)
            return [
                intersect_sop(
                    frozenset({frozenset({(id(node.tensor), it, False)})}), stencil_indices[it]
                )
                for it in range(ndim(node.tensor))
            ]

    def visit_PermuteIndices(self, node: LAIR.PermuteIndices) -> VisitorReturnType:
        operand_indices = self.visit(node.operand)
        if node.new_axes is None:
            return list(reversed(operand_indices))
        else:
            assert set(node.new_axes) == set(range(len(operand_indices)))
            return [operand_indices[dim] for dim in node.new_axes]

    def visit_VectorToColumnMatrix(self, node: LAIR.VectorToColumnMatrix):
        operand_indices = self.visit(node.operand)
        return [operand_indices[0], [frozenset({frozenset({(None, 0, False)})})]]

    def visit_VectorToRowMatrix(self, node: LAIR.VectorToRowMatrix):
        operand_indices = self.visit(node.operand)
        return [[frozenset({frozenset({(None, 0, False)})})], operand_indices[0]]

    def visit_VectorToDiagMatrix(self, node: LAIR.VectorToDiagMatrix):
        operand_indices = self.visit(node.operand)
        return [operand_indices[0], operand_indices[0]]

tensql/LAIR/visitors/IndexVisitor.py
sandialabs/TenSQL
from ... import LAIR
from ..Visitor import Visitor


class NodeReplacer(Visitor):
    def __init__(self, old_node: LAIR.Node, new_node: LAIR.Node):
        self.old_node = old_node
        self.new_node = new_node

    def visit(self, node: LAIR.Node) -> LAIR.Node:
        old_node_attributes = node.get_attributes()
        new_node_attributes = {}
        for name, attr in old_node_attributes.items():
            if attr is self.old_node:
                new_node_attributes[name] = self.new_node
            elif isinstance(attr, LAIR.Node):
                new_node_attributes[name] = self.visit(attr)
            else:
                new_node_attributes[name] = attr

        return type(node)(**new_node_attributes)

tensql/LAIR/visitors/NodeReplacer.py
sandialabs/TenSQL
from typing import Optional, Tuple, Set, FrozenSet, List, Dict
import copy

from ...util import GBTensor, ndim
from ..Visitor import Visitor
from ... import LAIR
from .index_algebra import intersect_sop, union_sop, complement_sop, SOPExpression


Memo = Tuple[SOPExpression]
VisitorReturnType = Memo

class ReverseIndexVisitor(Visitor):
    def __init__(self, forward_registry: Dict[LAIR.Node, List[SOPExpression]]):
        self.forward_registry = {
            k: copy.deepcopy(v) for k, v in forward_registry.items()
        }
        self.registry: Dict[LAIR.Node, List[SOPExpression]] = {}

    def visit(
        self, node: LAIR.Node, out_indices: Optional[Memo] = None
    ) -> VisitorReturnType:
        if out_indices is None:
            out_indices = tuple(frozenset() for _ in self.forward_registry[node])
        ret = super().visit(node, out_indices)
        self.registry[node] = ret
        return ret

    def visit_InnerBroadcast(
        self, node: LAIR.InnerBroadcast, out_indices: Memo
    ) -> VisitorReturnType:

        left_ret = [None for it_out in node.pattern[2]]
        for it_in, dim_in in enumerate(node.pattern[0]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    left_ret[it_in] = out_indices[it_out]
        self.visit(node.left, left_ret)

        right_ret = [None for it_out in node.pattern[2]]
        for it_in, dim_in in enumerate(node.pattern[1]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    right_ret[it_in] = out_indices[it_out]
        self.visit(node.right, right_ret)

        return out_indices

    def visit_InnerBroadcastMask(
        self, node: LAIR.InnerBroadcast, out_indices: Memo
    ) -> VisitorReturnType:

        left_ret = [None for it_out in node.pattern[2]]
        for it_in, dim_in in enumerate(node.pattern[0]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    left_ret[it_in] = out_indices[it_out]
        self.visit(node.left, left_ret)

        right_ret = [None for it_out in node.pattern[2]]
        for it_in, dim_in in enumerate(node.pattern[1]):
            for it_out, dim_out in enumerate(node.pattern[2]):
                if dim_in == dim_out:
                    right_ret[it_in] = out_indices[it_out]
        self.visit(node.right, right_ret)

        return out_indices

    def visit_CastLike(
        self, node: LAIR.CastLike, out_indices: Memo
    ) -> VisitorReturnType:
        self.visit(node.operand, out_indices)
        return out_indices

    def visit_DotMxM(self, node: LAIR.DotMxM, out_indices: Memo) -> VisitorReturnType:
        right_forward = self.forward_registry[node.right]
        left_forward = self.forward_registry[node.left]

        self.visit(node.left, (out_indices[0], right_forward[0]))
        self.visit(node.right, (left_forward[1], out_indices[1]))
        return out_indices

    def visit_DotMxV(self, node: LAIR.DotMxV, out_indices: Memo) -> VisitorReturnType:
        right_forward = self.forward_registry[node.right]
        left_forward = self.forward_registry[node.left]

        self.visit(
            node.left,
            (out_indices[0], right_forward[0]),
        )
        self.visit(
            node.right, (left_forward[1],)
        )
        return out_indices

    def visit_DotVxM(self, node: LAIR.DotVxM, out_indices: Memo) -> VisitorReturnType:
        right_forward = self.forward_registry[node.right]
        left_forward = self.forward_registry[node.left]
        self.visit(
            node.left, (right_forward[0],)
        )
        self.visit(
            node.right, (left_forward[0], out_indices[0]),
        )
        return out_indices

    def visit_ElementwiseAdd(
        self, node: LAIR.ElementwiseAdd, out_indices: Memo
    ) -> VisitorReturnType:
        right_forward = self.forward_registry[node.right]
        left_forward = self.forward_registry[node.left]
        self.visit(
            node.left,
            tuple(
                union_sop(out_indices[it], right_forward[it])
                for it in range(len(out_indices))
            ),
        )
        self.visit(
            node.right,
            tuple(
                union_sop(out_indices[it], left_forward[it])
                for it in range(len(out_indices))
            ),
        )
        return out_indices

    def visit_ElementwiseApply(
        self, node: LAIR.ElementwiseApply, out_indices: Memo
    ) -> VisitorReturnType:
        self.visit(node.operand, out_indices)
        return out_indices

    def visit_Cast(
        self, node: LAIR.Cast, out_indices: Memo
    ) -> VisitorReturnType:
        self.visit(node.operand, out_indices)
        return out_indices

    def visit_ElementwiseApplyBindFirst(
        self, node: LAIR.ElementwiseApplyBindFirst, out_indices: Memo
    ) -> VisitorReturnType:
        self.visit(node.operand, out_indices)
        return out_indices

    def visit_ElementwiseApplyBindSecond(
        self, node: LAIR.ElementwiseApplyBindSecond, out_indices: Memo
    ) -> VisitorReturnType:
        self.visit(node.operand, out_indices)
        return out_indices

    def visit_ElementwiseMultiply(
        self, node: LAIR.ElementwiseMultiply, out_indices: Memo
    ) -> VisitorReturnType:
        right_forward = self.forward_registry[node.right]
        left_forward = self.forward_registry[node.left]
        #print(f"{left_forward=}, {right_forward=}, {out_indices=}")
        if len(left_forward) == len(right_forward):
            self.visit(
                node.left,
                tuple(
                    intersect_sop(out_indices[it], right_forward[it])
                    for it in range(len(out_indices))
                ),
            )
            self.visit(
                node.right,
                tuple(
                    intersect_sop(out_indices[it], left_forward[it])
                    for it in range(len(out_indices))
                ),
            )
        elif len(left_forward) == 0:
            self.visit(node.right, out_indices)
        elif len(right_forward) == 0:
            self.visit(node.left, out_indices)
        return out_indices

    def visit_AntiMask(self, node: LAIR.AntiMask, out_indices: Memo) -> VisitorReturnType:
        operand_forward = self.forward_registry[node.operand]
        mask_forward = self.forward_registry[node.mask]
        self.visit(
            node.mask,
            tuple(
                intersect_sop(out_indices[it], complement_sop(operand_forward[it]))
                for it in range(len(out_indices))
            ),
        )
        self.visit(
            node.operand,
            tuple(
                intersect_sop(out_indices[it], complement_sop(mask_forward[it]))
                for it in range(len(out_indices))
            ),
        )
        return out_indices

    def visit_Mask(self, node: LAIR.Mask, out_indices: Memo) -> VisitorReturnType:
        operand_forward = self.forward_registry[node.operand]
        mask_forward = self.forward_registry[node.mask]
        self.visit(
            node.mask,
            tuple(
                intersect_sop(out_indices[it], operand_forward[it])
                for it in range(len(out_indices))
            ),
        )
        self.visit(
            node.operand,
            tuple(
                intersect_sop(out_indices[it], mask_forward[it])
                for it in range(len(out_indices))
            ),
        )
        return out_indices

    def visit_Output(self, node: LAIR.Output, out_indices: Memo) -> VisitorReturnType:
        self.visit(node.operand, out_indices)
        return out_indices

    def visit_Pattern(
        self, node: LAIR.Pattern, out_indices: Memo
    ) -> VisitorReturnType:
        self.visit(node.operand, out_indices)
        return out_indices

    def visit_PrimaryKey(
        self, node: LAIR.PrimaryKey, out_indices: Memo
    ) -> VisitorReturnType:
        self.visit(node.stencil, out_indices)
        return out_indices

    def visit_Promote(
        self, node: LAIR.Promote, out_indices: Memo
    ) -> VisitorReturnType:
        out_indices_lookup = {
            node.out_indices[it]: indices for it, indices in enumerate(out_indices)
        }
        self.visit(
            node.operand, tuple(out_indices_lookup[idx] for idx in node.operand_indices)
        )
        return out_indices

    def visit_Reduction(
        self, node: LAIR.Reduction, out_indices: Memo
    ) -> VisitorReturnType:
        operand_indices = []
        count = 0
        for it in range(len(out_indices) + len(node.axes)):
            if it in node.axes:
                operand_indices.append(frozenset())
            else:
                operand_indices.append(out_indices[count])
                count += 1
        self.visit(node.operand, tuple(operand_indices))
        return out_indices

    def visit_Tensor(self, node: LAIR.Tensor, out_indices: Memo) -> VisitorReturnType:
        if node.stencil is not None:
            self.visit(node.stencil, out_indices)
        return out_indices

    def visit_PermuteIndices(
        self, node: LAIR.PermuteIndices, out_indices: Memo
    ) -> VisitorReturnType:
        out_indices_lookup = {
            idx: indices for idx, indices in zip(node.new_axes, out_indices)
        }
        self.visit(
            node.operand,
            tuple(out_indices_lookup[idx] for idx in range(len(node.new_axes))),
        )
        return out_indices

    def visit_VectorToColumnMatrix(
        self, node: LAIR.VectorToColumnMatrix, out_indices: Memo
    ):
        self.visit(node.operand, (out_indices[0],))
        return out_indices

    def visit_VectorToRowMatrix(self, node: LAIR.VectorToRowMatrix, out_indices: Memo):
        self.visit(node.operand, (out_indices[1],))
        return out_indices

    def visit_VectorToDiagMatrix(
        self, node: LAIR.VectorToDiagMatrix, out_indices: Memo
    ):
        self.visit(node.operand, (union_sop(out_indices[0], out_indices[1]),))
        return out_indices

tensql/LAIR/visitors/ReverseIndexVisitor.py
sandialabs/TenSQL
from typing import Optional, Tuple, Set, FrozenSet, List, Dict
import itertools

from ...util import GBTensor, ndim, grb_shape

from ..Visitor import Visitor

from ... import LAIR

VisitorReturnType = Tuple[Optional[int], ...]

class ShapeVisitor(Visitor):
    def __init__(self):
        self.registry = {}

    def visit(self, node: LAIR.Node) -> VisitorReturnType:
        ret = super().visit(node)
        self.registry[node] = ret
        return ret

    def visit_InnerBroadcast(self, node: LAIR.InnerBroadcast) -> VisitorReturnType:
        dim_sizes = {}

        left_shape = self.visit(node.left)
        for idx, size in zip(node.left_indices, left_shape):
            assert idx not in dim_sizes
            dim_sizes[idx] = size

        right_shape = self.visit(node.right)
        for idx, size in zip(node.left_indices, left_shape):
            if idx in dim_sizes:
                assert dim_sizes[idx] == size
            else:
                dim_sizes[idx] = size

        return tuple(dim_sizes[oidx] for oidx in node.out_indices)

    def visit_InnerBroadcastMask(self, node: LAIR.InnerBroadcastMask) -> VisitorReturnType:
        dim_sizes = {}

        left_shape = self.visit(node.left)
        for idx, size in zip(node.left_indices, left_shape):
            assert idx not in dim_sizes
            dim_sizes[idx] = size

        right_shape = self.visit(node.right)
        for idx, size in zip(node.left_indices, left_shape):
            if idx in dim_sizes:
                assert dim_sizes[idx] == size
            else:
                dim_sizes[idx] = size

        return tuple(dim_sizes[oidx] for oidx in node.out_indices)

    def visit_CastLike(self, node: LAIR.CastLike) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_DotMxM(self, node: LAIR.DotMxM) -> VisitorReturnType:
        left_shape = self.visit(node.left)
        right_shape = self.visit(node.right)
        return [left_shape[0], right_shape[1]]

    def visit_DotMxV(self, node: LAIR.DotMxV) -> VisitorReturnType:
        left_shape = self.visit(node.left)
        self.visit(node.right)
        return (left_shape[0],)

    def visit_DotVxM(self, node: LAIR.DotVxM) -> VisitorReturnType:
        self.visit(node.left)
        right_shape = self.visit(node.right)
        return (right_shape[0],)

    def visit_ElementwiseAdd(self, node: LAIR.ElementwiseAdd) -> VisitorReturnType:
        left_shape = self.visit(node.left)
        right_shape = self.visit(node.right)
        assert len(left_shape) == len(right_shape)
        return left_shape

    def visit_ElementwiseApply(self, node: LAIR.ElementwiseApply) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_Cast(self, node: LAIR.Cast) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_ElementwiseApplyBindFirst(
        self, node: LAIR.ElementwiseApplyBindFirst
    ) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_ElementwiseApplyBindSecond(
        self, node: LAIR.ElementwiseApplyBindSecond
    ) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_ElementwiseMultiply(
        self, node: LAIR.ElementwiseMultiply
    ) -> VisitorReturnType:
        left_shape = self.visit(node.left)
        right_shape = self.visit(node.right)
        assert len(left_shape) == len(right_shape)
        return left_shape

    def visit_AntiMask(self, node: LAIR.AntiMask) -> VisitorReturnType:
        operand_shape = self.visit(node.operand)
        mask_shape = self.visit(node.mask)
        assert len(operand_shape) == len(mask_shape)
        return operand_shape

    def visit_Mask(self, node: LAIR.Mask) -> VisitorReturnType:
        operand_shape = self.visit(node.operand)
        mask_shape = self.visit(node.mask)
        assert len(operand_shape) == len(mask_shape)
        return operand_shape

    def visit_Output(self, node: LAIR.Output) -> VisitorReturnType:
        return self.visit(node.operand)

    def visit_Pattern(self, node: LAIR.Pattern) -> VisitorReturnType:
        return self.visit(node.operand)

    # def visit_PrimaryKey(self, node: LAIR.PrimaryKey) -> VisitorReturnType:
    #    return self.visit(node.)

    def visit_Promote(self, node: LAIR.Promote) -> VisitorReturnType:
        ret = []
        operand_shape = self.visit(node.operand)
        for oidx in node.out_indices:
            if oidx in node.operand_indices:
                ret.append(operand_shape[oidx])
            else:
                ret.append(None)
        return tuple(ret)

    def visit_Reduction(self, node: LAIR.Reduction) -> VisitorReturnType:
        operand_shape = self.visit(node.operand)
        return tuple(
            oidx
            for it_oidx, oidx in enumerate(operand_shape)
            if it_oidx not in node.axes
        )

    def visit_Tensor(self, node: LAIR.Tensor) -> VisitorReturnType:
        ret = grb_shape(node.tensor)
        if node.stencil is not None:
            assert grb_shape(node.stencil) == ret
        return tuple(ret)

    def visit_PermuteIndices(self, node: LAIR.PermuteIndices) -> VisitorReturnType:
        operand_shape = self.visit(node.operand)
        if node.new_axes is None:
            return tuple(reversed(operand_shape))
        else:
            assert set(node.new_axes) == set(range(len(operand_shape)))
            return tuple(operand_shape[dim] for dim in node.new_axes)

    def visit_VectorToColumnMatrix(self, node: LAIR.VectorToColumnMatrix):
        operand_shape = self.visit(node.operand)
        return (operand_shape[0], 1)

    def visit_VectorToRowMatrix(self, node: LAIR.VectorToRowMatrix):
        operand_shape = self.visit(node.operand)
        return (1, operand_shape[0])

    def visit_VectorToDiagMatrix(self, node: LAIR.VectorToDiagMatrix):
        operand_shape = self.visit(node.operand)
        return (operand_shape[0], operand_shape[0])

tensql/LAIR/visitors/ShapeVisitor.py
sandialabs/TenSQL
from .IndexVisitor import IndexVisitor
from .NodeReplacer import NodeReplacer
from .ReverseIndexVisitor import ReverseIndexVisitor
from .ShapeVisitor import ShapeVisitor
from . import index_algebra

tensql/LAIR/visitors/__init__.py
sandialabs/TenSQL
from typing import Optional, Tuple, Set, FrozenSet, Sequence

import operator
import functools
import itertools

AxisIdent = Tuple[Optional[int], int, bool]
POSExpression = FrozenSet[FrozenSet[AxisIdent]]
SOPExpression = FrozenSet[FrozenSet[AxisIdent]]

def complement_axis(ident: AxisIdent):
    return (ident[0], ident[1], not ident[2])

def intersect_sop(*args) -> SOPExpression:
    ret = set()
    args = [arg for arg in args if len(arg) > 0]
    for idx_tuple in itertools.product(*args):
        ret.add(functools.reduce(operator.ior, idx_tuple, frozenset()))
    return frozenset(ret)

def union_sop(*args: Sequence[SOPExpression]) -> SOPExpression:
    return functools.reduce(operator.ior, args, frozenset())

def complement_sop(sop: SOPExpression) -> SOPExpression:
    pos = frozenset(frozenset(complement_axis(ax) for ax in term) for term in sop)
    posop_terms = frozenset(frozenset((term,)) for term in pos)
    return intersect_sop(*posop_terms)

tensql/LAIR/visitors/index_algebra.py
sandialabs/TenSQL
from typing import Optional

import pygraphblas.types

from . import Types
from .util import GBType


class PrimaryKey:
    def __init__(
        self,
        column_name: str,
        type_: Types = Types.BigInt(),
        max_idx: Optional[int] = None,
    ) -> None:
        self._column_name = column_name

        valid_type_default_sizes = {
            Types.TinyInt: 2**7 - 1,
            Types.SmallInt: 2**15 - 1,
            Types.Integer: 2**31 - 1,
            Types.BigInt: 2**60
        }

        self._type = type_

        if isinstance(type_, Types.TinyInt):
            type_max_idx = 2 ** 7 - 1
        elif isinstance(type_, Types.SmallInt):
            type_max_idx = 2 ** 15 - 1
        elif isinstance(type_, Types.Integer):
            type_max_idx = 2 ** 31 - 1
        elif isinstance(type_, Types.BigInt):
            type_max_idx = 2 ** 60
        else:
            raise ValueError("Argument type_ to PrimaryKey() must be an instance of Types.TinyInt, Types.SmallInt, Types.Integer, or Types.BigInt")

        if max_idx is not None:
            assert max_idx <= type_max_idx
        else:
            max_idx = type_max_idx
        self._max_idx = max_idx

    @property
    def column_name(self) -> str:
        return self._column_name

    @property
    def gbtype(self) -> GBType:
        return self._type.as_pygrb

    @property
    def type_(self) -> Types.Type:
        return self._type

    @property
    def max_idx(self) -> int:
        return self._max_idx

tensql/PrimaryKey.py
sandialabs/TenSQL
import numpy as np
import numba
import pygraphblas.types
import pygraphblas.binaryop
from pygraphblas.base import lib as pygrb_lib, ffi as pygrb_ffi

from ._PyObjectHelper import lib as pyoh_lib, ffi as pyoh_ffi

class AddressWrapper:
  def __init__(self, name):
    self.address = pyoh_ffi.addressof(pyoh_lib, name)

class PyObject(pygraphblas.types.INT64):
#  default_one = 1
#  default_zero = 0
#  _gb_type = pygrb_lib.GrB_INT64
#  _c_type = "void*"
#  _typecode = "q"
#  _numba_t = numba.int64
#  _numpy_t = np.int64
#  _base_name = "UDT"
  pass

def load_helper_function(op, *, boolean):
  return pygraphblas.binaryop.BinaryOp(
    op, "PyObject", 
    AddressWrapper(f"PyObject_{op}"), 
    udt=PyObject,
    boolean=boolean
  )

PyObject.EQ = load_helper_function("EQ", boolean=True)
PyObject.NE = load_helper_function("NE", boolean=True)
PyObject.LT = load_helper_function("LT", boolean=True)
PyObject.GT = load_helper_function("GT", boolean=True)
PyObject.LE = load_helper_function("LE", boolean=True)
PyObject.GE = load_helper_function("GE", boolean=True)

PyObject.ANY = load_helper_function("ANY", boolean=False)
PyObject.FIRST = load_helper_function("FIRST", boolean=False)
PyObject.SECOND = load_helper_function("SECOND", boolean=False)

tensql/PyObjectHelper.py
sandialabs/TenSQL
from __future__ import annotations

from typing import TYPE_CHECKING

from .ABCExpression import ABCExpression

if TYPE_CHECKING:
    from .ArithmeticBinaryExpression import Add, Sub, Mul, Div
    from .ArithmeticConstant import ArithmeticConstant
    from .BooleanBinaryExpression import (
        GreaterThan,
        LessThan,
        GreaterThanOrEqual,
        LessThanOrEqual,
        Equal,
        NotEqual,
    )
    from .BooleanConstant import BooleanConstant
    from .ABCUnaryExpression import Negate, AbsoluteValue


class ABCArithmeticExpression(ABCExpression):
    def __add__(self, other: ABCArithmeticExpression) -> Add:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Add

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Add(self, other)

    def __radd__(self, other: ABCArithmeticExpression) -> Add:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Add

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Add(other, self)

    def __sub__(self, other: ABCArithmeticExpression) -> Sub:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Sub

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Sub(self, other)

    def __rsub__(self, other: ABCArithmeticExpression) -> Sub:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Sub

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Sub(other, self)

    def __mul__(self, other: ABCArithmeticExpression) -> Mul:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Mul

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Mul(self, other)

    def __rmul__(self, other: ABCArithmeticExpression) -> Mul:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Mul

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Mul(other, self)

    def __truediv__(self, other: ABCArithmeticExpression) -> Div:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Div

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Div(self, other)

    def __rtruediv__(self, other: ABCArithmeticExpression) -> Div:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Div

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Div(other, self)

    def __floordiv__(self, other: ABCArithmeticExpression) -> Div:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Div

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Div(self, other)

    def __rfloordiv__(self, other: ABCArithmeticExpression) -> Div:
        from .ArithmeticConstant import ArithmeticConstant
        from .ArithmeticBinaryExpression import Div

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Div(other, self)

    def __gt__(self, other: ABCArithmeticExpression) -> GreaterThan:
        from .ArithmeticConstant import ArithmeticConstant
        from .BooleanBinaryExpression import GreaterThan

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return GreaterThan(self, other)

    def __lt__(self, other: ABCArithmeticExpression) -> LessThan:
        from .ArithmeticConstant import ArithmeticConstant
        from .BooleanBinaryExpression import LessThan

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return LessThan(self, other)

    def __ge__(self, other: ABCArithmeticExpression) -> GreaterThanOrEqual:
        from .ArithmeticConstant import ArithmeticConstant
        from .BooleanBinaryExpression import GreaterThanOrEqual

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return GreaterThanOrEqual(self, other)

    def __le__(self, other: ABCArithmeticExpression) -> LessThanOrEqual:
        from .ArithmeticConstant import ArithmeticConstant
        from .BooleanBinaryExpression import LessThanOrEqual

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return LessThanOrEqual(self, other)

    def __pos__(self) -> ABCArithmeticExpression:
        from .ArithmeticConstant import ArithmeticConstant

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return self

    def __neg__(self) -> Negate:
        from .ArithmeticConstant import ArithmeticConstant
        from .ABCUnaryExpression import Negate

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return Negate(self)

    def __abs__(self) -> AbsoluteValue:
        from .ArithmeticConstant import ArithmeticConstant
        from .ABCUnaryExpression import AbsoluteValue

        if isinstance(other, (int, float)):
            other = ArithmeticConstant(other)
        return AbsoluteValue(self)

    def __eq__(self, other: ABCArithmeticExpression) -> Equal:
        from .ArithmeticConstant import ArithmeticConstant
        from .BooleanBinaryExpression import Equal

        if isinstance(other, (int, float, str)):
            other = ArithmeticConstant(other)
        elif isinstance(other, bool):
            other = BooleanConstant(other)
        return Equal(self, other)

    def __ne__(self, other: ABCArithmeticExpression) -> NotEqual:
        from .ArithmeticConstant import ArithmeticConstant
        from .BooleanBinaryExpression import NotEqual

        if isinstance(other, (int, float, str)):
            other = ArithmeticConstant(other)
        elif isinstance(other, bool):
            other = BooleanConstant(other)
        return NotEqual(self, other)

tensql/QIR/ABCArithmeticExpression.py
sandialabs/TenSQL
from __future__ import annotations

from typing import TYPE_CHECKING

from .ABCExpression import ABCExpression

if TYPE_CHECKING:
    from .BooleanBinaryExpression import LogicalOr, LogicalAnd, Equal, NotEqual
    from .BooleanConstant import BooleanConstant
    from .ABCUnaryExpression import LogicalNot


class ABCBooleanExpression(ABCExpression):
    def __or__(self, other: ABCBooleanExpression) -> LogicalOr:
        from .BooleanConstant import BooleanConstant
        from .BooleanBinaryExpression import LogicalOr

        if isinstance(other, bool):
            other = BooleanConstant(other)
        elif not isinstance(other, ABCBooleanExpression):
            raise NotImplementedError("Right hand side must be a boolean expression")
        return LogicalOr(self, other)

    def __ror__(self, other: ABCBooleanExpression) -> LogicalOr:
        from .BooleanConstant import BooleanConstant
        from .BooleanBinaryExpression import LogicalOr

        if isinstance(other, bool):
            other = BooleanConstant(other)
        elif not isinstance(other, ABCBooleanExpression):
            raise NotImplementedError("Left hand side must be a boolean expression")
        return LogicalOr(other, self)

    def __and__(self, other: ABCBooleanExpression) -> LogicalAnd:
        from .BooleanConstant import BooleanConstant
        from .BooleanBinaryExpression import LogicalAnd

        if isinstance(other, bool):
            other = BooleanConstant(other)
        elif not isinstance(other, ABCBooleanExpression):
            raise NotImplementedError("Right hand side must be a boolean expression")
        return LogicalAnd(self, other)

    def __rand__(self, other: ABCBooleanExpression) -> LogicalAnd:
        from .BooleanConstant import BooleanConstant
        from .BooleanBinaryExpression import LogicalAnd

        if isinstance(other, bool):
            other = BooleanConstant(other)
        elif not isinstance(other, ABCBooleanExpression):
            raise NotImplementedError("Left hand side must be a boolean expression")
        return LogicalAnd(other, self)

    def __not__(self) -> LogicalNot:
        from .BooleanConstant import BooleanConstant
        from .ABCUnaryExpression import LogicalNot

        return LogicalNot(self)

    def __eq__(self, other: ABCBooleanExpression) -> Equal:
        from .BooleanBinaryExpression import Equal

        if isinstance(other, bool):
            other = BooleanConstant(other)
        elif not isinstance(other, ABCBooleanExpression):
            raise NotImplementedError("Left hand side must be a boolean expression")
        return Equal(self, other)

    def __ne__(self, other: ABCBooleanExpression) -> NotEqual:
        from .BooleanBinaryExpression import NotEqual

        if isinstance(other, bool):
            other = BooleanConstant(other)
        elif not isinstance(other, ABCBooleanExpression):
            raise NotImplementedError("Left hand side must be a boolean expression")
        return NotEqual(self, other)

tensql/QIR/ABCBooleanExpression.py
sandialabs/TenSQL
from __future__ import annotations

import abc
from typing import Optional, Generator, TYPE_CHECKING

from .ABCNode import ABCNode
from .. import Types

if TYPE_CHECKING:
    from .QirColumn import QirColumn

class ABCExpression(ABCNode):
    def __init__(self, name: Optional[str]) -> None:
        super().__init__(name)

    @property
    @abc.abstractmethod
    def description(self) -> str:
        pass

    @property
    @abc.abstractmethod
    def columns(self) -> Generator[QirColumn]:
        pass

tensql/QIR/ABCExpression.py
sandialabs/TenSQL
from __future__ import annotations

import abc
import pathlib
import json

from typing import Generator, TYPE_CHECKING, Tuple, List, Optional

if TYPE_CHECKING:
    from .ABCExpression import ABCExpression


class ABCNode(abc.ABC):
    def __init__(self, name: Optional[str]):
        if name is not None:
            self._name = name
        else:
            self._name = self.description
        self._alias = self._name

    @property
    def name(self) -> str:
        return self._name

    @property
    def alias(self) -> str:
        return self._alias

    def aliased(self, name: str) -> ABCNode:
        self._alias = name
        return self
      
    @abc.abstractmethod
    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        pass

    @property
    def description(self) -> str:
        return f"{type(self).__name__}"

    def describe_as_edges(self) -> Generator[Tuple[ABCNode, ABCNode, str], None, None]:
        # print(type(self))
        for edge_name, child in self:
            # print(type(parent), type(self))
            yield (self, child, edge_name)
            yield from child.describe_as_edges()

    @staticmethod
    def save_graph(path: str, *ir_nodes: ABCNode) -> None:
        path = pathlib.Path(path)
        with path.open("w") as fout:
            all_edges: List[Tuple[ABCNode, ABCNode]] = []
            for node in ir_nodes:
                all_edges.extend(node.describe_as_edges())

            print("digraph G {", file=fout)
            nodes = set()
            for tail, head, edge_name in all_edges:
                nodes.add((head.description, id(head)))
                nodes.add((tail.description, id(tail)))
                print(
                    f"    n{id(head)} -> n{id(tail)} [label={json.dumps(edge_name)}]",
                    file=fout,
                )

            print("", file=fout)
            for node_type, node_id in nodes:
                print(f"    n{node_id} [label={json.dumps(node_type)}]", file=fout)

            print("}", file=fout)

tensql/QIR/ABCNode.py
sandialabs/TenSQL
import abc
from typing import Generator

from .ABCNode import ABCNode
from .Table import QirTable

from .. import Database


class ABCTableSequence(ABCNode):
    @property
    @abc.abstractmethod
    def tables(self) -> Generator[QirTable, None, None]:
        pass

    @property
    @abc.abstractmethod
    def db(self) -> Database:
        pass

tensql/QIR/ABCTableSequence.py
sandialabs/TenSQL
from typing import Generator, Optional, Tuple

from .ABCNode import ABCNode

from .ABCExpression import ABCExpression
from .ABCArithmeticExpression import ABCArithmeticExpression
from .QirColumn import QirColumn
from .ABCBooleanExpression import ABCBooleanExpression
from .. import Types

class ABCUnaryExpression(ABCExpression):
    def __init__(
        self, op: str, operand: ABCExpression, name: Optional[str] = None
    ) -> None:
        self._operand = operand
        self._op = op
        super().__init__(name)

    @property
    def operand(self) -> ABCExpression:
        return self._operand

    @property
    def op(self) -> str:
        return self._op

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        yield ("operand", self._operand)

    @property
    def columns(self) -> Generator[QirColumn, None, None]:
        yield from self._operand.columns


class LogicalNot(ABCUnaryExpression):
    def __init__(self, arg: ABCBooleanExpression, name: Optional[str] = None) -> None:
        super().__init__("NOT", arg, name)

    @property
    def operand(self) -> ABCBooleanExpression:
        return self._operand

    @property
    def description(self) -> str:
        return f"NOT {self.operand.description}"



tensql/QIR/ABCUnaryExpression.py
sandialabs/TenSQL
from typing import Optional

from .BinaryExpression import BinaryExpression
from .ABCArithmeticExpression import ABCArithmeticExpression


class ArithmeticBinaryExpression(BinaryExpression, ABCArithmeticExpression):
    valid_ops = {"+", "-", "*", "/", "//", "%"}

    def __init__(
        self,
        op: str,
        left: ABCArithmeticExpression,
        right: ABCArithmeticExpression,
        name: Optional[str] = None,
    ) -> None:
        if op not in self.valid_ops:
            raise ValueError(f"Parameter op must be one of {self.valid_ops!r}")
        super().__init__(op=op, first=left, second=right, name=name)

    @property
    def left(self) -> ABCArithmeticExpression:
        return self._first

    @property
    def right(self) -> ABCArithmeticExpression:
        return self._second


class Add(ArithmeticBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("+", *args, **kwargs)


class Sub(ArithmeticBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("-", *args, **kwargs)


class Div(ArithmeticBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("/", *args, **kwargs)


class Mul(ArithmeticBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("*", *args, **kwargs)


class QirModulus(ArithmeticBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("%", *args, **kwargs)


class QirFloorDiv(ArithmeticBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("//", *args, **kwargs)

tensql/QIR/ArithmeticBinaryExpression.py
sandialabs/TenSQL
from .QirColumn import QirColumn
from .ABCArithmeticExpression import ABCArithmeticExpression


class ArithmeticColumn(QirColumn, ABCArithmeticExpression):
    pass

tensql/QIR/ArithmeticColumn.py
sandialabs/TenSQL
from typing import Optional, Union

from .ABCArithmeticExpression import ABCArithmeticExpression
from .Constant import Constant


class ArithmeticConstant(Constant, ABCArithmeticExpression):
    def __init__(self, value: Union[int, float], name: Optional[str] = None) -> None:
        self._value = value
        super().__init__(value=value, name=name)

    @property
    def value(self) -> Union[int, float]:
        return self._value

tensql/QIR/ArithmeticConstant.py
sandialabs/TenSQL
from typing import Optional

from .ABCUnaryExpression import ABCUnaryExpression
from .ABCArithmeticExpression import ABCArithmeticExpression
from .. import Types

class ArithmeticUnaryExpression(ABCUnaryExpression, ABCArithmeticExpression):
  pass
    
class Negate(ArithmeticUnaryExpression):
    def __init__(
        self, arg: ABCArithmeticExpression, name: Optional[str] = None
    ) -> None:
        super().__init__("-", arg, name)

    @property
    def operand(self) -> ABCArithmeticExpression:
        return self._operand

    @property
    def description(self) -> str:
        return f"-{self.operand.description}"


class AbsoluteValue(ArithmeticUnaryExpression):
    def __init__(
        self, arg: ABCArithmeticExpression, name: Optional[str] = None
    ) -> None:
        super().__init__("ABS", arg, name)

    @property
    def operand(self) -> ABCArithmeticExpression:
        return self._operand

    @property
    def description(self) -> str:
        return f"ABS({self.operand.description})"


class Ones(ArithmeticUnaryExpression):
    def __init__(self, arg: ABCArithmeticExpression, name: Optional[str] = None) -> None:
        super().__init__("ONE", arg, name)

    @property
    def operand(self) -> ABCArithmeticExpression:
        return self._operand

    @property
    def description(self) -> str:
        return f"Ones({self.operand.description})"


class Cast(ArithmeticUnaryExpression):
    def __init__(self, arg: ABCArithmeticExpression, type_: Types.Type, name: Optional[str] = None) -> None:
        self.type_ = type_
        super().__init__("CAST", arg, name)

    @property
    def operand(self) -> ABCArithmeticExpression:
        return self._operand

    @property
    def description(self) -> str:
        return f"CAST({self.operand.description} AS {type(self.type_).__name__})"

tensql/QIR/ArithmeticUnaryExpression.py
sandialabs/TenSQL
from typing import Generator, Optional, Tuple

from .ABCNode import ABCNode
from .ABCExpression import ABCExpression
from .QirColumn import QirColumn


class BinaryExpression(ABCExpression):
    def __init__(
        self,
        op: str,
        first: ABCExpression,
        second: ABCExpression,
        name: Optional[str] = None,
    ) -> None:
        self._first = first
        self._second = second
        self._op = op
        if name is None:
            name = self.description
        super().__init__(name)

    @property
    def op(self) -> str:
        return self._op

    @property
    def description(self) -> str:
        return f"({self._first.description} {self._op} {self._second.description})"

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        yield ("first", self._first)
        yield ("second", self._second)

    @property
    def columns(self) -> Generator[QirColumn, None, None]:
        yield from self._first.columns
        yield from self._second.columns

tensql/QIR/BinaryExpression.py
sandialabs/TenSQL
from typing import Optional

from .BinaryExpression import BinaryExpression
from .ABCBooleanExpression import ABCBooleanExpression
from .ABCExpression import ABCExpression


class BooleanBinaryExpression(BinaryExpression, ABCBooleanExpression):
    valid_ops = {"==", "!=", ">", "<", "<=", ">=", "AND", "OR"}

    def __init__(
        self,
        op: str,
        left: ABCExpression,
        right: ABCExpression,
        name: Optional[str] = None,
    ) -> None:
        if op not in self.valid_ops:
            raise ValueError(f"Parameter op must be one of {self.valid_ops!r}")
        super().__init__(op=op, first=left, second=right, name=name)

    @property
    def left(self) -> ABCExpression:
        return self._first

    @property
    def right(self) -> ABCExpression:
        return self._second


class Equal(BooleanBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("==", *args, **kwargs)


class NotEqual(BooleanBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("!=", *args, **kwargs)


class GreaterThan(BooleanBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(">", *args, **kwargs)


class LessThan(BooleanBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("<", *args, **kwargs)


class GreaterThanOrEqual(BooleanBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(">=", *args, **kwargs)


class LessThanOrEqual(BooleanBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("<=", *args, **kwargs)


class LogicalOr(BooleanBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("OR", *args, **kwargs)


class LogicalAnd(BooleanBinaryExpression):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__("AND", *args, **kwargs)

tensql/QIR/BooleanBinaryExpression.py
sandialabs/TenSQL
from .ABCBooleanExpression import ABCBooleanExpression
from .QirColumn import QirColumn


class BooleanColumn(QirColumn, ABCBooleanExpression):
    pass

tensql/QIR/BooleanColumn.py
sandialabs/TenSQL
from typing import Optional
from .Constant import Constant
from .ABCBooleanExpression import ABCBooleanExpression


class BooleanConstant(Constant, ABCBooleanExpression):
    def __init__(self, value: bool, name: Optional[str] = None) -> None:
        self._value = value
        super().__init__(value=value, name=name)

    @property
    def value(self) -> bool:
        return self._value


class QirFalse(BooleanConstant):
    def __init__(self, name: Optional[str] = None) -> None:
        Constant.__init__(self=self, value=False, name=name)


class QirTrue(BooleanConstant):
    def __init__(self, name: Optional[str] = None) -> None:
        Constant.__init__(self=self, value=True, name=name)

tensql/QIR/BooleanConstant.py
sandialabs/TenSQL
from typing import Generator, Optional, Union, Tuple
from .ABCNode import ABCNode
from .ABCExpression import ABCExpression
from .QirColumn import QirColumn


class Constant(ABCExpression):
    def __init__(
        self, value: Union[int, float, bool], name: Optional[str] = None
    ) -> None:
        self._value = value

        super().__init__(name)

    @property
    def value(self) -> Union[int, float, bool]:
        return self._value

    @property
    def description(self) -> str:
        return repr(self._value)

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        yield from []

    @property
    def columns(self) -> Generator[QirColumn, None, None]:
        pass

tensql/QIR/Constant.py
sandialabs/TenSQL
from typing import Union, Optional

from ..Table import Table
from ..Database import Database
from .Query import Query
from .SelectQuery import SelectQuery
from .InsertQuery import InsertQuery
from .Table import QirTable

# from . import Query, SelectQuery, QirTable
from .LAIRTranslator import QueryVisitor


class QirDatabase:
    def __init__(self, db: Database):
        self._db = db
        self._executor = QueryVisitor(db)

    @property
    def db(self):
        return self._db

    def __getattr__(self, name: str) -> QirTable:
        table = self._db.tables.get(name)
        if table is None:
            raise AttributeError()
        else:
            return QirTable(self.db, name, table)

    def import_table(self, name, table) -> QirTable:
        ret = QirTable(self.db, name, table)
        self._db.add_table(name, ret.raw)
        return ret

    def drop_table(self, name) -> None:
        del self._db[name]

    def query(
      self,
      table: Union[QirTable, Table],
      alias: Optional[str] = None
    ):
        if isinstance(table, Table):
            table = QirTable(self, None, table, alias=alias)
        elif alias is not None:
            table = table.aliased(alias)

        return SelectQuery(self, table)

    def insert(self, table: QirTable, **kwargs):
        return InsertQuery(self, table, **kwargs)

    def run(self, node: Query) -> Table:
        return self._executor.visit(node)

tensql/QIR/Database.py
sandialabs/TenSQL
from typing import Generator, Tuple

from .ABCNode import ABCNode
from .ABCTableSequence import ABCTableSequence
from .Table import QirTable
from ..Database import Database


class From(ABCTableSequence):
    def __init__(self, table: QirTable):
        super().__init__(None)
        self._table = table

    @property
    def table(self) -> QirTable:
        return self._table

    @property
    def db(self) -> Database:
        return self._table.db

    @property
    def tables(self) -> Generator[QirTable, None, None]:
        yield self._table

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        yield ("table", self._table)

tensql/QIR/From.py
sandialabs/TenSQL
from __future__ import annotations
from typing import Optional, Tuple, Generator

from .ABCNode import ABCNode
from .QirColumn import QirColumn
from .BooleanColumn import BooleanColumn
from .ArithmeticColumn import ArithmeticColumn
from ..Table import Table
from ..Database import Database


class GroupBy(ABCNode):
    def __init__(self, *columns: ArithmeticColumn) -> None:
        assert all(isinstance(col, ArithmeticColumn) for col in columns)
        self._columns = tuple(columns)

    @property
    def columns(self) -> Tuple[QirColumn]:
        return self._columns

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        for col in self._columns:
            yield ("column", col)

tensql/QIR/GroupBy.py
sandialabs/TenSQL
from typing import Optional, Sequence, Union, Generator, Tuple, Any, TYPE_CHECKING, Dict

import numpy as np

from .ABCTableSequence import ABCTableSequence
from .Query import Query
from .ABCNode import ABCNode
from .Table import QirTable
from ..Table import Table

ValuesType = Dict[str, Any]
RowValuesSequenceType = Sequence[ValuesType]

class InsertQuery(Query):
  def __init__(
    self,
    qirdb: Any,
    table: QirTable,
    *,
#    values: Optional[RowValuesSequenceType] = None,
    if_not_exists: bool = False,
#    from_table: Optional[QirTable] = None
  ):
    self._db = qirdb
    self._table = table
    self._if_not_exists = if_not_exists
    self._values = []
    self._from_records = None
    self._from_table = None
#    self._from_table = from_table

#    if values is None:
#        self._values = []
#    else:
#        self._values = list(values)

  @property
  def db(self) -> Any:
    return self._db

  @property
  def table(self) -> Table:
    return self._table

  def values(
    self,
    args: Optional[RowValuesSequenceType] = None,
    **kwargs: ValuesType
  ):
    assert args is None or len(kwargs) == 0
    assert self._from_table is None
    assert self._from_records is None

    if args is not None:
      self._values.extend(args)
    else:
      self._values.append(kwargs)

    return self

  def from_table(
    self,
    table: QirTable
  ):
    assert len(self._values) == 0
    assert self._from_table is None
    assert self._from_records is None
    self._from_table = table
    return self

  def from_records(
    self,
    data: np.ndarray
  ):
    assert len(self._values) == 0
    assert self._from_table is None
    assert self._from_records is None

    self._from_records = data
    return self

  def run(self) -> Table:
    ret = self._db.run(self)
    return ret

  def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
    yield ("input", self._table)

tensql/QIR/InsertQuery.py
sandialabs/TenSQL
from typing import Optional, Generator, Tuple

from .ABCNode import ABCNode
from .ABCTableSequence import ABCTableSequence
from .Table import QirTable
from .ABCBooleanExpression import ABCBooleanExpression
from ..Database import Database


class Join(ABCTableSequence):
    join_types = {"inner", "outer", "left", "right"}

    def __init__(
        self,
        left: ABCTableSequence,
        right: QirTable,
        *,
        on_clause: Optional[ABCBooleanExpression],
        kind: Optional[str] = None,
    ) -> None:
        if kind is None:
            kind = "inner"

        if kind not in self.join_types:
            raise ValueError(f"Parameter kind must be one of {self.join_types!r}")

        self._left = left
        self._right = right
        self._kind = kind
        self._on_clause = on_clause

    @property
    def left(self) -> ABCTableSequence:
        return self._left

    @property
    def right(self) -> QirTable:
        return self._right

    @property
    def kind(self) -> str:
        return self._kind

    @property
    def db(self) -> Database:
        return self._right.db

    @property
    def tables(self) -> Generator[QirTable, None, None]:
        yield from self._left.tables
        yield self._right

    @property
    def on_clause(self) -> ABCBooleanExpression:
        return self._on_clause

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        yield ("left", self._left)
        yield ("table", self._right)
        yield ("on", self._on_clause)

tensql/QIR/Join.py
sandialabs/TenSQL
from __future__ import annotations
import collections
import itertools

from typing import Optional, Sequence, Callable, Tuple, Any, Dict, FrozenSet

import pygraphblas.types
from pygraphblas.types import BOOL as GBBool
from pygraphblas import Matrix, Vector, Scalar

from .. import Types
from ..PrimaryKey import PrimaryKey
from ..Column import Column
from ..Table import Table
from ..util import SequentialIdGenerator, EncoderDecoder
from .. import LAIR

Promotion = Callable[[LAIR.Node], LAIR.Node]
QualifiedColumnName = Tuple[str, str] # Qualified Name: (TableName, ColumnName)
JoinMatches = FrozenSet[QualifiedColumnName]
AxesSpecifier = Tuple[int, ...]

def PromoteStencil(
    pattern: Tuple[AxesSpecifier, AxesSpecifier, AxesSpecifier],
    A: LAIR.Node, 
    B: LAIR.Node, 
    kind: str
):
    NDIM = tuple()
    if kind == 'inner':
        if pattern == (NDIM, NDIM, NDIM):
            oStencil = LAIR.ElementwiseMultiply(A, B)
        elif pattern[0] == NDIM:
            if x.nvals > 0 and y.nvals > 0:
                oStencil = x.apply_second(grb_op, y)
            elif isinstance(x, Vector):
                oStencil = Vector.sparse(x.type, x.size)
            elif isinstance(x, Matrix):
                oStencil = Matrix.sparse(x.type, x.nrows, x.ncols)

class JoinPromotion:
    def __init__(self, idx_in: AxesSpecifier, idx_out: AxesSpecifier, stencil: LAIR.Node):
        self.idx_in = idx_in
        self.idx_out = idx_out
        self.stencil = stencil

    def __call__(self, node: LAIR.node) -> LAIR.Node:
        return LAIR.InnerBroadcastMask(
            node, self.idx_in,
            self.stencil, self.idx_out,
            self.idx_out
        )

class TableBuilder:
    def __init__(
        self,
        shape: Sequence[int],
        stencil_builder: LAIR.Node,
        column_builders: Dict[str, Optional[LAIR.Node]],
    ):
        self._shape = tuple(shape)
        self._stencil_builder = stencil_builder
        self._column_builders = column_builders

    @property
    def ndim(self) -> int:
        return len(self._shape)

    @property
    def shape(self) -> Tuple[int, ...]:
        return self._shape

    @property
    def stencil(self) -> LAIR.Node:
        return self._stencil_builder

    def get_column(self, name):
        #print(tuple(self._column_builders.keys()))
        return self._column_builders[name]

    def mask(self, mask: LAIR.Node) -> TableBuilder:
        out_cols = collections.OrderedDict()
        for cb_name, cb in self._column_builders.items():
            if cb is not None:
                out_cols[cb_name] = LAIR.Mask(cb, mask)
            else:
                out_cols[cb_name] = None
        return TableBuilder(self._shape, LAIR.Mask(self._stencil_builder, mask), out_cols)


class JoinerResult:
    def __init__(
        self,
        tables: Dict[str, TableBuilder],
        common_stencil: LAIR.Node,
        common_shape: Sequence[int],
        common_keys: Sequence[JoinMatches],
    ):
        self._tables = tables
        self._stencil = common_stencil
        self._shape = tuple(common_shape)
        self._keys = tuple(common_keys)

    @property
    def tables(self) -> Dict[str, TableBuilder]:
        return self._tables

    @property
    def stencil(self) -> LAIR.Node:
        return self._stencil

    @property
    def shape(self) -> Tuple[int, ...]:
        return self._shape

    @property
    def ndim(self) -> int:
        return len(self._shape)

    @property
    def primary_keys(self) -> Tuple[JoinMatches, ...]:
        return self._keys

    def mask(self, mask: LAIR.Node) -> JoinerResult:
        new_tables = {k: v.mask(mask) for k, v in self._tables.items()}
        new_stencil = LAIR.Mask(self._stencil, mask)
        return JoinerResult(new_tables, new_stencil, self._shape, self._keys)


class Joiner:
    INDEX_TYPE = Types.BigInt()

    def __init__(
        self,
        out_shape: Sequence[int],
        stencil: LAIR.Node,
        promote_left: Promotion,
        promote_right: Promotion,
        join_keys: Sequence[JoinMatches],
    ):
        join_keys = tuple(join_keys)
        out_shape = tuple(out_shape)
        assert len(join_keys) == len(out_shape)

        self._join_keys = join_keys
        self._out_shape = out_shape
        self._stencil = stencil
        self._promote_left = promote_left
        self._promote_right = promote_right

    @property
    def ndim(self) -> int:
        return len(self._out_shape)

    def promoteTable(self, alias: str, table: TableBuilder, *, is_left: bool) -> TableBuilder:
        if is_left:
            promote = self._promote_left
        else:
            promote = self._promote_right

        table_pk_qnames = []
        for itdim, (col_name, _) in enumerate(table._column_builders.items()):
            if itdim < table.ndim:
                table_pk_qnames.append((alias, col_name))

        out_cols = collections.OrderedDict()
        for itdim, dim_keys in enumerate(self._join_keys):
            match_cols = [qname for qname in table_pk_qnames if qname in dim_keys]
            if len(match_cols) == 0:
                col_name = f"__bc_pk_{itdim}__"
            elif len(match_cols) == 1:
                table_alias, col_name = match_cols[0]
            else:
                raise ValueError(
                    "Found multiple names for output primary key (join by tensor diagonal)"
                )
            if len(match_cols) == 0:
              pk = PrimaryKey(col_name, self.INDEX_TYPE)
              out_cols[col_name] = LAIR.PrimaryKey(self._stencil, itdim, pk.gbtype, name=col_name)

        for it, (col_name, in_col) in enumerate(table._column_builders.items()):
            #if it >= table.ndim:
            out_cols[col_name] = promote(in_col)

        return TableBuilder(self._out_shape, self._stencil, out_cols)

    def joinAdditionalTable(
        self, left: JoinerResult, right_alias: str, right: TableBuilder
    ) -> JoinerResult:
        tables = {}
        for table_alias, table in left.tables.items():
            assert table_alias not in tables
            tables[table_alias] = self.promoteTable(table_alias, table, is_left=True)

        assert right_alias not in tables
        tables[right_alias] = self.promoteTable(right_alias, right, is_left=False)

        return JoinerResult(tables, self._stencil, self._out_shape, self._join_keys)

    @classmethod
    def from_axes(
        cls,
        A: LAIR.Node,
        aCols: Sequence[QualifiedColumnName], #Qualified names of a's columns
        aShape: Sequence[int],
        B: LAIR.Node,
        bCols: Sequence[QualifiedColumnName], #Qualified names of b's columns
        bShape: Sequence[int],
        join_keys: Sequence[JoinMatches],
        kind: str = "inner",
    ) -> Joiner:
        aCols, aShape = tuple(aCols), tuple(aShape)
        bCols, bShape = tuple(bCols), tuple(bShape)

        #print(f"{aCols=}, {aShape=}, {bCols=}, {bShape=}, {join_keys=}")

        assert kind in {"inner", "outer", "left", "right"}
        if kind != "inner":
            raise NotImplementedError()  # Left, Right, and Outer broadcasts

        if len(aCols) < len(bCols):
            aCols, aShape, A, bCols, bShape, B = bCols, bShape, B, aCols, aShape, A
            reversed = True
            if kind == "left":
                kind = "right"
            elif kind == "right":
                kind = "left"
        else:
            reversed = False

        join_keys = tuple(join_keys)

        # Maps qualified names of relevant primary keys to integer ids
        all_idx = EncoderDecoder(SequentialIdGenerator())
        oIndices = []
        oCols = []
        for join_matches in join_keys:
            match_id: Optional[int] = None
            for qname in join_matches:
                if match_id is None:
                    match_id = all_idx[qname]
                    oCols.append(qname)
                else:
                    all_idx[qname] = match_id
            assert match_id is not None
            oIndices.append(match_id)

        aIndices = tuple(all_idx[col] for col in aCols)
        bIndices = tuple(all_idx[col] for col in bCols)
        oIndices = tuple(oIndices)

        # Does not support reductions or broadcasting to new axes
        assert set(oIndices) == set(aIndices) | set(bIndices)

        # Build the result shape
        oShape = [None for oidx in oIndices]
        for name, idx, dim in itertools.chain(zip(aCols, aIndices, aShape), zip(bCols, bIndices, bShape)):
            if idx >= len(oShape):
                continue
            elif oShape[idx] is None:
                oShape[idx] = dim
            elif dim != oShape[idx]:
                oname = oCols[idx]
                raise ValueError(
                    "Shape operand dimension {name} disagrees with expected dimension for output primary key {oname}"
                )
        #print(f"{oShape=}")
        assert None not in oShape

        if kind == "inner":
#            oStencil: LAIR.Node = LAIR.ElementwiseMultiply(
#                'LAND',
#                LAIR.Promote(A, aIndices, oIndices),
#                LAIR.Promote(B, bIndices, oIndices),
#            )
            oStencil: LAIR.Node = LAIR.InnerBroadcast("LAND", A, aIndices, B, bIndices, oIndices)
#        elif kind == "outer":
#            oStencil: LAIR.Node = LAIR.BroadcastOuter(A, aIndices, B, bIndices, oIndices)
        promote_left: JoinPromotion = JoinPromotion(aIndices, oIndices, oStencil)
        promote_right: JoinPromotion = JoinPromotion(bIndices, oIndices, oStencil)

        if reversed:
            promote_left, promote_right = promote_right, promote_left

        return Joiner(oShape, oStencil, promote_left, promote_right, join_keys)

tensql/QIR/Joiner.py
sandialabs/TenSQL
from __future__ import annotations

import collections
import contextlib
from typing import Any, List, Dict, Tuple, Sequence, FrozenSet, Optional
import itertools
import json
import sys

from .Table import QirTable
from .QirColumn import QirColumn

from .. import QIR, LAIR, Table, Database
from ..Column import Column
from ..util import SequentialIdGenerator, EncoderDecoder, grb_shape, tensor_to_numpy
from . import Joiner
from .Joiner import JoinerResult, JoinMatches
from .Visitor import Visitor
from ..LAIR.Optimizer import optimize
from .. import Types
from .._PyHashOCNT import PyHashOCNT

import pygraphblas
import pygraphblas.types
from pygraphblas import Scalar, Vector, Matrix

def _get_reduce_pattern(
    input_matches: Tuple[JoinMatches, ...], output_keys: Tuple[JoinMatches, ...]
) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
    NODIM = tuple()
    if len(input_matches) == 0:
        pattern = (NODIM, NODIM)
    elif len(input_matches) == 1:
        if len(output_keys) == 0:
            pattern = ((0,), NODIM)
        else:
            pattern = ((0,), (0,))
    elif len(input_matches) == 2:
        if len(output_keys) == 0:
            pattern = ((0, 1), NODIM)
        elif len(output_keys) == 2:
            pattern = ((0, 1), (0, 1))
        else:
            if any(qname in input_matches[0] for qname in output_keys[0]):
                pattern = ((0, 1), (0,))
            elif any(qname in input_matches[1] for qname in output_keys[0]):
                pattern = ((0, 1), (1,))
            else:
                assert False
    else:
        all_idx = {}
        for it_idx, idx_matches in enumerate(input_matches):
            for it, qname in enumerate(idx_matches):
                all_idx[qname] = it_idx

        ret = []
        for it_idx, idx_matches in enumerate(output_keys):
            matches = []
            match_ids = set()
            for it, qname in enumerate(idx_matches):
                match_id = all_idx.get(qname)
                if match_id is not None:
                    matches.append(qname)
                    match_ids.add(match_id)
                    break
            if len(match_ids) == 0:
                raise ValueError(
                    f"All output keys for dimension {it_idx} missing from input keys: {idx_matches}"
                )
            elif len(match_ids) > 1:
                raise ValueError(
                    f"Output keys for dimension {it_idx} matched input keys in multiple dimensions: {match_ids}"
                )
            else:
                ret.append(list(match_ids)[0])
        pattern = (tuple(range(len(input_matches))), tuple(ret))
    return pattern


class SelectVisitor(Visitor):
    def visit(self, node: QIR.ABCNode) -> LAIR.Node:
        return super().visit(node)

    def __init__(self, join_result: JoinerResult, expected_keys: Sequence[JoinMatches]):
        self._join_result = join_result
        self._expected_keys = tuple(expected_keys)
        self._refcounts = PyHashOCNT(16)

    @property
    def join_map(self) -> Dict[str, Table]:
        return self._join_result.tables

    def generic_visit(self, node: QIR.ABCNode) -> None:
        classes = [cls for cls in type(node).mro() if issubclass(cls, QIR.ABCNode)]

        raise NotImplementedError(
            f"{type(self).__name__} does not implement a visit method for any of {classes}"
        )

    def visit_QirColumn(self, node: QIR.QirColumn) -> LAIR.Node:
        return self.join_map[node.table.alias].get_column(node.raw.name)

    def visit_Constant(self, node: QIR.Constant) -> LAIR.Node:
        type_ = Types.from_native(type(node.value))
        ret = Scalar.from_type(type_.as_pygrb)
        if type_.sideload:
            ret[None], count = self._refcounts.insert(node.value)
        else:
            ret[None] = node.value

        out_idx = tuple(range(self._join_result.ndim))

        return LAIR.InnerBroadcastMask(
            LAIR.Tensor(ret, name=node.description, stencil=None), tuple(),
            self._join_result.stencil, out_idx,
            out_idx
        )
            
#        return LAIR.Tensor(ret, name=node.description, stencil=None)
          

#        stencil = Scalar.from_type(Types.Boolean._as_pygrb)
#        stencil[None] = True
#        stencil = LAIR.Tensor(
#          stencil, name=f"_constant.stencil", stencil=None
#        )
#
#        shape = tuple()
#        columns = collections.OrderedDict()
#        columns["value"] = LAIR.Tensor(
#            ret, name="_constant.value", stencil=stencil
#        )
#
#        table_builder = Joiner.TableBuilder(shape, stencil, columns)
#
#
#        return Joiner.TableBuilder(shape, stencil, columns)


    def visit_ABCUnaryExpression(self, node: QIR.ABCUnaryExpression) -> LAIR.Node:
        x = self.visit(node.operand)
        return LAIR.ElementwiseApply(node.op, x)

    def visit_Cast(self, node: QIR.Cast) -> LAIR.Node:
        x = self.visit(node.operand)
        return LAIR.Cast(x, node.type_.as_pygrb)

    def visit_ArithmeticBinaryExpression(
        self, node: QIR.ArithmeticBinaryExpression
    ) -> LAIR.Node:
        x = self.visit(node.left)
        y = self.visit(node.right)

        # (null + anything) = null, so everything equates to EWiseMult
        return LAIR.ElementwiseMultiply(node.op, x, y)

    def visit_ABCReductionExpression(
        self, node: QIR.ABCReductionExpression
    ) -> LAIR.Node:
        x = self.visit(node.operand)

        pattern = _get_reduce_pattern(
            self._join_result.primary_keys, self._expected_keys
        )
        #print(self._join_result.primary_keys, self._expected_keys, pattern)

        return LAIR.Reduction(node.op, x, set(pattern[0]) - set(pattern[1]))

    def visit_Count(self, node: QIR.Count) -> LAIR.Node:
        x = self.visit(node.operand)
        pattern = _get_reduce_pattern(
            self._join_result.primary_keys, self._expected_keys
        )
        x1 = LAIR.ElementwiseApply("ONE", x)

        return LAIR.Reduction(node.op, x1, set(pattern[0]) - set(pattern[1]))

    def visit_BooleanBinaryExpression(
        self, node: QIR.BooleanBinaryExpression
    ) -> LAIR.Node:
        x = self.visit(node.left)
        y = self.visit(node.right)

        # (null || anything) = null, so everything equates to EWiseMult
        return LAIR.ElementwiseMultiply(node.op, x, y)


from .Visitor import JoinPlan, JoinPlanOnClauseVisitor, JoinPlanVisitor


class FromVisitor(Visitor):
    def __init__(self, db: Database, outkeys: Sequence[QIR.ArithmeticColumn]):
        self.db = db
        self._outkeys = outkeys

    def visit_From(self, node: QIR.From) -> JoinerResult:
        return self.visit(node.table)

    def to_TableBuilder(self, node: QIR.QirTable) -> Joiner.TableBuilder:
        stencil = LAIR.Tensor(
            node.raw.stencil, name=f"{node.alias}.stencil", stencil=None
        )
        shape = node.raw.shape
        columns = collections.OrderedDict()
        for it, col in enumerate(node.columns()):
            if it < len(shape):
                columns[col.name] = LAIR.PrimaryKey(stencil, it, col.raw.type_, name=".".join(col.qualname))
            else:
                columns[col.name] = LAIR.Tensor(
                    col.raw.data, name=".".join(col.qualname), stencil=stencil
                )

        return Joiner.TableBuilder(shape, stencil, columns)

    def visit_QirTable(self, node: QIR.QirTable) -> JoinerResult:
        grb_table = self.to_TableBuilder(node)

        tables = {node.alias: grb_table}

        primary_keys = [{(pk.table.alias, pk.name)} for pk in node.primary_keys]
        #print(f"Building table {node.alias} with primary keys: {primary_keys}")

        return JoinerResult(
            tables=tables,
            common_stencil=grb_table.stencil,
            common_shape=grb_table.shape,
            common_keys=primary_keys,
        )

    def visit_Join(self, node: QIR.Join) -> JoinerResult:
        left_result = self.visit(node.left)
        join_plan = JoinPlanVisitor().visit(node)
        join_keys = join_plan.realize(self._outkeys)

        sA = left_result.stencil
        aidx = [tuple(ax)[0] for ax in left_result.primary_keys]
        ashape = left_result.shape
        B = self.to_TableBuilder(node.right)
        sB = B.stencil
        bidx = [(col.table.alias, col.name) for col in node.right.primary_keys]
        bshape = node.right.raw.shape

        joiner = Joiner.Joiner.from_axes(sA, aidx, ashape, sB, bidx, bshape, join_keys)
        return joiner.joinAdditionalTable(left_result, node.right.alias, B)

    def visit(self, node: QIR.ABCNode) -> JoinerResult:
        return super().visit(node)


class WhereVisitor(SelectVisitor):
    pass


class QueryVisitor(Visitor):
    def __init__(self, db: Database):
        self.db = db

    def visit_SelectQuery(self, node: QIR.SelectQuery) -> Table:
        ################################################
        # Construct the join plan
        ################################################

        QIR.ABCNode.save_graph("QIR-viz.dot", node)

        if node._group_by is not None:
            outkeys = node._group_by.columns
        else:
            outkeys = []
            for clause in node._selects:
                if isinstance(
                    clause, QIR.ArithmeticColumn
                ) and clause.table.raw.is_primary_key(clause.raw):
                    outkeys.append(clause)
                else:
                    break

        #print(f"{outkeys=}")
        join_result = FromVisitor(self.db, outkeys).visit(node._table_seq)

        if node._group_by is not None:
            expected_primary_keys = [{col.qualname} for col in node._group_by.columns]
        else:
            expected_primary_keys = join_result.primary_keys

        #print(f"{expected_primary_keys=}")

        if node._where is not None:
            where = WhereVisitor(join_result, expected_primary_keys).visit(node._where)
            join_result = join_result.mask(LAIR.Pattern(where))

        ############################################
        # Create pseudo-columns for primary keys
        ############################################
        columns = []
        column_names = []
        #print(f"{expected_primary_keys=}")
        #print(f"{outkeys=}")
        for pk_aliases, pk_select in itertools.zip_longest(
            expected_primary_keys, node._selects
        ):
            if pk_aliases is None:
                # Select clauses after primary keys
                if isinstance(
                    pk_select, QIR.ArithmeticColumn
                ) and pk_select.table.raw.is_primary_key(pk_select.raw):
                    select_qname = pk_select.qualname
                    raise ValueError(
                        f"Unexpected additional primary key discovered in select clause: {select_qname}."
                    )
            else:
                if (
                    pk_select is None
                    or not isinstance(pk_select, QIR.ArithmeticColumn)
                    or not pk_select.table.raw.is_primary_key(pk_select.raw)
                ):
                    #print(f"{pk_select is None=},{not isinstance(pk_select, QIR.ArithmeticColumn)=},{not pk_select.table.raw.is_primary_key(pk_select.raw)=},{pk_aliases=},{pk_select=}")
                    expected = set(pk_aliases)
                    if pk_select is None:
                        observed = None
                    else:
                        observed = pk_select.qualname
                    raise ValueError(
                        f"Query select statements must start with the correct primary keys in the correct order.  Expected one of: {expected}, but found {observed}."
                    )

                select_qname = pk_select.qualname
                if select_qname not in pk_aliases:
                    raise ValueError(
                        f"Query select statements must start with the correct primary keys in the correct order.  Found: {select_qname}, but expected one of: {set(pk_aliases)}"
                    )
                columns.append(
                    Column(pk_select.alias, None, type_=pk_select.raw.type_)
                )
            column_names.append(pk_select.name)

        ################################################
        # Generate the IR for the output columns
        ################################################
        if node._group_by is not None:
            # Handle Reductions
            reduce_pattern = _get_reduce_pattern(
                join_result.primary_keys, expected_primary_keys
            )

            axes = [x for x in reduce_pattern[0] if x not in reduce_pattern[1]]
            out_ndim = join_result.ndim - len(axes)

            #print(f"Reduction: {axes=}, {out_ndim=}")
            out_stencil = LAIR.Reduction("LOR", join_result.stencil, axes)
        else:
            out_ndim = join_result.ndim
            out_stencil = join_result.stencil

        column_builders = []
        for col_name, col in zip(column_names[out_ndim:], node._selects[out_ndim:]):
            column_builders.append(
                LAIR.Output(
                    SelectVisitor(join_result, expected_primary_keys).visit(col),
                    name = col.alias
                )
            )

        out_stencil = LAIR.Output(out_stencil, "stencil")

        ################################################
        # Optimize the query
        ################################################

        LAIR.Node.save_graph("IRViz-unoptimized.dot", *column_builders, out_stencil)
        column_builders = [optimize(column) for column in column_builders]
        out_stencil = optimize(out_stencil)
        LAIR.Node.save_graph("IRViz-optimized.dot", *column_builders, out_stencil)

        ################################################
        # Execute the intermediate representation
        ################################################
        from ..LAIR.Visitor import ExpressionVisitor as LAIRVisitor

        column_tensors = list(LAIRVisitor().run(*column_builders, out_stencil))
        out_stencil = column_tensors.pop()
        #print(f"{len(grb_shape(out_stencil))=},{out_ndim=}")
        assert len(grb_shape(out_stencil)) == out_ndim
        #print("Finished executing LAIR") ; sys.stdout.flush()
        #print("column_builder names", [col.name for col in column_builders])
        #print("select names", [qircol.alias for qircol in node._selects[out_ndim:]])
        #print("column tensors", [tensor.type for tensor in column_tensors])
        #print("column tensors from_grb", [Types.from_pygrb(tensor.type) for tensor in column_tensors])


        for qircol, tensor in zip(node._selects[out_ndim:], column_tensors):
            type_ = Types.from_pygrb(tensor.type)
            #print("Building column", qircol.alias, type_) ; sys.stdout.flush()
            col = Column(qircol.alias, tensor, type_=type_)
            if col.data is not None and col.is_sideloaded:
                as_lists = tensor_to_numpy(col.data)
                col._refcounts.insert_many_encoded_numpy(as_lists[-1])
            columns.append(col)

        return Table(grb_shape(out_stencil), columns, stencil_data=out_stencil)

    def visit_InsertQuery(self, node: QIR.InsertQuery) -> Table:
        ndim = node.table.raw.ndim
        from_table = node._from_table
        from_records = node._from_records

        if from_records is not None:
            table = node.table.raw.copy_schema()
            table.insert_from_numpy(from_records)
            from_table = QIR.QirTable(node.db, "__FROM_RECORDS__", table)

        outdata = {}
        required_keys = set()
        available_keys = set()
        for it, (cname, col) in enumerate(node.table.raw.columns):
            available_keys.add(cname)
            if it >= ndim:
              outdata[cname] = tuple([[] for _ in range(ndim)] + [[]])
            else:
              required_keys.add(cname)

        #print("Building tensors")  ; sys.stdout.flush()
        stencildata = [[] for _ in range(ndim)]
        nvals = 0
        for it_row, row in enumerate(node._values):
            assert len(row) >= len(required_keys) and all(k in row for k in required_keys)
            assert not any(k not in available_keys for k in row)

            nvals += 1

            for it, (cname, col) in enumerate(node.table.raw.columns):
                # Get the value or its default
                default = None
                value = row.get(cname, default)

                if it < ndim:
                    # Add the primary key value to the stencil data
                    assert isinstance(value, int) or isinstance(value, np.integer)
                    stencildata[it].append(value)
                elif value is not None:
                    # Add the value to the output lists for the column
                    for it in range(ndim): # Append primary keys
                        outdata[cname][it].append(stencildata[it][-1])
                    if col.is_sideloaded:
                        value = id(value)
                    outdata[cname][ndim].append(value) # Append the value
                else:
                    pass
        #print("Done building tensors") ; sys.stdout.flush()

        
        column_builders = []

        ################################################
        # Execute the intermediate representation
        ################################################
        from ..LAIR.Visitor import ExpressionVisitor as LAIRVisitor
        from ..LAIR.ElementwiseAdd import ElementwiseAdd
        from ..LAIR.Tensor import Tensor

        if from_table is None:
            stencildata.append([True] * nvals)
            addstencil = Table._initialize_tensor(
                Types.Boolean.as_pygrb,
                node.table.raw.shape,
                stencildata
            )
            # Ensure the primary key constraints aren't violated by the incoming data
            if (not node._if_not_exists) and (addstencil.nvals != nvals):
                raise ValueError("Primary key uniqueness constraint violated")
        else:
            addstencil = from_table.raw.stencil
        addstencil = Tensor(addstencil, name='INSERT_VALUES.stencil')

        table_stencil = Tensor(node.table.raw.stencil, name=f"{node.table.name}.stencil")
        if node._if_not_exists:
            # Don't upload stuff that is already in the database
            addstencil = LAIR.AntiMask(addstencil, table_stencil)

        refcount_sources = {}

        #print("Building LAIR")  ; sys.stdout.flush()
        for it, (cname, col) in enumerate(node.table.raw.columns):
            if it >= ndim:
                if from_table is None:
                    adddata = Table._initialize_tensor(
                        col.type_.as_pygrb,
                        node.table.raw.shape,
                        outdata[cname]
                    )
                else:
                    adddata = from_table[cname].raw.data
                adddata = Tensor(adddata, name=f"INSERT_VALUES.{cname}")

                if node._if_not_exists:
                    # Don't upload stuff that is already in the database
                    adddata = LAIR.Mask(adddata, addstencil)

                if col.is_sideloaded:
                    refcount_sources[cname] = adddata

                column_builders.append(
                    LAIR.Output(
                        ElementwiseAdd(
                            "FIRST",
                            Tensor(col.data, name=f"{node.table.name}.{cname}"),
                            adddata
                        ),
                        f"{node.table.name}.{cname}"
                    )
                )

        #Add the column builder for the output stencil
        column_builders.append(
            LAIR.Output(
                ElementwiseAdd(
                    "FIRST",
                    Tensor(node.table.raw.stencil, name=f"{node.table.name}.stencil"),
                    addstencil
                ),
                f"{node.table.name}.stencil"
            )
        )

        refcount_source_keys = list(refcount_sources.keys())
        refcount_source_values = [LAIR.Output(refcount_sources[k], f"refcounts[{k}]") for k in refcount_source_keys]
        column_builders.extend(refcount_source_values)

        #print("Optimizing LAIR")  ; sys.stdout.flush()
        LAIR.Node.save_graph("IRViz-unoptimized.dot", *column_builders)
        column_builders = [optimize(column) for column in column_builders]
        LAIR.Node.save_graph("IRViz-optimized.dot", *column_builders)
        #print("Running LAIR")  ; sys.stdout.flush()
        column_tensors = list(LAIRVisitor().run(*column_builders))
        #print("Finished executing LAIR") ; sys.stdout.flush()
        
        if len(refcount_sources) > 0:
            column_tensors, refcount_tensors = column_tensors[:-len(refcount_sources)], column_tensors[-len(refcount_sources):]
        else:
            refcount_tensors = []
        out_stencil = column_tensors.pop()
        
        # Ensure the primary key constraints aren't violated by the combination of incoming and existing data
        if (not node._if_not_exists) and (out_stencil.nvals != node.table.raw.stencil.nvals + addstencil.tensor.nvals):
            raise ValueError("Primary key uniqueness constraint violated")

#        for it, (cname, col) in enumerate(node.table.raw.columns):
#          if not col.is_sideloaded:
#            continue
#
#          print("sideloaded column", it, cname, col.data.nvals)
#          if col.data.nvals > 0:
#            continue
#
#          refcount_tensor = refcount_tensors[refcount_source_keys.index(cname)]
#          refcount_values = tensor_to_numpy(refcount_tensor)[-1]
#          result_tensor = column_tensors[it - ndim]
#          result_values = tensor_to_numpy(refcount_tensor)[-1]
#
#          setdiff = set(refcount_values) - set(result_values)
#          print(f"{setdiff=}")
#          assert len(setdiff) == 0


        # Build the reference counts for the output columns
#        print("Building reference counts") ; sys.stdout.flush()
        output_refcounts = {cname: col._refcounts.copy() for cname, col in node.table.raw.columns if col.is_sideloaded}
        assert len(output_refcounts) == len(refcount_source_keys)
        assert len(output_refcounts) == len(refcount_tensors)
        for cname, value_tensor in zip(refcount_source_keys, refcount_tensors):
            values = tensor_to_numpy(value_tensor)[-1]
#            if from_table is not None:
#              print(f"{sys.getrefcount(from_table)=}") ; sys.stdout.flush()
#              print(f"Showing refcounts._print_debug from {from_table.name}.{cname}") ; sys.stdout.flush()
#              from_table[cname].raw._refcounts._print_debug(True)
#              print("Output refcounts for", cname, f"{from_table[cname].raw.data.nvals=}, {values.size=}")
#            else:
#              print("Output refcounts for", cname, f"{nvals=}, {values.size=}")
            output_refcounts[cname].insert_many_encoded_numpy(values)
            #output_refcounts[cname]._print_debug(False)

        #print("Building columns")  ; sys.stdout.flush()
        columns = []
        for it, (cname, col) in enumerate(node.table.raw.columns):
            if col.is_sideloaded:
                refcounts = output_refcounts[cname]
            else:
                refcounts = None

            #print("Instantiating column", cname, col.type_) ; sys.stdout.flush()
            if it >= ndim:
                columns.append(Column(cname, column_tensors[it - ndim], type_=col.type_, refcounts = refcounts))
            else:
                columns.append(Column(cname, None, type_=col.type_, refcounts=refcounts))

            #print("Finished column", cname, col.type_) ; sys.stdout.flush()

        #print("Building table")  ; sys.stdout.flush()
        ret = Table(node.table.raw.shape, columns, stencil_data=out_stencil)
        #print("Replacing table in database")  ; sys.stdout.flush()
        node.db.db[node.table.name] = ret

        return ret

tensql/QIR/LAIRTranslator.py
sandialabs/TenSQL
from __future__ import annotations

from typing import Generator, Tuple, TYPE_CHECKING

from .ABCNode import ABCNode
from .ABCExpression import ABCExpression
from ..Column import Column

if TYPE_CHECKING:
    from .Table import QirTable


class QirColumn(ABCExpression):
    def __init__(self, table: QirTable, name: str, column: Column) -> None:
        self._table = table
        self._column = column

        super().__init__(name)

    @property
    def table(self) -> QirTable:
        return self._table

#    @property
#    def name(self) -> str:
#        return self._name

    @property
    def raw(self) -> Column:
        return self._column

    @property
    def qualname(self) -> Tuple[str, str]:
        return (self._table.alias, self._name)
        # return f"{self._table.alias!r}.{self._name!r}"

    @property
    def is_primary_key(self) -> bool:
        for it, (col_name, col) in self._table.raw.columns:
            if it < self._table.ndim:
                if col is self._column:
                    return True
            else:
                break
        return False

    @property
    def primary_key_size(self) -> int:
        for it, (col_name, col) in self._table.raw.columns:
            if it < self._table.ndim:
                if col is self._column:
                    return self._table.shape[it]
            else:
                break
        return None

    @property
    def description(self) -> str:
        if self._alias == self._name:
            return ".".join(self.qualname)
        else:
            return ".".join(self.qualname) + " AS " + self._alias

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        yield ("table", self._table)

    @property
    def columns(self) -> Generator[QirColumn, None, None]:
        yield self

tensql/QIR/QirColumn.py
sandialabs/TenSQL
import abc
from typing import Any

from .ABCNode import ABCNode
from ..Database import Database


class Query(ABCNode):
    def __init__(self, db: Database):
        self._db = db

    @property
    def db(self) -> Database:
        return self._db

    @abc.abstractmethod
    def run(self) -> Any:
        pass

tensql/QIR/Query.py
sandialabs/TenSQL
from typing import Generator, Optional, Tuple

from .ABCNode import ABCNode
from .ABCExpression import ABCExpression
from .ABCArithmeticExpression import ABCArithmeticExpression
from .QirColumn import QirColumn
from .ABCBooleanExpression import ABCBooleanExpression


class ABCReductionExpression(ABCExpression):
    def __init__(
        self, op: str, operand: ABCExpression, name: Optional[str] = None
    ) -> None:
        self._operand = operand
        self._op = op
        super().__init__(name)

    @property
    def op(self):
        return self._op

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        yield ("operand", self._operand)

    @property
    def columns(self) -> Generator[QirColumn, None, None]:
        yield from self._operand.columns

    @property
    def operand(self) -> ABCBooleanExpression:
        return self._operand


class Any(ABCReductionExpression, ABCArithmeticExpression):
    def __init__(self, arg: ABCArithmeticExpression, name: Optional[str] = None):
        super().__init__("ANY", arg, name)

    @property
    def description(self) -> str:
        return f"ANY({self.operand.description})"


class Sum(ABCReductionExpression, ABCArithmeticExpression):
    def __init__(self, arg: ABCArithmeticExpression, name: Optional[str] = None):
        super().__init__("+", arg, name)

    @property
    def description(self) -> str:
        return f"SUM({self.operand.description})"


class Count(ABCReductionExpression, ABCArithmeticExpression):
    def __init__(self, arg: ABCArithmeticExpression, name: Optional[str] = None):
        super().__init__("+", arg, name)

    @property
    def description(self) -> str:
        return f"COUNT({self.operand.description})"


class Min(ABCReductionExpression, ABCArithmeticExpression):
    def __init__(self, arg: ABCArithmeticExpression, name: Optional[str] = None):
        super().__init__("MIN", arg, name)

    @property
    def description(self) -> str:
        return f"MIN({self.operand.description})"


class Max(ABCReductionExpression, ABCArithmeticExpression):
    def __init__(self, arg: ABCArithmeticExpression, name: Optional[str] = None):
        super().__init__("MAX", arg, name)

    @property
    def description(self) -> str:
        return f"MAX({self.operand.description})"

tensql/QIR/ReductionExpression.py
sandialabs/TenSQL
from __future__ import annotations

import copy
from typing import Optional, Sequence, Union, Generator, Tuple, Any, TYPE_CHECKING

from .ArithmeticColumn import ArithmeticColumn

from .Query import Query
from .Table import QirTable
from .From import From
from .Join import Join
from .GroupBy import GroupBy
from .QirColumn import QirColumn
from .ABCNode import ABCNode
from .ABCTableSequence import ABCTableSequence
from .ABCBooleanExpression import ABCBooleanExpression
from .ABCExpression import ABCExpression
from .ABCBooleanExpression import ABCBooleanExpression

from ..Table import Table

if TYPE_CHECKING:
    from .ABCBooleanExpression import ABCBooleanExpression


class SelectQuery(Query, ABCTableSequence):
    def __init__(
        self,
        qirdb: Any,
        table: Union[QirTable, ABCTableSequence],
        selects: Optional[Sequence[ABCExpression]] = None,
        where: Optional[ABCBooleanExpression] = None,
        group_by: Optional[GroupBy] = None,
    ):
        self.qirdb = qirdb
        db = qirdb.db

        if isinstance(table, QirTable):
            self._table_seq = From(table)
        elif isinstance(table, ABCTableSequence):
            self._table_seq = table
        else:
            raise ValueError(
                "Parameter table must be an instance of QirTable or ABCQirTableSequence"
            )

        if selects is None:
            self._selects = []
        else:
            self._selects = [col for col in selects]

        self._where = where
        self._group_by = group_by

        super().__init__(db)

    def __copy__(self) -> SelectQuery:
        return SelectQuery(
            self.qirdb, self._table_seq, selects=self._selects, where=self._where
        )

    def copy(self) -> SelectQuery:
        return copy.copy(self)

    def tables(self) -> Generator[QirTable]:
        yield from self._table_seq.tables

    def columns(self) -> Generator[QirColumn]:
        for col in self._selects:
            yield from col._columns

    def join(
        self,
        table: QirTable,
        on_clause: ABCBooleanExpression = None,
        *,
        kind: Optional[str] = None,
    ) -> SelectQuery:
        ret = self.copy()
        ret._table_seq = Join(
            left=self._table_seq, right=table, on_clause=on_clause, kind=kind
        )
        return ret

    def select(self, *selects: ABCExpression) -> SelectQuery:
        for col in selects:
            if not isinstance(col, ABCExpression):
                raise ValueError(f"Expression is not an instance of ABCQirExpression: {col}")
        if self._selects != []:
            raise ValueError("select must only be called once on a given query")
        ret = self.copy()
        ret._selects.extend(selects)
        return ret

    def where(self, clause: ABCBooleanExpression) -> SelectQuery:
        if not isinstance(clause, ABCBooleanExpression):
            raise ValueError("Expression is not an instance of ABCQirBooleanExpression")
        if self._where is not None:
            raise ValueError("where must only be called once on a given query")
        ret = self.copy()
        ret._where = clause
        return ret

    def group_by(self, *columns: ArithmeticColumn):
        if not all(isinstance(column, ArithmeticColumn) for column in columns):
            raise ValueError("Arguments to group_by must instances of ArithmeticColumn")
        if self._group_by is not None:
            raise ValueError("group_by must only be called once on a given query")
        ret = self.copy()
        ret._group_by = GroupBy(*columns)
        return ret

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        for select in self._selects:
            yield ("output", select)
        yield ("inputs", self._table_seq)
        if self._where is not None:
            yield ("where", self._where)
        if self._group_by is not None:
            yield ("group_by", self._group_by)

    def run(self) -> Table:
        return self.qirdb.run(self)

    def run_subquery(self, alias) -> QirTable:
        result = self.qirdb.run(self)
        return QirTable(self.db, None, result, alias=alias)

tensql/QIR/SelectQuery.py
sandialabs/TenSQL
from __future__ import annotations
import collections
from typing import Optional, List, Generator, Tuple

import numpy as np

from .ABCNode import ABCNode
from .QirColumn import QirColumn
from .BooleanColumn import BooleanColumn
from .ArithmeticColumn import ArithmeticColumn
from ..Table import Table
from ..Database import Database


class QirTable(ABCNode):
    def __init__(
        self, db: Database, name: str, table: Table, alias: Optional[str] = None
    ) -> None:
        self._db = db
        self._table = table
        self._name = name

        if alias is not None:
            self._alias = alias
        else:
            self._alias = name

        columns = list(self._table.columns)
        self._namedtuple = collections.namedtuple(
            self.name, 
            [c_name for c_name, c in columns],
            defaults = [None] * len(columns)
        )
        rtype = []
        for cname, c in columns:
            rtype.append((cname, c.type_.as_numpy))
        self._record_type = np.dtype(rtype)

    @property
    def db(self) -> Database:
        return self._db

    @property
    def raw(self) -> Table:
        return self._table

    @property
    def namedtuple(self):
      return self._namedtuple

    @property
    def record_type(self) -> np.dtype:
      return self._record_type

    @property
    def alias(self) -> str:
        return self._alias

    @property
    def name(self) -> str:
        return self._name

    @property
    def primary_keys(self) -> List[QirColumn]:
        ret = []
        for it, (col_name, col) in enumerate(self._table.columns):
            if it < self._table.ndim:
                ret.append(ArithmeticColumn(self, col_name, col))
            else:
                break
        return ret

    def aliased(self, alias: str) -> QirTable:
        return QirTable(self._db, self._name, self._table, alias)

    def __getitem__(self, name: str) -> QirColumn:
        col = self._table._columns.get(name)
        if col is None or col.hidden:
            raise AttributeError()
        elif col.is_boolean:
            return BooleanColumn(self, name, col)
        else:
            return ArithmeticColumn(self, name, col)

    def __getattr__(self, name: str) -> QirColumn:
        col = self._table._columns.get(name)
        if col is None or col.hidden:
            raise AttributeError()
        elif col.is_boolean:
            return BooleanColumn(self, name, col)
        else:
            return ArithmeticColumn(self, name, col)

    def columns(self) -> Generator[QirColumn, None, None]:
        for col_name, col in self._table.columns:
            yield self[col_name]

    def __iter__(self) -> Generator[Tuple[str, ABCNode], None, None]:
        yield from []

    @property
    def description(self) -> str:
        return self.alias

tensql/QIR/Table.py
sandialabs/TenSQL
from typing import Any, List, Dict, Tuple, Sequence, FrozenSet, Optional
import itertools

from .Table import QirTable
from .QirColumn import QirColumn

from .. import QIR, Table, Database
from ..Column import Column
from ..util import SequentialIdGenerator, EncoderDecoder
from .Joiner import Joiner, JoinerResult, JoinMatches

import pygraphblas
import pygraphblas.types
from pygraphblas import Scalar


def _get_reduce_pattern(
    input_matches: Tuple[JoinMatches, ...], output_keys: Tuple[JoinMatches, ...]
) -> Tuple[Tuple[int, ...], Tuple[int, ...]]:
    NODIM = tuple()
    if len(input_matches) == 0:
        pattern = (NODIM, NODIM)
    elif len(input_matches) == 1:
        if len(output_keys) == 0:
            pattern = ((0,), NODIM)
        else:
            pattern = ((0,), (0,))
    elif len(input_matches) == 2:
        if len(output_keys) == 0:
            pattern = ((0, 1), NODIM)
        elif len(output_keys) == 2:
            pattern = ((0, 1), (0, 1))
        else:
            if any(qname in input_matches[0] for qname in output_keys[0]):
                pattern = ((0, 1), (0,))
            elif any(qname in input_matches[1] for qname in output_keys[0]):
                pattern = ((0, 1), (1,))
            else:
                assert False
    else:
        raise ValueError("Prototype only supports scalars, vectors, and matrices")
    return pattern


class Visitor:
    def visit(self, node: QIR.ABCNode) -> Any:
        for cls in type(node).mro():
            if issubclass(cls, QIR.ABCNode):
                visit_name = f"visit_{cls.__name__}"
                if hasattr(self, visit_name):
                    return getattr(self, visit_name)(node)
        return self.generic_visit(node)

    def generic_visit(self, node: QIR.ABCNode) -> None:
        classes = [cls for cls in type(node).mro() if issubclass(cls, QIR.ABCNode)]
        raise NotImplementedError(
            f"{type(self).__name__} does not implement a visit method for any of {classes}"
        )


class SelectVisitor(Visitor):
    UnaryPyOps = {"-": lambda x: -x, "NOT": lambda x: not x, "ABS": lambda x: abs(x)}

    UnaryGrbOps = {"-": "AINV", "NOT": "LNOT", "ABS": "ABS"}

    BinaryArithmeticPyOps = {
        "+": lambda x, y: x + y,
        "-": lambda x, y: x - y,
        "*": lambda x, y: x * y,
        "/": lambda x, y: x / y,
        "//": lambda x, y: x // y,
        "%": lambda x, y: x % y,
    }

    BinaryArithmeticGrbOps = {
        "+": "PLUS",
        "-": "MINUS",
        "*": "TIMES",
        "/": "DIV",
    }

    BinaryBooleanPyOps = {
        "==": lambda x, y: x == y,
        "!=": lambda x, y: x != y,
        ">": lambda x, y: x > y,
        "<": lambda x, y: x < y,
        "<=": lambda x, y: x <= y,
        ">=": lambda x, y: x >= y,
        "AND": lambda x, y: x and y,
        "OR": lambda x, y: x or y,
    }

    BinaryBooleanGrbOps = {
        "==": "EQ",
        "!=": "NE",
        ">": "GT",
        "<": "LT",
        "<=": "GE",
        ">=": "LE",
        "AND": "TIMES",
        "OR": "PLUS",
    }

    def __init__(self, join_result: JoinerResult, expected_keys: Sequence[JoinMatches]):
        self._join_result = join_result
        self._expected_keys = tuple(expected_keys)

    @property
    def join_map(self) -> Dict[str, Table]:
        return self._join_result.tables

    def generic_visit(self, node: QIR.ABCNode) -> None:
        classes = [cls for cls in type(node).mro() if issubclass(cls, QIR.ABCNode)]

        raise NotImplementedError(
            f"{type(self).__name__} does not implement a visit method for any of {classes}"
        )

    def visit_QirColumn(self, node: QIR.QirColumn) -> Column:
        return self.join_map[node.table.alias].get_column(node.name).copy()

    def visit_Constant(self, node: QIR.Constant) -> Column:
        ret = Scalar.from_value(node.value)
        return Column(node.name, ret, gbtype=ret.type)

    def visit_ABCUnaryExpression(self, node: QIR.ABCUnaryExpression) -> Column:
        x = self.visit(node.operand)
        if x.is_scalar:
            ret = Scalar.from_value(self.UnaryPyOps[node.op](x.data[0]))
        else:
            ret = x.apply(getattr(x.gbtype, self.UnaryGrbOps[node.op]))
        return Column(node.name, ret, gbtype=ret.type)

    def visit_ArithmeticBinaryExpression(
        self, node: QIR.ArithmeticBinaryExpression
    ) -> Column:
        x = self.visit(node.left)
        y = self.visit(node.right)

        assert x.gbtype == y.gbtype

        grb_op = getattr(x.gbtype, self.BinaryArithmeticGrbOps[node.op])
        if x.is_scalar and y.is_scalar:
            ret = Scalar.from_value(
                self.BinaryArithmeticPyOps[node.op](x.data[0], y.data[0])
            )
        elif (not x.is_scalar) and y.is_scalar:
            ret = x.data.apply_first(y.data, grb_op)
        elif x.is_scalar and (not y.is_scalar):
            ret = x.data.apply_second(y.data, grb_op)
        else:
            ret = x.data.eadd(y.data, grb_op)

        return Column(node.name, ret, gbtype=ret.type)

    def visit_ABCReductionExpression(
        self, node: QIR.ArithmeticBinaryExpression
    ) -> Column:
        x = self.visit(node.arg)

        NODIM = tuple()
        pattern = _get_reduce_pattern(
            self._join_result.primary_keys, self._expected_keys
        )
        grb_op = getattr(x.gbtype, f"{node.op}_MONOID")

        # Scalar Input
        if pattern == (NODIM, NODIM):
            ret = x

        # Vector Input
        elif pattern == ((0,), NODIM):
            ret = x.reduce(grb_op)
        elif pattern == ((0,), (0,)):
            ret = x

        # Matrix Input
        elif pattern == ((0, 1), NODIM):
            ret = x.reduce(grb_op)
        elif pattern == ((0, 1), (0,)):
            ret = x.T.reduce_vector(grb_op)
        elif pattern == ((0, 1), (1,)):
            ret = x.reduce_vector(grb_op)
        elif pattern == ((0, 1), (0, 1)):
            ret = x
        else:
            assert False

        return Column(node.name, ret, gbtype=ret.type)

    def visit_BooleanBinaryExpression(
        self, node: QIR.BooleanBinaryExpression
    ) -> Column:
        x = self.visit(node.left)
        y = self.visit(node.right)

        assert x.gbtype == y.gbtype

        if not x.is_scalar:
            if x.is_vector:
                ret = x.data.__class__.sparse(pygraphblas.types.BOOL, x.data.size)
            else:
                ret = x.data.__class__.sparse(
                    pygraphblas.types.BOOL, x.data.nrows, x.data.ncols
                )
        elif not y.is_scalar:
            if y.is_vector:
                ret = y.data.__class__.sparse(pygraphblas.types.BOOL, y.data.size)
            else:
                ret = y.data.__class__.sparse(
                    pygraphblas.types.BOOL, y.data.nrows, y.data.ncols
                )
        else:
            ret = None

        grb_op = getattr(x.gbtype, self.BinaryBooleanGrbOps[node.op])
        py_op = self.BinaryBooleanPyOps[node.op]
        if x.is_scalar and y.is_scalar:
            ret = Scalar.from_value(py_op(x.data[0], y.data[0]))
        elif (not x.is_scalar) and y.is_scalar:
            x.data.apply_second(grb_op, y.data, out=ret)
        elif x.is_scalar and (not y.is_scalar):
            y.data.apply_first(x.data, grb_op, out=ret)
        else:
            ret = x.data.emult(y.data, grb_op, out=ret)

        return Column(node.name, ret, gbtype=ret.type)


class JoinPlan:
    def __init__(self):
        self.index_constraints = []
        self.primary_keys = []
        self.tables_by_alias = {}

    def add_primary_key(self, pk1: QirColumn) -> None:
        ret = (pk1.table.alias, pk1.name)
        self.primary_keys.append(ret)
        return ret

    def add_constraint(self, pk1: QirColumn, pk2: QirColumn) -> None:
        # assert pk1.is_primary_key
        # assert pk2.is_primary_key
        # assert pk1.primary_key_size == pk2.primary_key_size

        t1 = pk1.table
        t2 = pk2.table.alias
        if self.tables_by_alias.setdefault(pk1.table.alias, pk1.table) is not pk1.table:
            raise ValueError(f"Duplicate table alias detected: {pk1.table.alias!r}")
        if self.tables_by_alias.setdefault(pk2.table.alias, pk2.table) is not pk2.table:
            raise ValueError(f"Duplicate table alias detected: {pk2.table.alias!r}")

        pk1_qname = self.add_primary_key(pk1)
        pk2_qname = self.add_primary_key(pk2)
        self.index_constraints.append((pk1_qname, pk2_qname))

    def add_constraints(
        self, constraints: Sequence[Tuple[QirColumn, QirColumn]]
    ) -> None:
        for pk1, pk2 in constraints:
            self.add_constraint(pk1, pk2)

    def realize(
        self, outkeys: Optional[Sequence[QIR.ArithmeticColumn]]
    ) -> Tuple[JoinMatches, ...]:
        "Groups input primary keys into output primary keys"
        if outkeys is None:
            outkeys = tuple()
        #print("&" * 80)
        #print(f"{outkeys=}")
        #print(f"{self.index_constraints=}")

        # Identify which primary keys need to be merged
        connected_components = []
        matched_pks = set()
        for pk1, pk2 in self.index_constraints:
            pk1_cc, pk2_cc = None, None
            matched_pks.add(pk1)
            matched_pks.add(pk2)
            for it_cc, cc in enumerate(connected_components):
                if pk1 in cc:
                    pk1_cc = it_cc
                if pk2 in cc:
                    pk2_cc = it_cc

            if pk1_cc is None and pk2_cc is None:
                # Create a new connected component
                connected_components.append({pk1, pk2})
            elif pk1_cc is not None and pk2_cc is not None:
                # Merge later connected component into earlier connected component
                if pk1_cc != pk2_cc:
                    target = connected_components[min(pk1_cc, pk2_cc)]
                    later = connected_components.pop(max(pk1_cc, pk2_cc))
                    target |= later
            elif pk1_cc is not None:
                # Add pk2 to pk1's connected component
                connected_components[pk1_cc].add(pk2)
            else:
                # Add pk1 to pk2's connected component
                connected_components[pk2_cc].add(pk1)

        #print("CC1", len(connected_components))
        #print(connected_components)

        # Add primary keys with no joins as stand-alone connected components
        for pk_qname in set(self.primary_keys) - matched_pks:
            connected_components.append([pk_qname])
        #print("CC2", len(connected_components))

        # Re-order the connected components based on outkeys.  Connected components
        # containing an output key should be moved to the front and have the same
        # order as outkeys.
        for outkey in reversed(outkeys):
            outkey_qname = (outkey.table.alias, outkey.name)
            for it_cc, cc in enumerate(connected_components):
                if outkey_qname in cc:
                    connected_components.insert(0, connected_components.pop(it_cc))
                    break
        #print("CC3", len(connected_components))

        return tuple(frozenset(x) for x in connected_components)


class JoinPlanOnClauseVisitor(Visitor):
    def __init__(self, target_table: QirTable):
        self._target_table = target_table

    def visit_Equal(self, node: QIR.Equal) -> List[Tuple[QirColumn, QirColumn]]:
        assert isinstance(node.left, QIR.QirColumn)
        assert isinstance(node.right, QIR.QirColumn)
        # print(node.left.table, node.right.table, self._target_table)
        # print(node.left.table.name, node.right.table.name, self._target_table.name)
        assert self._target_table.name in (node.left.table.name, node.right.table.name)

        assert node.left in node.left.table.primary_keys
        assert node.right in node.right.table.primary_keys
        return [(node.left, node.right)]

    def visit_LogicalAnd(
        self, node: QIR.LogicalAnd
    ) -> List[Tuple[QirColumn, QirColumn]]:
        ret = self.visit(node.left)
        ret.extend(self.visit(node.right))
        return ret


class JoinPlanVisitor(Visitor):
    def visit(self, node: QIR.ABCNode) -> JoinPlan:
        return super().visit(node)

    def visit_From(self, node: QIR.From) -> JoinPlan:
        return self.visit(node.table)

    def visit_QirTable(self, node: QIR.Table) -> JoinPlan:
        ret_map = JoinPlan()
        for pk in node.primary_keys:
            ret_map.add_primary_key(pk)
        return ret_map

    def visit_Join(self, node: QIR.Join) -> JoinPlan:
        ret = self.visit(node.left)
        for pk in node.right.primary_keys:
            ret.add_primary_key(pk)

        if node.on_clause is not None:
            ret.add_constraints(
                JoinPlanOnClauseVisitor(node.right).visit(node.on_clause)
            )
        return ret

tensql/QIR/Visitor.py
sandialabs/TenSQL
from .ABCNode import ABCNode
from .ABCExpression import ABCExpression
from .ABCArithmeticExpression import ABCArithmeticExpression
from .ABCBooleanExpression import ABCBooleanExpression
from .QirColumn import QirColumn
from .ArithmeticColumn import ArithmeticColumn
from .BooleanColumn import BooleanColumn
from .Constant import Constant
from .ArithmeticConstant import ArithmeticConstant
from .BooleanConstant import BooleanConstant
from .BooleanConstant import QirTrue
from .BooleanConstant import QirFalse
from .ABCUnaryExpression import ABCUnaryExpression
from .ABCUnaryExpression import LogicalNot
from .ArithmeticUnaryExpression import AbsoluteValue
from .ArithmeticUnaryExpression import Cast
from .ArithmeticUnaryExpression import Negate
from .ArithmeticUnaryExpression import Ones
from .BinaryExpression import BinaryExpression
from .ArithmeticBinaryExpression import (
    ArithmeticBinaryExpression,
    Add,
    Sub,
    Mul,
    Div,
    QirFloorDiv,
    QirModulus,
)
from .BooleanBinaryExpression import (
    BooleanBinaryExpression,
    Equal,
    NotEqual,
    GreaterThan,
    LessThan,
    GreaterThanOrEqual,
    LessThanOrEqual,
    LogicalOr,
    LogicalAnd,
)
from .Table import QirTable
from .ReductionExpression import Any, Sum, Min, Max, Count, ABCReductionExpression

from .ABCTableSequence import ABCTableSequence
from .From import From
from .Join import Join
from .Query import Query
from .SelectQuery import SelectQuery

from .Database import QirDatabase
from .LAIRTranslator import QueryVisitor

tensql/QIR/__init__.py
sandialabs/TenSQL
from __future__ import annotations

import collections
import json
import warnings
from typing import Union, Optional, Sequence, Tuple, Any, Dict

from pygraphblas import Matrix, Vector, Scalar
import pygraphblas.types
import pygraphblas.types
import h5py
import numpy as np

from .Column import Column
from .PrimaryKey import PrimaryKey
from . import util, Types
from .util import GBType, GBTensor, scalar_add, build_tensor, initialize_tensor, grb_shape, ndim, tensor_to_numpy, fill_tensor_from_numpy

class Table:
    """
    Represents a colleciton of tensors, each represented by a column
    """

    def __init__(
        self,
        shape: Sequence[int],
        columns: Sequence[Column],
        stencil_data: GBTensor = None,
    ) -> None:
        assert isinstance(shape, (list, tuple))
        assert len(shape) <= 2
        assert len(shape) <= len(columns)
        assert all(isinstance(d, int) for d in shape)

        assert isinstance(columns, (list, tuple))
        assert all(isinstance(col, Column) for col in columns)

        self._columns = collections.OrderedDict()
        self.shape = tuple(shape)
        # first_column = None

        first_column: Union[None, Column] = None

        for it_column, column in enumerate(columns):
            if it_column >= len(self.shape):
                first_column = column

            if column.name in self._columns:
                raise ValueError(f"Duplicate column name {column.name!r}")

            self._columns[column.name] = column

        # Create the stencil column representing the indices
        if stencil_data is None:
            stencil_data = self._build_tensor(pygraphblas.types.BOOL, self.shape)
            if first_column is not None:
                if not first_column.is_scalar:
                    #TODO: make this the union of all columns, instead of using the first one
                    mask_data = first_column.data.pattern()
                    stencil_data.assign_scalar(True, mask=mask_data)
                elif first_column.data.nvals > 0:
                    stencil_data[None] = True
        else:
            if isinstance(stencil_data, Scalar):
                new_stencil_data = Scalar.from_type(pygraphblas.types.BOOL)
                if stencil_data.nvals > 0 and stencil_data[None]:
                    new_stencil_data[None] = True
                stencil_data = new_stencil_data
            else:
                stencil_data = stencil_data.nonzero()

        name = "_stencil_"
        self._columns[name] = Column(
            name, stencil_data, hidden=True, type_ = Types.Boolean()
        )

        self.mask(stencil_data, inplace=True)

    def __getitem__(self, column_name: str) -> Optional[Column]:
      return self._columns.get(column_name)

    def copy_schema(self):
        new_cols = []
        for name, col in self.columns:
            if col.data is None:
                data = None
            else:
                data = self._build_tensor(col.gbtype, self.shape)
            new_cols.append(Column(name, data, hidden=col.hidden, type_=col.type_))
        return Table(self.shape, new_cols)

    @classmethod
    def from_definition(
        cls,
        column_types: Dict[str, Types.Type],
        primary_keys: Union[Sequence[PrimaryKey], Sequence[str]],
        shape: Optional[Sequence[int, ...]] = None
    ) -> Table:
        assert isinstance(column_types, dict)
        assert isinstance(primary_keys, (tuple, list))

        assert len(primary_keys) <= 2

        columns = []
        default_shape = []
        column_types = column_types.copy()

        for pk_or_name in primary_keys:
            if isinstance(pk_or_name, str):
                coltype = column_types.pop(pk_or_name)
                pk: PrimaryKey = PrimaryKey(pk_or_name, coltype)
            else:
                pk = pk_or_name

            default_shape.append(pk.max_idx)
            columns.append(Column(pk.column_name, None, type_=pk.type_))

        if shape is None:
            shape = default_shape

        shape = tuple(shape)

        for name, type_ in column_types.items():
            gbtype = type_.as_pygrb
            data = cls._build_tensor(gbtype, shape)
            columns.append(Column(name, data, type_=type_))

        return Table(shape, columns)

    @classmethod
    def from_tensor(
        cls, tensor: GBTensor, shape: Optional[Sequence[int]] = None, *, copy=False
    ) -> Table:
        # tensor_shape = util.shape(tensor)
        # if shape is not None:
        #   shape = tuple(shape)
        #   assert len(shape) == tensor_shape
        #   for new_dim, tensor_dim in

        if copy:
            tensor = tensor.dup()

        value_type = Types.from_pygrb(tensor.type)
        assert value_type is not None
        value_gbtype = value_type.as_pygrb

        if isinstance(tensor, Scalar):
            ret = Scalar.from_type(value_gbtype)
            if tensor.nvals > 0:
                ret[None] = tensor[None]

            columns = [Column("value", ret, type_=value_type)]
        elif isinstance(tensor, Vector):
            ret = tensor.apply(value_gbtype.IDENTITY)

            columns = [
                Column("idx", None, type_=Types.BigInt()),
                Column("value", ret, type_=value_type),
            ]
        elif isinstance(tensor, Matrix):
            ret = tensor.cast(value_gbtype)

            columns = [
                Column("ridx", None, type_=Types.BigInt()),
                Column("cidx", None, type_=Types.BigInt()),
                Column("value", ret, type_=value_type),
            ]
        else:
            raise NotImplementedError(
                "Prototype only supports scalars, vectors, and matrices"
            )

        return Table(util.shape(ret), columns)

    @property
    def schema(self):
        return {
            "nnz": self.count(),
            "shape": self.shape,
            "columns": [
                [c_name, c.schema]
                for c_name, c in self.columns
            ],
        }

    @property
    def ndim(self) -> int:
        "Returns the number of primary keys in the table"
        return len(self.shape)

    @property
    def is_scalar(self) -> bool:
        "Returns true iff the table has 0 primary keys"
        return len(self.shape) == 0

    @property
    def is_vector(self) -> bool:
        "Returns true iff the table has 1 primary keys"
        return len(self.shape) == 1

    @property
    def is_matrix(self) -> bool:
        "Returns true iff the table has 2 primary keys"
        return len(self.shape) == 2

    def is_primary_key(self, check_col: Column):
        "Checks whether a given column is a primary key column"
        for it_col, (col_name, col) in enumerate(self.columns):
            if it_col >= self.ndim:
                break
            if col is check_col:
                return True
        return False

    @property
    def primary_key_names(self) -> Tuple[str, ...]:
        "Returns the names of the table's primary keys"
        return tuple(col_name for it_col, (col_name, col) in enumerate(self.columns) if it_col < self.ndim)

    @property
    def stencil(self) -> GBTensor:
        "Returns the sparsity pattern of the table defining the rows"
        return self._columns["_stencil_"].data

    def get_column(self, name: str) -> Optional[Column]:
        "Returns a specific column associated with the given name (or None)"
        return self._columns.get(name)

    @property
    def columns(self) -> Sequence[Tuple[str, Column]]:
        "Generates a sequence of (column_name, column) pairs"
        for name, col in self._columns.items():
            if not col.hidden:
                yield name, col

    @classmethod
    def _build_tensor(cls, gbtype: pygraphblas.types.Type, shape: Sequence[int]) -> GBTensor:
        return build_tensor(gbtype, shape)

    @classmethod
    def _initialize_tensor(
        cls, gbtype: pygraphblas.types.Type, shape: Sequence[int], lists: Sequence[Sequence[Any]]
    ) -> GBTensor:
        return initialize_tensor(gbtype, shape, lists)

    def count(self) -> int:
        return self.stencil.nvals

    def add(self, other: Table, *, inplace=False, op: Any = None) -> Table:
        assert self.shape == other.shape
        assert len(self._columns) == len(other._columns)

        new_cols = []
        for self_col, other_col in zip(self._columns.values(), other._columns.values()):
            assert self_col.name == other_col.name
            assert (self_col.data is None) == (other_col.data is None)

            # if self_col.name == "_stencil_":
            #  continue

            if self_col.data is not None:
                if self.is_scalar:
                    data = scalar_add(self_col.data, other_col.data, op=op, inplace=inplace)
                elif not inplace:
                    data = self_col.data.eadd(other_col.data, add_op=op)
                elif inplace:
                        data = self_col.data.eadd(
                            other_col.data, add_op=op, out=self_col.data
                        )
            else:
                data = None

            if not inplace:
                new_cols.append(
                    Column(
                        self_col.name,
                        data,
                        type_=self_col.type_,
                        hidden=self_col.hidden,
                    )
                )

        if not inplace:
            return Table(self.shape, new_cols)
        else:
            return self

    def mask(self, mask: GBTensor, *, inplace=False):
        new_cols = []

        if not self.is_scalar:
            mask = self.stencil.emult(mask).nonzero()
        else:
            if self.stencil.nvals == 0:
                mask.clear()

        for it, (cname, col) in enumerate(self.columns):
            if col.data is not None:
                if not inplace:
                    data = self._build_tensor(col.type_, self.shape)
                else:
                    data = col.data

                if self.is_scalar:
                    if mask.nvals > 0 and col.data.nvals > 0:
                        data[0] = col.data[0]
                    else:
                        data.clear()
                else:
                    data.assign(col.data, mask=mask)
            else:
                data = None
            new_cols.append(Column(cname, data, type_=col.type_))

        if not inplace:
            return Table(self.shape, new_cols, stencil_data=mask)
        else:
            name = "_stencil_"
            self._columns[name] = Column(
                name, mask, hidden=True, type_=Types.Boolean()
            )
            return self

    def insert(self, **kwargs: Dict[str, Any]) -> None:
        assert set(kwargs.keys()) <= set(
            cname for cname, col in self._columns.items() if not col.hidden
        )
        pk = []
        for it, (cname, col) in enumerate(self.columns):
            if it >= self.ndim:
                break
            pk.append(kwargs[cname])
        
        if self.is_scalar:
            self.stencil[None] = True
            for cname, col in list(self.columns)[self.ndim :]:
                if cname in kwargs:
                    col._insert(None, kwargs[cname])
            return
        elif self.is_vector:
            # assign_scalar only works with slices on vectors
            pk = tuple(slice(idx, idx) for idx in pk)
        elif self.is_matrix:
            pk = tuple(pk)
        self.stencil.assign_scalar(True, *pk)
        for cname, col in list(self.columns)[self.ndim :]:
            if cname in kwargs:
                col._insert(pk, kwargs[cname])
        return

    def insert_from_numpy(self, data: np.ndarray) -> None:
        nptype = data.dtype
        colnames = set(c_name for c_name, column in self.columns)
        assert isinstance(nptype, np.dtype)
        assert all(nptype_name in colnames for nptype_name in nptype.names)
        assert all(nptype.fields[c_name][0] == column.nptype for c_name, column in self.columns)

        #TODO: rollback all changes to all columns and the stencil on failure
        pks = [data[c_name] for itdim, (c_name, column) in enumerate(self.columns) if itdim < self.ndim]
        stencil_data = pks + [[True] * data.size]
        add_stencil = self._initialize_tensor(self.stencil.type, grb_shape(self.stencil), stencil_data)
        self.stencil.eadd(add_stencil, out=self.stencil)

        for itdim, (c_name, column) in enumerate(self.columns):
            if itdim < self.ndim:
                pass
            else:
                column._insert_bulk(pks, data[c_name])


    def __len__(self) -> int:
        return self.stencil.nvals

    def to_records(self) -> np.ndarray:
        names, columns = zip(*self.columns)
        nvals, ndim = self.stencil.nvals, self.ndim

        #ret_dtype = [(name, column.type_.as_numpy) for name, column in self.columns]
        ret_dtype = [(name, object) for name, column in self.columns]
        ret = np.empty(shape=(nvals,), dtype=ret_dtype)

        stencil = self.stencil
        indices = tensor_to_numpy(stencil)[:-1]
        offsets = np.arange(stencil.nvals, dtype=np.int64)

        all_offsets_tensor = build_tensor(pygraphblas.types.INT64, self.shape)
        fill_tensor_from_numpy(all_offsets_tensor, *indices, offsets)

        if ndim == 1:
            ret[names[0]] = indices[0]
        elif ndim == 2:
            ret[names[0]], ret[names[1]] = indices[0], indices[1]

        for it_column, column in enumerate(columns):
            tensor = column.data
            if it_column < self.ndim:
                continue
            
            if ndim == 1:
                #TODO Cast mask to BOOL
                offset_tensor = all_offsets_tensor.extract(tensor)
            elif ndim == 2:
                bool_tensor = Matrix.sparse(typ=pygraphblas.types.BOOL, nrows=tensor.nrows, ncols=tensor.ncols)
                tensor.apply(tensor.type.IDENTITY, out=bool_tensor)
                offset_tensor = all_offsets_tensor.extract_matrix(mask=bool_tensor)

            offset_indices = list(tensor_to_numpy(offset_tensor))
            offset_values = offset_indices.pop()

            tensor_indices = list(tensor_to_numpy(tensor))
            tensor_values = tensor_indices.pop()

            #print(f"{offset_tensor.nvals=}, {tensor.nvals=}, {all_offsets_tensor.nvals=}, {self.stencil.nvals=}")

            #for off_idx, ten_idx in zip(offset_indices, tensor_indices):
            #    np.testing.assert_equal(off_idx, ten_idx)

            if column.is_sideloaded:
                tensor_values = column._refcounts.decode_many(tensor_values)

            ret[names[it_column]][offset_values] = tensor_values

        return ret


    def __iter__(self):
        if self.is_scalar:
            if self.stencil.nvals > 0:
                row = {}
                for it_col, (name, c) in enumerate(self.columns):
                    if c.data.nvals == 0:
                        value = None
                    elif c.is_sideloaded:
                        value = c._refcounts.decode(c.data[None])
                    else:
                        value = c.data[None]
                    row[name] = value
                yield row
        elif self.is_vector or self.is_matrix:
            ret = collections.OrderedDict()
            pk_names = [name for name, c in list(self.columns)[: self.ndim]]

            all_null = {name: None for it_col, (name, c) in enumerate(self.columns) if it_col >= self.ndim}

            stencil_lists = self.stencil.to_lists()
            for row in zip(*stencil_lists[:-1]):
                new_row = dict(zip(pk_names, row))
                new_row.update(all_null)
                ret[tuple(row)] = new_row

            for it_col, (name, c) in enumerate(self.columns):
                if it_col >= self.ndim:
                    col_lists = c.data.to_lists()
                    for col_row in zip(*col_lists):
                        pks, val = col_row[:-1], col_row[-1]
                        row = ret.get(pks)
                        if row is None:
                            pass
                        elif c.is_sideloaded:
                            row[name] = c._refcounts.decode(val)
                        else:
                            row[name] = val
            
            yield from ret.values()
        else:
            raise NotImplementedError(
                "Prototype only supports scalars, vectors, and matrices"
            )

    def save_hdf5(self, table_group, verbose=0) -> None:
        table_group.attrs.create(
            "table", json.dumps(self.schema), shape=tuple(), dtype=h5py.string_dtype()
        )

        if self.ndim == 0:
            idx_list = []
        else:
            idx_list = self.stencil.to_lists()[: self.ndim]

        data_group = table_group.create_group("data")
        stencil_group = table_group.create_group("stencil")
        pk_cols = []

        for it, (col_name, col) in enumerate(self.columns):

            if it < self.ndim:
                if verbose >= 1:
                    print("Saving", "primary key", col_name)
                values = np.array(idx_list[it], dtype=col.nptype)
                stencil_group.create_dataset(col_name, data=values)
                pk_cols.append(col)
            else:
                if verbose >= 1:
                    print("Saving", "column", col_name, "with", col.data.nvals, "nnz")
                col_group = data_group.create_group(col_name)
                col.save_hdf5(col_group, pk_cols, verbose=verbose-1)

    @classmethod
    def load_hdf5(cls, table_group) -> Table:
        metadata = json.loads(table_group.attrs["table"])
        shape = metadata["shape"]
        nnz = metadata['nnz']
        ndim = len(shape)
        columns = []

        idx_list = []
        data_group = table_group["data"]
        stencil_group = table_group["stencil"]
        for it, (col_name, col_info) in enumerate(metadata["columns"]):
            if it < ndim:
                column = Column.load_hdf5(None, col_name, shape, **col_info)
                idx_list.append(stencil_group[col_name])
            else:
                column = Column.load_hdf5(data_group[col_name], col_name, shape, **col_info)

            columns.append(column)

        if nnz > 0:
            stencil_data = cls._initialize_tensor(
                pygraphblas.types.BOOL, shape, idx_list + [[True] * nnz]
            )
        else:
            stencil_data = cls._build_tensor(pygraphblas.types.BOOL, shape)

        return Table(shape, columns, stencil_data)
    
    def __del__(self) -> None:
        pass
#        for it_col, (name, c) in enumerate(self.columns):
#            if it_col >= self.ndim:
#                if c.type_.sideload:
#                    col_lists = c.data.to_lists()
#                    for val in col_lists[-1]:
#                        c.type_.sideload_hash.remove_by_id(val)

tensql/Table.py
sandialabs/TenSQL
import abc
import re
import ast

from typing import Set

from pygraphblas.base import lib as pygrb_lib
from pygraphblas.base import ffi as pygrb_ffi
from pygraphblas import unary_op, binary_op
import h5py
import numba
import numpy as np
import pygraphblas.types
import numba
from ._PyHashOCNT import npbytes2npstr, npstr2npbytes

from . import PyObjectHelper

type_registry = {}

def register_type(cls):
  type_registry[cls.__name__.upper()] = cls
  return cls

class Type:
  _as_pygrb = None
  _as_numpy = None
  _as_python = None
  _as_h5py = None
  _unary_ops = set()
  _binary_ops = set()
  _sideload = False
  _is_string = False

  @property
  def name(self) -> str:
    return type(self).__name__ + "()"

  @classmethod
  @property
  def as_pygrb(cls):
    return cls._as_pygrb

  @classmethod
  @property
  def as_numpy(cls):
    return cls._as_numpy

  @classmethod
  @property
  def as_python(cls):
    return cls._as_python

  @classmethod
  @property
  def as_h5py(cls):
    return cls._as_h5py

  @classmethod
  @property
  def unary_ops(cls) -> Set[str]:
    return cls._unary_ops

  @classmethod
  @property
  def binary_ops(cls) -> Set[str]:
    return cls._binary_ops

  @classmethod
  @property
  def sideload(cls) -> bool:
    return cls._sideload

  @classmethod
  @property
  def is_string(cls) -> bool:
    return cls._is_string

#def boolean_binary_op(arg_type, nopython=True):
#    """ Based on code from PyGraphBLAS """
#    def inner(func):
#        func_name = func.__name__
#        sig = numba.void(
#            numba.types.CPointer(pygraphblas.types.INT8._numba_t),
#            numba.types.CPointer(arg_type._numba_t),
#            numba.types.CPointer(arg_type._numba_t),
#        )
#        inner_sig = numba.boolean(arg_type._numba_t, arg_type._numba_t)
#        jitfunc = numba.jit(inner_sig, nopython=nopython, forceobj=not nopython)(func)
#
#        @numba.cfunc(sig, nopython=True)
#        def wrapper(z, x, y):  # pragma: no cover
#            result = False
#            result = jitfunc(x[0], y[0])
#            z[0] = result
#
#        out = pygrb_ffi.new("GrB_BinaryOp*")
#        pygrb_lib.GrB_BinaryOp_new(
#            out,
#            pygrb_ffi.cast("GxB_binary_function", wrapper.address),
#            pygraphblas.types.BOOL._gb_type,
#            arg_type._gb_type,
#            arg_type._gb_type,
#        )
#
#        return pygraphblas.binaryop.BinaryOp(func_name, arg_type.__name__, out[0])
#
#    return inner

@register_type
class Boolean(Type):
  _as_pygrb = pygraphblas.types.BOOL
  _as_numpy = np.bool_
  _as_h5py = np.bool_
  _as_python = bool

  _unary_ops = {"NOT", "ONE"}
  _binary_ops = {"==", "!=", "AND", "OR"}

class NumericType(Type):
  _unary_ops = {"-", "NOT", "ABS", "ONE"}
  _binary_ops = {"+", "-", "*", "/", "==", "!=", "<", ">", "<=", ">=", "AND", "OR"}

@register_type
class TinyInt(NumericType):
  _as_pygrb = pygraphblas.types.INT8
  _as_numpy = np.int8
  _as_h5py = np.int8
  _as_python = int

@register_type
class SmallInt(NumericType):
  _as_pygrb = pygraphblas.types.INT16
  _as_numpy = np.int16
  _as_h5py = np.int16
  _as_python = int

@register_type
class Integer(NumericType):
  _as_pygrb = pygraphblas.types.INT32
  _as_numpy = np.int32
  _as_h5py = np.int32
  _as_python = int

@register_type
class BigInt(NumericType):
  _as_pygrb = pygraphblas.types.INT64
  _as_numpy = np.int64
  _as_h5py = np.int64
  _as_python = int

@register_type
class Real(NumericType):
  _as_pygrb = pygraphblas.types.FP32
  _as_numpy = np.float32
  _as_h5py = np.float32
  _as_python = float

@register_type
class DoublePrecision(NumericType):
  _as_pygrb = pygraphblas.types.FP64
  _as_numpy = np.float64
  _as_h5py = np.float64
  _as_python = float

class SideLoadType(Type):
  pass

class StringType(SideLoadType):
  _binary_ops = {"==", "!="}
  _sideload = True
  _is_string = True

  def serialize(self, value):
    if isinstance(value, str):
      return value.encode('utf-8')
    elif isinstance(value, np.ndarray):
      if value.size > 0:
        ret = np.empty_like(value)
        npstr2npbytes(value, ret)
      else:
        ret = np.empty_like(value, dtype=object)
      return ret
    else:
      raise ValueError(f"Unsupported type: {type(value)}")

#  def serialize_bulk(self, value_array: np.ndarray):
#    

  def deserialize(self, b_value):
    if isinstance(b_value, bytes):
      return b_value.decode('utf-8')
    elif isinstance(b_value, h5py.Dataset):
      return np.array(b_value.asstr())
    elif isinstance(b_value, np.ndarray):
      if b_value.size > 0:
        ret = np.empty_like(b_value)
        npbytes2npstr(b_value, ret)
        #decoder = np.vectorize(lambda x: x.decode('UTF-8'))
        #ret = decoder(b_value).astype(object)
      else:
        ret = np.zeros_like(b_value, dtype=object)
      return ret
    else:
      raise ValueError(f"Unsupported type: {type(value)}")

class StringGrbTypeBase(PyObjectHelper.PyObject):
#  default_one = 1
#  default_zero = 0
#  _gb_type = pygrb_lib.GrB_INT64
#  _c_type = "int64_t"
#  _typecode = "q"
#  _numba_t = numba.int64
#  _numpy_t = np.int64
#  _base_name = "INT64"
  _max_length = None

  @pygraphblas.unary_op(pygraphblas.INT64)
  def ONE(x):
    return 1

  @pygraphblas.unary_op(pygraphblas.INT64)
  def ZERO(x):
    return 0

  @pygraphblas.unary_op(pygraphblas.INT64)
  def IDENTITY(x):
    return x

  @pygraphblas.binary_op(pygraphblas.INT64)
  def FIRST(x, y):
    return x

  @pygraphblas.binary_op(pygraphblas.INT64)
  def SECOND(x, y):
    return y

  @pygraphblas.binary_op(pygraphblas.INT64)
  def ANY(x, y):
    return x

#  @classmethod
#  def _from_value(cls, value):
#    return pygraphblas.INT64._from_value(value)
#
#  @classmethod
#  def _to_value(cls, value):
#    return pygraphblas.INT64._to_value(value)

#MIN,MAX,PLUS,FIRSTI,FIRSTI1,FIRSTJ,FIRSTJ1,SECONDI,SECONDI1,SECONDJ,SECONDJ1,PAIR,LOR,LAND,LXOR

class CharGrbTypeBase(StringGrbTypeBase):
  pass

@register_type
class Char(StringType):
  _as_h5py = h5py.string_dtype('utf-8')
  _as_python = str
  _sideload = False

  def __init__(self, max_length: int):
    assert isinstance(max_length, int)
    assert max_length > 0
    self._max_length = max_length

    class CharGrbType(CharGrbTypeBase):
      _max_length = self._max_length
    
    self._as_pygrb = CharGrbType

  @property
  def name(self) -> str:
    return f"{type(self).__name__}({self._max_length})"

  @property
  def max_length(self) -> int:
    return self._max_length

  @property
  def as_numpy(self):
    return np.dtype(f'U{self.max_length}')

  @property
  def as_h5py(self):
    return h5py.string_dtype('utf-8', self.max_length * 4)

  @property
  def as_pygrb(self):
    return self._as_pygrb


class VarCharGrbTypeBase(StringGrbTypeBase):
  pass

@register_type
class VarChar(StringType):
  _as_numpy = np.object_
  _as_h5py = h5py.string_dtype('utf-8')
  _as_python = str

  def __init__(self, max_length: int):
    assert isinstance(max_length, int)
    assert max_length > 0
    self._max_length = max_length

  @property
  def name(self) -> str:
    return f"{type(self).__name__}({self._max_length})"

  @property
  def max_length(self) -> int:
    return self._max_length

  @property
  def as_pygrb(self):
    class VarCharGrbType(VarCharGrbTypeBase):
      _max_length = self._max_length

    return VarCharGrbType

class TextGrbType(StringGrbTypeBase):
  pass

@register_type
class Text(StringType):
  _as_numpy = np.object_
  _as_h5py = h5py.string_dtype('utf-8')
  _as_python = str

  @property
  def as_pygrb(self):
    return TextGrbType

def from_pygrb(grb_type):
  if (type(grb_type) == pygraphblas.types.BOOL) or (grb_type is pygraphblas.types.BOOL):
    return Boolean()
  elif (type(grb_type) == pygraphblas.types.INT8) or (grb_type is pygraphblas.types.INT8):
    return TinyInt()
  elif (type(grb_type) == pygraphblas.types.INT16) or (grb_type is pygraphblas.types.INT16):
    return SmallInt()
  elif (type(grb_type) == pygraphblas.types.INT32) or (grb_type is pygraphblas.types.INT32):
    return Integer()
  elif (type(grb_type) == pygraphblas.types.INT64) or (grb_type is pygraphblas.types.INT64):
    return BigInt()
  elif (type(grb_type) == pygraphblas.types.FP32) or (grb_type is pygraphblas.types.FP32):
    return Real()
  elif (type(grb_type) == pygraphblas.types.FP64) or (grb_type is pygraphblas.types.FP64):
    return DoublePrecision()
  elif isinstance(grb_type, CharGrbTypeBase) or issubclass(grb_type, CharGrbTypeBase):
    return Char(grb_type.max_length)
  elif isinstance(grb_type, VarCharGrbTypeBase) or issubclass(grb_type, VarCharGrbTypeBase):
    return VarChar(grb_type.max_length)
  elif isinstance(grb_type, TextGrbType) or issubclass(grb_type, TextGrbType):
    return Text()
  else:
    raise NotImplementedError("Invalid grb_type")

def from_numpy(np_type):
  if isinstance(grb_type, np.bool_) or issubclass(grb_type, np.bool_):
    return Boolean()
  elif isinstance(grb_type, np.int8) or issubclass(grb_type, np.int8):
    return TinyInt()
  elif isinstance(grb_type, np.int16) or issubclass(grb_type, np.int16):
    return SmallInt()
  elif isinstance(grb_type, np.int32) or issubclass(grb_type, np.int32):
    return Integer()
  elif isinstance(grb_type, np.int64) or issubclass(grb_type, np.int64):
    return BigInt()
  elif isinstance(grb_type, np.float32) or issubclass(grb_type, np.float32):
    return Real()
  elif isinstance(grb_type, np.float64) or issubclass(grb_type, np.float64):
    return DoublePrecision()
  else:
    raise NotImplementedError(f"Unsupported np_type: {np_type!s}")

def from_native(py_type):
  if issubclass(py_type, int):
    return BigInt()
  elif issubclass(py_type, float):
    return DoublePrecision()
  elif issubclass(py_type, str):
    return Text()
  else:
    raise NotImplementedError(f"Unsupported py_type: {py_type!s}")

def from_string(type_decl: str):
  re_typename = r'(\w+)(\(.*\))'
  match = re.fullmatch(re_typename, type_decl)
  if match is None:
    raise ValueError(f"Invalid type string: {type_decl!r}")
  typename = match.group(1).upper()
  args = ast.literal_eval(match.group(2))
  if not isinstance(args, tuple):
    args = (args,)
  if typename not in type_registry:
    raise LookupError(f"Invalid type name: {typename!r}")
  return type_registry[typename](*args)

tensql/Types.py
sandialabs/TenSQL
from .Column import Column
from .PrimaryKey import PrimaryKey
from .Table import Table
from .Database import Database
#from . import LAIR
#from . import QIR
from . import _version

__version__ = _version.get_versions()["version"]

tensql/__init__.py
sandialabs/TenSQL
# This file helps to compute a version number in source trees obtained from
# git-archive tarball (such as those provided by githubs download-from-tag
# feature). Distribution tarballs (built by setup.py sdist) and build
# directories (produced by setup.py build) will contain a much shorter file
# that just contains the computed version number.

# This file is released into the public domain.
# Generated by versioneer-0.28
# https://github.com/python-versioneer/python-versioneer

"""Git implementation of _version.py."""

import errno
import os
import re
import subprocess
import sys
from typing import Callable, Dict
import functools


def get_keywords():
    """Get the keywords needed to look up the version information."""
    # these strings will be replaced by git during git-archive.
    # setup.py/versioneer.py will grep for the variable names, so they must
    # each be defined on a line of their own. _version.py will just call
    # get_keywords().
    git_refnames = "$Format:%d$"
    git_full = "$Format:%H$"
    git_date = "$Format:%ci$"
    keywords = {"refnames": git_refnames, "full": git_full, "date": git_date}
    return keywords


class VersioneerConfig:
    """Container for Versioneer configuration parameters."""


def get_config():
    """Create, populate and return the VersioneerConfig() object."""
    # these strings are filled in when 'setup.py versioneer' creates
    # _version.py
    cfg = VersioneerConfig()
    cfg.VCS = "git"
    cfg.style = ""
    cfg.tag_prefix = "v"
    cfg.parentdir_prefix = "None"
    cfg.versionfile_source = "tensql/_version.py"
    cfg.verbose = False
    return cfg


class NotThisMethod(Exception):
    """Exception raised if a method is not valid for the current scenario."""


LONG_VERSION_PY: Dict[str, str] = {}
HANDLERS: Dict[str, Dict[str, Callable]] = {}


def register_vcs_handler(vcs, method):  # decorator
    """Create decorator to mark a method as the handler of a VCS."""

    def decorate(f):
        """Store f in HANDLERS[vcs][method]."""
        if vcs not in HANDLERS:
            HANDLERS[vcs] = {}
        HANDLERS[vcs][method] = f
        return f

    return decorate


def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=None):
    """Call the given command(s)."""
    assert isinstance(commands, list)
    process = None

    popen_kwargs = {}
    if sys.platform == "win32":
        # This hides the console window if pythonw.exe is used
        startupinfo = subprocess.STARTUPINFO()
        startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        popen_kwargs["startupinfo"] = startupinfo

    for command in commands:
        try:
            dispcmd = str([command] + args)
            # remember shell=False, so use git.cmd on windows, not just git
            process = subprocess.Popen(
                [command] + args,
                cwd=cwd,
                env=env,
                stdout=subprocess.PIPE,
                stderr=(subprocess.PIPE if hide_stderr else None),
                **popen_kwargs,
            )
            break
        except OSError:
            e = sys.exc_info()[1]
            if e.errno == errno.ENOENT:
                continue
            if verbose:
                print("unable to run %s" % dispcmd)
                print(e)
            return None, None
    else:
        if verbose:
            print("unable to find command, tried %s" % (commands,))
        return None, None
    stdout = process.communicate()[0].strip().decode()
    if process.returncode != 0:
        if verbose:
            print("unable to run %s (error)" % dispcmd)
            print("stdout was %s" % stdout)
        return None, process.returncode
    return stdout, process.returncode


def versions_from_parentdir(parentdir_prefix, root, verbose):
    """Try to determine the version from the parent directory name.

    Source tarballs conventionally unpack into a directory that includes both
    the project name and a version string. We will also support searching up
    two directory levels for an appropriately named parent directory
    """
    rootdirs = []

    for _ in range(3):
        dirname = os.path.basename(root)
        if dirname.startswith(parentdir_prefix):
            return {
                "version": dirname[len(parentdir_prefix) :],
                "full-revisionid": None,
                "dirty": False,
                "error": None,
                "date": None,
            }
        rootdirs.append(root)
        root = os.path.dirname(root)  # up a level

    if verbose:
        print(
            "Tried directories %s but none started with prefix %s"
            % (str(rootdirs), parentdir_prefix)
        )
    raise NotThisMethod("rootdir doesn't start with parentdir_prefix")


@register_vcs_handler("git", "get_keywords")
def git_get_keywords(versionfile_abs):
    """Extract version information from the given file."""
    # the code embedded in _version.py can just fetch the value of these
    # keywords. When used from setup.py, we don't want to import _version.py,
    # so we do it with a regexp instead. This function is not used from
    # _version.py.
    keywords = {}
    try:
        with open(versionfile_abs, "r") as fobj:
            for line in fobj:
                if line.strip().startswith("git_refnames ="):
                    mo = re.search(r'=\s*"(.*)"', line)
                    if mo:
                        keywords["refnames"] = mo.group(1)
                if line.strip().startswith("git_full ="):
                    mo = re.search(r'=\s*"(.*)"', line)
                    if mo:
                        keywords["full"] = mo.group(1)
                if line.strip().startswith("git_date ="):
                    mo = re.search(r'=\s*"(.*)"', line)
                    if mo:
                        keywords["date"] = mo.group(1)
    except OSError:
        pass
    return keywords


@register_vcs_handler("git", "keywords")
def git_versions_from_keywords(keywords, tag_prefix, verbose):
    """Get version information from git keywords."""
    if "refnames" not in keywords:
        raise NotThisMethod("Short version file found")
    date = keywords.get("date")
    if date is not None:
        # Use only the last line.  Previous lines may contain GPG signature
        # information.
        date = date.splitlines()[-1]

        # git-2.2.0 added "%cI", which expands to an ISO-8601 -compliant
        # datestamp. However we prefer "%ci" (which expands to an "ISO-8601
        # -like" string, which we must then edit to make compliant), because
        # it's been around since git-1.5.3, and it's too difficult to
        # discover which version we're using, or to work around using an
        # older one.
        date = date.strip().replace(" ", "T", 1).replace(" ", "", 1)
    refnames = keywords["refnames"].strip()
    if refnames.startswith("$Format"):
        if verbose:
            print("keywords are unexpanded, not using")
        raise NotThisMethod("unexpanded keywords, not a git-archive tarball")
    refs = {r.strip() for r in refnames.strip("()").split(",")}
    # starting in git-1.8.3, tags are listed as "tag: foo-1.0" instead of
    # just "foo-1.0". If we see a "tag: " prefix, prefer those.
    TAG = "tag: "
    tags = {r[len(TAG) :] for r in refs if r.startswith(TAG)}
    if not tags:
        # Either we're using git < 1.8.3, or there really are no tags. We use
        # a heuristic: assume all version tags have a digit. The old git %d
        # expansion behaves like git log --decorate=short and strips out the
        # refs/heads/ and refs/tags/ prefixes that would let us distinguish
        # between branches and tags. By ignoring refnames without digits, we
        # filter out many common branch names like "release" and
        # "stabilization", as well as "HEAD" and "master".
        tags = {r for r in refs if re.search(r"\d", r)}
        if verbose:
            print("discarding '%s', no digits" % ",".join(refs - tags))
    if verbose:
        print("likely tags: %s" % ",".join(sorted(tags)))
    for ref in sorted(tags):
        # sorting will prefer e.g. "2.0" over "2.0rc1"
        if ref.startswith(tag_prefix):
            r = ref[len(tag_prefix) :]
            # Filter out refs that exactly match prefix or that don't start
            # with a number once the prefix is stripped (mostly a concern
            # when prefix is '')
            if not re.match(r"\d", r):
                continue
            if verbose:
                print("picking %s" % r)
            return {
                "version": r,
                "full-revisionid": keywords["full"].strip(),
                "dirty": False,
                "error": None,
                "date": date,
            }
    # no suitable tags, so version is "0+unknown", but full hex is still there
    if verbose:
        print("no suitable tags, using unknown + full revision id")
    return {
        "version": "0+unknown",
        "full-revisionid": keywords["full"].strip(),
        "dirty": False,
        "error": "no suitable tags",
        "date": None,
    }


@register_vcs_handler("git", "pieces_from_vcs")
def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):
    """Get version from 'git describe' in the root of the source tree.

    This only gets called if the git-archive 'subst' keywords were *not*
    expanded, and _version.py hasn't already been rewritten with a short
    version string, meaning we're inside a checked out source tree.
    """
    GITS = ["git"]
    if sys.platform == "win32":
        GITS = ["git.cmd", "git.exe"]

    # GIT_DIR can interfere with correct operation of Versioneer.
    # It may be intended to be passed to the Versioneer-versioned project,
    # but that should not change where we get our version from.
    env = os.environ.copy()
    env.pop("GIT_DIR", None)
    runner = functools.partial(runner, env=env)

    _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, hide_stderr=not verbose)
    if rc != 0:
        if verbose:
            print("Directory %s not under git control" % root)
        raise NotThisMethod("'git rev-parse --git-dir' returned error")

    # if there is a tag matching tag_prefix, this yields TAG-NUM-gHEX[-dirty]
    # if there isn't one, this yields HEX[-dirty] (no NUM)
    describe_out, rc = runner(
        GITS,
        [
            "describe",
            "--tags",
            "--dirty",
            "--always",
            "--long",
            "--match",
            f"{tag_prefix}[[:digit:]]*",
        ],
        cwd=root,
    )
    # --long was added in git-1.5.5
    if describe_out is None:
        raise NotThisMethod("'git describe' failed")
    describe_out = describe_out.strip()
    full_out, rc = runner(GITS, ["rev-parse", "HEAD"], cwd=root)
    if full_out is None:
        raise NotThisMethod("'git rev-parse' failed")
    full_out = full_out.strip()

    pieces = {}
    pieces["long"] = full_out
    pieces["short"] = full_out[:7]  # maybe improved later
    pieces["error"] = None

    branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"], cwd=root)
    # --abbrev-ref was added in git-1.6.3
    if rc != 0 or branch_name is None:
        raise NotThisMethod("'git rev-parse --abbrev-ref' returned error")
    branch_name = branch_name.strip()

    if branch_name == "HEAD":
        # If we aren't exactly on a branch, pick a branch which represents
        # the current commit. If all else fails, we are on a branchless
        # commit.
        branches, rc = runner(GITS, ["branch", "--contains"], cwd=root)
        # --contains was added in git-1.5.4
        if rc != 0 or branches is None:
            raise NotThisMethod("'git branch --contains' returned error")
        branches = branches.split("\n")

        # Remove the first line if we're running detached
        if "(" in branches[0]:
            branches.pop(0)

        # Strip off the leading "* " from the list of branches.
        branches = [branch[2:] for branch in branches]
        if "master" in branches:
            branch_name = "master"
        elif not branches:
            branch_name = None
        else:
            # Pick the first branch that is returned. Good or bad.
            branch_name = branches[0]

    pieces["branch"] = branch_name

    # parse describe_out. It will be like TAG-NUM-gHEX[-dirty] or HEX[-dirty]
    # TAG might have hyphens.
    git_describe = describe_out

    # look for -dirty suffix
    dirty = git_describe.endswith("-dirty")
    pieces["dirty"] = dirty
    if dirty:
        git_describe = git_describe[: git_describe.rindex("-dirty")]

    # now we have TAG-NUM-gHEX or HEX

    if "-" in git_describe:
        # TAG-NUM-gHEX
        mo = re.search(r"^(.+)-(\d+)-g([0-9a-f]+)$", git_describe)
        if not mo:
            # unparsable. Maybe git-describe is misbehaving?
            pieces["error"] = "unable to parse git-describe output: '%s'" % describe_out
            return pieces

        # tag
        full_tag = mo.group(1)
        if not full_tag.startswith(tag_prefix):
            if verbose:
                fmt = "tag '%s' doesn't start with prefix '%s'"
                print(fmt % (full_tag, tag_prefix))
            pieces["error"] = "tag '%s' doesn't start with prefix '%s'" % (
                full_tag,
                tag_prefix,
            )
            return pieces
        pieces["closest-tag"] = full_tag[len(tag_prefix) :]

        # distance: number of commits since tag
        pieces["distance"] = int(mo.group(2))

        # commit: short hex revision ID
        pieces["short"] = mo.group(3)

    else:
        # HEX: no tags
        pieces["closest-tag"] = None
        out, rc = runner(GITS, ["rev-list", "HEAD", "--left-right"], cwd=root)
        pieces["distance"] = len(out.split())  # total number of commits

    # commit date: see ISO-8601 comment in git_versions_from_keywords()
    date = runner(GITS, ["show", "-s", "--format=%ci", "HEAD"], cwd=root)[0].strip()
    # Use only the last line.  Previous lines may contain GPG signature
    # information.
    date = date.splitlines()[-1]
    pieces["date"] = date.strip().replace(" ", "T", 1).replace(" ", "", 1)

    return pieces


def plus_or_dot(pieces):
    """Return a + if we don't already have one, else return a ."""
    if "+" in pieces.get("closest-tag", ""):
        return "."
    return "+"


def render_pep440(pieces):
    """Build up version string, with post-release "local version identifier".

    Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you
    get a tagged build and then dirty it, you'll get TAG+0.gHEX.dirty

    Exceptions:
    1: no tags. git_describe was just HEX. 0+untagged.DISTANCE.gHEX[.dirty]
    """
    if pieces["closest-tag"]:
        rendered = pieces["closest-tag"]
        if pieces["distance"] or pieces["dirty"]:
            rendered += plus_or_dot(pieces)
            rendered += "%d.g%s" % (pieces["distance"], pieces["short"])
            if pieces["dirty"]:
                rendered += ".dirty"
    else:
        # exception #1
        rendered = "0+untagged.%d.g%s" % (pieces["distance"], pieces["short"])
        if pieces["dirty"]:
            rendered += ".dirty"
    return rendered


def render_pep440_branch(pieces):
    """TAG[[.dev0]+DISTANCE.gHEX[.dirty]] .

    The ".dev0" means not master branch. Note that .dev0 sorts backwards
    (a feature branch will appear "older" than the master branch).

    Exceptions:
    1: no tags. 0[.dev0]+untagged.DISTANCE.gHEX[.dirty]
    """
    if pieces["closest-tag"]:
        rendered = pieces["closest-tag"]
        if pieces["distance"] or pieces["dirty"]:
            if pieces["branch"] != "master":
                rendered += ".dev0"
            rendered += plus_or_dot(pieces)
            rendered += "%d.g%s" % (pieces["distance"], pieces["short"])
            if pieces["dirty"]:
                rendered += ".dirty"
    else:
        # exception #1
        rendered = "0"
        if pieces["branch"] != "master":
            rendered += ".dev0"
        rendered += "+untagged.%d.g%s" % (pieces["distance"], pieces["short"])
        if pieces["dirty"]:
            rendered += ".dirty"
    return rendered


def pep440_split_post(ver):
    """Split pep440 version string at the post-release segment.

    Returns the release segments before the post-release and the
    post-release version number (or -1 if no post-release segment is present).
    """
    vc = str.split(ver, ".post")
    return vc[0], int(vc[1] or 0) if len(vc) == 2 else None


def render_pep440_pre(pieces):
    """TAG[.postN.devDISTANCE] -- No -dirty.

    Exceptions:
    1: no tags. 0.post0.devDISTANCE
    """
    if pieces["closest-tag"]:
        if pieces["distance"]:
            # update the post release segment
            tag_version, post_version = pep440_split_post(pieces["closest-tag"])
            rendered = tag_version
            if post_version is not None:
                rendered += ".post%d.dev%d" % (post_version + 1, pieces["distance"])
            else:
                rendered += ".post0.dev%d" % (pieces["distance"])
        else:
            # no commits, use the tag as the version
            rendered = pieces["closest-tag"]
    else:
        # exception #1
        rendered = "0.post0.dev%d" % pieces["distance"]
    return rendered


def render_pep440_post(pieces):
    """TAG[.postDISTANCE[.dev0]+gHEX] .

    The ".dev0" means dirty. Note that .dev0 sorts backwards
    (a dirty tree will appear "older" than the corresponding clean one),
    but you shouldn't be releasing software with -dirty anyways.

    Exceptions:
    1: no tags. 0.postDISTANCE[.dev0]
    """
    if pieces["closest-tag"]:
        rendered = pieces["closest-tag"]
        if pieces["distance"] or pieces["dirty"]:
            rendered += ".post%d" % pieces["distance"]
            if pieces["dirty"]:
                rendered += ".dev0"
            rendered += plus_or_dot(pieces)
            rendered += "g%s" % pieces["short"]
    else:
        # exception #1
        rendered = "0.post%d" % pieces["distance"]
        if pieces["dirty"]:
            rendered += ".dev0"
        rendered += "+g%s" % pieces["short"]
    return rendered


def render_pep440_post_branch(pieces):
    """TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] .

    The ".dev0" means not master branch.

    Exceptions:
    1: no tags. 0.postDISTANCE[.dev0]+gHEX[.dirty]
    """
    if pieces["closest-tag"]:
        rendered = pieces["closest-tag"]
        if pieces["distance"] or pieces["dirty"]:
            rendered += ".post%d" % pieces["distance"]
            if pieces["branch"] != "master":
                rendered += ".dev0"
            rendered += plus_or_dot(pieces)
            rendered += "g%s" % pieces["short"]
            if pieces["dirty"]:
                rendered += ".dirty"
    else:
        # exception #1
        rendered = "0.post%d" % pieces["distance"]
        if pieces["branch"] != "master":
            rendered += ".dev0"
        rendered += "+g%s" % pieces["short"]
        if pieces["dirty"]:
            rendered += ".dirty"
    return rendered


def render_pep440_old(pieces):
    """TAG[.postDISTANCE[.dev0]] .

    The ".dev0" means dirty.

    Exceptions:
    1: no tags. 0.postDISTANCE[.dev0]
    """
    if pieces["closest-tag"]:
        rendered = pieces["closest-tag"]
        if pieces["distance"] or pieces["dirty"]:
            rendered += ".post%d" % pieces["distance"]
            if pieces["dirty"]:
                rendered += ".dev0"
    else:
        # exception #1
        rendered = "0.post%d" % pieces["distance"]
        if pieces["dirty"]:
            rendered += ".dev0"
    return rendered


def render_git_describe(pieces):
    """TAG[-DISTANCE-gHEX][-dirty].

    Like 'git describe --tags --dirty --always'.

    Exceptions:
    1: no tags. HEX[-dirty]  (note: no 'g' prefix)
    """
    if pieces["closest-tag"]:
        rendered = pieces["closest-tag"]
        if pieces["distance"]:
            rendered += "-%d-g%s" % (pieces["distance"], pieces["short"])
    else:
        # exception #1
        rendered = pieces["short"]
    if pieces["dirty"]:
        rendered += "-dirty"
    return rendered


def render_git_describe_long(pieces):
    """TAG-DISTANCE-gHEX[-dirty].

    Like 'git describe --tags --dirty --always -long'.
    The distance/hash is unconditional.

    Exceptions:
    1: no tags. HEX[-dirty]  (note: no 'g' prefix)
    """
    if pieces["closest-tag"]:
        rendered = pieces["closest-tag"]
        rendered += "-%d-g%s" % (pieces["distance"], pieces["short"])
    else:
        # exception #1
        rendered = pieces["short"]
    if pieces["dirty"]:
        rendered += "-dirty"
    return rendered


def render(pieces, style):
    """Render the given version pieces into the requested style."""
    if pieces["error"]:
        return {
            "version": "unknown",
            "full-revisionid": pieces.get("long"),
            "dirty": None,
            "error": pieces["error"],
            "date": None,
        }

    if not style or style == "default":
        style = "pep440"  # the default

    if style == "pep440":
        rendered = render_pep440(pieces)
    elif style == "pep440-branch":
        rendered = render_pep440_branch(pieces)
    elif style == "pep440-pre":
        rendered = render_pep440_pre(pieces)
    elif style == "pep440-post":
        rendered = render_pep440_post(pieces)
    elif style == "pep440-post-branch":
        rendered = render_pep440_post_branch(pieces)
    elif style == "pep440-old":
        rendered = render_pep440_old(pieces)
    elif style == "git-describe":
        rendered = render_git_describe(pieces)
    elif style == "git-describe-long":
        rendered = render_git_describe_long(pieces)
    else:
        raise ValueError("unknown style '%s'" % style)

    return {
        "version": rendered,
        "full-revisionid": pieces["long"],
        "dirty": pieces["dirty"],
        "error": None,
        "date": pieces.get("date"),
    }


def get_versions():
    """Get version information or return default if unable to do so."""
    # I am in _version.py, which lives at ROOT/VERSIONFILE_SOURCE. If we have
    # __file__, we can work backwards from there to the root. Some
    # py2exe/bbfreeze/non-CPython implementations don't do __file__, in which
    # case we can only use expanded keywords.

    cfg = get_config()
    verbose = cfg.verbose

    try:
        return git_versions_from_keywords(get_keywords(), cfg.tag_prefix, verbose)
    except NotThisMethod:
        pass

    try:
        root = os.path.realpath(__file__)
        # versionfile_source is the relative path from the top of the source
        # tree (where the .git directory might live) to this file. Invert
        # this to find the root from __file__.
        for _ in cfg.versionfile_source.split("/"):
            root = os.path.dirname(root)
    except NameError:
        return {
            "version": "0+unknown",
            "full-revisionid": None,
            "dirty": None,
            "error": "unable to find root of source tree",
            "date": None,
        }

    try:
        pieces = git_pieces_from_vcs(cfg.tag_prefix, root, verbose)
        return render(pieces, cfg.style)
    except NotThisMethod:
        pass

    try:
        if cfg.parentdir_prefix:
            return versions_from_parentdir(cfg.parentdir_prefix, root, verbose)
    except NotThisMethod:
        pass

    return {
        "version": "0+unknown",
        "full-revisionid": None,
        "dirty": None,
        "error": "unable to compute version",
        "date": None,
    }

tensql/_version.py
sandialabs/TenSQL
import unittest

import numpy as np
import pygraphblas
import pygraphblas.types

from ..util import export, ndim
from .test_util import register
from .. import Column, Types

__all__ = []

@export
@register()
class TestColumnScalarInit(unittest.TestCase):
    def runTest(self):
        type_ = Types.BigInt()
        data = pygraphblas.Scalar.from_type(type_.as_pygrb)
        name = "Test"
        column = Column(name, data, type_=type_)
        self.assertEqual(column.name, name)
        self.assertEqual(column.hidden, False)
        self.assertIs(column.data, data)
        self.assertIs(column.gbtype, type_.as_pygrb)
        self.assertTrue(column.is_scalar)
        self.assertFalse(column.is_vector)
        self.assertFalse(column.is_matrix)
        self.assertFalse(column.is_boolean)
        self.assertEqual(column.nptype, np.int64)
        self.assertEqual(ndim(column.data), 0)

@export
@register()
class TestColumnVectorInit(unittest.TestCase):
    def runTest(self):
        type_ = Types.BigInt()
        data = pygraphblas.Vector.sparse(type_.as_pygrb, size=10)
        name = "Test"
        column = Column(name, data, type_=type_)
        self.assertEqual(column.name, name)
        self.assertEqual(column.hidden, False)
        self.assertIs(column.data, data)
        self.assertIs(column.gbtype, type_.as_pygrb)
        self.assertFalse(column.is_scalar)
        self.assertTrue(column.is_vector)
        self.assertFalse(column.is_matrix)
        self.assertFalse(column.is_boolean)
        self.assertEqual(column.nptype, np.int64)
        self.assertEqual(ndim(column.data), 1)

@export
@register()
class TestColumnMatrixInit(unittest.TestCase):
    def runTest(self):
        type_ = Types.BigInt()
        data = pygraphblas.Matrix.sparse(type_.as_pygrb, nrows=10, ncols=10)
        name = "Test"
        column = Column(name, data, type_=type_)
        self.assertEqual(column.name, name)
        self.assertEqual(column.hidden, False)
        self.assertIs(column.data, data)
        self.assertIs(column.gbtype, type_.as_pygrb)
        self.assertFalse(column.is_scalar)
        self.assertFalse(column.is_vector)
        self.assertTrue(column.is_matrix)
        self.assertFalse(column.is_boolean)
        self.assertEqual(column.nptype, np.int64)
        self.assertEqual(ndim(column.data), 2)

@export
@register()
class TestColumnCopy(unittest.TestCase):
    def runTest(self):
        type_ = Types.BigInt()
        data = pygraphblas.Matrix.sparse(type_.as_pygrb, nrows=10, ncols=10)
        name = "Test"
        type_ = Types.BigInt()
        column = Column(name, data, type_=type_)
        column2 = column.copy()
        self.assertEqual(column2.name, name)
        self.assertEqual(column2.hidden, False)
        self.assertEqual(column2.data, data)
        self.assertIsNot(column2.data, column.data)
        self.assertIs(column2.gbtype, type_.as_pygrb)

@export
@register()
class TestColumnScalarInitText(unittest.TestCase):
    def runTest(self):
        type_ = Types.Text()
        name = "Test"
        data = pygraphblas.Scalar.from_type(type_.as_pygrb)
        column = Column(name, data, type_=type_)
        self.assertEqual(column.name, name)
        self.assertEqual(column.hidden, False)
        self.assertIs(column.data, data)
        self.assertIs(column.gbtype, type_.as_pygrb)
        self.assertTrue(column.is_scalar)
        self.assertFalse(column.is_vector)
        self.assertFalse(column.is_matrix)
        self.assertFalse(column.is_boolean)
        self.assertEqual(column.nptype, type_.as_numpy)
        self.assertEqual(ndim(column.data), 0)

tensql/test/Column.py
sandialabs/TenSQL
import h5py
import pygraphblas
import pygraphblas.types

from ..util import export
from .test_util import register, TmpDirMixin
from .. import Column, Table, Database, Types

__all__ = []


@export
@register()
class TestDBTableInitNoStencil(TmpDirMixin):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(col_type.as_pygrb, size=shape[0])
        col_name = "Data"
        column = Column(col_name, col_data, type_=col_type)

        pk_name, pk_data, pk_gbtype = "PK", None, Types.BigInt()
        pk_column = Column(pk_name, pk_data, type_=pk_gbtype)

        db = Database()
        table = db.create_table("tmp", shape, [pk_column, column])

        self.assertEqual(table.shape, shape)
        self.assertEqual(
            list(table.columns), [(pk_name, pk_column), (col_name, column)]
        )

@export
@register()
class TestTableFromScalar(TmpDirMixin):
    def runTest(self):
        col_type = Types.BigInt()
        col_data = pygraphblas.Scalar.from_type(col_type.as_pygrb)

        db = Database()
        table = db.create_table_from_tensor("tmp", col_data)

        [[col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(table.shape, tuple())
        self.assertEqual(table.ndim, 0)
        self.assertTrue(table.is_scalar)
        self.assertFalse(table.is_vector)
        self.assertFalse(table.is_matrix)
        self.assertEqual(table.primary_key_names, tuple())
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)


@export
@register()
class TestTableFromVector(TmpDirMixin):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(typ=col_type.as_pygrb, size=shape[0])

        db = Database()
        table = db.create_table_from_tensor("tmp", col_data)

        [[pk_name, pk], [col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(pk_name, "idx")
        self.assertEqual(table.shape, shape)
        self.assertEqual(table.ndim, 1)
        self.assertFalse(table.is_scalar)
        self.assertTrue(table.is_vector)
        self.assertFalse(table.is_matrix)
        self.assertTrue(table.is_primary_key(pk))
        self.assertEqual(table.primary_key_names, (pk_name,))
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)

@export
@register()
class TestTableFromMatrix(TmpDirMixin):
    def runTest(self):
        shape = (10, 7)
        col_type = Types.BigInt()
        col_data = pygraphblas.Matrix.sparse(typ=col_type.as_pygrb, nrows=shape[0], ncols=shape[1])

        db = Database()
        table = db.create_table_from_tensor("tmp", col_data)

        [[pk1_name, pk1], [pk2_name, pk2], [col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(pk1_name, "ridx")
        self.assertEqual(pk2_name, "cidx")
        self.assertEqual(table.shape, shape)
        self.assertEqual(table.ndim, 2)
        self.assertFalse(table.is_scalar)
        self.assertFalse(table.is_vector)
        self.assertTrue(table.is_matrix)
        self.assertEqual(table.primary_key_names, (pk1_name, pk2_name))
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)

@export
@register()
class TestTableFromDefinitionScalar(TmpDirMixin):
    def runTest(self):
        col_type = Types.BigInt()

        db = Database()
        table = db.create_table_from_definition(
            "tmp",
            {
                "value": col_type
            },
            []
        )

        [[col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(table.shape, tuple())
        self.assertEqual(table.ndim, 0)
        self.assertTrue(table.is_scalar)
        self.assertFalse(table.is_vector)
        self.assertFalse(table.is_matrix)
        self.assertEqual(table.primary_key_names, tuple())
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)

@export
@register()
class TestTableFromDefinitionVector(TmpDirMixin):
    def runTest(self):
        shape = (2 ** 60,)
        col_type = Types.BigInt()

        db = Database()
        table = db.create_table_from_definition(
            "tmp",
            {
                "idx": col_type,
                "value": col_type
            },
            ["idx"]
        )

        [[pk_name, pk], [col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(pk_name, "idx")
        self.assertEqual(table.shape, shape)
        self.assertEqual(table.ndim, 1)
        self.assertFalse(table.is_scalar)
        self.assertTrue(table.is_vector)
        self.assertFalse(table.is_matrix)
        self.assertTrue(table.is_primary_key(pk))
        self.assertEqual(table.primary_key_names, (pk_name,))
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)

@export
@register()
class TestTableFromDefinitionMatrix(TmpDirMixin):
    def runTest(self):
        shape = (2 ** 60, 2 ** 60)
        col_type = Types.BigInt()

        db = Database()
        table = db.create_table_from_definition(
            "tmp",
            {
                "ridx": col_type,
                "cidx": col_type,
                "value": col_type
            },
            ["ridx", "cidx"]
        )

        [[pk1_name, pk1], [pk2_name, pk2], [col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(pk1_name, "ridx")
        self.assertEqual(pk2_name, "cidx")
        self.assertEqual(table.shape, shape)
        self.assertEqual(table.ndim, 2)
        self.assertFalse(table.is_scalar)
        self.assertFalse(table.is_vector)
        self.assertTrue(table.is_matrix)
        self.assertEqual(table.primary_key_names, (pk1_name, pk2_name))
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)


@export
@register()
class TestTableSaveLoadHDF5Scalar(TmpDirMixin):
    def runTest(self):
        col_type = Types.BigInt()
        col_data = pygraphblas.Scalar.from_type(typ=col_type.as_pygrb)

        db = Database()
        table = db.create_table_from_tensor("tmp", col_data)
        
        data = [
            (4,),
        ]

        def test(db, N):
            db.save("tmp.h5")
            db2 = Database.load("tmp.h5")
            table2 = db2['tmp']

            self.assertEqual(table2.count(), N)
            self.assertEqual(len(table2), N)
            observed = list(table2)
            expected = [{"value":val} for (val,) in data[:N]]
            self.assertEqual(expected, observed)

        test(db, 0)
        table.insert(value=data[0][0])
        test(db, 1)

@export
@register()
class TestTableSaveLoadHDF5Vector(TmpDirMixin):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(typ=col_type.as_pygrb, size=shape[0])

        db = Database()
        table = db.create_table_from_tensor("tmp", col_data)
        
        (N,) = shape
        data = [
            (0, 2),
            (N-1, 4)
        ]

        def test(db, N):
            db.save("tmp.h5")
            db2 = Database.load("tmp.h5")
            table2 = db2['tmp']

            self.assertEqual(table2.count(), N)
            self.assertEqual(len(table2), N)

            observed = list(table2)
            expected = [{"idx":idx, "value":val} for idx, val in data[:N]]
            self.assertEqual(expected, observed)

        test(db, 0)
        for it_count, (idx, val) in enumerate(data):
            table.insert(idx=idx, value=val)
            test(db, it_count+1)

@export
@register()
class TestTableSaveLoadHDF5Matrix(TmpDirMixin):
    def runTest(self):
        shape = (10, 7)
        col_type = Types.BigInt()
        col_data = pygraphblas.Matrix.sparse(typ=col_type.as_pygrb, nrows=shape[0], ncols=shape[1])

        db = Database()
        table = db.create_table_from_tensor("tmp", col_data)
        
        M, N = shape
        data = [
            (0, 0, 2),
            (0, N-1, 4),
            (M-1, 0, 8),
            (M-1, N-1, 16)
        ]

        def test(db, N):
            db.save("tmp.h5")
            db2 = Database.load("tmp.h5")
            table2 = db2['tmp']

            self.assertEqual(table2.count(), N)
            self.assertEqual(len(table2), N)

            observed = list(table2)
            expected = [{"ridx":ridx, "cidx":cidx, "value":val} for ridx, cidx, val in data[:N]]
            self.assertEqual(expected, observed)

        test(db, 0)
        for it_count, (ridx, cidx, val) in enumerate(data):
            table.insert(ridx=ridx, cidx=cidx, value=val)
            test(db, it_count+1)

tensql/test/Database.py
sandialabs/TenSQL
from .._PyHashOCNT import PyHashOCNT
import unittest
from ..util import export
from .test_util import register, TmpDirMixin
import sys

__all__ = []

@export
@register()
class TestOCNT_init(unittest.TestCase):
    def runTest(self):
        ocnt = PyHashOCNT(16)
        self.assertEqual(ocnt.size, 0)
        del ocnt

@export
@register()
class TestOCNT_insert_delete(unittest.TestCase):
    def runTest(self):
        x = "asdfeasdfebcxcqse"
        x_refcount = sys.getrefcount(x)
        ocnt = PyHashOCNT(16)
        ocnt.insert(x)
        self.assertEqual(ocnt.size, 1)
        self.assertEqual(x_refcount + 1, sys.getrefcount(x))
        del ocnt
        self.assertEqual(x_refcount, sys.getrefcount(x))

@export
@register()
class TestOCNT_insert_remove_delete(unittest.TestCase):
    def runTest(self):
        x = "asdfeasdfeewafeasef"
        x_refcount = sys.getrefcount(x)
        ocnt = PyHashOCNT(16)
        self.assertEqual(ocnt.size, 0)

        ocnt.insert(x)
        self.assertEqual(ocnt.size, 1)
        self.assertEqual(x_refcount + 1, sys.getrefcount(x))

        ocnt.remove(x)
        self.assertEqual(x_refcount, sys.getrefcount(x))

        del ocnt
        self.assertEqual(x_refcount, sys.getrefcount(x))

@export
@register()
class TestOCNT_insert_decode_remove_delete(unittest.TestCase):
    def runTest(self):
        x = "asdfeasdfesdfesdf"
        x_refcount = sys.getrefcount(x)
        ocnt = PyHashOCNT(16)

        ocnt.insert(x)
        self.assertEqual(ocnt.size, 1)
        self.assertEqual(x_refcount + 1, sys.getrefcount(x))

        x_id = id(x)
        self.assertEqual(x_refcount + 1, sys.getrefcount(x))

        y = ocnt.decode(x_id)
        self.assertEqual(x_refcount+2, sys.getrefcount(x))

        del y
        self.assertEqual(x_refcount+1, sys.getrefcount(x))

        ocnt.remove(x)
        self.assertEqual(x_refcount, sys.getrefcount(x))

        del ocnt
        self.assertEqual(x_refcount, sys.getrefcount(x))

tensql/test/PyHashOCNT.py
sandialabs/TenSQL
import os
import unittest

import h5py
import numpy as np
import pygraphblas
import pygraphblas.types

from ..util import export
from .test_util import register, TmpDirMixin
from .. import Column, Table, Types

__all__ = []


@export
@register()
class TestTableInitNoStencil(unittest.TestCase):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(col_type.as_pygrb, size=shape[0])
        col_name = "Data"
        column = Column(col_name, col_data, type_=col_type)

        pk_name, pk_data, pk_type = "PK", None, Types.BigInt()
        pk_column = Column(pk_name, pk_data, type_=pk_type)

        table = Table(shape, [pk_column, column])

        self.assertEqual(table.shape, shape)
        self.assertEqual(
            list(table.columns), [(pk_name, pk_column), (col_name, column)]
        )

@export
@register()
class TestTableFromScalar(unittest.TestCase):
    def runTest(self):
        col_type = Types.BigInt()
        col_data = pygraphblas.Scalar.from_type(col_type.as_pygrb)

        table = Table.from_tensor(col_data)

        [[col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(table.shape, tuple())
        self.assertEqual(table.ndim, 0)
        self.assertTrue(table.is_scalar)
        self.assertFalse(table.is_vector)
        self.assertFalse(table.is_matrix)
        self.assertEqual(table.primary_key_names, tuple())
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)


@export
@register()
class TestTableFromVector(unittest.TestCase):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(typ=col_type.as_pygrb, size=shape[0])

        table = Table.from_tensor(col_data)

        [[pk_name, pk], [col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(pk_name, "idx")
        self.assertEqual(table.shape, shape)
        self.assertEqual(table.ndim, 1)
        self.assertFalse(table.is_scalar)
        self.assertTrue(table.is_vector)
        self.assertFalse(table.is_matrix)
        self.assertTrue(table.is_primary_key(pk))
        self.assertEqual(table.primary_key_names, (pk_name,))
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)

@export
@register()
class TestTableFromMatrix(unittest.TestCase):
    def runTest(self):
        shape = (10, 7)
        col_type = Types.BigInt()
        col_data = pygraphblas.Matrix.sparse(typ=col_type.as_pygrb, nrows=shape[0], ncols=shape[1])

        table = Table.from_tensor(col_data)

        [[pk1_name, pk1], [pk2_name, pk2], [col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(pk1_name, "ridx")
        self.assertEqual(pk2_name, "cidx")
        self.assertEqual(table.shape, shape)
        self.assertEqual(table.ndim, 2)
        self.assertFalse(table.is_scalar)
        self.assertFalse(table.is_vector)
        self.assertTrue(table.is_matrix)
        self.assertEqual(table.primary_key_names, (pk1_name, pk2_name))
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)

@export
@register()
class TestTableInsertIntoScalar(unittest.TestCase):
    def runTest(self):
        col_type = Types.BigInt()
        col_data = pygraphblas.Scalar.from_type(typ=col_type.as_pygrb)

        table = Table.from_tensor(col_data)
        
        table = Table.from_definition(
            {
                "value": col_type,
                "value2": col_type
            },
            []
        )
        
        data = [
            (4,),
        ]

        table.insert(value=data[0][0])

        self.assertEqual(table.count(), 1)
        self.assertEqual(len(table), 1)

        observed = list(table)
        expected = [{"value":val, "value2": None} for (val,) in data]
        self.assertEqual(expected, observed)

@export
@register()
class TestTableInsertIntoVector(unittest.TestCase):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(typ=col_type.as_pygrb, size=shape[0])

        table = Table.from_definition(
            {
                "idx": col_type,
                "value": col_type,
                "value2": col_type
            },
            ["idx"]
        )
        
        (N,) = shape
        data = [
            (0, 2),
            (5, 4),
            (N-1, 8)
        ]

        for it_count, (idx, val) in enumerate(data):
            table.insert(idx=idx, value=val)

            self.assertEqual(table.count(), it_count+1)
            self.assertEqual(len(table), it_count+1)
            self.assertEqual(table._columns['value2'].data.nvals, 0)

            observed = list(table)
            expected = [{"idx":idx, "value":val, "value2": None} for idx, val in data[:it_count+1]]
            self.assertEqual(expected, observed)

@export
@register()
class TestTableInsertIntoVector2(unittest.TestCase):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(typ=col_type.as_pygrb, size=shape[0])

        table = Table.from_definition(
            {
                "idx": col_type,
                "value": col_type,
                "value2": col_type
            },
            ["idx"]
        )
        
        (N,) = shape
        data = [
            (0, 2),
            (5, 4),
            (N-1, 8)
        ]

        for it_count, (idx, val) in enumerate(data):
            table.insert(idx=idx, value=val)

            self.assertEqual(table.count(), it_count+1)
            self.assertEqual(len(table), it_count+1)
            self.assertEqual(table._columns['value2'].data.nvals, 0)

            observed = table.to_records()
            self.assertEqual(("idx", "value", "value2"), observed.dtype.names)

            expected = np.empty(shape=(it_count+1,), dtype=observed.dtype)
            for off, (idx, val) in enumerate(data[:it_count+1]):
                expected[off]["idx"] = idx
                expected[off]["value"] = val
                expected[off]["value2"] = None
            np.testing.assert_equal(expected, observed)


@export
@register()
class TestTableInsertIntoMatrix(unittest.TestCase):
    def runTest(self):
        shape = (10, 7)
        col_type = Types.BigInt()
        col_data = pygraphblas.Matrix.sparse(typ=col_type.as_pygrb, nrows=shape[0], ncols=shape[1])

        table = Table.from_definition(
            {
                "ridx": col_type,
                "cidx": col_type,
                "value": col_type,
                "value2": col_type
            },
            ["ridx", "cidx"]
        )
        
        M, N = shape
        data = [
            (0, 0, 2),
            (0, N-1, 4),
            (M-1, 0, 8),
            (M-1, N-1, 16)
        ]

        for it_count, (ridx, cidx, val) in enumerate(data):
            table.insert(ridx=ridx, cidx=cidx, value=val)

            self.assertEqual(table.count(), it_count+1)
            self.assertEqual(len(table), it_count+1)

            observed = list(table)
            expected = [{"ridx":ridx, "cidx":cidx, "value":val, "value2":None} for ridx, cidx, val in data[:it_count+1]]
            self.assertEqual(expected, observed)



@export
@register()
class TestTableFromDefinitionScalar(unittest.TestCase):
    def runTest(self):
        col_type = Types.BigInt()

        table = Table.from_definition(
            {
                "value": col_type
            },
            []
        )

        [[col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(table.shape, tuple())
        self.assertEqual(table.ndim, 0)
        self.assertTrue(table.is_scalar)
        self.assertFalse(table.is_vector)
        self.assertFalse(table.is_matrix)
        self.assertEqual(table.primary_key_names, tuple())
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)

@export
@register()
class TestTableFromDefinitionVector(unittest.TestCase):
    def runTest(self):
        shape = (2 ** 60,)
        col_type = Types.BigInt()

        table = Table.from_definition(
            {
                "idx": col_type,
                "value": col_type
            },
            ["idx"]
        )

        [[pk_name, pk], [col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(pk_name, "idx")
        self.assertEqual(table.shape, shape)
        self.assertEqual(table.ndim, 1)
        self.assertFalse(table.is_scalar)
        self.assertTrue(table.is_vector)
        self.assertFalse(table.is_matrix)
        self.assertTrue(table.is_primary_key(pk))
        self.assertEqual(table.primary_key_names, (pk_name,))
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)

@export
@register()
class TestTableFromDefinitionMatrix(unittest.TestCase):
    def runTest(self):
        shape = (2 ** 60, 2 ** 60)
        col_type = Types.BigInt()

        table = Table.from_definition(
            {
                "ridx": col_type,
                "cidx": col_type,
                "value": col_type
            },
            ["ridx", "cidx"]
        )

        [[pk1_name, pk1], [pk2_name, pk2], [col_name, col]] = list(table.columns)
        self.assertEqual(col_name, "value")
        #self.assertEqual(col.data, col_data)
        self.assertEqual(pk1_name, "ridx")
        self.assertEqual(pk2_name, "cidx")
        self.assertEqual(table.shape, shape)
        self.assertEqual(table.ndim, 2)
        self.assertFalse(table.is_scalar)
        self.assertFalse(table.is_vector)
        self.assertTrue(table.is_matrix)
        self.assertEqual(table.primary_key_names, (pk1_name, pk2_name))
        self.assertIs(table.get_column(col_name), col)
        self.assertEqual(table.count(), 0)
        self.assertEqual(len(table), 0)


@export
@register()
class TestTableSaveLoadHDF5Scalar(TmpDirMixin):
    def runTest(self):
        col_type = Types.BigInt()
        col_data = pygraphblas.Scalar.from_type(typ=col_type.as_pygrb)

        table = Table.from_definition(
            {
                "value": col_type,
                "value2": col_type
            },
            []
        )
        
        data = [
            (4,),
        ]

        def test(table, N):
            with h5py.File("tmp.h5", "w") as fout:
                table.save_hdf5(fout.create_group("tmp"))
            
            with h5py.File("tmp.h5", "r") as fin:
                table2 = Table.load_hdf5(fin['tmp'])

            self.assertEqual(table2.count(), N)
            self.assertEqual(len(table2), N)
            observed = list(table2)
            expected = [{"value":val, "value2": None} for (val,) in data[:N]]
            self.assertEqual(expected, observed)

        test(table, 0)
        table.insert(value=data[0][0])
        test(table, 1)

@export
@register()
class TestTableSaveLoadHDF5Vector(TmpDirMixin):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(typ=col_type.as_pygrb, size=shape[0])

        table = Table.from_definition(
            {
                "idx": col_type,
                "value": col_type,
                "value2": col_type
            },
            ["idx"]
        )
        
        (N,) = shape
        data = [
            (0, 2),
            (N-1, 4)
        ]

        def test(table, N):
            with h5py.File("tmp.h5", "w") as fout:
                table.save_hdf5(fout.create_group("tmp"))
        
            with h5py.File("tmp.h5", "r") as fin:
                table2 = Table.load_hdf5(fin['tmp'])

            self.assertEqual(table2.count(), N)
            self.assertEqual(len(table2), N)

            observed = list(table2)
            expected = [{"idx":idx, "value":val, "value2":None} for idx, val in data[:N]]
            self.assertEqual(expected, observed)

        test(table, 0)
        for it_count, (idx, val) in enumerate(data):
            table.insert(idx=idx, value=val)
            test(table, it_count+1)

@export
@register()
class TestTableSaveLoadHDF5Matrix(TmpDirMixin):
    def runTest(self):
        shape = (10, 7)
        col_type = Types.BigInt()
        col_data = pygraphblas.Matrix.sparse(typ=col_type.as_pygrb, nrows=shape[0], ncols=shape[1])

        table = Table.from_definition(
            {
                "ridx": col_type,
                "cidx": col_type,
                "value": col_type,
                "value2": col_type
            },
            ["ridx", "cidx"]
        )
        
        M, N = shape
        data = [
            (0, 0, 2),
            (0, N-1, 4),
            (M-1, 0, 8),
            (M-1, N-1, 16)
        ]

        def test(table, N):
            with h5py.File("tmp.h5", "w") as fout:
                table.save_hdf5(fout.create_group("tmp"))
        
            with h5py.File("tmp.h5", "r") as fin:
                table2 = Table.load_hdf5(fin['tmp'])

            self.assertEqual(table2.count(), N)
            self.assertEqual(len(table2), N)

            observed = list(table2)
            expected = [{"ridx":ridx, "cidx":cidx, "value":val, "value2": None} for ridx, cidx, val in data[:N]]
            self.assertEqual(expected, observed)

        test(table, 0)
        for it_count, (ridx, cidx, val) in enumerate(data):
            table.insert(ridx=ridx, cidx=cidx, value=val)
            test(table, it_count+1)

@export
@register()
class TestTableAddScalar(TmpDirMixin):
    def runTest(self):
        col_type = Types.BigInt()
        col_data = pygraphblas.Scalar.from_type(typ=col_type.as_pygrb)

        table = Table.from_tensor(col_data)
        
        data = [
            (4,),
            (8,),
            (16,),
        ]

        def test(table, N):
            table2 = Table.from_tensor(col_data)
            for it, (val,) in enumerate(data[:N]):
                table.insert(value=val)
                table2.add(table, inplace=True)

            observed = list(table2)
            if N == 0:
                self.assertEqual(table2.count(), 0)
                self.assertEqual(len(table2), 0)
                expected = []
            else:
                self.assertEqual(table2.count(), 1)
                self.assertEqual(len(table2), 1)
                expected = [{"value":sum(val for (val,) in data[:N])}]
            self.assertEqual(expected, observed)

        test(table, 0)
        test(table, 1)
        test(table, 2)
        test(table, 3)



@export
@register()
class TestTableAddVector(TmpDirMixin):
    def runTest(self):
        shape = (10,)
        col_type = Types.BigInt()
        col_data = pygraphblas.Vector.sparse(typ=col_type.as_pygrb, size=shape[0])

        table = Table.from_tensor(col_data)
        
        (N,) = shape
        data = [
            (0, 2),
            (2, 4),
            (N-1, 8)
        ]

        def test(table, N):
            table2 = Table.from_tensor(col_data)
            table2.add(table, inplace=True)

            self.assertEqual(table2.count(), N)
            self.assertEqual(len(table2), N)

            observed = list(table2)
            expected = [{"idx":idx, "value":val} for idx, val in data[:N]]
            self.assertEqual(expected, observed)

        test(table, 0)
        for it_count, (idx, val) in enumerate(data):
            table.insert(idx=idx, value=val)
            test(table, it_count+1)

@export
@register()
class TestTableAddMatrix(TmpDirMixin):
    def runTest(self):
        shape = (10, 7)
        col_type = Types.BigInt()
        col_data = pygraphblas.Matrix.sparse(typ=col_type.as_pygrb, nrows=shape[0], ncols=shape[1])

        table = Table.from_tensor(col_data, copy=True)
        
        M, N = shape
        data = [
            (0, 0, 2),
            (0, N-1, 4),
            (M-1, 0, 8),
            (M-1, N-1, 16)
        ]

        def test(table, n):
            table2 = Table.from_tensor(col_data, copy=True)
            table2.add(table, inplace=True)

            self.assertEqual(table2.count(), n)
            self.assertEqual(len(table2), n)
            observed = list(table2)
            expected = [{"ridx":ridx, "cidx":cidx, "value":val} for ridx, cidx, val in data[:n]]
            self.assertEqual(expected, observed)

        test(table, 0)
        for it_count, (ridx, cidx, val) in enumerate(data):
            table.insert(ridx=ridx, cidx=cidx, value=val)
            test(table, it_count+1)

@export
@register()
class TestTableSaveLoadHDF5ScalarText(TmpDirMixin):
    def runTest(self):
        col_type = Types.Text()
        col_data = pygraphblas.Scalar.from_type(typ=col_type.as_pygrb)

        table = Table.from_definition(
            {
                "value": col_type,
                "value2": col_type
            },
            []
        )
        
        data = [
            ("Hello World",),
        ]

        def test(table, N):
            with h5py.File("tmp.h5", "w") as fout:
                table.save_hdf5(fout.create_group("tmp"))
            
            with h5py.File("tmp.h5", "r") as fin:
                table2 = Table.load_hdf5(fin['tmp'])

            self.assertEqual(table2.count(), N)
            self.assertEqual(len(table2), N)
            observed = list(table2)
            expected = [{"value":val, "value2": None} for (val,) in data[:N]]
            self.assertEqual(expected, observed)

        test(table, 0)
        table.insert(value=data[0][0])
        test(table, 1)

tensql/test/Table.py
sandialabs/TenSQL
from .Column import *
from .Database import *
from .PyHashOCNT import *
from .Table import *
from .test_queries import *

tensql/test/__init__.py
sandialabs/TenSQL
import numpy as np

from .common import CommonDatabaseUnitTest
from ..test_util import register
from ...util import export
from ... import QIR

__all__ = []

@export
@register()
class TestQuery_AA(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        A = qdb.Matrix.aliased("A")
        B = qdb.Matrix.aliased("B")

        result = (
            qdb.query(A)
                .join(B, A.cidx == B.ridx)
                .select(A.ridx, B.cidx, QIR.Sum(A.value * B.value).aliased("Result"))
                .group_by(A.ridx, B.cidx)
                .run()
        )

        observed = result.get_column('Result').data[:99, :99].to_numpy()
        expected = self.np_matrix.dot(self.np_matrix)

        np.testing.assert_allclose(expected, observed, rtol=1e-4, atol=1e-4)

tensql/test/test_queries/AA.py
sandialabs/TenSQL
import base64
import gc
import os
import random
import sys
import time
import ctypes

import numpy as np
import psutil

from .common import CommonDatabaseUnitTest
from ..test_util import register
from ...util import export
from ... import QIR
from ... import Types

__all__ = []

def get_memory_usage():
  return psutil.Process(os.getpid()).memory_info().rss

@export
@register()
class TestMemoryLeaks_insert_float_values(CommonDatabaseUnitTest):
    @staticmethod
    def get_random_rows(num=1024*1024*4):
      return [dict(idx=it, value=random.random()) for it in range(num)]

    def runTest(self):
        qdb = QIR.QirDatabase(self.db)

        vec = self.db.create_table_from_definition(
            "vec",
            {
                "idx": Types.BigInt(),
                "value": Types.Real()
            },
            ["idx"]
        )
        
        gc.collect()
        ctypes.CDLL('libc.so.6').malloc_trim(0) 

        starting_memory = get_memory_usage()

        input_rows = self.get_random_rows()

        result = (
            qdb.insert(qdb.vec)
              .values(input_rows)
              .run()
        )
        self.assertEqual(len(input_rows), result.count())

        del input_rows

        self.assertGreater(get_memory_usage(), starting_memory + 1024*1024*4)
        del self.db['vec']
        del result
        gc.collect()
        ctypes.CDLL('libc.so.6').malloc_trim(0) 
        self.assertLess(get_memory_usage(), starting_memory + 1024*1024*4)

@export
@register()
class TestMemoryLeaks_insert_string_values(CommonDatabaseUnitTest):
    @staticmethod
    def get_random_strings(num_strings=1024*32, string_size = 1024*64):
      return [dict(idx=it, value=base64.b64encode(random.randbytes(string_size))) for it in range(num_strings)]

    def runTest(self):
        qdb = QIR.QirDatabase(self.db)

        strvec = self.db.create_table_from_definition(
            "StrVec",
            {
                "idx": Types.BigInt(),
                "value": Types.Text()
            },
            ["idx"]
        )

        gc.collect()
        ctypes.CDLL('libc.so.6').malloc_trim(0) 

        starting_memory = get_memory_usage()

        input_rows = self.get_random_strings()

        result = (
            qdb.insert(qdb.StrVec)
              .values(input_rows)
              .run()
        )
        self.assertEqual(len(input_rows), result.count())

        del input_rows

        self.assertGreater(get_memory_usage(), starting_memory + 1024*1024*1024)
        del self.db["StrVec"]
        del result
        gc.collect()
        ctypes.CDLL('libc.so.6').malloc_trim(0) 
        self.assertLess(get_memory_usage(), starting_memory + 1024*1024*4)


#@export
#@register()
#class TestMemoryLeaks_string_query_result(TestMemoryLeaks_string_insert):
#    def setUp(self):
#        super().setUp()
#        super().runTest()
#
#    def runTest(self):
#        qdb = QIR.QirDatabase(self.db)
#        result = (
#            qdb.query(qdb.StrVec)
#              .select(qdb.StrVec.idx, qdb.StrVec.value)
#              .run()
#        )
#        self.assertEqual(
#          len(self.data),
#          result.count()
#        )
#        self.assertEqual(
#            self.data,
#            list(result)
#        )
#

tensql/test/test_queries/MemoryLeaks.py
sandialabs/TenSQL
from .common import CommonDatabaseUnitTest
from ..test_util import register
from ...util import export
from ... import QIR

import numpy as np

__all__ = []

@export
@register()
class TestQuery_PrimaryKeyFilter(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        A = qdb.Matrix.aliased("A")

        result = (
            qdb.query(A)
                .select(A.ridx, A.cidx, A.value)
                .group_by(A.ridx, A.cidx)
                .where(A.ridx >= A.cidx)
                .run()
        )

        observed = result.get_column('value').data[:99, :99].to_numpy()
        expected = np.tril(self.np_matrix)

        np.testing.assert_allclose(expected, observed, rtol=1e-4, atol=1e-4)

@export
@register()
class TestQuery_PrimaryKeyFilterPlusOne(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        A = qdb.Matrix.aliased("A")

        result = (
            qdb.query(A)
                .select(A.ridx, A.cidx, A.value)
                .group_by(A.ridx, A.cidx)
                .where(A.ridx >= (A.cidx + 1))
                .run()
        )

        observed = result.get_column('value').data[:99, :99].to_numpy()
        expected = np.tril(self.np_matrix, k=-1)

        np.testing.assert_allclose(expected, observed, rtol=1e-4, atol=1e-4)
tensql/test/test_queries/PrimaryKeyFilter.py
sandialabs/TenSQL
import numpy as np

from .common import CommonDatabaseUnitTest
from ..test_util import register
from ...util import export
from ... import QIR
from ... import Types

__all__ = []

@export
@register()
class TestInsert_simple(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        A = qdb.Matrix

        result = (
            qdb.insert(A).values([
                dict(ridx=99, cidx=0, value=25),
                dict(ridx=0, cidx=99, value=16)
            ]).run()
        )

        observed = result.get_column('value').data[:99, :99].to_numpy()
        expected = np.copy(self.np_matrix)
        expected[99, 0] = 25
        expected[0, 99] = 16

        np.testing.assert_allclose(expected, observed, rtol=1e-4, atol=1e-4)

tensql/test/test_queries/SimpleInsert.py
sandialabs/TenSQL
import sys
import gc

import numpy as np

from .common import CommonDatabaseUnitTest
from ..test_util import register
from ...util import export
from ... import QIR
from ... import Types

__all__ = []

@export
@register()
class TestString_insert(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)

        strvec = self.db.create_table_from_definition(
            "StrVec",
            {
                "idx": Types.BigInt(),
                "value": Types.Text()
            },
            ["idx"]
        )

        input_rows = [
            dict(idx=1, value="sdkfjkel"),
            dict(idx=2, value="sadfewewqw"),
            dict(idx=3, value="asdfewqwwe")
        ]
        self.data = input_rows

        result = (
            qdb.insert(qdb.StrVec)
              .values(input_rows)
              .run()
        )

        self.assertEqual(len(input_rows), result.count())

        rows = list(result)

        self.assertEqual(input_rows, rows)

@export
@register()
class TestString_query_result(TestString_insert):
    def setUp(self):
        super().setUp()
        super().runTest()

    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        result = (
            qdb.query(qdb.StrVec)
              .select(qdb.StrVec.idx, qdb.StrVec.value)
              .run()
        )
        self.assertEqual(
          len(self.data),
          result.count()
        )
        self.assertEqual(
            self.data,
            list(result)
        )

@export
@register()
class TestString_where_expression_equals(TestString_insert):
    def setUp(self):
        super().setUp()
        super().runTest()

    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        data = self.data
        for row in data:
          idx, value = row['idx'], row['value']
          result = (
              qdb.query(qdb.StrVec)
                .select(qdb.StrVec.idx, qdb.StrVec.value)
                .where(qdb.StrVec.value == value)
                .run()
          )
          self.assertEqual(
              sorted([tuple(row.items()) for row in data if row['value'] == value]),
              sorted(tuple(row.items()) for row in result)
          )

@export
@register()
class TestString_where_expression_notequals(TestString_insert):
    def setUp(self):
        super().setUp()
        super().runTest()

    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        data = self.data
        for row in data:
          idx, value = row['idx'], row['value']
          result = (
              qdb.query(qdb.StrVec)
                .select(qdb.StrVec.idx, qdb.StrVec.value)
                .where(qdb.StrVec.value != value)
                .run()
          )
          self.assertEqual(
              sorted([tuple(row.items()) for row in data if row['value'] != value]),
              sorted(tuple(row.items()) for row in result)
          )
          

tensql/test/test_queries/Strings.py
sandialabs/TenSQL
from .AA import *
from .MemoryLeaks import *
from .PrimaryKeyFilter import *
from .SimpleInsert import *
from .Strings import *
from .xAy import *

tensql/test/test_queries/__init__.py
sandialabs/TenSQL
import unittest

import numpy as np

from ..test_util import TmpDirMixin
from ... import Database, Types

class CommonDatabaseUnitTest(unittest.TestCase):
    def setUp(self):
        super().setUp()
        self.db = Database()

        self.matrix = self.db.create_table_from_definition(
            "Matrix",
            {
                "ridx": Types.BigInt(),
                "cidx": Types.BigInt(),
                "value": Types.Real()
            },
            ["ridx", "cidx"]
        )

        self.np_matrix = np.zeros(shape=(100, 100))
        for it in range(100):
            val = 1.0 / (it+1)
            self.matrix.insert(ridx=it, cidx=it, value=val)
            self.np_matrix[it, it] = 1.0 / (it+1)
            if(it > 0):
                val = 7.0 / (it + 1)
                self.matrix.insert(ridx=it, cidx=it-1, value=7.0 / (it + 1))
                self.np_matrix[it, it-1] = val

                val = 5.0 / (it + 1)
                self.matrix.insert(ridx=it-1, cidx=it, value=5.0 / (it + 1))
                self.np_matrix[it-1, it] = val

        self.left_vector = self.db.create_table_from_definition(
            "LeftVector",
            {
                "idx": Types.BigInt(),
                "value": Types.Real()
            },
            ["idx"]
        )
        self.np_left_vector = np.zeros(shape=(100,), dtype=np.float32)
        for it in range(100):
            if it % 3 != 0:
                val = 3.0 / (it+1)
                self.np_left_vector[it] = val
                self.left_vector.insert(idx=it, value=val)

        self.right_vector = self.db.create_table_from_definition(
            "RightVector",
            {
                "idx": Types.BigInt(),
                "value": Types.Real()
            },
            ["idx"]
        )
        self.np_right_vector = np.zeros(shape=(100,), dtype=np.float32)
        for it in range(100):
            if it % 5 != 0:
                val = 2.0 / (it+1)
                self.np_right_vector[it] = val
                self.right_vector.insert(idx=it, value=val)


tensql/test/test_queries/common.py
sandialabs/TenSQL
from .common import CommonDatabaseUnitTest
from ..test_util import register
from ...util import export
from ... import QIR
import numpy as np

__all__ = []

np.set_printoptions(suppress=True)

@export
@register()
class TestQuery_x_ewisemult_A(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        A = qdb.Matrix.aliased("A")
        x = qdb.LeftVector.aliased("x")

        result = (
            qdb.query(A)
                .join(x, A.ridx == x.idx)
                .select(A.ridx, A.cidx, (A.value * x.value).aliased("Result"))
                .group_by(A.ridx, A.cidx)
                .run()
        )


        expected = np.expand_dims(self.np_left_vector, 1) * self.np_matrix
        observed = np.zeros_like(expected)
        result = result.get_column('Result').data
        observed[result.npI, result.npJ] = result.npV

        np.testing.assert_allclose(expected, observed, rtol=1e-4)

@export
@register()
class TestQuery_y_ewisemult_A(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        A = qdb.Matrix.aliased("A")
        y = qdb.RightVector.aliased("y")

        result = (
            qdb.query(A)
                .join(y, A.cidx == y.idx)
                .select(A.ridx, A.cidx, (A.value * y.value).aliased("Result"))
                .group_by(A.ridx, A.cidx)
                .run()
        )


        expected = np.expand_dims(self.np_right_vector, 0) * self.np_matrix
        observed = np.zeros_like(expected)
        result = result.get_column('Result').data
        observed[result.npI, result.npJ] = result.npV

        np.testing.assert_allclose(expected, observed, rtol=1e-4)

@export
@register()
class TestQuery_xy_ewisemult_A(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        A = qdb.Matrix.aliased("A")
        x = qdb.LeftVector.aliased("x")
        y = qdb.RightVector.aliased("y")

        result = (
            qdb.query(A)
                .join(x, A.ridx == x.idx)
                .join(y, A.cidx == y.idx)
                .select(A.ridx, A.cidx, (A.value * x.value * y.value).aliased("Result"))
                .group_by(A.ridx, A.cidx)
                .run()
        )


        expected = np.expand_dims(self.np_left_vector, 1) * np.expand_dims(self.np_right_vector, 0) * self.np_matrix
        observed = np.zeros_like(expected)
        result_col = result.get_column('Result')
        O = result.get_column('Result').data
        observed[O.npI, O.npJ] = O.npV

        self.assertEqual(result_col.data.nvals, result.stencil.nvals)
        np.testing.assert_allclose(expected, observed, rtol=1e-4)

@export
@register()
class TestQuery_xAy(CommonDatabaseUnitTest):
    def runTest(self):
        qdb = QIR.QirDatabase(self.db)
        A = qdb.Matrix.aliased("A")
        x = qdb.LeftVector.aliased("x")
        y = qdb.RightVector.aliased("y")

        result = (
            qdb.query(A)
                .join(x, A.ridx == x.idx)
                .join(y, A.cidx == y.idx)
                .select(QIR.Sum(A.value * x.value * y.value).aliased("Result"))
                .group_by()
                .run()
        )

        result_col = result.get_column('Result')
        observed = result_col.data[None]
        expected = self.np_left_vector.dot(self.np_matrix).dot(self.np_right_vector)

        self.assertEqual(result.stencil.nvals, 1)
        self.assertEqual(result_col.data.nvals, 1)
        self.assertAlmostEqual(expected, observed, places=4)

tensql/test/test_queries/xAy.py
sandialabs/TenSQL
"""Helpful tools for creating automated tests"""

import collections
import itertools
import os
import pathlib
import sys
import tempfile
import unittest

from ..util import export

__all__ = ["test_registry"]

test_registry = collections.defaultdict(lambda: [])

@export
def register(*tags):
    """Registers a unit test for the custom test harness"""
    tags = list(tags)
    def register_inner(test_case):
        for tag in itertools.chain(("all",), tags):
            test_registry[tag].append(test_case)
        return test_case
    return register_inner

@export
class TmpDirMixin(unittest.TestCase):
    """Unittest mixin that creates and manages a temporary directory"""
    def setUp(self):
        """Create the temporary directory"""
        super().setUp()
        self.old_cwd = os.getcwd()
        self.tmpdir = tempfile.TemporaryDirectory()
        os.chdir(self.tmpdir.__enter__())

    def tearDown(self):
        """Delete the temporary directory"""
        os.chdir(self.old_cwd)
        self.tmpdir.__exit__(None, None, None)
        self.tmpdir.cleanup()
        super().tearDown()


tensql/test/test_util.py
sandialabs/TenSQL
import abc
import hashlib
import collections
import inspect
import threading
import json
from tkinter import N

from typing import List, Optional, Union, Any, Type, Sequence, Tuple

import numpy as np
from pygraphblas.types import Type as GBType
from pygraphblas import Matrix, Vector, Scalar
import pygraphblas
import pygraphblas.types
from pygraphblas.base import NULL, ffi

from . import Types

class Panic(pygraphblas.base.GraphBLASException):
  pass

class OutOfMemory(pygraphblas.base.GraphBLASException):
  pass

class InsufficientSpace(pygraphblas.base.GraphBLASException):
  pass

class InvalidObject(pygraphblas.base.GraphBLASException):
  pass

class IndexOutOfBounds(pygraphblas.base.GraphBLASException):
  pass

class EmptyObject(pygraphblas.base.GraphBLASException):
  pass

class UninitializedObject(pygraphblas.base.GraphBLASException):
  pass

class NullPointer(pygraphblas.base.GraphBLASException):
  pass

class InvalidValue(pygraphblas.base.GraphBLASException):
  pass

class InvalidIndex(pygraphblas.base.GraphBLASException):
  pass

class DomainMismatch(pygraphblas.base.GraphBLASException):
  pass

class DimensionMismatch(pygraphblas.base.GraphBLASException):
  pass

class OutputNotEmpty(pygraphblas.base.GraphBLASException):
  pass

class GrbNotImplemented(pygraphblas.base.GraphBLASException):
  pass

pygraphblas.base._error_codes.update({
  -1: UninitializedObject,
  -2: NullPointer,
  -3: InvalidValue,
  -4: InvalidIndex,
  -5: DomainMismatch,
  -6: DimensionMismatch,
  -7: OutputNotEmpty,
  -8: GrbNotImplemented,
  -101: Panic,
  -102: OutOfMemory,
  -103: InsufficientSpace,
  -104: InvalidObject,
  -105: IndexOutOfBounds,
  -106: EmptyObject
})

class DummyTensor:
    def __init__(self, shape: Tuple[int, ...]):
        self._shape = shape

    @property
    def shape(self) -> int:
        return self._shape

    @property
    def ndim(self) -> int:
        return len(self._shape)


GBTensor = Union[Matrix, Vector, Scalar, DummyTensor]
Axes = Sequence[int]

UnaryPyOps = {
    "-": lambda x: -x,
    "NOT": lambda x: not x,
    "ABS": lambda x: abs(x),
    "ONE": lambda x: 1,
}

UnaryGrbOps = {
    "-": "AINV",
    "NOT": "LNOT",
    "ABS": "ABS",
    "ONE": "ONE"
}

BinaryArithmeticPyOps = {
    "+": lambda x, y: x + y,
    "-": lambda x, y: x - y,
    "*": lambda x, y: x * y,
    "/": lambda x, y: x / y,
    "//": lambda x, y: x // y,
    "%": lambda x, y: x % y,
}

BinaryArithmeticGrbOps = {
    "+": "PLUS",
    "-": "MINUS",
    "*": "TIMES",
    "/": "DIV",
}

BinaryBooleanPyOps = {
    "==": lambda x, y: x == y,
    "!=": lambda x, y: x != y,
    ">": lambda x, y: x > y,
    "<": lambda x, y: x < y,
    "<=": lambda x, y: x <= y,
    ">=": lambda x, y: x >= y,
    "AND": lambda x, y: x and y,
    "OR": lambda x, y: x or y,
}

BinaryBooleanGrbOps = {
    "==": "EQ",
    "!=": "NE",
    ">": "GT",
    "<": "LT",
    "<=": "GE",
    ">=": "LE",
    "AND": "TIMES",
    "OR": "PLUS",
}


def grb_shape(obj: GBTensor) -> Sequence[int]:
    if isinstance(obj, Scalar):
        return tuple()
    elif isinstance(obj, Vector):
        return (obj.size,)
    elif isinstance(obj, Matrix):
        return (obj.nrows, obj.ncols)
    elif isinstance(obj, DummyTensor):
        return obj.shape
    else:
        raise ValueError(
            f"Prototype is only implemented for matrices, vectors, and scalars, not {type(obj)=}"
        )


shape = grb_shape


def ndim(obj: GBTensor) -> int:
    if isinstance(obj, Scalar):
        return 0
    elif isinstance(obj, Vector):
        return 1
    elif isinstance(obj, Matrix):
        return 2
    elif isinstance(obj, DummyTensor):
        return obj.ndim
    else:
        raise ValueError(
            f"Prototype is only implemented for matrices, vectors, and scalars, not {type(obj)=}"
        )

def scalar_add(first: Scalar, second: Scalar, *, op: Optional[str]=None, inplace: bool=False):
    if inplace:
        data = first
    else:
        data = Scalar.from_type(first.type)

    if op is None:
        op = "+"

    if first.nvals == 0 and second.nvals == 0:
        pass
    elif first.nvals > 0 and second.nvals == 0:
        data[None] = first[None]
    elif first.nvals == 0 and second.nvals > 0:
        data[None] = second[None]
    elif first.nvals > 0 and second.nvals > 0:
        data[None] = BinaryArithmeticPyOps[op](
            first[None], second[None]
        )

    return data


class ABCIdGenerator(abc.ABC):
    def __init__(self):
        pass

    @abc.abstractmethod
    def __call__(self, value):
        pass


class SequentialIdGenerator(ABCIdGenerator):
    def __init__(self):
        self.counter = 0
        self.lock = threading.Lock()

    def increment(self, amount=1):
        with self.lock:
            ret = self.counter
            self.counter += amount
        return ret

    def __call__(self, value):
        return self.increment()

class HashIdGenerator(ABCIdGenerator):
    KEYSPACE_NIBBLES = 15  # 60 bits of hash

    def __init__(self, algo="sha256", salt=None):
        self.hasher = hashlib.new(algo)
        if salt is not None:
            self.hasher.update(salt)

    def __call__(self, value):
        hasher = self.hasher.copy()
        if isinstance(value, str):
            hasher.update(value.encode("utf-8"))
        else:
            hasher.update(value)
        return int(hasher.hexdigest()[: self.KEYSPACE_NIBBLES], 16)


BinaryPyOps = {
    "+": lambda x, y: x + y,
    "-": lambda x, y: x - y,
    "*": lambda x, y: x * y,
    "/": lambda x, y: x / y,
    "//": lambda x, y: x // y,
    "%": lambda x, y: x % y,
    "==": lambda x, y: x == y,
    "!=": lambda x, y: x != y,
    ">": lambda x, y: x > y,
    "<": lambda x, y: x < y,
    "<=": lambda x, y: x <= y,
    ">=": lambda x, y: x >= y,
    "AND": lambda x, y: x and y,
    "OR": lambda x, y: x or y,
    "FIRST": lambda x, y: x,
    "SECOND": lambda x, y: y,
}

BinaryGrbOps = {
    "+": "PLUS",
    "-": "MINUS",
    "*": "TIMES",
    "/": "DIV",
    "==": "EQ",
    "!=": "NE",
    ">": "GT",
    "<": "LT",
    ">=": "GE",
    "<=": "LE",
    "AND": "TIMES",
    "OR": "PLUS",
    "FIRST": "FIRST",
    "SECOND": "SECOND",
    "LAND": "LAND",
    "LOR": "LOR",
}


class EncoderDecoder:
    def __init__(self, id_generator: SequentialIdGenerator):
        self.id_generator = id_generator
        self.encoder = collections.OrderedDict()
        self.decoder = {}
        self.lock = threading.Lock()
        self._np_encode = np.vectorize(lambda x: self[x], otypes=[np.uint64])

    def __getitem__(self, name: str) -> Any:
        return self.get(name)[0]

    def get(self, name: str) -> Any:
        ret = self.encoder.get(name)
        if ret is None:
            with self.lock:
                ret = self.encoder.get(name)
                if ret is None:
                    ret = self.id_generator(name)
                    self.encoder[name] = ret
                    self.decoder[ret] = name

            return ret, True
        else:
            return ret, False

    def __setitem__(self, name: str, value: Any) -> None:
        with self.lock:
            prev_value = self.encoder.get(name)
            if prev_value is None:
                self.encoder[name] = value
                self.decoder.setdefault(value, name)
            else:
                raise ValueError(f"Key {name} was already set")

    def __contains__(self, name: str) -> bool:
        return name in self.encoder

    def encode(self, name: str) -> Any:
        return self.encoder.get(name)

    def encode_many(self, names: Sequence[str]) -> np.ndarray:
        return self._np_encode(names)

    def decode(self, ident: Any) -> str:
        return self.decoder.get(ident)

    def items(self):
        for k, v in self.encoder.items():
            yield k, v

    def __len__(self):
        return len(self.encoder)

    def to_json(self):
        return json.dumps(self.encoder)


def export(obj):
    """
    A helpful decorator used to control __all__
    Args:
        obj: the object whose name should be added to __all__
    Returns:
        obj
    """
    inspect.getmodule(obj).__all__.append(obj.__name__)
    return obj

def vector_grb_check_success(vector, res):
    if res != pygraphblas.lib.GrB_SUCCESS:
        error_string = ffi.new("char**")
        pygraphblas.lib.GrB_Vector_error(error_string, vector._vector[0])
        raise pygraphblas.base._error_codes[res](ffi.string(error_string[0]))

def matrix_grb_check_success(vector, res):
    if res != pygraphblas.lib.GrB_SUCCESS:
        error_string = ffi.new("char**")
        pygraphblas.lib.GrB_Matrix_error(error_string, vector._vector[0])
        raise pygraphblas.base._error_codes[res](ffi.string(error_string[0]))

def grb_check_success(res):
    pygraphblas.base._check(res)

def extract_indices(obj, axis, typ=None):
    if typ is None:
        typ = Types.BigInt()

    if not isinstance(typ, Types.BigInt):
        raise NotImplementedError("Primary key types other than BigInt are not yet supported.")
    
    assert isinstance(typ, Types.BigInt)

    if isinstance(obj, Scalar):
        raise ValueError("Scalar has no indices")
    elif isinstance(obj, Vector):
        if axis != 0:
          raise ValueError("axis must be 0 when obj argument is a vector")
        ret = Vector.sparse(pygraphblas.types.INT64, size=obj.size)

        nvals = obj.nvals
        _nvals = ffi.new("GrB_Index[1]", [nvals])
        I = ffi.new("GrB_Index[%s]" % nvals)
        X = NULL
        vector_grb_check_success(
          obj,
          obj.type._Vector_extractTuples(I, X, _nvals, obj._vector[0])
        )

        X = ffi.cast("int64_t[%s]" % nvals, I)

        grb_check_success(
          pygraphblas.lib.GrB_Vector_build_INT64(
            ret._vector[0],
            I, X,
            nvals,
            typ.as_pygrb.FIRST.get_op()
          )
        )
        return ret
    elif isinstance(obj, Matrix):
        if axis not in (0,1):
          raise ValueError("axis must be 0 or 1 when obj argument is a matrix")
        ret = Matrix.sparse(pygraphblas.types.INT64, nrows=obj.nrows, ncols=obj.ncols)

        nvals = obj.nvals
        _nvals = ffi.new("GrB_Index[1]", [nvals])
        I = ffi.new("GrB_Index[%s]" % nvals)
        J = ffi.new("GrB_Index[%s]" % nvals)
        X = NULL
        matrix_grb_check_success(
          obj,
          obj.type._Matrix_extractTuples(I, J, X, _nvals, obj._matrix[0])
        )

        if axis == 0:
          X = I
        elif axis == 1:
          X = J

        X = ffi.cast("int64_t[%s]" % nvals, X)

        grb_check_success(
          pygraphblas.lib.GrB_Matrix_build_INT64(
            ret._matrix[0],
            I, J, X,
            nvals,
            typ.as_pygrb.FIRST.get_op()
          )
        )
        return ret
    else:
        raise ValueError(
            "Prototype is only implemented for matrices, vectors, and scalars"
        )

def build_tensor(gbtype: pygraphblas.types.Type, shape: Sequence[int]) -> GBTensor:
    assert issubclass(gbtype, pygraphblas.types.Type)

    if len(shape) == 0:
        return Scalar.from_type(gbtype)
    elif len(shape) == 1:
        return Vector.sparse(gbtype, *shape)
    elif len(shape) == 2:
        return Matrix.sparse(gbtype, *shape)
    else:
        raise NotImplementedError(
            "Prototype only supports scalars, vectors, and matrices"
        )

def tensor_to_numpy(obj: GBTensor) -> Tuple[np.ndarray]:
    nvals, typ = obj.nvals, obj.type

    if ndim(obj) == 0:
        ret = np.zeros(shape=(obj.nvals,), dtype=typ._numpy_t)
        if obj.nvals > 0:
            ret[0] = obj[None]
        return (ret,)
    elif ndim(obj) == 1:
        _nvals = ffi.new("GrB_Index[1]", [nvals])
        I = ffi.new(f"GrB_Index[{nvals}]")
        X = ffi.new(f"{typ._c_type}[{nvals}]")
        vector_grb_check_success(
            obj,
            obj.type._Vector_extractTuples(I, X, _nvals, obj._vector[0])
        )
        return (
            np.frombuffer(ffi.buffer(I, ffi.sizeof(I)), dtype=np.uint64),
            np.frombuffer(ffi.buffer(X, ffi.sizeof(X)), dtype=typ._numpy_t)
        )
    elif ndim(obj) == 2:
        _nvals = ffi.new("GrB_Index[1]", [nvals])
        I = ffi.new(f"GrB_Index[{nvals}]")
        J = ffi.new(f"GrB_Index[{nvals}]")
        X = ffi.new(f"{typ._c_type}[{nvals}]")
        matrix_grb_check_success(
            obj,
            obj.type._Matrix_extractTuples(I, J, X, _nvals, obj._matrix[0])
        )
        return (
            np.frombuffer(ffi.buffer(I, ffi.sizeof(I)), dtype=np.uint64),
            np.frombuffer(ffi.buffer(J, ffi.sizeof(J)), dtype=np.uint64),
            np.frombuffer(ffi.buffer(X, ffi.sizeof(X)), dtype=typ._numpy_t)
        )
    else:
        raise NotImplementedError(
            "Prototype only supports scalars, vectors, and matrices"
        )

def fill_tensor_from_numpy(obj: GBTensor, *lists: Sequence[np.ndarray]) -> None:
    assert len(lists) == ndim(obj) + 1

    if ndim(obj) == 0:
        if lists[0].size > 0:
            obj[None] = lists[0][0]
        return
    elif ndim(obj) not in (1, 2):
        raise NotImplementedError(
            "Prototype only supports scalars, vectors, and matrices"
        )

    build_func_name_parts = ["GrB"]
    if ndim(obj) == 1:
        build_func_name_parts.append("Vector")
        grb_tensor = obj._vector[0]
    elif ndim(obj) == 2:
        build_func_name_parts.append("Matrix")
        grb_tensor = obj._matrix[0]

    build_func_name_parts.append("build")

    typ = obj.type
    build_func_name_parts.append(typ._base_name)

    build_func_name = "_".join(build_func_name_parts)
    assert hasattr(pygraphblas.lib, build_func_name)
    build_func = getattr(pygraphblas.lib, build_func_name)

    indices, value = lists[:-1], lists[-1]
    nvals = value.size
    assert all(idx.size == value.size for idx in indices)
    ffi_lists = [ffi.cast(f"uint64_t[{nvals}]", idx.ctypes.data) for idx in indices]
    ffi_lists.append(
        ffi.cast(f"{typ._c_type}[{nvals}]", value.ctypes.data)
    )

    #print([x for x in ffi_lists], nvals)

    grb_check_success(
        build_func(
            grb_tensor,
            *ffi_lists,
            nvals,
            typ.FIRST.get_op()
        )
    )
    

def initialize_tensor(
    gbtype: pygraphblas.types.Type, shape: Sequence[int], lists: Sequence[Sequence[Any]]
) -> GBTensor:
    try:
        is_gbtype_subclass = issubclass(gbtype, pygraphblas.types.Type)
    except TypeError:
        is_gbtype_subclass = False
    
    if not (isinstance(gbtype, pygraphblas.types.Type) or is_gbtype_subclass):
        raise ValueError(f"Expected instance of pygraphblas.types.Type, not {gbtype}")
    assert len(lists) == len(shape) + 1

    #lists = [np.array(l).tolist() for l in lists]

    if len(shape) == 0:
        ret = Scalar.from_type(gbtype)
        if len(lists[0]) == 0:
            pass
        elif len(lists[0]) == 1:
            ret[None] = lists[0][0]
        else:
            raise ValueError("Scalars should only have zero or one entries!")
        return ret
    elif len(shape) == 1:
        ret = Vector.sparse(gbtype, size=shape[0])
        if len(lists[0]) > 0:
            fill_tensor_from_numpy(
                ret,
                np.array(lists[0], dtype=np.uint64, copy=False),
                np.array(lists[1], dtype=gbtype._numpy_t, copy=False),
            )
    elif len(shape) == 2:
        ret = Matrix.sparse(
            gbtype,
            nrows=shape[0],
            ncols=shape[1]
        )
        if len(lists[0]) > 0:
            fill_tensor_from_numpy(
                ret,
                np.array(lists[0], dtype=np.uint64, copy=False),
                np.array(lists[1], dtype=np.uint64, copy=False),
                np.array(lists[2], dtype=gbtype._numpy_t, copy=False)
            )
#            ret = Matrix.from_lists(
#                I=lists[0],
#                J=lists[1],
#                V=lists[2],
#                nrows=shape[0],
#                ncols=shape[1],
#                typ=gbtype,
#            )
    else:
        raise ValueError("Prototype only supports scalars, vectors, and matrices")

    return ret

tensql/util.py
sandialabs/TenSQL