pgai：将维基百科导入 PostgreSQL

飞象数据 2025-06-09

121

原文地址：

https://www.cybertec-postgresql.com/en/pgai-importing-wikipedia-into-postgresql/

在处理人工智能时，关键问题之一是：在哪里可以找到用于实验的数据？虽然互联网上有很多地方可以满足对测试数据的需求，但维基百科无疑是最有趣的入门选择之一。那么，我们如何将维基百科导入 PostgreSQL 并使用它呢？以下是它的工作原理。

使用 pgai 加载维基百科

迄今为止最有用的扩展之一是pgai，它是 PostgreSQL 中一个强大的人工智能工具箱。幸运的是，它还包含一种将维基百科加载到数据库的简便方法。

https://github.com/timescale/pgai

一旦您在系统上安装了 pgai 或启动了 Github 上所示的现成容器，您就可以在数据库中启用该扩展：

cybertec=# CREATE EXTENSION IF NOT EXISTS ai CASCADE;
NOTICE:  installing required extension "vector"
NOTICE:  installing required extension "plpython3u"
CREATE EXTENSION

我们在这里要做的就是调用 ai.load_dataset 函数，它有大量我们可以利用的有用参数：

cybertec=# \x
Expanded display is on.
cybertec=# \df+ ai.load_dataset
List of functions
-[ RECORD 1 ]-------+-------------------------------------------------
Schema              | ai
Name                | load_dataset
Result data type    | bigint
Argument data types | name text, config_name text DEFAULT NULL::text, 
		      split text DEFAULT NULL::text, 
		      schema_name name DEFAULT 'public'::name, 
                      table_name name DEFAULT NULL::name, 
                      if_table_exists text DEFAULT 'error'::text, 
                      field_types jsonb DEFAULT NULL::jsonb, 
 		      integer DEFAULT 5000, 
                      max_batches integer DEFAULT NULL::integer, 
		      kwargs jsonb DEFAULT '{}'::jsonb
Type                | func
Volatility          | volatile
Parallel            | unsafe
Owner               | hs
Security            | invoker
Access privileges   | hs=X*/hs +
                    | pg_database_owner=X*/hs
Language            | plpython3u
Internal name       | 
Description         |

第一个参数告诉函数要加载哪个数据集。在本例中，我们直接加载维基百科及其相关版本。前两个参数非常重要。此外，我们还可以定义要将数据加载到哪个目标表，或者当要使用的表已存在时该如何操作。另外：通常情况下，我们不需要大量的行，因此我们基本上可以定义要加载多少数据。显然，获取所有这些数据需要相当长的时间，而且您可能并不总是需要一个庞大的数据集来使用维基百科和 PostgreSQL。

那么，让我们尝试一下，看看会发生什么。以下代码片段包含一个已导入部分数据的示例调用：

cybertec=#  SELECT ai.load_dataset(
'wikimedia/wikipedia', '20231101.en', 
		table_name=>'wiki', batch_size=>1000, 
		max_batches=>100000, if_table_exists=>'append'
);

如果此命令成功执行，我们就可以使用新创建的表了：

cybertec=# SELECT count(*) FROM wiki;
  count  
---------
 6407814
(1 row)


cybertec=# \d wiki
              Table "public.wiki"
 Column | Type | Collation | Nullable | Default 
--------+------+-----------+----------+---------
 id     | text |           | not null | 
 url    | text |           |          | 
 title  | text |           |          | 
 text   | text |           |          | 
Indexes:
    "wiki_pkey" PRIMARY KEY, btree (id)

如您所见，我们有大约 640 万行，包含 URL、标题和网站文本。

现在，所有这些数据对于各种用途来说都很有趣。然而，我们远未完成。如果你想参与机器学习，下一步就是使用一种叫做“嵌入”的技术，pgai 也很好地处理了这个问题。

嵌入文本以进行语义搜索

正如我们之前关于语义搜索的文章中所讨论的，将这些文本转换为向量非常重要，这可以通过著名且备受推崇的 pgvector 扩展来理解。pgai 为我们提供了一种很好的方法。嵌入数百万个文档并非几秒钟就能完成，因此 pgai 的做法是分块运行。

