PostgreSQL 指南¶

使用用户名和密码进行身份验证¶

要连接到 PostgreSQL 数据库，必须在 URI 中提供用户名和密码。例如，

postgresql://username:password@hostname:port/dbname

有关详细信息，请参阅 PostgreSQL 文档。

import os

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

with conn.cursor() as cur:
    cur.execute("SELECT 1")
    print(cur.fetchone())
    # Output: (1,)

conn.close()

标准输出¶

(1,)

从 Arrow 数据集创建/追加到表¶

指南来源：postgresql_create_dataset_table.py

ADBC 可以轻松地将 PyArrow 数据集加载到您的数据存储中。

import os
import tempfile
from pathlib import Path

import pyarrow
import pyarrow.csv
import pyarrow.dataset
import pyarrow.feather
import pyarrow.parquet

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

为了测试目的，我们首先确保将要使用的表不存在。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS csvtable")
    cur.execute("DROP TABLE IF EXISTS ipctable")
    cur.execute("DROP TABLE IF EXISTS pqtable")
    cur.execute("DROP TABLE IF EXISTS csvdataset")
    cur.execute("DROP TABLE IF EXISTS ipcdataset")
    cur.execute("DROP TABLE IF EXISTS pqdataset")

conn.commit()

生成示例数据¶

tempdir = tempfile.TemporaryDirectory(
    prefix="adbc-docs-",
    ignore_cleanup_errors=True,
)
root = Path(tempdir.name)
table = pyarrow.table(
    [
        [1, 1, 2],
        ["foo", "bar", "baz"],
    ],
    names=["ints", "strs"],
)

首先我们将写入单个文件。

csv_file = root / "example.csv"
pyarrow.csv.write_csv(table, csv_file)

ipc_file = root / "example.arrow"
pyarrow.feather.write_feather(table, ipc_file)

parquet_file = root / "example.parquet"
pyarrow.parquet.write_table(table, parquet_file)

我们还将生成一些分区数据集。

csv_dataset = root / "csv_dataset"
pyarrow.dataset.write_dataset(
    table,
    csv_dataset,
    format="csv",
    partitioning=["ints"],
)

ipc_dataset = root / "ipc_dataset"
pyarrow.dataset.write_dataset(
    table,
    ipc_dataset,
    format="feather",
    partitioning=["ints"],
)

parquet_dataset = root / "parquet_dataset"
pyarrow.dataset.write_dataset(
    table,
    parquet_dataset,
    format="parquet",
    partitioning=["ints"],
)

将 CSV 文件加载到 PostgreSQL¶

我们可以直接将 pyarrow.RecordBatchReader（来自 open_csv）传递给 adbc_ingest。我们还可以传递 pyarrow.dataset.Dataset 或 pyarrow.dataset.Scanner。

with conn.cursor() as cur:
    reader = pyarrow.csv.open_csv(csv_file)
    cur.adbc_ingest("csvtable", reader, mode="create")

    reader = pyarrow.dataset.dataset(
        csv_dataset,
        format="csv",
        partitioning=["ints"],
    )
    cur.adbc_ingest("csvdataset", reader, mode="create")

conn.commit()

