feat: Support S3 Vectors (#3330)

* feat: Support S3 Vectors

Bring Amazon S3 Vectors — AWS's native, cost-optimised vector store
for similarity search, RAG, and AI agents — to AWS SDK for pandas.
With this change, indexing a DataFrame and running an approximate-
nearest-neighbour query is a one-liner; on-the-fly embedding via
Amazon Bedrock is built in.

Implementation lives in a private subpackage `awswrangler/s3/_vectors/`
and is re-exported flat on `wr.s3.*`, matching the s3_tables convention.

Public surface (14 functions on `wr.s3`):
  - Buckets:  create_vector_bucket, delete_vector_bucket,
              list_vector_buckets, get_vector_bucket
  - Indexes:  create_vector_index, delete_vector_index,
              list_vector_indexes, get_vector_index
  - Data:     put_vectors, put_vectors_from_df, get_vectors,
              delete_vectors, list_vectors, query_vectors

Highlights:
  - End-to-end RAG: pass `text_column` + `bedrock_model_id` to
    put_vectors_from_df and awswrangler embeds each row via Bedrock
    (Titan / Cohere) and writes the resulting vectors plus all other
    columns as filterable metadata. query_vectors mirrors this with
    `query_text` / `query_vector`.
  - MongoDB-style metadata filters ($eq, $ne, $gt, $gte, $lt, $lte,
    $in, $nin, $exists, $and, $or) evaluated server-side during search.
  - Automatic chunking to AWS API limits (500/put, 100/get, 500/delete)
    and parallel-segment list_vectors (up to 16 segments).
  - Float32 coercion + non-finite rejection; NaN / pd.NA / None
    metadata cells dropped per row.

* test: add mocked unit tests and live integration tests for S3 Vectors

- tests/unit/test_s3_vectors_mocked.py — 46 tests using unittest.mock,
  no AWS required. Covers chunking, target resolution, float32 coercion,
  NaN/pd.NA metadata drop, parallel-segment list, query top-k bounds,
  Bedrock Titan/Cohere request and response shapes, and aliasing identity.
- tests/unit/test_s3_vectors.py — 10 live integration tests using new
  vector_bucket (session-scope) and vector_index (function-scope)
  fixtures in tests/conftest.py. Fixtures self-bootstrap via
  create_vector_bucket / create_vector_index; no CDK stack required.

* docs: api reference and tutorial for S3 Vectors

- docs/source/api.rst: new "Amazon S3 Vectors" section after
  "Amazon S3 Tables", listing all 14 public functions.
- tutorials/043 - Amazon S3 Vectors.ipynb: end-to-end walkthrough
  covering bucket/index lifecycle, discovery, Bedrock-embedded writes,
  semantic queries with metadata filters, per-key CRUD, bulk export,
  and cleanup.
- README.md: tutorial 043 entry.

* fix: type-parametrize np.ndarray for mypy on Python 3.10

* fix: use hyphen in vector index fixture name (underscore is invalid)

* fix: avoid zero-norm vectors in live integration tests (cosine rejects them)

* feat: add `chunked` parameter to list_vectors

Memory-friendly streaming for indexes too large to materialise in one
DataFrame. Mirrors the `chunked: bool | int` convention used by
`s3.read_parquet` and `athena.read_sql_query`:

  - chunked=False (default) — unchanged; returns a DataFrame and keeps
    the parallel-segment fan-out (up to 16 segments).
  - chunked=True — yields one DataFrame per underlying API page.
  - chunked=INTEGER — yields DataFrames of exactly N rows (final frame
    may be shorter).

Chunked streaming is single-segment and sequential by design;
`use_threads` is ignored in that mode, since lazy iteration across
parallel segments would require buffering and defeat the memory win.

Internally, `_list_segment` is now a thin materialiser around a new
`_iter_list_pages` generator, so the per-segment pagination logic
(including `max_items` enforcement across pages) is shared by both the
eager and streaming paths.

Tests: 4 new cases cover per-page yield, exact-size chunking, laziness
(no API call before the first `next()`), and `max_items` truncation of
the chunked stream.

Author: kukushking

README.md

@@ -153,6 +153,7 @@ Read our [docs](https://aws-sdk-pandas.readthedocs.io/en/3.16.1/scale.html) or h
   - [039 - Athena Iceberg](https://github.com/aws/aws-sdk-pandas/blob/main/tutorials/039%20-%20Athena%20Iceberg.ipynb)
   - [040 - EMR Serverless](https://github.com/aws/aws-sdk-pandas/blob/main/tutorials/040%20-%20EMR%20Serverless.ipynb)
   - [041 - Apache Spark on Amazon Athena](https://github.com/aws/aws-sdk-pandas/blob/main/tutorials/041%20-%20Apache%20Spark%20on%20Amazon%20Athena.ipynb)
+  - [043 - Amazon S3 Vectors](https://github.com/aws/aws-sdk-pandas/blob/main/tutorials/043%20-%20Amazon%20S3%20Vectors.ipynb)
 - [**API Reference**](https://aws-sdk-pandas.readthedocs.io/en/3.16.1/api.html)
   - [Amazon S3](https://aws-sdk-pandas.readthedocs.io/en/3.16.1/api.html#amazon-s3)
   - [AWS Glue Catalog](https://aws-sdk-pandas.readthedocs.io/en/3.16.1/api.html#aws-glue-catalog)

awswrangler/_utils.py

@@ -67,6 +67,7 @@
 
     ServiceName = Literal[
         "athena",
+        "bedrock-runtime",
         "cleanrooms",
         "dynamodb",
         "ec2",
@@ -85,6 +86,7 @@
         "secretsmanager",
         "sts",
         "s3tables",
+        "s3vectors",
         "timestream-query",
         "timestream-write",
     ]

awswrangler/s3/__init__.py

@@ -21,6 +21,22 @@
 )
 from awswrangler.s3._select import select_query
 from awswrangler.s3._upload import upload
+from awswrangler.s3._vectors import (
+    create_vector_bucket,
+    create_vector_index,
+    delete_vector_bucket,
+    delete_vector_index,
+    delete_vectors,
+    get_vector_bucket,
+    get_vector_index,
+    get_vectors,
+    list_vector_buckets,
+    list_vector_indexes,
+    list_vectors,
+    put_vectors,
+    put_vectors_from_df,
+    query_vectors,
+)
 from awswrangler.s3._wait import wait_objects_exist, wait_objects_not_exist
 from awswrangler.s3._write_deltalake import to_deltalake, to_deltalake_streaming
 from awswrangler.s3._write_excel import to_excel
@@ -71,4 +87,18 @@
     "delete_table",
     "from_iceberg",
     "to_iceberg",
+    "create_vector_bucket",
+    "delete_vector_bucket",
+    "list_vector_buckets",
+    "get_vector_bucket",
+    "create_vector_index",
+    "delete_vector_index",
+    "list_vector_indexes",
+    "get_vector_index",
+    "put_vectors",
+    "put_vectors_from_df",
+    "get_vectors",
+    "delete_vectors",
+    "list_vectors",
+    "query_vectors",
 ]

awswrangler/s3/_vectors/__init__.py

@@ -0,0 +1,43 @@
+"""Amazon S3 Vectors (PRIVATE subpackage).
+
+Public functions are re-exported through ``awswrangler.s3``. This module is private —
+its layout may change without notice.
+"""
+
+from awswrangler.s3._vectors._mgmt import (
+    create_vector_bucket,
+    create_vector_index,
+    delete_vector_bucket,
+    delete_vector_index,
+    get_vector_bucket,
+    get_vector_index,
+    list_vector_buckets,
+    list_vector_indexes,
+)
+from awswrangler.s3._vectors._read import (
+    get_vectors,
+    list_vectors,
+    query_vectors,
+)
+from awswrangler.s3._vectors._write import (
+    delete_vectors,
+    put_vectors,
+    put_vectors_from_df,
+)
+
+__all__ = [
+    "create_vector_bucket",
+    "delete_vector_bucket",
+    "list_vector_buckets",
+    "get_vector_bucket",
+    "create_vector_index",
+    "delete_vector_index",
+    "list_vector_indexes",
+    "get_vector_index",
+    "put_vectors",
+    "put_vectors_from_df",
+    "get_vectors",
+    "delete_vectors",
+    "list_vectors",
+    "query_vectors",
+]

awswrangler/s3/_vectors/_bedrock.py

@@ -0,0 +1,73 @@
+"""Amazon S3 Vectors - Bedrock embedding helper (PRIVATE)."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any, Callable
+
+import boto3
+
+from awswrangler import _utils, exceptions
+from awswrangler._executor import _get_executor
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+_logger: logging.Logger = logging.getLogger(__name__)
+
+
+def embed_texts(
+    texts: list[str],
+    model_id: str,
+    model_kwargs: dict[str, Any] | None = None,
+    use_threads: bool | int = True,
+    boto3_session: boto3.Session | None = None,
+) -> list[list[float]]:
+    """Embed a list of strings via Amazon Bedrock, optionally in parallel.
+
+    Supported model id prefixes: ``amazon.titan-embed-text-*``, ``cohere.embed-*``.
+    """
+    if not texts:
+        return []
+
+    extra = dict(model_kwargs or {})
+    build_body: Callable[[str], dict[str, Any]]
+    parse_response: Callable[[dict[str, Any]], list[float]]
+
+    if model_id.startswith("amazon.titan-embed-text"):
+
+        def build_body(text: str) -> dict[str, Any]:
+            return {"inputText": text, **extra}
+
+        def parse_response(payload: dict[str, Any]) -> list[float]:
+            return list(payload["embedding"])
+
+    elif model_id.startswith("cohere.embed"):
+        extra.setdefault("input_type", "search_document")
+
+        def build_body(text: str) -> dict[str, Any]:
+            return {"texts": [text], **extra}
+
+        def parse_response(payload: dict[str, Any]) -> list[float]:
+            return list(payload["embeddings"][0])
+
+    else:
+        raise exceptions.InvalidArgument(
+            f"Unsupported Bedrock embedding model_id '{model_id}'. "
+            "Pre-compute embeddings and pass them via `vector_column` instead. "
+            "Supported model id prefixes: 'amazon.titan-embed-text', 'cohere.embed'."
+        )
+
+    def embed_one(client: "BaseClient", text: str) -> list[float]:
+        response = client.invoke_model(  # type: ignore[attr-defined]
+            modelId=model_id,
+            accept="application/json",
+            contentType="application/json",
+            body=json.dumps(build_body(text)).encode("utf-8"),
+        )
+        return parse_response(json.loads(response["body"].read()))
+
+    client = _utils.client(service_name="bedrock-runtime", session=boto3_session)
+    executor = _get_executor(use_threads=use_threads)
+    return list(executor.map(embed_one, client, texts))