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")

但是，Dataset 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()

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

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,)

注意

如果您习惯于格式字符串风格 %s 语法（例如 psycopg 等库用于绑定参数），请注意，这不受支持 - 仅支持 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()

获取表的 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()

默认情况下，假设“active”目录/模式。

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 已集成到流行的 dataframe 库 pandas 中。 Pandas 可以使用 ADBC 与 PostgreSQL 和其他数据库交换数据。与使用 SQLAlchemy 或其他选项相比，将 ADBC 与 pandas 结合使用可以获得更好的性能，例如避免过度转换为 Python 对象和从 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 一起使用，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