def test_greedy_add_item_quantity():
collection2 = collections_pb2.Collection(replication_factor=2)
collection3 = collections_pb2.Collection(replication_factor=3)
shard2 = collections_pb2.Shard()
shard3 = collections_pb2.Shard()
item_2x = Item(collection2, shard2)
item_3x = Item(collection3, shard3)
items = [item_2x, item_3x]
knapsacks = [
Knapsack(collections_pb2.Node()),
Knapsack(collections_pb2.Node()),
Knapsack(collections_pb2.Node()),
]
algo = GreedyAlgorithm()
algo.add(items, knapsacks)
for knapsack in knapsacks:
assert item_3x in knapsack.items
count = 0
for knapsack in knapsacks:
if item_2x in knapsack.items:
count += 1
assert count == 2
def collection_2shards_2d(tmpdir):
X, metadatas = gen_random_vectors_and_metadatas(dimension=2,
size=20,
dtype="float32",
id_prefix="shard_1")
index_proto_shard_1 = gen_faiss_index_proto(tmpdir, X, metadatas,
"shard_1")
X, metadatas = gen_random_vectors_and_metadatas(dimension=2,
size=25,
dtype="float32",
id_prefix="shard_2")
index_proto_shard_2 = gen_faiss_index_proto(tmpdir, X, metadatas,
"shard_2")
proto = collections_pb2.Collection(
name="test_name",
replication_factor=1,
enable_id_to_vector=True,
shards=[
collections_pb2.Shard(name="shard_1",
weight=20.0,
index=index_proto_shard_1),
collections_pb2.Shard(name="shard_2",
weight=25.0,
index=index_proto_shard_2),
],
)
yield Collection.from_proto(proto)
def test_greedy_add_two_knapsack():
collection = collections_pb2.Collection()
shard1 = collections_pb2.Shard()
shard2 = collections_pb2.Shard()
node1 = collections_pb2.Node()
node2 = collections_pb2.Node()
items = [Item(collection, shard1), Item(collection, shard2)]
knapsacks = [Knapsack(node1), Knapsack(node2)]
algo = GreedyAlgorithm()
algo.add(items, knapsacks)
for knapsack in knapsacks:
assert len(knapsack.items) == 1
def create_shard_proto(collection_name: str,
name: str) -> collections_pb2.Shard:
"""Creates a shard protobuf that points to a local data source file"""
return collections_pb2.Shard(
name=name,
index=indices_pb2.BaseIndex(faiss_index=indices_pb2.FaissIndex(
data_source=data_sources_pb2.DataSource(
local_data_source=data_sources_pb2.LocalDataSource(
filename=get_shard_filename(collection_name, name))))))
def test_knapsack_add_item():
shard = collections_pb2.Shard(weight=3.14)
collection = collections_pb2.Collection(shards=[shard])
node = collections_pb2.Node()
item = Item(collection, shard)
knapsack = Knapsack(node)
knapsack.add_item(item)
assert item in knapsack.items
assert knapsack.current_weight == pytest.approx(3.14)
def test_knapsack_add_duplicate_item():
shard = collections_pb2.Shard(weight=1.12)
collection = collections_pb2.Collection(shards=[shard])
node = collections_pb2.Node()
item = Item(collection, shard)
knapsack = Knapsack(node)
knapsack.add_item(item)
with pytest.raises(KnapsackItemException) as excinfo:
knapsack.add_item(item)
assert "Item already exists in this knapsack" == str(excinfo.value)
def test_greedy_add_one_knapsack():
shard = collections_pb2.Shard()
collection = collections_pb2.Collection()
node = collections_pb2.Node()
item = Item(collection, shard)
knapsack = Knapsack(node)
algo = GreedyAlgorithm()
algo.add([item], [knapsack])
assert item in knapsack.items
def test_greedy_add_item_quantity_to_much():
collection = collections_pb2.Collection(replication_factor=2)
shard = collections_pb2.Shard()
node = collections_pb2.Node()
item_2x = Item(collection, shard)
knapsack = Knapsack(node)
algo = GreedyAlgorithm()
algo.add([item_2x], [knapsack])
assert len(knapsack.items) == 1
def shard_3d(tmpdir):
X, metadatas = gen_random_vectors_and_metadatas(dimension=3,
size=10,
dtype="float32")
index = gen_faiss_index_proto(tmpdir, X, metadatas)
proto = collections_pb2.Shard(
name="shard_1",
weight=20.0,
index=index,
)
yield Shard.from_proto(proto)