FEATURE: Try to automatically handle larger embedding indexes (#403)
* FEATURE: Try to automatically handle larger embedding indexes * linteeeeeeeer
This commit is contained in:
parent
dd42a4e47b
commit
6fc1c9f7a6
|
@ -80,20 +80,33 @@ module DiscourseAi
|
||||||
end
|
end
|
||||||
|
|
||||||
def create_index!(table_name, memory, lists, probes)
|
def create_index!(table_name, memory, lists, probes)
|
||||||
|
tries = 0
|
||||||
index_name = index_name(table_name)
|
index_name = index_name(table_name)
|
||||||
DB.exec("SET work_mem TO '#{memory}';")
|
DB.exec("SET work_mem TO '#{memory}';")
|
||||||
DB.exec("SET maintenance_work_mem TO '#{memory}';")
|
DB.exec("SET maintenance_work_mem TO '#{memory}';")
|
||||||
DB.exec(<<~SQL)
|
begin
|
||||||
DROP INDEX IF EXISTS #{index_name};
|
DB.exec(<<~SQL)
|
||||||
CREATE INDEX IF NOT EXISTS
|
DROP INDEX IF EXISTS #{index_name};
|
||||||
#{index_name}
|
CREATE INDEX IF NOT EXISTS
|
||||||
ON
|
#{index_name}
|
||||||
#{table_name}
|
ON
|
||||||
USING
|
#{table_name}
|
||||||
ivfflat (embeddings #{pg_index_type})
|
USING
|
||||||
WITH
|
ivfflat (embeddings #{pg_index_type})
|
||||||
(lists = #{lists});
|
WITH
|
||||||
SQL
|
(lists = #{lists});
|
||||||
|
SQL
|
||||||
|
rescue PG::ProgramLimitExceeded => e
|
||||||
|
parsed_error = e.message.match(/memory required is (\d+ [A-Z]{2}), ([a-z_]+)/)
|
||||||
|
if parsed_error[1].present? && parsed_error[2].present?
|
||||||
|
DB.exec("SET #{parsed_error[2]} TO '#{parsed_error[1].tr(" ", "")}';")
|
||||||
|
tries += 1
|
||||||
|
retry if tries < 3
|
||||||
|
else
|
||||||
|
raise e
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
DB.exec("COMMENT ON INDEX #{index_name} IS '#{Time.now.to_i}';")
|
DB.exec("COMMENT ON INDEX #{index_name} IS '#{Time.now.to_i}';")
|
||||||
DB.exec("RESET work_mem;")
|
DB.exec("RESET work_mem;")
|
||||||
DB.exec("RESET maintenance_work_mem;")
|
DB.exec("RESET maintenance_work_mem;")
|
||||||
|
|
Loading…
Reference in New Issue