https://www.cybertec-postgresql.com/semantic-search-in-postgresql-an-overview/

为了做到这一点，我们必须运行一个命令：

cybertec=# SELECT ai.create_vectorizer(
     'wiki'::regclass,
     destination => 'demo_wiki_emb',
     embedding => ai.embedding_ollama('all-minilm', 384),
     chunking => ai.chunking_recursive_character_text_splitter('text')
);
 create_vectorizer 
-------------------
                 6
(1 row)


Time: 24135.802 ms (00:24.136)

我们所做的就是调用 ai.create_vectorizer 函数。它被要求嵌入“wiki”表的内容。向量数据应该可以通过名为“demo_wiki_emb”的视图访问。下一步是告诉系统使用哪个模型进行嵌入。在我们的例子中，我使用了一个名为“all-minilm”的相当小的模型。但是，有很多类似的模型可以调用：

更多信息请访问Ollama 网站。只需确定最适合您的方案，然后在 Ollama 服务器上本地运行即可。

https://ollama.com/blog/embedding-models

接下来我们要做的就是跟踪嵌入过程的进度。具体操作如下：

cybertec=# select * from ai.vectorizer_status;
 id | source_table |         target_table         |          view          |    pending_items    | disabled 
----+--------------+------------------------------+------------------------+---------------------+----------
  6 | public.wiki  | public.wiki_embeddings_store | public.wiki_embeddings | 9223372036854775807 | f
(1 row)




cybertec=# select ai.vectorizer_queue_pending(6, exact_count=>true);
 vectorizer_queue_pending 
--------------------------
                  6407814
(1 row)

有两种方法可以实际检查 PostgreSQL 内部的情况。我个人更喜欢精确版本，但两种方法都可以。

然而，这里存在一个问题：什么也没发生。没有任何进展。原因是我们必须启动一个 Python 程序来真正完成大部分繁重的工作。最简单的方法如下：

$ /usr/local/bin/pgai vectorizer worker -d "dbname=cybertec host=localhost"
2025-03-21 20:16:04 [info     ] running vectorizer             vectorizer_id=6

这是启动该过程最简单的方法，我们已经看到了一些进展。对于少量数据来说，这应该没问题，但运行庞大的数据集可能需要很长时间——真的很长时间。

解决这个问题的一种方法是并行执行：

$ /usr/local/bin/pgai vectorizer worker --help
Usage: pgai vectorizer worker [OPTIONS]


