GeezerVison is showing github.com — pgvector-python/tests/test_sqlmodel.py at master · pgvector/pgvector-python

Expand file tree

test_sqlmodel.py

Copy path

More file actions

Latest commit

History

243 lines (204 loc) · 10.3 KB

test_sqlmodel.py

Copy path

File metadata and controls

243 lines (204 loc) · 10.3 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

from pgvector import SparseVector

from pgvector.sqlalchemy import VECTOR, HALFVEC, BIT, SPARSEVEC, avg, sum

from pydantic import ConfigDict

import pytest

from sqlalchemy.exc import StatementError

from sqlmodel import Field, Index, Session, SQLModel, col, create_engine, delete, select, text

engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test')

with Session(engine) as session:

# typing issue

# https://github.com/fastapi/sqlmodel/issues/909

# https://github.com/fastapi/sqlmodel/pull/1657

session.exec(text('CREATE EXTENSION IF NOT EXISTS vector')) # type: ignore

class Item(SQLModel, table=True):

__tablename__ = 'sqlmodel_item' # type: ignore

model_config = ConfigDict(arbitrary_types_allowed=True) # type: ignore

id: int | None = Field(default=None, primary_key=True)

# typing issue

# https://github.com/fastapi/sqlmodel/discussions/1228

# https://github.com/fastapi/sqlmodel/pull/1345

embedding: list[float] | None = Field(default=None, sa_type=VECTOR(3)) # type: ignore

half_embedding: list[float] | None = Field(default=None, sa_type=HALFVEC(3)) # type: ignore

binary_embedding: str | None = Field(default=None, sa_type=BIT(3)) # type: ignore

sparse_embedding: SparseVector | list[float] | None = Field(default=None, sa_type=SPARSEVEC(3)) # type: ignore

SQLModel.metadata.drop_all(engine)

SQLModel.metadata.create_all(engine)

index = Index(

'sqlmodel_index',

col(Item.embedding), # col() for typing

postgresql_using='hnsw',

postgresql_with={'m': 16, 'ef_construction': 64},

postgresql_ops={'embedding': 'vector_l2_ops'}

)

index.create(engine)

def create_items() -> None:

with Session(engine) as session:

session.add(Item(id=1, embedding=[1, 1, 1], half_embedding=[1, 1, 1], binary_embedding='000', sparse_embedding=SparseVector([1, 1, 1])))

session.add(Item(id=2, embedding=[2, 2, 2], half_embedding=[2, 2, 2], binary_embedding='101', sparse_embedding=SparseVector([2, 2, 2])))

session.add(Item(id=3, embedding=[1, 1, 2], half_embedding=[1, 1, 2], binary_embedding='111', sparse_embedding=SparseVector([1, 1, 2])))

session.commit()

class TestSqlmodel:

def setup_method(self) -> None:

with Session(engine) as session:

session.exec(delete(Item))

session.commit()

def test_orm(self) -> None:

item = Item(embedding=[1.5, 2, 3])

item2 = Item(embedding=[4, 5, 6])

item3 = Item()

with Session(engine) as session:

session.add(item)

session.add(item2)

session.add(item3)

session.commit()

stmt = select(Item)

with Session(engine) as session:

items = session.exec(stmt).all()

assert items[0].id == 1

assert items[1].id == 2

assert items[2].id == 3

assert items[0].embedding == [1.5, 2, 3]

assert items[1].embedding == [4, 5, 6]

assert items[2].embedding is None

def test_vector(self) -> None:

with Session(engine) as session:

session.add(Item(id=1, embedding=[1, 2, 3]))

session.commit()

item = session.get_one(Item, 1)

assert item.embedding == [1, 2, 3]

def test_vector_l2_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.embedding).l2_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 3, 2]

def test_vector_max_inner_product(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.embedding).max_inner_product([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [2, 3, 1]

def test_vector_cosine_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.embedding).cosine_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 2, 3]

def test_vector_l1_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.embedding).l1_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 3, 2]

