Zayden-Jung/foundry-core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Read the published research corpus and avoid expensive paging plans [skip ci]

Browse Source
This commit is contained in:
Codex Bot
2026-11-07 15:51:17 +08:00
parent 913bf4292b
commit 1363792bb0
10 changed files with 355 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/research/postgres.go
+61 -17
View File
@@ -15,8 +15,10 @@ import (
// and BIGINT and creates the English title/abstract GIN expression below.
// No schema creation, indexing or other writes belong to this repository.
type PostgresRepository struct {
db *sql.DB
prefix string
db *sql.DB
prefix string
schema string
encoded bool
}
var schemaName = regexp.MustCompile(`^[a-z][a-z0-9_]{0,35}$`)
@@ -25,7 +27,7 @@ func NewPostgresRepository(db *sql.DB, schema string) (*PostgresRepository, erro
if db == nil || !schemaName.MatchString(schema) {
return nil, errors.New("invalid research database configuration")
}
return &PostgresRepository{db: db, prefix: `"` + schema + `".`}, nil
return &PostgresRepository{db: db, prefix: `"` + schema + `".`, schema: schema}, nil
}
func (r *PostgresRepository) query(value string) string {
@@ -39,7 +41,7 @@ func (r *PostgresRepository) begin(ctx context.Context) (*sql.Tx, error) {
}
_, err = tx.ExecContext(ctx, `SELECT set_config('statement_timeout','8000',true),
set_config('lock_timeout','1000',true), set_config('work_mem','2MB',true),
set_config('max_parallel_workers_per_gather','0',true)`)
set_config('max_parallel_workers_per_gather','0',true), set_config('jit','off',true)`)
if err != nil {
_ = tx.Rollback()
return nil, err
@@ -55,6 +57,16 @@ func (r *PostgresRepository) Check(ctx context.Context) error {
return err
}
defer tx.Rollback()
var annotation sql.NullString
err = tx.QueryRowContext(ctx, `SELECT pg_catalog.obj_description(oid,'pg_namespace')
FROM pg_catalog.pg_namespace WHERE nspname=$1`, r.schema).Scan(&annotation)
if err != nil {
return err
}
encoded, err := encodedSchema(annotation)
if err != nil {
return err
}
rows, err := tx.QueryContext(ctx, r.query(`SELECT p.id,p.title,p.abstract,p.authors_json,p.first_published,p.title_key,
r.external_id,r.url,r.metadata_json,so.name,i.scheme,i.value,t.namespace,t.value,
s.id,e.year,e.track,v.acronym,d.status,d.evidence_url,d.observed_at,d.decision_date
@@ -69,7 +81,11 @@ func (r *PostgresRepository) Check(ctx context.Context) error {
if err = rows.Close(); err != nil {
return err
}
return tx.Commit()
if err = tx.Commit(); err != nil {
return err
}
r.encoded = encoded
return nil
}
func whereClause(filters Filters) (string, []any) {
@@ -98,12 +114,12 @@ func whereClause(filters Filters) (string, []any) {
type scanner interface{ Scan(...any) error }
func scanPaper(row scanner) (Paper, error) {
func (r *PostgresRepository) scanPaper(row scanner) (Paper, error) {
p := Paper{Authors: []string{}, Topics: []string{}, Sources: []Source{}}
var id int64
var abstract, published sql.NullString
var authors string
if err := row.Scan(&id, &p.Title, &abstract, &authors, &published); err != nil {
if err := row.Scan(&id, r.text(&p.Title), r.text(&abstract), r.text(&authors), r.text(&published)); err != nil {
return p, err
}
p.ID = strconv.FormatInt(id, 10)
@@ -133,9 +149,14 @@ func decodeAuthors(raw string) ([]string, error) {
continue
}
var tuple []string
if json.Unmarshal(entry, &tuple) != nil || len(tuple) == 0 || len(tuple) > 3 {
if json.Unmarshal(entry, &tuple) != nil || len(tuple) == 0 || len(tuple) > 4 {
return nil, errors.New("unsupported stored author")
}
// arXiv may append an affiliation after family, given and suffix.
// It is source metadata, not part of the display name.
if len(tuple) == 4 {
tuple = tuple[:3]
}
if len(tuple) >= 2 {
tuple[0], tuple[1] = tuple[1], tuple[0]
}
@@ -159,6 +180,10 @@ func (r *PostgresRepository) List(ctx context.Context, filters Filters) (Page, e
return result, err
}
defer tx.Rollback()
if r.encoded {
filters.Query = encodeText(filters.Query)
filters.Topic = encodeText(filters.Topic)
}
where, args := whereClause(filters)
if err = tx.QueryRowContext(ctx, r.query(`SELECT count(*) FROM {schema}papers p`+where), args...).Scan(&result.Total); err != nil {
return result, err
@@ -168,14 +193,33 @@ func (r *PostgresRepository) List(ctx context.Context, filters Filters) (Page, e
order = "p.title_key,p.id"
}
args = append(args, filters.PageSize, (filters.Page-1)*filters.PageSize)
query := fmt.Sprintf(`SELECT p.id,p.title,p.abstract,p.authors_json,p.first_published
FROM {schema}papers p%s ORDER BY %s LIMIT $%d OFFSET $%d`, where, order, len(args)-1, len(args))
// Select IDs before loading potentially TOASTed titles/abstracts. With
// sparse categories, selecting full rows here can make the planner walk
// hundreds of thousands of wide heap rows before finding eight matches.
candidateQuery := `SELECT p.id FROM {schema}papers p`
candidateOrder := order
prefix := ""
if where != "" {
columns := "p.id"
if filters.Sort == "title" {
columns += ",p.title_key"
}
// Apply search before LIMIT/ORDER BY so a sparse term cannot turn
// into a backward heap scan that recomputes vectors for millions.
prefix = fmt.Sprintf(`matches AS MATERIALIZED (SELECT %s FROM {schema}papers p%s), `, columns, where)
candidateQuery = "SELECT id FROM matches"
candidateOrder = strings.ReplaceAll(order, "p.", "")
}
query := fmt.Sprintf(`WITH %spage_ids AS MATERIALIZED (
%s ORDER BY %s LIMIT $%d OFFSET $%d
) SELECT p.id,p.title,p.abstract,p.authors_json,p.first_published
FROM page_ids JOIN {schema}papers p ON p.id=page_ids.id ORDER BY %s`, prefix, candidateQuery, candidateOrder, len(args)-1, len(args), order)
rows, err := tx.QueryContext(ctx, r.query(query), args...)
if err != nil {
return result, err
}
for rows.Next() {
p, err := scanPaper(rows)
p, err := r.scanPaper(rows)
if err != nil {
_ = rows.Close()
return result, err
@@ -196,7 +240,7 @@ func (r *PostgresRepository) List(ctx context.Context, filters Filters) (Page, e
}
for rows.Next() {
var topic string
if err = rows.Scan(&topic); err != nil {
if err = rows.Scan(r.text(&topic)); err != nil {
_ = rows.Close()
return result, err
}
@@ -216,7 +260,7 @@ func (r *PostgresRepository) Get(ctx context.Context, id int64) (Paper, error) {
return Paper{}, err
}
defer tx.Rollback()
p, err := scanPaper(tx.QueryRowContext(ctx, r.query(`SELECT id,title,abstract,authors_json,first_published FROM {schema}papers WHERE id=$1`), id))
p, err := r.scanPaper(tx.QueryRowContext(ctx, r.query(`SELECT id,title,abstract,authors_json,first_published FROM {schema}papers WHERE id=$1`), id))
if err != nil {
return p, err
}
@@ -252,7 +296,7 @@ func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Pa
for rows.Next() {
var id int64
var topic string
if err = rows.Scan(&id, &topic); err != nil {
if err = rows.Scan(&id, r.text(&topic)); err != nil {
_ = rows.Close()
return err
}
@@ -273,7 +317,7 @@ func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Pa
var id int64
var source Source
var metadata string
if err = rows.Scan(&id, &source.Label, &source.ExternalID, &source.URL, &metadata); err != nil {
if err = rows.Scan(&id, r.text(&source.Label), r.text(&source.ExternalID), r.text(&source.URL), r.text(&metadata)); err != nil {
_ = rows.Close()
return err
}
@@ -309,7 +353,7 @@ func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Pa
for rows.Next() {
var id int64
var value string
if err = rows.Scan(&id, &value); err != nil {
if err = rows.Scan(&id, r.text(&value)); err != nil {
_ = rows.Close()
return err
}
@@ -333,7 +377,7 @@ func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Pa
for rows.Next() {
var id int64
var publication Publication
if err = rows.Scan(&id, &publication.Venue, &publication.Year, &publication.Track, &publication.Status, &publication.EvidenceURL); err != nil {
if err = rows.Scan(&id, r.text(&publication.Venue), &publication.Year, r.text(&publication.Track), r.text(&publication.Status), r.text(&publication.EvidenceURL)); err != nil {
_ = rows.Close()
return err
}