Options:
  --version                       Show the version and exit.
  -d, --db-url TEXT               The database URL to connect to  [default:
                                  postgres://postgres@localhost:5432/postgres]
  -i, --vectorizer-id INTEGER     Only fetch work from the given vectorizer
                                  ids. If not provided, all vectorizers will
                                  be fetched.
  --log-level [DEBUG|INFO|WARN|ERROR|FATAL|CRITICAL]
  --poll-interval TIME DURATION   The interval, in duration string or integer
                                  (seconds), to wait before checking for new
                                  work after processing all available work in
                                  the queue.  [default: 5m]
  --once                          Exit after processing all available work
                                  (implies --exit-on-error).
  -c, --concurrency INTEGER RANGE
                                  [x>=1]
  --exit-on-error BOOLEAN         Exit immediately when an error occurs.
  --help                          Show this message and exit.

问题的解决方案在于“-c”选项，它为场景增加了一些并行性。如果你想更快地运行此类操作，这绝对很有帮助。

一段时间后我们的队列应该是空的：

cybertec=# select ai.vectorizer_queue_pending(6, exact_count=>true);
 vectorizer_queue_pending 
--------------------------
                        0
(1 row)

检查结果

数据导入后，我们可以检查结果：

cybertec=# \d+
                                         List of relations
 Schema |        Name         | Type  | Owner | Persistence | Access method |  Size   | Description 
--------+---------------------+-------+-------+-------------+---------------+---------+-------------
 public | demo_wiki_emb       | view  | hs    | permanent   |               | 0 bytes | 
 public | demo_wiki_emb_store | table | hs    | permanent   | heap          | 92 GB   | 
 public | wiki                | table | hs    | permanent   | heap          | 13 GB   | 
(3 rows)

我们在这里看到，pgai 生成了一个表和一个视图，让我们可以方便地读取数据（=原始数据 + 向量）。我们还可以看到，我们的模型生成的数据量非常庞大。

进程返回的数据可以在该表中找到（如之前的函数调用所定义）：

cybertec=# \d demo_wiki_emb_store
                   Table "public.demo_wiki_emb_store"
     Column     |    Type     | Collation | Nullable |      Default      
----------------+-------------+-----------+----------+-------------------
 embedding_uuid | uuid        |           | not null | gen_random_uuid()
 id             | text        |           | not null | 
 chunk_seq      | integer     |           | not null | 
 chunk          | text        |           | not null | 
 embedding      | vector(384) |           | not null | 
Indexes:
    "demo_wiki_emb_store_pkey" PRIMARY KEY, btree (embedding_uuid)
    "demo_wiki_emb_store_id_chunk_seq_key" UNIQUE CONSTRAINT, btree (id, chunk_seq)
Foreign-key constraints:
    "demo_wiki_emb_store_id_fkey" FOREIGN KEY (id) REFERENCES wiki(id) ON DELETE CASCADE

我们可以看到，该表引用了 wiki 表。为了更轻松地访问数据，pgai 创建了一个易于使用的视图：

cybertec=# \d+ demo_wiki_emb
                              View "public.demo_wiki_emb"
     Column     |    Type     | Collation | Nullable | Default | Storage  | Description 
----------------+-------------+-----------+----------+---------+----------+-------------
 embedding_uuid | uuid        |           |          |         | plain    | 
 chunk_seq      | integer     |           |          |         | plain    | 
 chunk          | text        |           |          |         | extended | 
 embedding      | vector(384) |           |          |         | external | 
 id             | text        |           |          |         | extended | 
 url            | text        |           |          |         | extended | 
 title          | text        |           |          |         | extended | 
 text           | text        |           |          |         | extended | 
View definition:
 SELECT t.embedding_uuid,
    t.chunk_seq,
    t.chunk,
    t.embedding,
    t.id,
    s.url,
    s.title,
    s.text
   FROM demo_wiki_emb_store t
     LEFT JOIN wiki s ON t.id = s.id;

它的作用是将原始数据与这些嵌入连接起来，并将所有内容显示为单一视图，这使得交互变得非常简单。但让我在这里重点关注一些其他细节：

cybertec=# explain SELECT count(*) FROM demo_wiki_emb;
                                                QUERY PLAN
---------------------------------------------------------------------------------------------------------------
 Finalize Aggregate  (cost=1105977.00..1105977.01 rows=1 width=8)
   ->  Gather  (cost=1105976.78..1105976.99 rows=2 width=8)
         Workers Planned: 2
         ->  Partial Aggregate  (cost=1104976.78..1104976.79 rows=1 width=8)
               ->  Parallel Index Only Scan using demo_wiki_emb_store_id_chunk_seq_key on demo_wiki_emb_store t  
			(cost=0.56..1062313.83 rows=17065183 width=0)
(5 rows)

我们在这里看到了什么？在前面的清单中，视图显示了一个连接。但为什么在这个执行计划中我们看不到连接呢？答案是：“连接修剪”。优化器已经确定它可以安全地从此操作中移除“wiki”表，而不会带来任何风险。如果您想了解更多信息，我们之前的一篇博客文章解释了这个概念。

https://www.cybertec-postgresql.com/en/join-pruning-cool-stuff-in-postgresql/

进一步阅读

在本文中，您学习了如何快速将维基百科数据导入 PostgreSQL，以及如何将这些维基百科数据转换为向量。接下来的文章将在此基础上，展示我们有哪些索引选项、如何快速索引以及如何实际利用这些数据。

postgresql

文章转载自飞象数据，如果涉嫌侵权，请发送邮件至：contact@modb.pro进行举报，并提供相关证据，一经查实，墨天轮将立刻删除相关内容。

pgai：将维基百科导入 PostgreSQL

使用 pgai 加载维基百科

嵌入文本以进行语义搜索

检查结果

进一步阅读

评论