This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1a83f11e7a8 Add `S3VectorsCreateIndexOperator` (#66121)
1a83f11e7a8 is described below
commit 1a83f11e7a8f768756bf6721c5ee3a0aad10e469
Author: John Jackson <[email protected]>
AuthorDate: Mon May 4 06:48:21 2026 -0700
Add `S3VectorsCreateIndexOperator` (#66121)
---
providers/amazon/docs/operators/s3_vectors.rst | 24 ++++--
.../providers/amazon/aws/operators/s3_vectors.py | 83 ++++++++++++++++++++
.../tests/system/amazon/aws/example_s3_vectors.py | 27 ++++++-
.../unit/amazon/aws/operators/test_s3_vectors.py | 89 ++++++++++++++++++++++
4 files changed, 217 insertions(+), 6 deletions(-)
diff --git a/providers/amazon/docs/operators/s3_vectors.rst
b/providers/amazon/docs/operators/s3_vectors.rst
index 88db11ce778..00ef5f88ab6 100644
--- a/providers/amazon/docs/operators/s3_vectors.rst
+++ b/providers/amazon/docs/operators/s3_vectors.rst
@@ -25,7 +25,7 @@ Use the operators below to manage S3 Vectors resources.
.. _howto/operator:S3VectorsCreateVectorBucketOperator:
Create a Vector Bucket
-~~~~~~~~~~~~~~~~~~~~~~
+----------------------
To create an Amazon S3 Vectors vector bucket, use
:class:`~airflow.providers.amazon.aws.operators.s3_vectors.S3VectorsCreateVectorBucketOperator`.
@@ -36,15 +36,24 @@ To create an Amazon S3 Vectors vector bucket, use
:start-after: [START howto_operator_s3vectors_create_vector_bucket]
:end-before: [END howto_operator_s3vectors_create_vector_bucket]
-Reference
-~~~~~~~~~
+.. _howto/operator:S3VectorsCreateIndexOperator:
-* `AWS boto3 Library Documentation for S3 Vectors
<https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3vectors.html>`__
+Create an Index
+---------------
+
+To create an index in an Amazon S3 Vectors vector bucket, use
+:class:`~airflow.providers.amazon.aws.operators.s3_vectors.S3VectorsCreateIndexOperator`.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3_vectors.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_s3vectors_create_index]
+ :end-before: [END howto_operator_s3vectors_create_index]
.. _howto/operator:S3VectorsDeleteVectorBucketOperator:
Delete a Vector Bucket
-~~~~~~~~~~~~~~~~~~~~~~
+----------------------
To delete an Amazon S3 Vectors vector bucket, use
:class:`~airflow.providers.amazon.aws.operators.s3_vectors.S3VectorsDeleteVectorBucketOperator`.
@@ -54,3 +63,8 @@ To delete an Amazon S3 Vectors vector bucket, use
:dedent: 4
:start-after: [START howto_operator_s3vectors_delete_vector_bucket]
:end-before: [END howto_operator_s3vectors_delete_vector_bucket]
+
+Reference
+---------
+
+* `AWS boto3 Library Documentation for S3 Vectors
<https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3vectors.html>`__
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py
index 7708c56a58f..c1990ad65f1 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py
@@ -126,3 +126,86 @@ class
S3VectorsDeleteVectorBucketOperator(AwsBaseOperator[AwsBaseHook]):
def execute(self, context: Context) -> None:
self.hook.conn.delete_vector_bucket(vectorBucketName=self.vector_bucket_name)
self.log.info("Deleted vector bucket %s", self.vector_bucket_name)
+
+
+class S3VectorsCreateIndexOperator(AwsBaseOperator[AwsBaseHook]):
+ """
+ Create an index in an Amazon S3 Vectors vector bucket.
+
+ An index stores vectors and supports similarity search queries.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:S3VectorsCreateIndexOperator`
+
+ :param vector_bucket_name: The name of the vector bucket. (templated)
+ :param index_name: The name of the index to create. (templated)
+ :param data_type: The data type for vectors (e.g. ``float32``). (templated)
+ :param dimension: The number of dimensions for each vector.
+ :param distance_metric: The distance metric for similarity search (e.g.
``cosine``, ``euclidean``).
+ :param metadata_configuration: Optional metadata configuration dict.
+ :param if_exists: Behavior when the index already exists.
+ ``"fail"`` raises an error, ``"skip"`` returns the existing index ARN.
+ """
+
+ aws_hook_class = AwsBaseHook
+ template_fields: tuple[str, ...] = (
+ *AwsBaseOperator.template_fields,
+ "vector_bucket_name",
+ "index_name",
+ "data_type",
+ "distance_metric",
+ )
+ template_fields_renderers = {"metadata_configuration": "json"}
+
+ def __init__(
+ self,
+ *,
+ vector_bucket_name: str,
+ index_name: str,
+ data_type: str,
+ dimension: int,
+ distance_metric: str = "cosine",
+ metadata_configuration: dict[str, Any] | None = None,
+ if_exists: Literal["fail", "skip"] = "skip",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.vector_bucket_name = vector_bucket_name
+ self.index_name = index_name
+ self.data_type = data_type
+ self.dimension = dimension
+ self.distance_metric = distance_metric
+ self.metadata_configuration = metadata_configuration
+ self.if_exists = if_exists
+
+ @property
+ def _hook_parameters(self) -> dict[str, Any]:
+ return {**super()._hook_parameters, "client_type": "s3vectors"}
+
+ def execute(self, context: Context) -> str:
+ self.log.info("Creating index %s in vector bucket %s",
self.index_name, self.vector_bucket_name)
+ kwargs: dict[str, Any] = prune_dict(
+ {
+ "vectorBucketName": self.vector_bucket_name,
+ "indexName": self.index_name,
+ "dataType": self.data_type,
+ "dimension": self.dimension,
+ "distanceMetric": self.distance_metric,
+ "metadataConfiguration": self.metadata_configuration,
+ }
+ )
+ try:
+ response = self.hook.conn.create_index(**kwargs)
+ index_arn = response["indexArn"]
+ except ClientError as e:
+ if e.response["Error"]["Code"] == "ConflictException" and
self.if_exists == "skip":
+ self.log.info("Index %s already exists, skipping.",
self.index_name)
+ response = self.hook.conn.get_index(
+ vectorBucketName=self.vector_bucket_name,
indexName=self.index_name
+ )
+ index_arn = response["index"]["indexArn"]
+ else:
+ raise
+ self.log.info("Index %s: %s", self.index_name, index_arn)
+ return index_arn
diff --git a/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py
b/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py
index 454ae72e995..54e5e5c2599 100644
--- a/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py
+++ b/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py
@@ -19,6 +19,7 @@ from __future__ import annotations
from datetime import datetime
from airflow.providers.amazon.aws.operators.s3_vectors import (
+ S3VectorsCreateIndexOperator,
S3VectorsCreateVectorBucketOperator,
S3VectorsDeleteVectorBucketOperator,
)
@@ -28,14 +29,24 @@ from system.amazon.aws.utils import ENV_ID_KEY,
SystemTestContextBuilder
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
if AIRFLOW_V_3_0_PLUS:
- from airflow.sdk import TriggerRule
+ from airflow.sdk import TriggerRule, task
else:
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
from airflow.utils.trigger_rule import TriggerRule # type:
ignore[no-redef,attr-defined]
DAG_ID = "example_s3_vectors"
sys_test_context_task = SystemTestContextBuilder().build()
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_index(vector_bucket_name: str, index_name: str):
+ """Delete the index."""
+ import boto3
+
+
boto3.client("s3vectors").delete_index(vectorBucketName=vector_bucket_name,
indexName=index_name)
+
+
with DAG(
dag_id=DAG_ID,
schedule=None,
@@ -45,6 +56,7 @@ with DAG(
test_context = sys_test_context_task()
env_id = test_context[ENV_ID_KEY]
bucket_name = f"{env_id}-test-vectors"
+ index_name = f"{env_id}-test-index"
# [START howto_operator_s3vectors_create_vector_bucket]
create_vector_bucket = S3VectorsCreateVectorBucketOperator(
@@ -53,6 +65,17 @@ with DAG(
)
# [END howto_operator_s3vectors_create_vector_bucket]
+ # [START howto_operator_s3vectors_create_index]
+ create_index = S3VectorsCreateIndexOperator(
+ task_id="create_index",
+ vector_bucket_name=bucket_name,
+ index_name=index_name,
+ data_type="float32",
+ dimension=128,
+ distance_metric="cosine",
+ )
+ # [END howto_operator_s3vectors_create_index]
+
# [START howto_operator_s3vectors_delete_vector_bucket]
delete_vector_bucket = S3VectorsDeleteVectorBucketOperator(
task_id="delete_vector_bucket",
@@ -64,6 +87,8 @@ with DAG(
chain(
test_context,
create_vector_bucket,
+ create_index,
+ delete_index(vector_bucket_name=bucket_name, index_name=index_name),
delete_vector_bucket,
)
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py
index f155bb11367..109bec8303d 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py
@@ -18,9 +18,11 @@ from __future__ import annotations
from unittest.mock import MagicMock
+import pytest
from botocore.exceptions import ClientError
from airflow.providers.amazon.aws.operators.s3_vectors import (
+ S3VectorsCreateIndexOperator,
S3VectorsCreateVectorBucketOperator,
S3VectorsDeleteVectorBucketOperator,
)
@@ -112,3 +114,90 @@ class TestS3VectorsDeleteVectorBucketOperator:
vector_bucket_name=BUCKET_NAME,
)
validate_template_fields(op)
+
+
+INDEX_NAME = "test-index"
+INDEX_ARN =
"arn:aws:s3vectors:us-east-1:123456789012:vector-bucket/test-bucket/index/test-index"
+
+
+class TestS3VectorsCreateIndexOperator:
+ def setup_method(self):
+ self.operator = S3VectorsCreateIndexOperator(
+ task_id="create_index",
+ vector_bucket_name=BUCKET_NAME,
+ index_name=INDEX_NAME,
+ data_type="float32",
+ dimension=128,
+ )
+
+ def test_execute(self):
+ mock_conn = MagicMock()
+ mock_conn.create_index.return_value = {"indexArn": INDEX_ARN}
+ self.operator.hook.conn = mock_conn
+
+ result = self.operator.execute({})
+
+ mock_conn.create_index.assert_called_once_with(
+ vectorBucketName=BUCKET_NAME,
+ indexName=INDEX_NAME,
+ dataType="float32",
+ dimension=128,
+ distanceMetric="cosine",
+ )
+ assert result == INDEX_ARN
+
+ def test_execute_with_metadata_config(self):
+ meta_config = {"nonFilterableMetadataKeys": ["key1"]}
+ op = S3VectorsCreateIndexOperator(
+ task_id="create_index",
+ vector_bucket_name=BUCKET_NAME,
+ index_name=INDEX_NAME,
+ data_type="float32",
+ dimension=128,
+ distance_metric="euclidean",
+ metadata_configuration=meta_config,
+ )
+ mock_conn = MagicMock()
+ mock_conn.create_index.return_value = {"indexArn": INDEX_ARN}
+ op.hook.conn = mock_conn
+
+ result = op.execute({})
+
+ call_kwargs = mock_conn.create_index.call_args[1]
+ assert call_kwargs["distanceMetric"] == "euclidean"
+ assert call_kwargs["metadataConfiguration"] == meta_config
+ assert result == INDEX_ARN
+
+ def test_execute_skip_existing(self):
+ mock_conn = MagicMock()
+ mock_conn.create_index.side_effect = ClientError(
+ {"Error": {"Code": "ConflictException", "Message": "Already
exists"}},
+ "CreateIndex",
+ )
+ mock_conn.get_index.return_value = {"index": {"indexArn": INDEX_ARN}}
+ self.operator.hook.conn = mock_conn
+
+ result = self.operator.execute({})
+ assert result == INDEX_ARN
+
+ def test_execute_fail_on_conflict(self):
+ op = S3VectorsCreateIndexOperator(
+ task_id="create_index",
+ vector_bucket_name=BUCKET_NAME,
+ index_name=INDEX_NAME,
+ data_type="float32",
+ dimension=128,
+ if_exists="fail",
+ )
+ mock_conn = MagicMock()
+ mock_conn.create_index.side_effect = ClientError(
+ {"Error": {"Code": "ConflictException", "Message": "Already
exists"}},
+ "CreateIndex",
+ )
+ op.hook.conn = mock_conn
+
+ with pytest.raises(ClientError):
+ op.execute({})
+
+ def test_template_fields(self):
+ validate_template_fields(self.operator)