with conn.cursor() as cur:
    cur.execute("SELECT ints, strs FROM csvtable ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

    cur.execute("SELECT ints, strs FROM csvdataset ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

将 Arrow IPC (Feather) 文件加载到 PostgreSQL¶

with conn.cursor() as cur:
    reader = pyarrow.ipc.RecordBatchFileReader(ipc_file)

由于 PyArrow API 的怪癖，我们必须将文件读入内存。

    cur.adbc_ingest("ipctable", reader.read_all(), mode="create")

不过，数据集 API 会将数据流式传输到内存中，然后传输到 PostgreSQL。

    reader = pyarrow.dataset.dataset(
        ipc_dataset,
        format="feather",
        partitioning=["ints"],
    )
    cur.adbc_ingest("ipcdataset", reader, mode="create")

conn.commit()

with conn.cursor() as cur:
    cur.execute("SELECT ints, strs FROM ipctable ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

    cur.execute("SELECT ints, strs FROM ipcdataset ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

将 Parquet 文件加载到 PostgreSQL¶

with conn.cursor() as cur:
    reader = pyarrow.parquet.ParquetFile(parquet_file)
    cur.adbc_ingest("pqtable", reader.iter_batches(), mode="create")

    reader = pyarrow.dataset.dataset(
        parquet_dataset,
        format="parquet",
        partitioning=["ints"],
    )
    cur.adbc_ingest("pqdataset", reader, mode="create")

conn.commit()

with conn.cursor() as cur:
    cur.execute("SELECT ints, strs FROM pqtable ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

    cur.execute("SELECT ints, strs FROM pqdataset ORDER BY ints, strs ASC")
    assert cur.fetchall() == [(1, "bar"), (1, "foo"), (2, "baz")]

清理¶

conn.close()
tempdir.cleanup()

从 Arrow 表创建/追加到表¶

指南来源：postgresql_create_append_table.py

ADBC 允许使用 Arrow 表创建和追加到数据库表。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

为了测试目的，我们首先确保将要使用的表不存在。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("DROP TABLE IF EXISTS example2")

现在我们可以创建表了。

with conn.cursor() as cur:
    data = pyarrow.table(
        [
            [1, 2, None, 4],
        ],
        schema=pyarrow.schema(
            [
                ("ints", "int32"),
            ]
        ),
    )
    cur.adbc_ingest("example", data, mode="create")

conn.commit()

摄取后，我们可以获取结果。

with conn.cursor() as cur:
    cur.execute("SELECT * FROM example")
    assert cur.fetchone() == (1,)
    assert cur.fetchone() == (2,)

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (4,)

如果我们尝试再次摄取，它会失败，因为表已经存在。

with conn.cursor() as cur:
    try:
        cur.adbc_ingest("example", data, mode="create")
    except conn.ProgrammingError:
        pass
    else:
        raise RuntimeError("Should have failed!")

conn.rollback()

相反，我们可以追加到表中。

with conn.cursor() as cur:
    cur.adbc_ingest("example", data, mode="append")

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (8,)

我们还可以选择在表不存在时创建表，否则追加。

with conn.cursor() as cur:
    cur.adbc_ingest("example2", data, mode="create_append")

    cur.execute("SELECT COUNT(*) FROM example2")
    assert cur.fetchone() == (4,)

    cur.adbc_ingest("example2", data, mode="create_append")

    cur.execute("SELECT COUNT(*) FROM example2")
    assert cur.fetchone() == (8,)

最后，我们可以替换表。

with conn.cursor() as cur:
    cur.adbc_ingest("example", data.slice(0, 2), mode="replace")

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (2,)

conn.close()

创建/追加到临时表¶

指南来源：postgresql_create_temp_table.py

ADBC 也允许创建和追加到临时表。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

为了测试目的，我们首先确保将要使用的表不存在。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")

要创建临时表，只需指定选项“temporary”。

data = pyarrow.table(
    [
        [1, 2, None, 4],
    ],
    schema=pyarrow.schema(
        [
            ("ints", "int32"),
        ]
    ),
)

with conn.cursor() as cur:
    cur.adbc_ingest("example", data, mode="create", temporary=True)

conn.commit()

摄取后，我们可以获取结果。

with conn.cursor() as cur:
    cur.execute("SELECT * FROM example")
    assert cur.fetchone() == (1,)
    assert cur.fetchone() == (2,)

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (4,)

临时表与常规表是分开的，即使它们具有相同的名称。

with conn.cursor() as cur:
    cur.adbc_ingest("example", data.slice(0, 2), mode="create", temporary=False)

conn.commit()

with conn.cursor() as cur:

由于我们有两个同名表，我们必须在此处明确引用普通的临时表。

    cur.execute("SELECT COUNT(*) FROM public.example")
    assert cur.fetchone() == (2,)

    cur.execute("SELECT COUNT(*) FROM example")
    assert cur.fetchone() == (4,)

conn.close()

关闭连接后，临时表被隐式删除。如果重新连接，表将不存在；我们将只看到“普通”表。

with adbc_driver_postgresql.dbapi.connect(uri) as conn:
    with conn.cursor() as cur:
        cur.execute("SELECT COUNT(*) FROM example")
        assert cur.fetchone() == (2,)

所有常规摄取选项也适用于临时表。更多示例请参阅从 Arrow 数据集创建/追加到表。

执行带绑定参数的语句¶

指南来源：postgresql_execute_bind.py

ADBC 允许使用 Python 和 Arrow 值作为绑定参数。目前，PostgreSQL 驱动程序仅支持不生成结果集的查询的绑定参数。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将创建一个示例表进行测试。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("CREATE TABLE example (ints INT, bigints BIGINT)")

conn.commit()

我们可以绑定 Python 值

with conn.cursor() as cur:
    cur.executemany("INSERT INTO example VALUES ($1, $2)", [(1, 2), (3, 4)])

    cur.execute("SELECT SUM(ints) FROM example")
    assert cur.fetchone() == (4,)

注意

如果您习惯于像 psycopg 等库用于绑定参数的格式字符串样式 %s 语法，请注意，此语法不支持——只支持 PostgreSQL 原生的 $1 语法。

我们还可以绑定 Arrow 值

with conn.cursor() as cur:
    data = pyarrow.record_batch(
        [
            [5, 6],
            [7, 8],
        ],
        names=["$1", "$2"],
    )
    cur.executemany("INSERT INTO example VALUES ($1, $2)", data)

    cur.execute("SELECT SUM(ints) FROM example")
    assert cur.fetchone() == (15,)

conn.close()

执行不带 COPY 的语句¶

指南来源：postgresql_execute_nocopy.py

ADBC 驱动程序默认尝试使用 COPY 执行查询，因为它对于大型结果集更快。但是，PostgreSQL 不支持所有类型的查询的 COPY。例如，SHOW 查询将不起作用。在这种情况下，您可以显式禁用 COPY 优化。

import os

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

可以在创建游标时设置该选项

with conn.cursor(
    adbc_stmt_kwargs={
        adbc_driver_postgresql.StatementOptions.USE_COPY.value: False,
    }
) as cur:
    cur.execute("SHOW ALL")
    print(cur.fetch_arrow_table().schema)

或者可以在之后设置

with conn.cursor() as cur:
    cur.adbc_statement.set_options(
        **{
            adbc_driver_postgresql.StatementOptions.USE_COPY.value: False,
        }
    )
    cur.execute("SHOW ALL")
    print(cur.fetch_arrow_table().schema)

如果没有该选项，查询将失败，因为驱动程序尝试使用 COPY 执行查询

with conn.cursor() as cur:
    try:
        cur.execute("SHOW ALL")
    except conn.Error:
        pass
    else:
        raise RuntimeError("Expected error")

conn.close()

stdout¶

name: string
setting: string
description: string
name: string
setting: string
description: string

获取表的 Arrow 模式¶

指南来源：postgresql_get_table_schema.py

ADBC 允许您将表的模式作为 Arrow 模式获取。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将创建一些示例表进行测试。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("CREATE TABLE example (ints INT, bigints BIGINT)")

    cur.execute("CREATE SCHEMA IF NOT EXISTS other_schema")
    cur.execute("DROP TABLE IF EXISTS other_schema.example")
    cur.execute("CREATE TABLE other_schema.example (strings TEXT, values INT)")

conn.commit()

默认情况下，假定“活动”目录/模式。

assert conn.adbc_get_table_schema("example") == pyarrow.schema(
    [
        ("ints", "int32"),
        ("bigints", "int64"),
    ]
)

我们可以显式指定 PostgreSQL 模式，以获取不同命名空间中表的 Arrow 模式。

注意

在 PostgreSQL 中，您只能查询连接到的数据库（目录）。因此我们不能在此处指定目录（或者说，这样做没有意义）。

请注意，NUMERIC 列被读取为字符串，因为 PostgreSQL 小数不映射到 Arrow 小数。

assert conn.adbc_get_table_schema(
    "example",
    db_schema_filter="other_schema",
) == pyarrow.schema(
    [
        ("strings", "string"),
        ("values", "int32"),
    ]
)

conn.close()

获取查询的 Arrow 模式¶

指南来源：postgresql_get_query_schema.py

ADBC 允许您获取结果集的模式，而无需执行查询。

import os

import pyarrow

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将创建一个示例表进行测试。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("CREATE TABLE example (ints INT, bigints BIGINT)")

conn.commit()

expected = pyarrow.schema(
    [
        ("ints", "int32"),
        ("bigints", "int64"),
    ]
)

with conn.cursor() as cur:
    assert cur.adbc_execute_schema("SELECT * FROM example") == expected

PostgreSQL 在这里不知道类型，所以它只是返回一个猜测。

    assert cur.adbc_execute_schema("SELECT $1 AS res") == pyarrow.schema(
        [
            ("res", "string"),
        ]
    )

conn.close()

列出目录、模式和表¶

指南来源：postgresql_list_catalogs.py

ADBC 允许列出数据库中的表、目录和模式。

import os

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将创建一个示例表来查找。

with conn.cursor() as cur:
    cur.execute("DROP TABLE IF EXISTS example")
    cur.execute("CREATE TABLE example (ints INT, bigints BIGINT)")

conn.commit()

数据以 PyArrow RecordBatchReader 的形式给出。

objects = conn.adbc_get_objects(depth="all").read_all()

为了方便起见，我们将其转换为纯 Python 数据。

objects = objects.to_pylist()
catalog = objects[0]
assert catalog["catalog_name"] == "postgres"

db_schema = catalog["catalog_db_schemas"][0]
assert db_schema["db_schema_name"] == "public"

tables = db_schema["db_schema_tables"]
example = [table for table in tables if table["table_name"] == "example"]
assert len(example) == 1
example = example[0]

assert example["table_columns"][0]["column_name"] == "ints"
assert example["table_columns"][1]["column_name"] == "bigints"

conn.close()

使用 SQLAlchemy 进行连接池¶

指南来源：postgresql_pool.py

ADBC 不实现连接池，因为这通常不是 DBAPI 驱动程序的特性。相反，请使用第三方连接池，例如 SQLAlchemy 中内置的连接池。

import os

import sqlalchemy.pool

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]

source = adbc_driver_postgresql.dbapi.connect(uri)

adbc_driver_manager.dbapi.Connection.adbc_clone() 从现有连接打开新连接，尽可能共享内部资源。例如，PostgreSQL 驱动程序将共享内部 OID 缓存，从而节省一些连接开销。

pool = sqlalchemy.pool.QueuePool(source.adbc_clone, max_overflow=1, pool_size=2)

现在我们可以从连接池中获取连接；SQLAlchemy 覆盖了 close() 以将连接返回到连接池。

注意

与底层的 ADBC 连接不同，SQLAlchemy 的包装器不支持上下文管理器协议。

conn = pool.connect()

assert pool.checkedin() == 0
assert pool.checkedout() == 1

with conn.cursor() as cur:
    cur.execute("SELECT 1")
    assert cur.fetchone() == (1,)

conn.close()

assert pool.checkedin() == 1
assert pool.checkedout() == 0

source.close()

使用 Pandas 和 ADBC¶

指南来源：postgresql_pandas.py

ADBC 集成到 pandas 中，这是一个流行的 DataFrame 库。Pandas 可以使用 ADBC 与 PostgreSQL 和其他数据库交换数据。与使用 SQLAlchemy 或其他选项相比，将 ADBC 与 pandas 结合使用可以获得更好的性能，例如避免过多地转换为和从 Python 对象。

import os

import pandas as pd

import adbc_driver_postgresql.dbapi

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]
conn = adbc_driver_postgresql.dbapi.connect(uri)

我们将使用 pd.DataFrame.to_sql 创建一个示例表。

data = pd.DataFrame(
    {
        "ints": [1, 2, None, 4],
        "strs": ["a", "b", "c", "d"],
    }
)
data.to_sql("example", conn, if_exists="replace")
conn.commit()

创建表后，我们可以将 ADBC 连接和 SQL 查询传递给 pd.read_sql，以将结果集作为 pandas DataFrame 获取。

df = pd.read_sql("SELECT * FROM example WHERE ints > 1", conn)

assert len(df) == 2

conn.close()

与 ADBC 接口相比，pandas 提供了更方便、更高级别的 API，特别是对于那些已经使用 pandas 的人。

使用 Polars 和 ADBC¶

指南来源：postgresql_polars.py

ADBC 可以与 Polars 配合使用，这是一个用 Rust 编写的 DataFrame 库。根据其文档

如果后端支持直接返回 Arrow 数据，则将使用此功能高效地实例化 DataFrame；否则，DataFrame 将从行式数据初始化。

显然，ADBC 直接返回 Arrow 数据，这使得 ADBC 和 Polars 自然而然地相互匹配。

import os

import polars as pl

uri = os.environ["ADBC_POSTGRESQL_TEST_URI"]

我们将使用 Polars 创建一个示例表，使用 polars.DataFrame.write_database()。我们不需要自己用 Polars 打开 ADBC 连接。

data = pl.DataFrame(
    {
        "ints": [1, 2, None, 4],
        "strs": ["a", "b", "c", "d"],
    }
)
data.write_database("example", uri, engine="adbc", if_table_exists="replace")

创建表后，我们可以使用 polars.read_database_uri() 获取结果。同样，我们只需传递 URI 并告诉 Polars 为我们管理 ADBC。

df = pl.read_database_uri("SELECT * FROM example WHERE ints > 1", uri, engine="adbc")

assert len(df) == 2