def test_halfvec(self) -> None:

with Session(engine) as session:

session.add(Item(id=1, half_embedding=[1, 2, 3]))

session.commit()

item = session.get_one(Item, 1)

assert item.half_embedding == [1, 2, 3]

def test_halfvec_l2_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.half_embedding).l2_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 3, 2]

def test_halfvec_max_inner_product(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.half_embedding).max_inner_product([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [2, 3, 1]

def test_halfvec_cosine_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.half_embedding).cosine_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 2, 3]

def test_halfvec_l1_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.half_embedding).l1_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 3, 2]

def test_bit(self) -> None:

with Session(engine) as session:

session.add(Item(id=1, binary_embedding='101'))

session.commit()

item = session.get_one(Item, 1)

assert item.binary_embedding == '101'

def test_bit_hamming_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.binary_embedding).hamming_distance('101'))) # type: ignore

assert [v.id for v in items] == [2, 3, 1]

def test_bit_jaccard_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.binary_embedding).jaccard_distance('101'))) # type: ignore

assert [v.id for v in items] == [2, 3, 1]

def test_sparsevec(self) -> None:

with Session(engine) as session:

session.add(Item(id=1, sparse_embedding=[1, 2, 3]))

session.commit()

item = session.get_one(Item, 1)

assert item.sparse_embedding == SparseVector([1, 2, 3])

def test_sparsevec_l2_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.sparse_embedding).l2_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 3, 2]

def test_sparsevec_max_inner_product(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.sparse_embedding).max_inner_product([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [2, 3, 1]

def test_sparsevec_cosine_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.sparse_embedding).cosine_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 2, 3]

def test_sparsevec_l1_distance(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).order_by(col(Item.sparse_embedding).l1_distance([1, 1, 1]))) # type: ignore

assert [v.id for v in items] == [1, 3, 2]

def test_filter(self) -> None:

create_items()

with Session(engine) as session:

items = session.exec(select(Item).filter(col(Item.embedding).l2_distance([1, 1, 1]) < 1)) # type: ignore

assert [v.id for v in items] == [1]

def test_select(self) -> None:

with Session(engine) as session:

session.add(Item(embedding=[2, 3, 3]))

items = session.exec(select(col(Item.embedding).l2_distance([1, 1, 1]))).all() # type: ignore

assert items[0] == 3

def test_vector_avg(self) -> None:

with Session(engine) as session:

res = session.exec(select(avg(Item.embedding))).first()

assert res is None

session.add(Item(embedding=[1, 2, 3]))

session.add(Item(embedding=[4, 5, 6]))

res = session.exec(select(avg(Item.embedding))).first()

assert res == [2.5, 3.5, 4.5]

def test_vector_sum(self) -> None:

with Session(engine) as session:

res = session.exec(select(sum(Item.embedding))).first()

assert res is None

session.add(Item(embedding=[1, 2, 3]))

session.add(Item(embedding=[4, 5, 6]))

res = session.exec(select(sum(Item.embedding))).first()

assert res == [5, 7, 9]

def test_halfvec_avg(self) -> None:

with Session(engine) as session:

res = session.exec(select(avg(Item.half_embedding))).first()

assert res is None

session.add(Item(half_embedding=[1, 2, 3]))

session.add(Item(half_embedding=[4, 5, 6]))

res = session.exec(select(avg(Item.half_embedding))).first()

assert res == [2.5, 3.5, 4.5]

def test_halfvec_sum(self) -> None:

with Session(engine) as session:

res = session.exec(select(sum(Item.half_embedding))).first()

assert res is None

session.add(Item(half_embedding=[1, 2, 3]))

session.add(Item(half_embedding=[4, 5, 6]))

res = session.exec(select(sum(Item.half_embedding))).first()

assert res == [5, 7, 9]

def test_bad_dimensions(self) -> None:

item = Item(embedding=[1, 2])

with Session(engine) as session:

session.add(item)

with pytest.raises(StatementError, match='expected 3 dimensions, not 2'):

session.commit()