From f4de5fd316d640e148634724232c0b9baf9d0bfe Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 28 Nov 2023 17:22:04 +0530 Subject: [PATCH 01/34] Add CRUD operations arounf schema and class objects --- airflow/providers/weaviate/hooks/weaviate.py | 5 +++++ tests/providers/weaviate/hooks/test_weaviate.py | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 6c9b1b67872b5..db0d47eaabbdb 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -25,11 +25,15 @@ import requests from tenacity import Retrying, retry, retry_if_exception, retry_if_exception_type, stop_after_attempt + from weaviate import Client as WeaviateClient +from weaviate import UnexpectedStatusCodeException from weaviate.auth import AuthApiKey, AuthBearerToken, AuthClientCredentials, AuthClientPassword + from weaviate.exceptions import ObjectAlreadyExistsException from weaviate.util import generate_uuid5 + from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.hooks.base import BaseHook @@ -157,6 +161,7 @@ def test_connection(self) -> tuple[bool, str]: self.log.error("Error testing Weaviate connection: %s", e) return False, str(e) + @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.exceptions.RequestException)) def create_class(self, class_json: dict[str, Any]) -> None: """Create a new class.""" client = self.conn diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index fc6d7db6aed8f..968cb75674c15 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -24,8 +24,13 @@ import pytest import requests import weaviate + from weaviate import ObjectAlreadyExistsException +import requests +import weaviate + + from airflow.models import Connection from airflow.providers.weaviate.hooks.weaviate import WeaviateHook From 4ae25d0e46de7bcef5524172b33147dbfa48b8b1 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Thu, 30 Nov 2023 09:49:18 +0530 Subject: [PATCH 02/34] Handle casees when the properties are not in same order --- airflow/providers/weaviate/hooks/weaviate.py | 1 + tests/providers/weaviate/hooks/test_weaviate.py | 1 + 2 files changed, 2 insertions(+) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index db0d47eaabbdb..3216313ac6301 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -33,6 +33,7 @@ from weaviate.exceptions import ObjectAlreadyExistsException from weaviate.util import generate_uuid5 +from tenacity import Retrying, retry, retry_if_exception_type, stop_after_attempt from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.hooks.base import BaseHook diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 968cb75674c15..4bcc19d5b60dd 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -675,3 +675,4 @@ def test_contains_schema(get_schema, classes_to_test, expected_result, weaviate_ ] } assert weaviate_hook.check_subset_of_schema(classes_to_test) == expected_result + From 6f1ca74a1016a6a72b95f26d52825a181f39a56c Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Thu, 30 Nov 2023 09:53:56 +0530 Subject: [PATCH 03/34] Change the methods name and docstring --- tests/providers/weaviate/hooks/test_weaviate.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 4bcc19d5b60dd..968cb75674c15 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -675,4 +675,3 @@ def test_contains_schema(get_schema, classes_to_test, expected_result, weaviate_ ] } assert weaviate_hook.check_subset_of_schema(classes_to_test) == expected_result - From 9f94fea4354a3c8ebe42fb56d8c11f7430ed5ead Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Thu, 30 Nov 2023 19:02:37 +0530 Subject: [PATCH 04/34] Resolve conflicts --- airflow/providers/weaviate/hooks/weaviate.py | 1 + .../providers/weaviate/operators/weaviate.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 3216313ac6301..f3a653df8a4f5 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -23,6 +23,7 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, Dict, List, cast +import pandas as pd import requests from tenacity import Retrying, retry, retry_if_exception, retry_if_exception_type, stop_after_attempt diff --git a/airflow/providers/weaviate/operators/weaviate.py b/airflow/providers/weaviate/operators/weaviate.py index 4e07a59edba99..496a85f0a3397 100644 --- a/airflow/providers/weaviate/operators/weaviate.py +++ b/airflow/providers/weaviate/operators/weaviate.py @@ -44,10 +44,16 @@ class WeaviateIngestOperator(BaseOperator): :param conn_id: The Weaviate connection. :param class_name: The Weaviate class to be used for storing the data objects into. +<<<<<<< HEAD :param input_data: The list of dicts or pandas dataframe representing Weaviate data objects to generate embeddings on (or provides custom vectors) and store them in the Weaviate class. :param input_json: (Deprecated) The JSON representing Weaviate data objects to generate embeddings on (or provides custom vectors) and store them in the Weaviate class. +======= + :param input_data: The list of dicts or pandas dataframe representing Weaviate data objects to generate\ + embeddings on (or provides custom vectors) and store them in the Weaviate class. Either input_json + or input_callable should be provided. +>>>>>>> e9508bb3b4 (Resolve conflicts) :param vector_col: key/column name in which the vectors are stored. """ @@ -57,8 +63,12 @@ def __init__( self, conn_id: str, class_name: str, +<<<<<<< HEAD input_json: list[dict[str, Any]] | pd.DataFrame | None = None, input_data: list[dict[str, Any]] | pd.DataFrame | None = None, +======= + input_data: list[dict[str, Any]] | pd.DataFrame, +>>>>>>> e9508bb3b4 (Resolve conflicts) vector_col: str = "Vector", **kwargs: Any, ) -> None: @@ -68,6 +78,7 @@ def __init__( super().__init__(**kwargs) self.class_name = class_name self.conn_id = conn_id +<<<<<<< HEAD self.vector_col = vector_col if input_data is not None: @@ -81,6 +92,10 @@ def __init__( self.input_data = input_json else: raise TypeError("Either input_json or input_data is required") +======= + self.input_data = input_data + self.vector_col = vector_col +>>>>>>> e9508bb3b4 (Resolve conflicts) @cached_property def hook(self) -> WeaviateHook: @@ -90,8 +105,12 @@ def hook(self) -> WeaviateHook: def execute(self, context: Context) -> None: self.log.debug("Input data: %s", self.input_data) self.hook.batch_data( +<<<<<<< HEAD self.class_name, self.input_data, **self.batch_params, vector_col=self.vector_col, +======= + self.class_name, self.input_data, vector_col=self.vector_col, **self.batch_params +>>>>>>> e9508bb3b4 (Resolve conflicts) ) From c766814c274d145156d9dd64bd69b9b0335bd8e6 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 1 Dec 2023 18:12:20 +0530 Subject: [PATCH 05/34] Make sure the retrying is working as expected --- airflow/providers/weaviate/hooks/weaviate.py | 5 +++++ tests/providers/weaviate/hooks/test_weaviate.py | 1 + 2 files changed, 6 insertions(+) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index f3a653df8a4f5..c14d8ca33ee9f 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -29,6 +29,7 @@ from weaviate import Client as WeaviateClient from weaviate import UnexpectedStatusCodeException + from weaviate.auth import AuthApiKey, AuthBearerToken, AuthClientCredentials, AuthClientPassword from weaviate.exceptions import ObjectAlreadyExistsException @@ -386,6 +387,10 @@ def check_subset_of_schema(self, classes_objects: list) -> bool: return False return True + @staticmethod + def check_http_error_is_retryable(exc: BaseException): + return isinstance(exc, requests.HTTPError) and not exc.response.ok + def batch_data( self, class_name: str, diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 968cb75674c15..1194bd4d88695 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -29,6 +29,7 @@ import requests import weaviate +from requests import HTTPError from airflow.models import Connection From 7a02b8ef45f140c7c7c732a171eb10cb5deda12d Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 5 Dec 2023 18:50:08 +0530 Subject: [PATCH 06/34] Address PR comments --- airflow/providers/weaviate/hooks/weaviate.py | 68 +++------------- .../providers/weaviate/hooks/test_weaviate.py | 81 +++++++------------ 2 files changed, 40 insertions(+), 109 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index c14d8ca33ee9f..06d43bcf02551 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -23,13 +23,8 @@ from functools import cached_property from typing import TYPE_CHECKING, Any, Dict, List, cast -import pandas as pd import requests -from tenacity import Retrying, retry, retry_if_exception, retry_if_exception_type, stop_after_attempt - from weaviate import Client as WeaviateClient -from weaviate import UnexpectedStatusCodeException - from weaviate.auth import AuthApiKey, AuthBearerToken, AuthClientCredentials, AuthClientPassword from weaviate.exceptions import ObjectAlreadyExistsException @@ -41,7 +36,7 @@ from airflow.hooks.base import BaseHook if TYPE_CHECKING: - from typing import Literal, Sequence + from typing import Sequence, Literal import pandas as pd from weaviate import ConsistencyLevel @@ -49,23 +44,6 @@ ExitingSchemaOptions = Literal["replace", "fail", "ignore"] -HTTP_RETRY_STATUS_CODE = [429, 500, 503, 504] -REQUESTS_EXCEPTIONS_TYPES = ( - requests.RequestException, - requests.exceptions.ConnectionError, - requests.exceptions.HTTPError, - requests.exceptions.ConnectTimeout, -) - - -def check_http_error_is_retryable(exc: BaseException): - return ( - isinstance(exc, requests.exceptions.RequestException) - and exc.response - and exc.response.status_code - and exc.response.status_code in HTTP_RETRY_STATUS_CODE - ) - class WeaviateHook(BaseHook): """ @@ -79,13 +57,7 @@ class WeaviateHook(BaseHook): conn_type = "weaviate" hook_name = "Weaviate" - def __init__( - self, - conn_id: str = default_conn_name, - retry_status_codes: list[int] | None = None, - *args: Any, - **kwargs: Any, - ) -> None: + def __init__(self, conn_id: str = default_conn_name, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.conn_id = conn_id @@ -170,14 +142,7 @@ def create_class(self, class_json: dict[str, Any]) -> None: client = self.conn client.schema.create_class(class_json) - @retry( - reraise=True, - stop=stop_after_attempt(3), - retry=( - retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) - | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) - ), - ) + @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.exceptions.RequestException)) def create_schema(self, schema_json: dict[str, Any] | str) -> None: """ Create a new Schema. @@ -203,14 +168,7 @@ def _convert_dataframe_to_list(data: list[dict[str, Any]] | pd.DataFrame) -> lis data = json.loads(data.to_json(orient="records")) return cast(List[Dict[str, Any]], data) - @retry( - reraise=True, - stop=stop_after_attempt(3), - retry=( - retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) - | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) - ), - ) + @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.exceptions.RequestException)) def get_schema(self, class_name: str | None = None): """Get the schema from Weaviate. @@ -237,13 +195,9 @@ def delete_classes(self, class_names: list[str] | str, if_error: str = "stop") - try: for attempt in Retrying( stop=stop_after_attempt(3), - retry=( - retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) - | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) - ), + retry=retry_if_exception_type(requests.exceptions.RequestException), ): with attempt: - print(attempt) client.schema.delete_class(class_name) except Exception as e: if if_error == "continue": @@ -261,16 +215,15 @@ def delete_all_schema(self): client = self.get_client() return client.schema.delete_all() + @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.ConnectionError)) def update_config(self, class_name: str, config: dict): """Update a schema configuration for a specific class.""" client = self.get_client() client.schema.update_config(class_name=class_name, config=config) - def create_or_replace_classes( - self, schema_json: dict[str, Any] | str, existing: ExitingSchemaOptions = "ignore" - ): + def upsert_classes(self, schema_json: dict[str, Any] | str, existing: ExitingSchemaOptions = "ignore"): """ - Create or replace the classes in schema of Weaviate database. + Create or update the classes in schema of Weaviate database. :param schema_json: Json containing the schema. Format {"class_name": "class_dict"} .. seealso:: `example of class_dict `_. @@ -419,10 +372,7 @@ def batch_data( for index, data_obj in enumerate(data): for attempt in Retrying( stop=stop_after_attempt(retry_attempts_per_object), - retry=( - retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) - | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) - ), + retry=retry_if_exception_type(requests.exceptions.RequestException), ): with attempt: self.log.debug( diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 1194bd4d88695..fce7bc1735eb8 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -17,14 +17,11 @@ from __future__ import annotations from contextlib import ExitStack -from unittest import mock -from unittest.mock import MagicMock, Mock + +from unittest.mock import MagicMock, Mock, patch, call import pandas as pd import pytest -import requests -import weaviate - from weaviate import ObjectAlreadyExistsException import requests @@ -46,32 +43,32 @@ def weaviate_hook(): mock_conn = Mock() # Patch the WeaviateHook get_connection method to return the mock connection - with mock.patch.object(WeaviateHook, "get_connection", return_value=mock_conn): + with patch.object(WeaviateHook, "get_connection", return_value=mock_conn): hook = WeaviateHook(conn_id=TEST_CONN_ID) return hook @pytest.fixture def mock_auth_api_key(): - with mock.patch("airflow.providers.weaviate.hooks.weaviate.AuthApiKey") as m: + with patch("airflow.providers.weaviate.hooks.weaviate.AuthApiKey") as m: yield m @pytest.fixture def mock_auth_bearer_token(): - with mock.patch("airflow.providers.weaviate.hooks.weaviate.AuthBearerToken") as m: + with patch("airflow.providers.weaviate.hooks.weaviate.AuthBearerToken") as m: yield m @pytest.fixture def mock_auth_client_credentials(): - with mock.patch("airflow.providers.weaviate.hooks.weaviate.AuthClientCredentials") as m: + with patch("airflow.providers.weaviate.hooks.weaviate.AuthClientCredentials") as m: yield m @pytest.fixture def mock_auth_client_password(): - with mock.patch("airflow.providers.weaviate.hooks.weaviate.AuthClientPassword") as m: + with patch("airflow.providers.weaviate.hooks.weaviate.AuthClientPassword") as m: yield m @@ -132,7 +129,7 @@ def setup_method(self, monkeypatch): for conn in conns: monkeypatch.setenv(f"AIRFLOW_CONN_{conn.conn_id.upper()}", conn.get_uri()) - @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_api_key_in_extra(self, mock_client, mock_auth_api_key): hook = WeaviateHook(conn_id=self.weaviate_api_key1) hook.get_conn() @@ -141,7 +138,7 @@ def test_get_conn_with_api_key_in_extra(self, mock_client, mock_auth_api_key): url=self.host, auth_client_secret=mock_auth_api_key(api_key=self.api_key), additional_headers={} ) - @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_token_in_extra(self, mock_client, mock_auth_api_key): # when token is passed in extra hook = WeaviateHook(conn_id=self.weaviate_api_key2) @@ -151,7 +148,7 @@ def test_get_conn_with_token_in_extra(self, mock_client, mock_auth_api_key): url=self.host, auth_client_secret=mock_auth_api_key(api_key=self.api_key), additional_headers={} ) - @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_access_token_in_extra(self, mock_client, mock_auth_bearer_token): hook = WeaviateHook(conn_id=self.client_bearer_token) hook.get_conn() @@ -166,7 +163,7 @@ def test_get_conn_with_access_token_in_extra(self, mock_client, mock_auth_bearer additional_headers={}, ) - @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_client_secret_in_extra(self, mock_client, mock_auth_client_credentials): hook = WeaviateHook(conn_id=self.weaviate_client_credentials) hook.get_conn() @@ -179,7 +176,7 @@ def test_get_conn_with_client_secret_in_extra(self, mock_client, mock_auth_clien additional_headers={}, ) - @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_client_password_in_extra(self, mock_client, mock_auth_client_password): hook = WeaviateHook(conn_id=self.client_password) hook.get_conn() @@ -190,7 +187,7 @@ def test_get_conn_with_client_password_in_extra(self, mock_client, mock_auth_cli additional_headers={}, ) - @mock.patch("airflow.providers.weaviate.hooks.weaviate.generate_uuid5") + @patch("airflow.providers.weaviate.hooks.weaviate.generate_uuid5") def test_create_object(self, mock_gen_uuid, weaviate_hook): """ Test the create_object method of WeaviateHook. @@ -236,7 +233,7 @@ def test_get_of_get_or_create_object(self, weaviate_hook): tenant=None, ) - @mock.patch("airflow.providers.weaviate.hooks.weaviate.generate_uuid5") + @patch("airflow.providers.weaviate.hooks.weaviate.generate_uuid5") def test_create_of_get_or_create_object(self, mock_gen_uuid, weaviate_hook): """ Test the create part of get_or_create_object method of WeaviateHook. @@ -286,8 +283,8 @@ def test_get_all_objects(self, weaviate_hook): return_value = weaviate_hook.get_all_objects(class_name="TestClass") assert weaviate_hook.get_object.call_args_list == [ - mock.call(after=None, class_name="TestClass"), - mock.call(after=3, class_name="TestClass"), + call(after=None, class_name="TestClass"), + call(after=3, class_name="TestClass"), ] assert return_value == [{"name": "Test1", "id": 2}, {"name": "Test2", "id": 3}] @@ -307,8 +304,8 @@ def test_get_all_objects_returns_dataframe(self, weaviate_hook): return_value = weaviate_hook.get_all_objects(class_name="TestClass", as_dataframe=True) assert weaviate_hook.get_object.call_args_list == [ - mock.call(after=None, class_name="TestClass"), - mock.call(after=3, class_name="TestClass"), + call(after=None, class_name="TestClass"), + call(after=3, class_name="TestClass"), ] import pandas @@ -442,7 +439,7 @@ def test_batch_data(data, expected_length, weaviate_hook): assert mock_batch_context.add_data_object.call_count == expected_length -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_conn") +@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_conn") def test_batch_data_retry(get_conn, weaviate_hook): """Test to ensure retrying working as expected""" data = [{"name": "chandler"}, {"name": "joey"}, {"name": "ross"}] @@ -466,9 +463,9 @@ def test_batch_data_retry(get_conn, weaviate_hook): ], ids=["ignore", "replace", "fail", "invalid_option"], ) -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.delete_classes") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.create_schema") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") +@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.delete_classes") +@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.create_schema") +@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") def test_upsert_schema_scenarios( get_schema, create_schema, delete_classes, get_schema_value, existing, expected_value, weaviate_hook ): @@ -481,16 +478,16 @@ def test_upsert_schema_scenarios( if existing in ["fail", "invalid_option"]: stack.enter_context(pytest.raises(ValueError)) get_schema.return_value = get_schema_value - weaviate_hook.create_or_replace_classes(schema_json=schema_json, existing=existing) + weaviate_hook.upsert_classes(schema_json=schema_json, existing=existing) create_schema.assert_called_once_with({"classes": expected_value}) if existing == "replace": delete_classes.assert_called_once_with(class_names=["B"]) -@mock.patch("builtins.open") -@mock.patch("json.load") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.create_schema") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") +@patch("builtins.open") +@patch("json.load") +@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.create_schema") +@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") def test_upsert_schema_json_file_param(get_schema, create_schema, load, open, weaviate_hook): """Test if schema_json is path to a json file""" get_schema.return_value = {"classes": [{"class": "A"}, {"class": "B"}]} @@ -498,12 +495,12 @@ def test_upsert_schema_json_file_param(get_schema, create_schema, load, open, we "B": {"class": "B"}, "C": {"class": "C"}, } - weaviate_hook.create_or_replace_classes(schema_json="/tmp/some_temp_file.json", existing="ignore") + weaviate_hook.upsert_classes(schema_json="/tmp/some_temp_file.json", existing="ignore") create_schema.assert_called_once_with({"classes": [{"class": "C"}]}) -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_client") -def test_delete_classes(get_client, weaviate_hook): +@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_client") +def test_delete_schema(get_client, weaviate_hook): class_names = ["class_a", "class_b"] get_client.return_value.schema.delete_class.side_effect = [ weaviate.UnexpectedStatusCodeException("something failed", requests.Response()), @@ -519,22 +516,6 @@ def test_delete_classes(get_client, weaviate_hook): weaviate_hook.delete_classes("class_a", if_error="stop") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_client") -def test_http_errors_of_delete_classes(get_client, weaviate_hook): - class_names = ["class_a", "class_b"] - resp = requests.Response() - resp.status_code = 429 - get_client.return_value.schema.delete_class.side_effect = [ - requests.exceptions.HTTPError(response=resp), - None, - requests.exceptions.ConnectionError, - None, - ] - error_list = weaviate_hook.delete_classes(class_names, if_error="continue") - assert error_list == [] - assert get_client.return_value.schema.delete_class.call_count == 4 - - @pytest.mark.parametrize( argnames=["classes_to_test", "expected_result"], argvalues=[ @@ -633,7 +614,7 @@ def test_http_errors_of_delete_classes(get_client, weaviate_hook): "swapped_properties", ), ) -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") +@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") def test_contains_schema(get_schema, classes_to_test, expected_result, weaviate_hook): get_schema.return_value = { "classes": [ From 9d367bfb2b5714c029f78e89082c1d1596ee1fd4 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 5 Dec 2023 20:53:47 +0530 Subject: [PATCH 07/34] Remove retry logic from --- tests/providers/weaviate/hooks/test_weaviate.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index fce7bc1735eb8..00541841669a0 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -26,7 +26,6 @@ import requests import weaviate -from requests import HTTPError from airflow.models import Connection From d547a22ca9a39a0fadf573ae9126131726a7249f Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 5 Dec 2023 21:41:27 +0530 Subject: [PATCH 08/34] Remove vector_col params and dataframe support --- airflow/providers/weaviate/hooks/weaviate.py | 2 +- .../providers/weaviate/operators/weaviate.py | 21 +------------------ .../operators/weaviate.rst | 2 +- 3 files changed, 3 insertions(+), 22 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 06d43bcf02551..0db59142ebd86 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -347,7 +347,7 @@ def check_http_error_is_retryable(exc: BaseException): def batch_data( self, class_name: str, - data: list[dict[str, Any]] | pd.DataFrame, + data: list[dict[str, Any]], batch_config_params: dict[str, Any] | None = None, vector_col: str = "Vector", retry_attempts_per_object: int = 5, diff --git a/airflow/providers/weaviate/operators/weaviate.py b/airflow/providers/weaviate/operators/weaviate.py index 496a85f0a3397..9d28790ac0733 100644 --- a/airflow/providers/weaviate/operators/weaviate.py +++ b/airflow/providers/weaviate/operators/weaviate.py @@ -27,7 +27,6 @@ if TYPE_CHECKING: import pandas as pd - from airflow.utils.context import Context @@ -44,16 +43,10 @@ class WeaviateIngestOperator(BaseOperator): :param conn_id: The Weaviate connection. :param class_name: The Weaviate class to be used for storing the data objects into. -<<<<<<< HEAD :param input_data: The list of dicts or pandas dataframe representing Weaviate data objects to generate embeddings on (or provides custom vectors) and store them in the Weaviate class. :param input_json: (Deprecated) The JSON representing Weaviate data objects to generate embeddings on (or provides custom vectors) and store them in the Weaviate class. -======= - :param input_data: The list of dicts or pandas dataframe representing Weaviate data objects to generate\ - embeddings on (or provides custom vectors) and store them in the Weaviate class. Either input_json - or input_callable should be provided. ->>>>>>> e9508bb3b4 (Resolve conflicts) :param vector_col: key/column name in which the vectors are stored. """ @@ -63,12 +56,8 @@ def __init__( self, conn_id: str, class_name: str, -<<<<<<< HEAD input_json: list[dict[str, Any]] | pd.DataFrame | None = None, input_data: list[dict[str, Any]] | pd.DataFrame | None = None, -======= - input_data: list[dict[str, Any]] | pd.DataFrame, ->>>>>>> e9508bb3b4 (Resolve conflicts) vector_col: str = "Vector", **kwargs: Any, ) -> None: @@ -78,7 +67,7 @@ def __init__( super().__init__(**kwargs) self.class_name = class_name self.conn_id = conn_id -<<<<<<< HEAD + self.vector_col = vector_col if input_data is not None: @@ -92,10 +81,6 @@ def __init__( self.input_data = input_json else: raise TypeError("Either input_json or input_data is required") -======= - self.input_data = input_data - self.vector_col = vector_col ->>>>>>> e9508bb3b4 (Resolve conflicts) @cached_property def hook(self) -> WeaviateHook: @@ -105,12 +90,8 @@ def hook(self) -> WeaviateHook: def execute(self, context: Context) -> None: self.log.debug("Input data: %s", self.input_data) self.hook.batch_data( -<<<<<<< HEAD self.class_name, self.input_data, **self.batch_params, vector_col=self.vector_col, -======= - self.class_name, self.input_data, vector_col=self.vector_col, **self.batch_params ->>>>>>> e9508bb3b4 (Resolve conflicts) ) diff --git a/docs/apache-airflow-providers-weaviate/operators/weaviate.rst b/docs/apache-airflow-providers-weaviate/operators/weaviate.rst index 5ec262ab7a2d3..05063376f4d95 100644 --- a/docs/apache-airflow-providers-weaviate/operators/weaviate.rst +++ b/docs/apache-airflow-providers-weaviate/operators/weaviate.rst @@ -28,7 +28,7 @@ into the database. Using the Operator ^^^^^^^^^^^^^^^^^^ -The WeaviateIngestOperator requires the ``input_data`` as an input to the operator. Use the ``conn_id`` parameter to specify the Weaviate connection to use to +The WeaviateIngestOperator requires the ``input_text`` as an input to the operator. Use the ``conn_id`` parameter to specify the Weaviate connection to use to connect to your account. An example using the operator to ingest data with custom vectors retrieved from XCOM: From ed9501cd80463c87caf96d5ae0454a7d1434030e Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Thu, 7 Dec 2023 17:55:10 +0530 Subject: [PATCH 09/34] Remove unwanted retry logic --- airflow/providers/weaviate/hooks/weaviate.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 0db59142ebd86..680b9c29db4dc 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -136,7 +136,6 @@ def test_connection(self) -> tuple[bool, str]: self.log.error("Error testing Weaviate connection: %s", e) return False, str(e) - @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.exceptions.RequestException)) def create_class(self, class_json: dict[str, Any]) -> None: """Create a new class.""" client = self.conn @@ -215,7 +214,6 @@ def delete_all_schema(self): client = self.get_client() return client.schema.delete_all() - @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.ConnectionError)) def update_config(self, class_name: str, config: dict): """Update a schema configuration for a specific class.""" client = self.get_client() From 410ee8e635f20b05ca58ab16bbd508a65d4221ca Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 8 Dec 2023 16:03:47 +0530 Subject: [PATCH 10/34] Address PR comments --- airflow/providers/weaviate/hooks/weaviate.py | 9 +++++---- tests/providers/weaviate/hooks/test_weaviate.py | 11 ++++------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 680b9c29db4dc..af7957edb13a5 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -24,13 +24,12 @@ from typing import TYPE_CHECKING, Any, Dict, List, cast import requests +from tenacity import Retrying, retry, retry_if_exception_type, stop_after_attempt from weaviate import Client as WeaviateClient from weaviate.auth import AuthApiKey, AuthBearerToken, AuthClientCredentials, AuthClientPassword - from weaviate.exceptions import ObjectAlreadyExistsException from weaviate.util import generate_uuid5 -from tenacity import Retrying, retry, retry_if_exception_type, stop_after_attempt from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.hooks.base import BaseHook @@ -219,9 +218,11 @@ def update_config(self, class_name: str, config: dict): client = self.get_client() client.schema.update_config(class_name=class_name, config=config) - def upsert_classes(self, schema_json: dict[str, Any] | str, existing: ExitingSchemaOptions = "ignore"): + def create_or_replace_classes( + self, schema_json: dict[str, Any] | str, existing: ExitingSchemaOptions = "ignore" + ): """ - Create or update the classes in schema of Weaviate database. + Create or replace the classes in schema of Weaviate database. :param schema_json: Json containing the schema. Format {"class_name": "class_dict"} .. seealso:: `example of class_dict `_. diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 00541841669a0..0f16c3548cde5 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -17,16 +17,13 @@ from __future__ import annotations from contextlib import ExitStack - -from unittest.mock import MagicMock, Mock, patch, call +from unittest.mock import MagicMock, Mock, call, patch import pandas as pd import pytest -from weaviate import ObjectAlreadyExistsException - import requests import weaviate - +from weaviate import ObjectAlreadyExistsException from airflow.models import Connection from airflow.providers.weaviate.hooks.weaviate import WeaviateHook @@ -477,7 +474,7 @@ def test_upsert_schema_scenarios( if existing in ["fail", "invalid_option"]: stack.enter_context(pytest.raises(ValueError)) get_schema.return_value = get_schema_value - weaviate_hook.upsert_classes(schema_json=schema_json, existing=existing) + weaviate_hook.create_or_replace_classes(schema_json=schema_json, existing=existing) create_schema.assert_called_once_with({"classes": expected_value}) if existing == "replace": delete_classes.assert_called_once_with(class_names=["B"]) @@ -494,7 +491,7 @@ def test_upsert_schema_json_file_param(get_schema, create_schema, load, open, we "B": {"class": "B"}, "C": {"class": "C"}, } - weaviate_hook.upsert_classes(schema_json="/tmp/some_temp_file.json", existing="ignore") + weaviate_hook.create_or_replace_classes(schema_json="/tmp/some_temp_file.json", existing="ignore") create_schema.assert_called_once_with({"classes": [{"class": "C"}]}) From c656fc88ea27bf9f8923f5a26f252c21f8e6ba83 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 8 Dec 2023 17:53:29 +0530 Subject: [PATCH 11/34] Resolve ruff-lint issues --- airflow/providers/weaviate/hooks/weaviate.py | 15 +++++++++++---- airflow/providers/weaviate/operators/weaviate.py | 1 + 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index af7957edb13a5..5ab22872b711f 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -30,12 +30,11 @@ from weaviate.exceptions import ObjectAlreadyExistsException from weaviate.util import generate_uuid5 - from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.hooks.base import BaseHook if TYPE_CHECKING: - from typing import Sequence, Literal + from typing import Literal, Sequence import pandas as pd from weaviate import ConsistencyLevel @@ -140,7 +139,11 @@ def create_class(self, class_json: dict[str, Any]) -> None: client = self.conn client.schema.create_class(class_json) - @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.exceptions.RequestException)) + @retry( + reraise=True, + stop=stop_after_attempt(3), + retry=retry_if_exception_type(requests.exceptions.RequestException), + ) def create_schema(self, schema_json: dict[str, Any] | str) -> None: """ Create a new Schema. @@ -166,7 +169,11 @@ def _convert_dataframe_to_list(data: list[dict[str, Any]] | pd.DataFrame) -> lis data = json.loads(data.to_json(orient="records")) return cast(List[Dict[str, Any]], data) - @retry(reraise=True, stop=stop_after_attempt(3), retry=retry_if_exception_type(requests.exceptions.RequestException)) + @retry( + reraise=True, + stop=stop_after_attempt(3), + retry=retry_if_exception_type(requests.exceptions.RequestException), + ) def get_schema(self, class_name: str | None = None): """Get the schema from Weaviate. diff --git a/airflow/providers/weaviate/operators/weaviate.py b/airflow/providers/weaviate/operators/weaviate.py index 9d28790ac0733..a39f95fcab1c2 100644 --- a/airflow/providers/weaviate/operators/weaviate.py +++ b/airflow/providers/weaviate/operators/weaviate.py @@ -27,6 +27,7 @@ if TYPE_CHECKING: import pandas as pd + from airflow.utils.context import Context From 81ba6b79d5b05f88f01946225bf82351112e389f Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 8 Dec 2023 18:03:20 +0530 Subject: [PATCH 12/34] Remove unwanted changes --- airflow/providers/weaviate/hooks/weaviate.py | 2 +- airflow/providers/weaviate/operators/weaviate.py | 1 - docs/apache-airflow-providers-weaviate/operators/weaviate.rst | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 5ab22872b711f..26349bf90f81c 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -353,7 +353,7 @@ def check_http_error_is_retryable(exc: BaseException): def batch_data( self, class_name: str, - data: list[dict[str, Any]], + data: list[dict[str, Any]] | pd.DataFrame, batch_config_params: dict[str, Any] | None = None, vector_col: str = "Vector", retry_attempts_per_object: int = 5, diff --git a/airflow/providers/weaviate/operators/weaviate.py b/airflow/providers/weaviate/operators/weaviate.py index a39f95fcab1c2..4e07a59edba99 100644 --- a/airflow/providers/weaviate/operators/weaviate.py +++ b/airflow/providers/weaviate/operators/weaviate.py @@ -68,7 +68,6 @@ def __init__( super().__init__(**kwargs) self.class_name = class_name self.conn_id = conn_id - self.vector_col = vector_col if input_data is not None: diff --git a/docs/apache-airflow-providers-weaviate/operators/weaviate.rst b/docs/apache-airflow-providers-weaviate/operators/weaviate.rst index 05063376f4d95..5ec262ab7a2d3 100644 --- a/docs/apache-airflow-providers-weaviate/operators/weaviate.rst +++ b/docs/apache-airflow-providers-weaviate/operators/weaviate.rst @@ -28,7 +28,7 @@ into the database. Using the Operator ^^^^^^^^^^^^^^^^^^ -The WeaviateIngestOperator requires the ``input_text`` as an input to the operator. Use the ``conn_id`` parameter to specify the Weaviate connection to use to +The WeaviateIngestOperator requires the ``input_data`` as an input to the operator. Use the ``conn_id`` parameter to specify the Weaviate connection to use to connect to your account. An example using the operator to ingest data with custom vectors retrieved from XCOM: From 1533537fbac9f5095bf103477365101ae4a4d392 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 8 Dec 2023 18:18:51 +0530 Subject: [PATCH 13/34] Remove unwanted changes --- .../providers/weaviate/hooks/test_weaviate.py | 55 ++++++++++--------- 1 file changed, 28 insertions(+), 27 deletions(-) diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 0f16c3548cde5..09e31c70af5a0 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -17,7 +17,8 @@ from __future__ import annotations from contextlib import ExitStack -from unittest.mock import MagicMock, Mock, call, patch +from unittest import mock +from unittest.mock import MagicMock, Mock import pandas as pd import pytest @@ -39,32 +40,32 @@ def weaviate_hook(): mock_conn = Mock() # Patch the WeaviateHook get_connection method to return the mock connection - with patch.object(WeaviateHook, "get_connection", return_value=mock_conn): + with mock.patch.object(WeaviateHook, "get_connection", return_value=mock_conn): hook = WeaviateHook(conn_id=TEST_CONN_ID) return hook @pytest.fixture def mock_auth_api_key(): - with patch("airflow.providers.weaviate.hooks.weaviate.AuthApiKey") as m: + with mock.patch("airflow.providers.weaviate.hooks.weaviate.AuthApiKey") as m: yield m @pytest.fixture def mock_auth_bearer_token(): - with patch("airflow.providers.weaviate.hooks.weaviate.AuthBearerToken") as m: + with mock.patch("airflow.providers.weaviate.hooks.weaviate.AuthBearerToken") as m: yield m @pytest.fixture def mock_auth_client_credentials(): - with patch("airflow.providers.weaviate.hooks.weaviate.AuthClientCredentials") as m: + with mock.patch("airflow.providers.weaviate.hooks.weaviate.AuthClientCredentials") as m: yield m @pytest.fixture def mock_auth_client_password(): - with patch("airflow.providers.weaviate.hooks.weaviate.AuthClientPassword") as m: + with mock.patch("airflow.providers.weaviate.hooks.weaviate.AuthClientPassword") as m: yield m @@ -125,7 +126,7 @@ def setup_method(self, monkeypatch): for conn in conns: monkeypatch.setenv(f"AIRFLOW_CONN_{conn.conn_id.upper()}", conn.get_uri()) - @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_api_key_in_extra(self, mock_client, mock_auth_api_key): hook = WeaviateHook(conn_id=self.weaviate_api_key1) hook.get_conn() @@ -134,7 +135,7 @@ def test_get_conn_with_api_key_in_extra(self, mock_client, mock_auth_api_key): url=self.host, auth_client_secret=mock_auth_api_key(api_key=self.api_key), additional_headers={} ) - @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_token_in_extra(self, mock_client, mock_auth_api_key): # when token is passed in extra hook = WeaviateHook(conn_id=self.weaviate_api_key2) @@ -144,7 +145,7 @@ def test_get_conn_with_token_in_extra(self, mock_client, mock_auth_api_key): url=self.host, auth_client_secret=mock_auth_api_key(api_key=self.api_key), additional_headers={} ) - @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_access_token_in_extra(self, mock_client, mock_auth_bearer_token): hook = WeaviateHook(conn_id=self.client_bearer_token) hook.get_conn() @@ -159,7 +160,7 @@ def test_get_conn_with_access_token_in_extra(self, mock_client, mock_auth_bearer additional_headers={}, ) - @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_client_secret_in_extra(self, mock_client, mock_auth_client_credentials): hook = WeaviateHook(conn_id=self.weaviate_client_credentials) hook.get_conn() @@ -172,7 +173,7 @@ def test_get_conn_with_client_secret_in_extra(self, mock_client, mock_auth_clien additional_headers={}, ) - @patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") + @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateClient") def test_get_conn_with_client_password_in_extra(self, mock_client, mock_auth_client_password): hook = WeaviateHook(conn_id=self.client_password) hook.get_conn() @@ -183,7 +184,7 @@ def test_get_conn_with_client_password_in_extra(self, mock_client, mock_auth_cli additional_headers={}, ) - @patch("airflow.providers.weaviate.hooks.weaviate.generate_uuid5") + @mock.patch("airflow.providers.weaviate.hooks.weaviate.generate_uuid5") def test_create_object(self, mock_gen_uuid, weaviate_hook): """ Test the create_object method of WeaviateHook. @@ -229,7 +230,7 @@ def test_get_of_get_or_create_object(self, weaviate_hook): tenant=None, ) - @patch("airflow.providers.weaviate.hooks.weaviate.generate_uuid5") + @mock.patch("airflow.providers.weaviate.hooks.weaviate.generate_uuid5") def test_create_of_get_or_create_object(self, mock_gen_uuid, weaviate_hook): """ Test the create part of get_or_create_object method of WeaviateHook. @@ -279,8 +280,8 @@ def test_get_all_objects(self, weaviate_hook): return_value = weaviate_hook.get_all_objects(class_name="TestClass") assert weaviate_hook.get_object.call_args_list == [ - call(after=None, class_name="TestClass"), - call(after=3, class_name="TestClass"), + mock.call(after=None, class_name="TestClass"), + mock.call(after=3, class_name="TestClass"), ] assert return_value == [{"name": "Test1", "id": 2}, {"name": "Test2", "id": 3}] @@ -300,8 +301,8 @@ def test_get_all_objects_returns_dataframe(self, weaviate_hook): return_value = weaviate_hook.get_all_objects(class_name="TestClass", as_dataframe=True) assert weaviate_hook.get_object.call_args_list == [ - call(after=None, class_name="TestClass"), - call(after=3, class_name="TestClass"), + mock.call(after=None, class_name="TestClass"), + mock.call(after=3, class_name="TestClass"), ] import pandas @@ -435,7 +436,7 @@ def test_batch_data(data, expected_length, weaviate_hook): assert mock_batch_context.add_data_object.call_count == expected_length -@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_conn") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_conn") def test_batch_data_retry(get_conn, weaviate_hook): """Test to ensure retrying working as expected""" data = [{"name": "chandler"}, {"name": "joey"}, {"name": "ross"}] @@ -459,9 +460,9 @@ def test_batch_data_retry(get_conn, weaviate_hook): ], ids=["ignore", "replace", "fail", "invalid_option"], ) -@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.delete_classes") -@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.create_schema") -@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.delete_classes") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.create_schema") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") def test_upsert_schema_scenarios( get_schema, create_schema, delete_classes, get_schema_value, existing, expected_value, weaviate_hook ): @@ -480,10 +481,10 @@ def test_upsert_schema_scenarios( delete_classes.assert_called_once_with(class_names=["B"]) -@patch("builtins.open") -@patch("json.load") -@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.create_schema") -@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") +@mock.patch("builtins.open") +@mock.patch("json.load") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.create_schema") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") def test_upsert_schema_json_file_param(get_schema, create_schema, load, open, weaviate_hook): """Test if schema_json is path to a json file""" get_schema.return_value = {"classes": [{"class": "A"}, {"class": "B"}]} @@ -495,7 +496,7 @@ def test_upsert_schema_json_file_param(get_schema, create_schema, load, open, we create_schema.assert_called_once_with({"classes": [{"class": "C"}]}) -@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_client") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_client") def test_delete_schema(get_client, weaviate_hook): class_names = ["class_a", "class_b"] get_client.return_value.schema.delete_class.side_effect = [ @@ -610,7 +611,7 @@ def test_delete_schema(get_client, weaviate_hook): "swapped_properties", ), ) -@patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_schema") def test_contains_schema(get_schema, classes_to_test, expected_result, weaviate_hook): get_schema.return_value = { "classes": [ From c4cf2dbbe9196bec1433ba798b93b7cec0763723 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 8 Dec 2023 18:59:10 +0530 Subject: [PATCH 14/34] Change the exception to rety on --- airflow/providers/weaviate/hooks/weaviate.py | 52 +++++++++++++++---- .../providers/weaviate/hooks/test_weaviate.py | 18 ++++++- 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 26349bf90f81c..6c9b1b67872b5 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, cast import requests -from tenacity import Retrying, retry, retry_if_exception_type, stop_after_attempt +from tenacity import Retrying, retry, retry_if_exception, retry_if_exception_type, stop_after_attempt from weaviate import Client as WeaviateClient from weaviate.auth import AuthApiKey, AuthBearerToken, AuthClientCredentials, AuthClientPassword from weaviate.exceptions import ObjectAlreadyExistsException @@ -42,6 +42,23 @@ ExitingSchemaOptions = Literal["replace", "fail", "ignore"] +HTTP_RETRY_STATUS_CODE = [429, 500, 503, 504] +REQUESTS_EXCEPTIONS_TYPES = ( + requests.RequestException, + requests.exceptions.ConnectionError, + requests.exceptions.HTTPError, + requests.exceptions.ConnectTimeout, +) + + +def check_http_error_is_retryable(exc: BaseException): + return ( + isinstance(exc, requests.exceptions.RequestException) + and exc.response + and exc.response.status_code + and exc.response.status_code in HTTP_RETRY_STATUS_CODE + ) + class WeaviateHook(BaseHook): """ @@ -55,7 +72,13 @@ class WeaviateHook(BaseHook): conn_type = "weaviate" hook_name = "Weaviate" - def __init__(self, conn_id: str = default_conn_name, *args: Any, **kwargs: Any) -> None: + def __init__( + self, + conn_id: str = default_conn_name, + retry_status_codes: list[int] | None = None, + *args: Any, + **kwargs: Any, + ) -> None: super().__init__(*args, **kwargs) self.conn_id = conn_id @@ -142,7 +165,10 @@ def create_class(self, class_json: dict[str, Any]) -> None: @retry( reraise=True, stop=stop_after_attempt(3), - retry=retry_if_exception_type(requests.exceptions.RequestException), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), ) def create_schema(self, schema_json: dict[str, Any] | str) -> None: """ @@ -172,7 +198,10 @@ def _convert_dataframe_to_list(data: list[dict[str, Any]] | pd.DataFrame) -> lis @retry( reraise=True, stop=stop_after_attempt(3), - retry=retry_if_exception_type(requests.exceptions.RequestException), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), ) def get_schema(self, class_name: str | None = None): """Get the schema from Weaviate. @@ -200,9 +229,13 @@ def delete_classes(self, class_names: list[str] | str, if_error: str = "stop") - try: for attempt in Retrying( stop=stop_after_attempt(3), - retry=retry_if_exception_type(requests.exceptions.RequestException), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), ): with attempt: + print(attempt) client.schema.delete_class(class_name) except Exception as e: if if_error == "continue": @@ -346,10 +379,6 @@ def check_subset_of_schema(self, classes_objects: list) -> bool: return False return True - @staticmethod - def check_http_error_is_retryable(exc: BaseException): - return isinstance(exc, requests.HTTPError) and not exc.response.ok - def batch_data( self, class_name: str, @@ -378,7 +407,10 @@ def batch_data( for index, data_obj in enumerate(data): for attempt in Retrying( stop=stop_after_attempt(retry_attempts_per_object), - retry=retry_if_exception_type(requests.exceptions.RequestException), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), ): with attempt: self.log.debug( diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 09e31c70af5a0..fc6d7db6aed8f 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -497,7 +497,7 @@ def test_upsert_schema_json_file_param(get_schema, create_schema, load, open, we @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_client") -def test_delete_schema(get_client, weaviate_hook): +def test_delete_classes(get_client, weaviate_hook): class_names = ["class_a", "class_b"] get_client.return_value.schema.delete_class.side_effect = [ weaviate.UnexpectedStatusCodeException("something failed", requests.Response()), @@ -513,6 +513,22 @@ def test_delete_schema(get_client, weaviate_hook): weaviate_hook.delete_classes("class_a", if_error="stop") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.get_client") +def test_http_errors_of_delete_classes(get_client, weaviate_hook): + class_names = ["class_a", "class_b"] + resp = requests.Response() + resp.status_code = 429 + get_client.return_value.schema.delete_class.side_effect = [ + requests.exceptions.HTTPError(response=resp), + None, + requests.exceptions.ConnectionError, + None, + ] + error_list = weaviate_hook.delete_classes(class_names, if_error="continue") + assert error_list == [] + assert get_client.return_value.schema.delete_class.call_count == 4 + + @pytest.mark.parametrize( argnames=["classes_to_test", "expected_result"], argvalues=[ From 5d794c8f9149523f3f5a464e8920d34f3a938c65 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 12 Dec 2023 02:00:57 +0530 Subject: [PATCH 15/34] [WIP] Add ingest methods --- airflow/providers/weaviate/hooks/weaviate.py | 245 ++++++++++++++++-- .../providers/weaviate/hooks/test_weaviate.py | 161 +++++++++++- 2 files changed, 389 insertions(+), 17 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 6c9b1b67872b5..db9d247cf335b 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -19,11 +19,13 @@ import contextlib import json +import logging import warnings from functools import cached_property -from typing import TYPE_CHECKING, Any, Dict, List, cast +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, cast import requests +import weaviate.exceptions from tenacity import Retrying, retry, retry_if_exception, retry_if_exception_type, stop_after_attempt from weaviate import Client as WeaviateClient from weaviate.auth import AuthApiKey, AuthBearerToken, AuthClientCredentials, AuthClientPassword @@ -80,6 +82,7 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) + self.logger = logging.getLogger("airflow.task") self.conn_id = conn_id @classmethod @@ -385,8 +388,10 @@ def batch_data( data: list[dict[str, Any]] | pd.DataFrame, batch_config_params: dict[str, Any] | None = None, vector_col: str = "Vector", + uuid_col: str = "id", retry_attempts_per_object: int = 5, - ) -> None: + tenant: str | None = None, + ) -> list: """ Add multiple objects or object references at once into weaviate. @@ -396,28 +401,46 @@ def batch_data( .. seealso:: `batch_config_params options `__ :param vector_col: name of the column containing the vector. :param retry_attempts_per_object: number of time to try in case of failure before giving up. + :param tenant: The tenant to which the object will be added. + :param uuid_col: Name of the column containing the UUID. """ client = self.conn if not batch_config_params: batch_config_params = {} client.batch.configure(**batch_config_params) data = self._convert_dataframe_to_list(data) + insertion_errors = [] with client.batch as batch: # Batch import all data - for index, data_obj in enumerate(data): - for attempt in Retrying( - stop=stop_after_attempt(retry_attempts_per_object), - retry=( - retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) - | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) - ), - ): - with attempt: - self.log.debug( - "Attempt %s of importing data: %s", attempt.retry_state.attempt_number, index + 1 - ) - vector = data_obj.pop(vector_col, None) - batch.add_data_object(data_obj, class_name, vector=vector) + try: + for index, data_obj in enumerate(data): + for attempt in Retrying( + stop=stop_after_attempt(retry_attempts_per_object), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), + ): + with attempt: + vector = data_obj.pop(vector_col, None) + uuid = data_obj.pop(uuid_col, None) + self.log.debug( + "Attempt %s of inserting object with uuid: %s", + attempt.retry_state.attempt_number, + uuid, + ) + batch.add_data_object( + data_object=data_obj, + class_name=class_name, + vector=vector, + uuid=uuid, + tenant=tenant, + ) + self.log.debug("Inserted object with uuid: %s", uuid) + except Exception as e: + insertion_errors.append({"uuid": uuid, "result": {"errors": str(e)}}) + self.logger.error(f"Failed to add object with UUID {uuid}. Error: {e}") + return insertion_errors def query_with_vector( self, @@ -606,3 +629,193 @@ def object_exists(self, uuid: str | UUID, **kwargs) -> bool: """ client = self.conn return client.data_object.exists(uuid, **kwargs) + + def _generate_uuids( + self, + df: pd.DataFrame, + class_name: str, + unique_columns: list[str] | None = None, + vector_column: str | None = None, + uuid_column: str | None = None, + ) -> tuple[pd.DataFrame, str]: + """ + Adds UUIDs to a DataFrame, useful for upsert operations where UUIDs must be known before ingestion. + + By default, UUIDs are generated using a custom function if 'uuid_column' is not specified. + The function can potentially ingest the same data multiple times with different UUIDs. + + :param df: A dataframe with data to generate a UUID from. + :param class_name: The name of the class use as part of the uuid namespace. + :param uuid_column: Name of the column to create. Default is 'id'. + :param unique_columns: A list of columns to use for UUID generation. By default, all columns except + vector_column will be used. + :param vector_column: Name of the column containing the vector data. If specified the vector will be + removed prior to generating the uuid. + """ + column_names = df.columns.to_list() + + unique_columns = unique_columns or column_names + unique_columns.sort() + + difference_columns = set(unique_columns).difference(set(df.columns.to_list())) + if difference_columns: + raise ValueError(f"Columns {', '.join(difference_columns)} don't exist in dataframe") + + if uuid_column is None: + self.logger.info("No uuid_column provided. Generating UUIDs as column name `id`.") + if "id" in column_names: + raise ValueError( + "Property 'id' already in dataset. Consider renaming or specify 'uuid_column'." + ) + else: + uuid_column = "id" + + if uuid_column in column_names: + raise ValueError( + f"Property {uuid_column} already in dataset. Consider renaming or specify a different" + f" 'uuid_column'." + ) + + df[uuid_column] = ( + df[unique_columns] + .drop(columns=[vector_column], inplace=False, errors="ignore") + .apply(lambda row: generate_uuid5(identifier=row.to_dict(), namespace=class_name), axis=1) + ) + + return df, uuid_column + + def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_name: str, existing: str): + """ + Check if the objects with uuid exist or not. + + :param data: A single pandas DataFrame. + :param uuid_column: Column with pre-generated UUIDs. + :param class_name: Name of the class in Weaviate schema where data is to be ingested. + """ + existing_uuid = set() + non_existing_uuid = set() + + if existing == "replace": + existing_uuid = set(data[uuid_column].to_list()) + else: + self.logger.info(f"checking if {data.shape[0]} objects exists.") + for uuid in data[uuid_column]: + for attempt in Retrying( + stop=stop_after_attempt(5), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), + ): + with attempt: + if self.object_exists(uuid=uuid, class_name=class_name): + existing_uuid.add(uuid) + self.logger.debug("object with uuid %s exists.", uuid) + else: + non_existing_uuid.add(uuid) + self.logger.debug("object with uuid %s don't exists.", uuid) + + self.logger.info( + f"Objects to override {len(existing_uuid)} and {len(non_existing_uuid)} " f"objects to create" + ) + return existing_uuid, non_existing_uuid + + def _delete_objects(self, uuids: Iterable, class_name: str, retry_attempts_per_object: int = 5): + """Helper function for `create_or_replace_objects()` to delete multiple objects.""" + for uuid in uuids: + for attempt in Retrying( + stop=stop_after_attempt(retry_attempts_per_object), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), + ): + with attempt: + try: + self.delete_object(uuid=uuid, class_name=class_name) + self.logger.debug("Deleted object with uuid %s", uuid) + except weaviate.exceptions.UnexpectedStatusCodeException as e: + if e.status_code == 404: + self.logger.debug("Tried to delete a non existent object with uuid %s", uuid) + else: + self.logger.debug( + "Error occurred while trying to delete object with uuid %s", uuid + ) + raise e + + self.logger.info(f"Deleted {len(uuids)} objects.") + + def create_or_replace_objects( + self, + data: pd.DataFrame | list[dict[str, Any]], + class_name: str, + existing: str = "skip", + unique_columns: str = None, + uuid_column: str | None = None, + vector_column: str = None, + batch_config_params: dict = None, + tenant: str | None = None, + ): + """ + create or replace objects. + + :param data: A single pandas DataFrame or a list of dicts to be ingested. + :param class_name: Name of the class in Weaviate schema where data is to be ingested. + :param existing: Strategy for handling existing data: 'skip', or 'replace'. Default is 'skip'. + :param unique_columns: Columns in DataFrame or keys in dict uniquely identifying each document, + required for 'upsert' operations. + :param uuid_column: Column with pre-generated UUIDs. If not provided, UUIDs will be generated. + :param vector_column: Column with embedding vectors for pre-embedded data. + :param batch_config_params: Additional parameters for Weaviate batch configuration. + :param tenant: The tenant to which the object will be added. + """ + import pandas as pd + + if existing not in ["skip", "replace", "error"]: + raise ValueError( + "Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'upsert', 'error'." + ) + + if isinstance(data, list): + data = pd.json_normalize(data) + + self.logger.info(f"Inserting {data.shape[0]} objects.") + + if uuid_column is None or uuid_column not in data.columns: + data, uuid_column = self._generate_uuids( + df=data, + class_name=class_name, + unique_columns=unique_columns, + vector_column=vector_column, + uuid_column=uuid_column, + ) + uuids_to_create = set() + existing_uuid, non_existing_uuid = self._check_existing_objects( + data=data, uuid_column=uuid_column, class_name=class_name, existing=existing + ) + if existing == "error" and len(existing_uuid): + self.logger.info(f"Found duplicate UUIDs {' ,'.join(existing_uuid)}") + raise ValueError( + f"Found {len(existing_uuid)} object with duplicate UUIDs. You can either ignore or replace" + f" them by passing 'existing=skip' or 'existing=replace' respectively." + ) + elif existing == "replace": + uuids_to_create = existing_uuid.union(non_existing_uuid) + self._delete_objects(existing_uuid, class_name=class_name) + elif existing == "skip": + uuids_to_create = non_existing_uuid + data = data[data[uuid_column].isin(uuids_to_create)] + if data.shape[0]: + self.logger.info(f"Batch inserting {data.shape[0]} objects.") + insertion_errors = self.batch_data( + class_name=class_name, + data=data, + batch_config_params=batch_config_params, + vector_col=vector_column, + uuid_col=uuid_column, + tenant=tenant, + ) + if insertion_errors: + self.logger.info(f"Failed to insert {len(insertion_errors)} objects.") + # Rollback object that were not created properly + self._delete_objects([item["uuid"] for item in insertion_errors], class_name=class_name) diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index fc6d7db6aed8f..41facd90e44cf 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -18,7 +18,7 @@ from contextlib import ExitStack from unittest import mock -from unittest.mock import MagicMock, Mock +from unittest.mock import ANY, MagicMock, Mock import pandas as pd import pytest @@ -670,3 +670,162 @@ def test_contains_schema(get_schema, classes_to_test, expected_result, weaviate_ ] } assert weaviate_hook.check_subset_of_schema(classes_to_test) == expected_result + + +@mock.patch("weaviate.util.generate_uuid5") +def test___generate_uuids(generate_uuid5, weaviate_hook): + df = pd.DataFrame.from_dict({"name": ["ross", "bob"], "age": ["12", "22"], "gender": ["m", "m"]}) + with pytest.raises(ValueError, match=r"Columns last_name don't exist in dataframe"): + weaviate_hook._generate_uuids( + df=df, class_name="test", unique_columns=["name", "age", "gender", "last_name"] + ) + + df = pd.DataFrame.from_dict( + {"id": [1, 2], "name": ["ross", "bob"], "age": ["12", "22"], "gender": ["m", "m"]} + ) + with pytest.raises( + ValueError, match=r"Property 'id' already in dataset. Consider renaming or specify" r" 'uuid_column'" + ): + weaviate_hook._generate_uuids(df=df, class_name="test", unique_columns=["name", "age", "gender"]) + + with pytest.raises( + ValueError, + match=r"Property age already in dataset. Consider renaming or specify" r" a different 'uuid_column'.", + ): + weaviate_hook._generate_uuids( + df=df, uuid_column="age", class_name="test", unique_columns=["name", "age", "gender"] + ) + + +def test_replace__check_existing_objects(weaviate_hook): + df = pd.DataFrame.from_dict( + { + "id": ["1", "2", "3"], + "name": ["ross", "bob", "joy"], + "age": ["12", "22", "15"], + "gender": ["m", "m", "f"], + } + ) + existing_uuid, non_existing_uuid = weaviate_hook._check_existing_objects( + data=df, uuid_column="id", class_name="test", existing="replace" + ) + assert existing_uuid == {"1", "2", "3"} + + +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.object_exists") +def test_skip__check_existing_objects(object_exists, weaviate_hook): + df = pd.DataFrame.from_dict( + { + "id": ["1", "2", "3"], + "name": ["ross", "bob", "joy"], + "age": ["12", "22", "15"], + "gender": ["m", "m", "f"], + } + ) + resp = requests.Response() + resp.status_code = 429 + requests.exceptions.HTTPError(response=resp) + http_exception = requests.exceptions.HTTPError(response=resp) + object_exists.side_effect = [True, http_exception, http_exception, True, False] + existing_uuid, non_existing_uuid = weaviate_hook._check_existing_objects( + data=df, uuid_column="id", class_name="test", existing="skip" + ) + assert existing_uuid == {"1", "2"} + assert non_existing_uuid == {"3"} + assert object_exists.call_count == 5 + + +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.delete_object") +def test__delete_objects(delete_object, weaviate_hook): + resp = requests.Response() + resp.status_code = 429 + requests.exceptions.HTTPError(response=resp) + http_429_exception = requests.exceptions.HTTPError(response=resp) + + resp = requests.Response() + resp.status_code = 404 + not_found_exception = weaviate.exceptions.UnexpectedStatusCodeException( + message="object not found", response=resp + ) + + delete_object.side_effect = [not_found_exception, None, http_429_exception, http_429_exception, None] + weaviate_hook._delete_objects(uuids=["1", "2", "3"], class_name="test") + assert delete_object.call_count == 5 + + +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") +def test_error_create_or_replace_objects(_generate_uuids, _check_existing_objects, weaviate_hook): + df = pd.DataFrame.from_dict( + { + "id": ["1", "2", "3"], + "name": ["ross", "bob", "joy"], + "age": ["12", "22", "15"], + "gender": ["m", "m", "f"], + } + ) + + _check_existing_objects.return_value = ({"1"}, {"2", "3"}) + _generate_uuids.return_value = (None, None) + with pytest.raises(ValueError, match=f"Found {len({'1'})} object with duplicate UUIDs. You can either"): + weaviate_hook.create_or_replace_objects(data=df, class_name="test", existing="error") + + +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._delete_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.batch_data") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") +def test_skip_create_or_replace_objects( + _generate_uuids, _check_existing_objects, batch_data, _delete_objects, weaviate_hook +): + df = pd.DataFrame.from_dict( + { + "id": ["1", "2", "3"], + "name": ["ross", "bob", "joy"], + "age": ["12", "22", "15"], + "gender": ["m", "m", "f"], + } + ) + + class_name = "test" + existing_uuid, non_existing_uuid = ({"1"}, {"2", "3"}) + batch_data.return_value = [{"uuid": i} for i in non_existing_uuid] + _check_existing_objects.return_value = (existing_uuid, non_existing_uuid) + _generate_uuids.return_value = (df, "id") + weaviate_hook.create_or_replace_objects(data=df, class_name=class_name, existing="skip") + batch_data.assert_called_with( + class_name="test", data=ANY, batch_config_params=None, vector_col=None, uuid_col="id", tenant=None + ) + assert set(batch_data.call_args_list[0].kwargs["data"]["id"].to_list()) == non_existing_uuid + assert sorted(_delete_objects.call_args.args[0]) == sorted(list(non_existing_uuid)) + + +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._delete_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.batch_data") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") +def test_replace_create_or_replace_objects( + _generate_uuids, _check_existing_objects, batch_data, _delete_objects, weaviate_hook +): + df = pd.DataFrame.from_dict( + { + "id": ["1", "2", "3"], + "name": ["ross", "bob", "joy"], + "age": ["12", "22", "15"], + "gender": ["m", "m", "f"], + } + ) + + class_name = "test" + existing_uuid, non_existing_uuid = ({"1"}, {"2", "3"}) + batch_data.return_value = [] + _check_existing_objects.return_value = (existing_uuid, non_existing_uuid) + _generate_uuids.return_value = (df, "id") + weaviate_hook.create_or_replace_objects(data=df, class_name=class_name, existing="replace") + _delete_objects.assert_called_with(existing_uuid, class_name=class_name) + batch_data.assert_called_with( + class_name="test", data=ANY, batch_config_params=None, vector_col=None, uuid_col="id", tenant=None + ) + assert set(batch_data.call_args_list[0].kwargs["data"]["id"].to_list()) == existing_uuid.union( + non_existing_uuid + ) From b1a4aa3c33865022adbe5c504385f3a4acd64597 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 12 Dec 2023 03:59:06 +0530 Subject: [PATCH 16/34] Fix static checks --- airflow/providers/weaviate/hooks/weaviate.py | 45 ++++++++++---------- 1 file changed, 22 insertions(+), 23 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index db9d247cf335b..c9195f407268d 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -19,10 +19,9 @@ import contextlib import json -import logging import warnings from functools import cached_property -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, cast +from typing import TYPE_CHECKING, Any, Dict, List, cast import requests import weaviate.exceptions @@ -82,7 +81,6 @@ def __init__( **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) - self.logger = logging.getLogger("airflow.task") self.conn_id = conn_id @classmethod @@ -439,7 +437,7 @@ def batch_data( self.log.debug("Inserted object with uuid: %s", uuid) except Exception as e: insertion_errors.append({"uuid": uuid, "result": {"errors": str(e)}}) - self.logger.error(f"Failed to add object with UUID {uuid}. Error: {e}") + self.log.error("Failed to add object with UUID %s. Error: %s", uuid, e) return insertion_errors def query_with_vector( @@ -662,7 +660,7 @@ def _generate_uuids( raise ValueError(f"Columns {', '.join(difference_columns)} don't exist in dataframe") if uuid_column is None: - self.logger.info("No uuid_column provided. Generating UUIDs as column name `id`.") + self.log.info("No uuid_column provided. Generating UUIDs as column name `id`.") if "id" in column_names: raise ValueError( "Property 'id' already in dataset. Consider renaming or specify 'uuid_column'." @@ -698,7 +696,7 @@ def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_na if existing == "replace": existing_uuid = set(data[uuid_column].to_list()) else: - self.logger.info(f"checking if {data.shape[0]} objects exists.") + self.log.info("checking if %s objects exists.", data.shape[0]) for uuid in data[uuid_column]: for attempt in Retrying( stop=stop_after_attempt(5), @@ -710,17 +708,17 @@ def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_na with attempt: if self.object_exists(uuid=uuid, class_name=class_name): existing_uuid.add(uuid) - self.logger.debug("object with uuid %s exists.", uuid) + self.log.debug("object with uuid %s exists.", uuid) else: non_existing_uuid.add(uuid) - self.logger.debug("object with uuid %s don't exists.", uuid) + self.log.debug("object with uuid %s don't exists.", uuid) - self.logger.info( + self.log.info( f"Objects to override {len(existing_uuid)} and {len(non_existing_uuid)} " f"objects to create" ) return existing_uuid, non_existing_uuid - def _delete_objects(self, uuids: Iterable, class_name: str, retry_attempts_per_object: int = 5): + def _delete_objects(self, uuids: Sequence, class_name: str, retry_attempts_per_object: int = 5): """Helper function for `create_or_replace_objects()` to delete multiple objects.""" for uuid in uuids: for attempt in Retrying( @@ -733,27 +731,25 @@ def _delete_objects(self, uuids: Iterable, class_name: str, retry_attempts_per_o with attempt: try: self.delete_object(uuid=uuid, class_name=class_name) - self.logger.debug("Deleted object with uuid %s", uuid) + self.log.debug("Deleted object with uuid %s", uuid) except weaviate.exceptions.UnexpectedStatusCodeException as e: if e.status_code == 404: - self.logger.debug("Tried to delete a non existent object with uuid %s", uuid) + self.log.debug("Tried to delete a non existent object with uuid %s", uuid) else: - self.logger.debug( - "Error occurred while trying to delete object with uuid %s", uuid - ) + self.log.debug("Error occurred while trying to delete object with uuid %s", uuid) raise e - self.logger.info(f"Deleted {len(uuids)} objects.") + self.log.info("Deleted %s objects.", len(uuids)) def create_or_replace_objects( self, data: pd.DataFrame | list[dict[str, Any]], class_name: str, existing: str = "skip", - unique_columns: str = None, + unique_columns: list[str] | str | None = None, uuid_column: str | None = None, - vector_column: str = None, - batch_config_params: dict = None, + vector_column: str = "Vector", + batch_config_params: dict | None = None, tenant: str | None = None, ): """ @@ -771,6 +767,9 @@ def create_or_replace_objects( """ import pandas as pd + if isinstance(unique_columns, str): + unique_columns = [unique_columns] + if existing not in ["skip", "replace", "error"]: raise ValueError( "Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'upsert', 'error'." @@ -779,7 +778,7 @@ def create_or_replace_objects( if isinstance(data, list): data = pd.json_normalize(data) - self.logger.info(f"Inserting {data.shape[0]} objects.") + self.log.info("Inserting %s objects.", data.shape[0]) if uuid_column is None or uuid_column not in data.columns: data, uuid_column = self._generate_uuids( @@ -794,7 +793,7 @@ def create_or_replace_objects( data=data, uuid_column=uuid_column, class_name=class_name, existing=existing ) if existing == "error" and len(existing_uuid): - self.logger.info(f"Found duplicate UUIDs {' ,'.join(existing_uuid)}") + self.log.info("Found duplicate UUIDs %s", " ,".join(existing_uuid)) raise ValueError( f"Found {len(existing_uuid)} object with duplicate UUIDs. You can either ignore or replace" f" them by passing 'existing=skip' or 'existing=replace' respectively." @@ -806,7 +805,7 @@ def create_or_replace_objects( uuids_to_create = non_existing_uuid data = data[data[uuid_column].isin(uuids_to_create)] if data.shape[0]: - self.logger.info(f"Batch inserting {data.shape[0]} objects.") + self.log.info("Batch inserting %s objects.", data.shape[0]) insertion_errors = self.batch_data( class_name=class_name, data=data, @@ -816,6 +815,6 @@ def create_or_replace_objects( tenant=tenant, ) if insertion_errors: - self.logger.info(f"Failed to insert {len(insertion_errors)} objects.") + self.log.info("Failed to insert %s objects.", len(insertion_errors)) # Rollback object that were not created properly self._delete_objects([item["uuid"] for item in insertion_errors], class_name=class_name) From 967ac86ec28a8b1e52443d17e204e735f9e295c2 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 12 Dec 2023 15:36:11 +0530 Subject: [PATCH 17/34] Optimize for error case --- airflow/providers/weaviate/hooks/weaviate.py | 10 +++-- .../providers/weaviate/hooks/test_weaviate.py | 38 ++++++++++++------- 2 files changed, 31 insertions(+), 17 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index c9195f407268d..8bb64c4b1bd7c 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -35,7 +35,7 @@ from airflow.hooks.base import BaseHook if TYPE_CHECKING: - from typing import Literal, Sequence + from typing import Collection, Literal, Sequence import pandas as pd from weaviate import ConsistencyLevel @@ -690,8 +690,8 @@ def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_na :param uuid_column: Column with pre-generated UUIDs. :param class_name: Name of the class in Weaviate schema where data is to be ingested. """ - existing_uuid = set() - non_existing_uuid = set() + existing_uuid: set = set() + non_existing_uuid: set = set() if existing == "replace": existing_uuid = set(data[uuid_column].to_list()) @@ -708,6 +708,8 @@ def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_na with attempt: if self.object_exists(uuid=uuid, class_name=class_name): existing_uuid.add(uuid) + if existing == "error": + return existing_uuid, non_existing_uuid self.log.debug("object with uuid %s exists.", uuid) else: non_existing_uuid.add(uuid) @@ -718,7 +720,7 @@ def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_na ) return existing_uuid, non_existing_uuid - def _delete_objects(self, uuids: Sequence, class_name: str, retry_attempts_per_object: int = 5): + def _delete_objects(self, uuids: Collection, class_name: str, retry_attempts_per_object: int = 5): """Helper function for `create_or_replace_objects()` to delete multiple objects.""" for uuid in uuids: for attempt in Retrying( diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 41facd90e44cf..eff103429315b 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -18,7 +18,7 @@ from contextlib import ExitStack from unittest import mock -from unittest.mock import ANY, MagicMock, Mock +from unittest.mock import MagicMock, Mock import pandas as pd import pytest @@ -697,6 +697,23 @@ def test___generate_uuids(generate_uuid5, weaviate_hook): ) +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.object_exists") +def test_error__check_existing_objects(object_exists, weaviate_hook): + df = pd.DataFrame.from_dict( + { + "id": ["1", "2", "3"], + "name": ["ross", "bob", "joy"], + "age": ["12", "22", "15"], + "gender": ["m", "m", "f"], + } + ) + object_exists.return_value = True + existing_uuid, non_existing_uuid = weaviate_hook._check_existing_objects( + data=df, uuid_column="id", class_name="test", existing="error" + ) + assert existing_uuid == {"1"} + + def test_replace__check_existing_objects(weaviate_hook): df = pd.DataFrame.from_dict( { @@ -755,7 +772,7 @@ def test__delete_objects(delete_object, weaviate_hook): @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") -def test_error_create_or_replace_objects(_generate_uuids, _check_existing_objects, weaviate_hook): +def test_error_option_of_create_or_replace_objects(_generate_uuids, _check_existing_objects, weaviate_hook): df = pd.DataFrame.from_dict( { "id": ["1", "2", "3"], @@ -775,7 +792,7 @@ def test_error_create_or_replace_objects(_generate_uuids, _check_existing_object @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.batch_data") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") -def test_skip_create_or_replace_objects( +def test_skip_option_of_create_or_replace_objects( _generate_uuids, _check_existing_objects, batch_data, _delete_objects, weaviate_hook ): df = pd.DataFrame.from_dict( @@ -793,18 +810,16 @@ def test_skip_create_or_replace_objects( _check_existing_objects.return_value = (existing_uuid, non_existing_uuid) _generate_uuids.return_value = (df, "id") weaviate_hook.create_or_replace_objects(data=df, class_name=class_name, existing="skip") - batch_data.assert_called_with( - class_name="test", data=ANY, batch_config_params=None, vector_col=None, uuid_col="id", tenant=None + pd.testing.assert_frame_equal( + batch_data.call_args_list[0].kwargs["data"], df[df["id"].isin(non_existing_uuid)] ) - assert set(batch_data.call_args_list[0].kwargs["data"]["id"].to_list()) == non_existing_uuid - assert sorted(_delete_objects.call_args.args[0]) == sorted(list(non_existing_uuid)) @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._delete_objects") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.batch_data") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") -def test_replace_create_or_replace_objects( +def test_replace_option_of_create_or_replace_objects( _generate_uuids, _check_existing_objects, batch_data, _delete_objects, weaviate_hook ): df = pd.DataFrame.from_dict( @@ -823,9 +838,6 @@ def test_replace_create_or_replace_objects( _generate_uuids.return_value = (df, "id") weaviate_hook.create_or_replace_objects(data=df, class_name=class_name, existing="replace") _delete_objects.assert_called_with(existing_uuid, class_name=class_name) - batch_data.assert_called_with( - class_name="test", data=ANY, batch_config_params=None, vector_col=None, uuid_col="id", tenant=None - ) - assert set(batch_data.call_args_list[0].kwargs["data"]["id"].to_list()) == existing_uuid.union( - non_existing_uuid + pd.testing.assert_frame_equal( + batch_data.call_args_list[0].kwargs["data"], df[df["id"].isin(existing_uuid.union(non_existing_uuid))] ) From 26527d6c8248c9a80add795a3758bee583dfd41a Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 12 Dec 2023 16:14:45 +0530 Subject: [PATCH 18/34] Fix test docstring and handle missing case --- airflow/providers/weaviate/hooks/weaviate.py | 37 +++++++++++++++---- .../providers/weaviate/hooks/test_weaviate.py | 10 ++++- 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 8bb64c4b1bd7c..8d0bcb14ecfc0 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -637,7 +637,7 @@ def _generate_uuids( uuid_column: str | None = None, ) -> tuple[pd.DataFrame, str]: """ - Adds UUIDs to a DataFrame, useful for upsert operations where UUIDs must be known before ingestion. + Adds UUIDs to a DataFrame, useful for replace operations where UUIDs must be known before ingestion. By default, UUIDs are generated using a custom function if 'uuid_column' is not specified. The function can potentially ingest the same data multiple times with different UUIDs. @@ -684,11 +684,12 @@ def _generate_uuids( def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_name: str, existing: str): """ - Check if the objects with uuid exist or not. + Helper function to check if the objects with uuid exist or not. :param data: A single pandas DataFrame. :param uuid_column: Column with pre-generated UUIDs. :param class_name: Name of the class in Weaviate schema where data is to be ingested. + :param existing: Strategy for handling existing data: 'skip', or 'replace'. """ existing_uuid: set = set() non_existing_uuid: set = set() @@ -721,7 +722,13 @@ def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_na return existing_uuid, non_existing_uuid def _delete_objects(self, uuids: Collection, class_name: str, retry_attempts_per_object: int = 5): - """Helper function for `create_or_replace_objects()` to delete multiple objects.""" + """ + Helper function for `create_or_replace_objects()` to delete multiple objects. + + :param uuids: Collection of uuids. + :param class_name: Name of the class in Weaviate schema where data is to be ingested. + :param retry_attempts_per_object: number of time to try in case of failure before giving up. + """ for uuid in uuids: for attempt in Retrying( stop=stop_after_attempt(retry_attempts_per_object), @@ -753,10 +760,18 @@ def create_or_replace_objects( vector_column: str = "Vector", batch_config_params: dict | None = None, tenant: str | None = None, - ): + ) -> list: """ create or replace objects. + Provides users with multiple ways of dealing with existing values. + 1. replace: replace the existing object with new object. This option requires to identify the rows + uniquely, which by default is done by using all columns(except `vector`) to create a uuid. + User can modify this behaviour by providing `unique_columns` params. + 2. skip: skip the existing objects. + 3. error: raise an error if an existing object is found. + + :param data: A single pandas DataFrame or a list of dicts to be ingested. :param class_name: Name of the class in Weaviate schema where data is to be ingested. :param existing: Strategy for handling existing data: 'skip', or 'replace'. Default is 'skip'. @@ -766,6 +781,7 @@ def create_or_replace_objects( :param vector_column: Column with embedding vectors for pre-embedded data. :param batch_config_params: Additional parameters for Weaviate batch configuration. :param tenant: The tenant to which the object will be added. + :return: list of uuids which failed to create """ import pandas as pd @@ -773,9 +789,7 @@ def create_or_replace_objects( unique_columns = [unique_columns] if existing not in ["skip", "replace", "error"]: - raise ValueError( - "Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'upsert', 'error'." - ) + raise ValueError("Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'error'.") if isinstance(data, list): data = pd.json_normalize(data) @@ -783,13 +797,17 @@ def create_or_replace_objects( self.log.info("Inserting %s objects.", data.shape[0]) if uuid_column is None or uuid_column not in data.columns: - data, uuid_column = self._generate_uuids( + ( + data, + uuid_column, + ) = self._generate_uuids( df=data, class_name=class_name, unique_columns=unique_columns, vector_column=vector_column, uuid_column=uuid_column, ) + uuids_to_create = set() existing_uuid, non_existing_uuid = self._check_existing_objects( data=data, uuid_column=uuid_column, class_name=class_name, existing=existing @@ -806,6 +824,7 @@ def create_or_replace_objects( elif existing == "skip": uuids_to_create = non_existing_uuid data = data[data[uuid_column].isin(uuids_to_create)] + insertion_errors = [] if data.shape[0]: self.log.info("Batch inserting %s objects.", data.shape[0]) insertion_errors = self.batch_data( @@ -820,3 +839,5 @@ def create_or_replace_objects( self.log.info("Failed to insert %s objects.", len(insertion_errors)) # Rollback object that were not created properly self._delete_objects([item["uuid"] for item in insertion_errors], class_name=class_name) + + return insertion_errors diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index eff103429315b..be38ca47686bc 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -806,10 +806,16 @@ def test_skip_option_of_create_or_replace_objects( class_name = "test" existing_uuid, non_existing_uuid = ({"1"}, {"2", "3"}) - batch_data.return_value = [{"uuid": i} for i in non_existing_uuid] + batch_data_return_value = [{"uuid": i} for i in non_existing_uuid] + batch_data.return_value = batch_data_return_value _check_existing_objects.return_value = (existing_uuid, non_existing_uuid) _generate_uuids.return_value = (df, "id") - weaviate_hook.create_or_replace_objects(data=df, class_name=class_name, existing="skip") + + insertion_errors = weaviate_hook.create_or_replace_objects( + data=df, class_name=class_name, existing="skip" + ) + + assert batch_data_return_value == insertion_errors pd.testing.assert_frame_equal( batch_data.call_args_list[0].kwargs["data"], df[df["id"].isin(non_existing_uuid)] ) From 8c833e813f52d5a3468a5d8bb94f99cb3f4af70d Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 12 Dec 2023 17:09:36 +0530 Subject: [PATCH 19/34] Remove duplicate rows from dataframe --- airflow/providers/weaviate/hooks/weaviate.py | 19 +++++++++++-------- .../providers/weaviate/hooks/test_weaviate.py | 2 +- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 8d0bcb14ecfc0..ded51fc443459 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -632,7 +632,7 @@ def _generate_uuids( self, df: pd.DataFrame, class_name: str, - unique_columns: list[str] | None = None, + unique_columns: list[str], vector_column: str | None = None, uuid_column: str | None = None, ) -> tuple[pd.DataFrame, str]: @@ -652,9 +652,6 @@ def _generate_uuids( """ column_names = df.columns.to_list() - unique_columns = unique_columns or column_names - unique_columns.sort() - difference_columns = set(unique_columns).difference(set(df.columns.to_list())) if difference_columns: raise ValueError(f"Columns {', '.join(difference_columns)} don't exist in dataframe") @@ -781,19 +778,21 @@ def create_or_replace_objects( :param vector_column: Column with embedding vectors for pre-embedded data. :param batch_config_params: Additional parameters for Weaviate batch configuration. :param tenant: The tenant to which the object will be added. - :return: list of uuids which failed to create + :return: list of UUID which failed to create """ import pandas as pd - if isinstance(unique_columns, str): - unique_columns = [unique_columns] - if existing not in ["skip", "replace", "error"]: raise ValueError("Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'error'.") if isinstance(data, list): data = pd.json_normalize(data) + if isinstance(unique_columns, str): + unique_columns = [unique_columns] + elif unique_columns is None: + unique_columns = sorted(data.columns.to_list()) + self.log.info("Inserting %s objects.", data.shape[0]) if uuid_column is None or uuid_column not in data.columns: @@ -807,6 +806,10 @@ def create_or_replace_objects( vector_column=vector_column, uuid_column=uuid_column, ) + # drop duplicate rows, using uuid_column and unique_columns + data = data.drop_duplicates( + subset=list({*unique_columns, uuid_column} - {vector_column, None}), keep="first" + ) uuids_to_create = set() existing_uuid, non_existing_uuid = self._check_existing_objects( diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index be38ca47686bc..0b415649de3c7 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -783,7 +783,7 @@ def test_error_option_of_create_or_replace_objects(_generate_uuids, _check_exist ) _check_existing_objects.return_value = ({"1"}, {"2", "3"}) - _generate_uuids.return_value = (None, None) + _generate_uuids.return_value = (df, "id") with pytest.raises(ValueError, match=f"Found {len({'1'})} object with duplicate UUIDs. You can either"): weaviate_hook.create_or_replace_objects(data=df, class_name="test", existing="error") From dd4d5984b2a9dbbb4dd80dcb922225a6dd843080 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Tue, 19 Dec 2023 23:57:55 +0530 Subject: [PATCH 20/34] Updated the logic for object creation based on documents --- airflow/providers/weaviate/hooks/weaviate.py | 323 +++++++++++------- .../providers/weaviate/hooks/test_weaviate.py | 127 +++---- 2 files changed, 248 insertions(+), 202 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index ded51fc443459..6d58c804c1ede 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -20,7 +20,7 @@ import contextlib import json import warnings -from functools import cached_property +from functools import cached_property, partial from typing import TYPE_CHECKING, Any, Dict, List, cast import requests @@ -384,6 +384,7 @@ def batch_data( self, class_name: str, data: list[dict[str, Any]] | pd.DataFrame, + insertion_errors: list, batch_config_params: dict[str, Any] | None = None, vector_col: str = "Vector", uuid_col: str = "id", @@ -401,43 +402,51 @@ def batch_data( :param retry_attempts_per_object: number of time to try in case of failure before giving up. :param tenant: The tenant to which the object will be added. :param uuid_col: Name of the column containing the UUID. + :param insertion_errors: list to hold errors while inserting. """ client = self.conn if not batch_config_params: batch_config_params = {} + + # configuration for context manager for __exit__ method to callback on errors for weaviate + # batch ingestion. + if not batch_config_params.get("callback"): + batch_config_params.update({"callback": partial(self.process_batch_errors, insertion_errors)}) + + if not batch_config_params.get("timeout_retries"): + batch_config_params.update({"timeout_retries": 5}) + + if not batch_config_params.get("connection_error_retries"): + batch_config_params.update({"connection_error_retries": 5}) + client.batch.configure(**batch_config_params) data = self._convert_dataframe_to_list(data) - insertion_errors = [] with client.batch as batch: # Batch import all data - try: - for index, data_obj in enumerate(data): - for attempt in Retrying( - stop=stop_after_attempt(retry_attempts_per_object), - retry=( - retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) - | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) - ), - ): - with attempt: - vector = data_obj.pop(vector_col, None) - uuid = data_obj.pop(uuid_col, None) - self.log.debug( - "Attempt %s of inserting object with uuid: %s", - attempt.retry_state.attempt_number, - uuid, - ) - batch.add_data_object( - data_object=data_obj, - class_name=class_name, - vector=vector, - uuid=uuid, - tenant=tenant, - ) - self.log.debug("Inserted object with uuid: %s", uuid) - except Exception as e: - insertion_errors.append({"uuid": uuid, "result": {"errors": str(e)}}) - self.log.error("Failed to add object with UUID %s. Error: %s", uuid, e) + for index, data_obj in enumerate(data): + for attempt in Retrying( + stop=stop_after_attempt(retry_attempts_per_object), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), + ): + with attempt: + vector = data_obj.pop(vector_col, None) + uuid = data_obj.pop(uuid_col, None) + self.log.debug( + "Attempt %s of inserting object with uuid: %s", + attempt.retry_state.attempt_number, + uuid, + ) + batch.add_data_object( + data_object=data_obj, + class_name=class_name, + vector=vector, + uuid=uuid, + tenant=tenant, + ) + self.log.debug("Inserted object with uuid: %s", uuid) return insertion_errors def query_with_vector( @@ -628,6 +637,35 @@ def object_exists(self, uuid: str | UUID, **kwargs) -> bool: client = self.conn return client.data_object.exists(uuid, **kwargs) + def _delete_objects(self, uuids: Collection, class_name: str, retry_attempts_per_object: int = 5): + """ + Helper function for `create_or_replace_objects()` to delete multiple objects. + + :param uuids: Collection of uuids. + :param class_name: Name of the class in Weaviate schema where data is to be ingested. + :param retry_attempts_per_object: number of time to try in case of failure before giving up. + """ + for uuid in uuids: + for attempt in Retrying( + stop=stop_after_attempt(retry_attempts_per_object), + retry=( + retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) + | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) + ), + ): + with attempt: + try: + self.delete_object(uuid=uuid, class_name=class_name) + self.log.debug("Deleted object with uuid %s", uuid) + except weaviate.exceptions.UnexpectedStatusCodeException as e: + if e.status_code == 404: + self.log.debug("Tried to delete a non existent object with uuid %s", uuid) + else: + self.log.debug("Error occurred while trying to delete object with uuid %s", uuid) + raise e + + self.log.info("Deleted %s objects.", len(uuids)) + def _generate_uuids( self, df: pd.DataFrame, @@ -679,87 +717,125 @@ def _generate_uuids( return df, uuid_column - def _check_existing_objects(self, data: pd.DataFrame, uuid_column: str, class_name: str, existing: str): + def _check_existing_documents( + self, data: pd.DataFrame, document_column: str, class_name: str, uuid_column: str + ) -> tuple[set, set]: """ - Helper function to check if the objects with uuid exist or not. + Get all object uuids belonging to a document. :param data: A single pandas DataFrame. - :param uuid_column: Column with pre-generated UUIDs. - :param class_name: Name of the class in Weaviate schema where data is to be ingested. - :param existing: Strategy for handling existing data: 'skip', or 'replace'. - """ - existing_uuid: set = set() - non_existing_uuid: set = set() + :param document_column: The name of the property to query. + :param class_name: The name of the class to query. + :param uuid_column: The name of the column containing the UUID. + """ + offset = 0 + limit = 2000 + documents_to_uuid: dict = {} + existing_documents = set() + document_keys = set(data[document_column]) + non_existing_documents = document_keys.copy() + while True: + data_objects = ( + self.conn.query.get(properties=[document_column], class_name=class_name) + .with_additional([uuid_column]) + .with_where( + { + "operator": "Or", + "operands": [ + {"valueText": key, "path": document_column, "operator": "Equal"} + for key in document_keys + ], + } + ) + .with_offset(offset) + .with_limit(limit) + .do()["data"]["Get"][class_name] + ) + if len(data_objects) == 0: + break + for data_object in data_objects: + document_url = data_object[document_column] - if existing == "replace": - existing_uuid = set(data[uuid_column].to_list()) - else: - self.log.info("checking if %s objects exists.", data.shape[0]) - for uuid in data[uuid_column]: - for attempt in Retrying( - stop=stop_after_attempt(5), - retry=( - retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) - | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) - ), - ): - with attempt: - if self.object_exists(uuid=uuid, class_name=class_name): - existing_uuid.add(uuid) - if existing == "error": - return existing_uuid, non_existing_uuid - self.log.debug("object with uuid %s exists.", uuid) - else: - non_existing_uuid.add(uuid) - self.log.debug("object with uuid %s don't exists.", uuid) + if document_url not in documents_to_uuid: + documents_to_uuid[document_url] = set() + existing_documents.add(document_url) + non_existing_documents.remove(document_url) - self.log.info( - f"Objects to override {len(existing_uuid)} and {len(non_existing_uuid)} " f"objects to create" - ) - return existing_uuid, non_existing_uuid + documents_to_uuid[document_url].add(data_object["_additional"][uuid_column]) + offset = offset + limit + return existing_documents, non_existing_documents - def _delete_objects(self, uuids: Collection, class_name: str, retry_attempts_per_object: int = 5): - """ - Helper function for `create_or_replace_objects()` to delete multiple objects. + def _delete_all_documents_objects( + self, + document_keys: list[str], + document_column: str, + class_name: str, + batch_delete_error: list | None = None, + tenant: str | None = None, + batch_config_params: dict[str, Any] | None = None, + ): + if not batch_config_params: + batch_config_params = {} - :param uuids: Collection of uuids. - :param class_name: Name of the class in Weaviate schema where data is to be ingested. - :param retry_attempts_per_object: number of time to try in case of failure before giving up. + # configuration for context manager for __exit__ method to callback on errors for weaviate + # batch ingestion. + if not batch_config_params.get("callback"): + batch_config_params.update({"callback": partial(self.process_batch_errors, batch_delete_error)}) + + self.conn.batch.configure(**batch_config_params) + + with self.conn.batch as batch: + batch.consistency_level = weaviate.data.replication.ConsistencyLevel.ALL + batch.delete_objects( + class_name=class_name, + # same where operator as in the GraphQL API + where={ + "operator": "Or", + "operands": [ + { + "path": [document_column], + "operator": "Equal", + "valueText": key, + } + for key in document_keys + ], + }, + output="verbose", + dry_run=False, + tenant=tenant, + ) + return batch_delete_error + + def process_batch_errors(self, batch_errors: list, results: list, verbose: bool = True) -> None: """ - for uuid in uuids: - for attempt in Retrying( - stop=stop_after_attempt(retry_attempts_per_object), - retry=( - retry_if_exception(lambda exc: check_http_error_is_retryable(exc)) - | retry_if_exception_type(REQUESTS_EXCEPTIONS_TYPES) - ), - ): - with attempt: - try: - self.delete_object(uuid=uuid, class_name=class_name) - self.log.debug("Deleted object with uuid %s", uuid) - except weaviate.exceptions.UnexpectedStatusCodeException as e: - if e.status_code == 404: - self.log.debug("Tried to delete a non existent object with uuid %s", uuid) - else: - self.log.debug("Error occurred while trying to delete object with uuid %s", uuid) - raise e + Processes the results from batch operation and collects any errors. - self.log.info("Deleted %s objects.", len(uuids)) + :param batch_errors: list to populate in case of error + :param results: Results from the batch operation. + :param verbose: Flag to enable verbose logging. + """ + for item in results: + if "errors" in item["result"]: + item_error = {"uuid": item["id"], "errors": item["result"]["errors"]} + if verbose: + self.log.info( + f"Error occurred in batch process for {item['id']} with error {item['result']['errors']}" + ) + batch_errors.append(item_error) - def create_or_replace_objects( + def create_or_replace_document_objects( self, data: pd.DataFrame | list[dict[str, Any]], class_name: str, + document_column: str, existing: str = "skip", - unique_columns: list[str] | str | None = None, uuid_column: str | None = None, vector_column: str = "Vector", batch_config_params: dict | None = None, tenant: str | None = None, - ) -> list: + ): """ - create or replace objects. + create or replace objects belonging to documents. Provides users with multiple ways of dealing with existing values. 1. replace: replace the existing object with new object. This option requires to identify the rows @@ -768,12 +844,11 @@ def create_or_replace_objects( 2. skip: skip the existing objects. 3. error: raise an error if an existing object is found. - :param data: A single pandas DataFrame or a list of dicts to be ingested. :param class_name: Name of the class in Weaviate schema where data is to be ingested. :param existing: Strategy for handling existing data: 'skip', or 'replace'. Default is 'skip'. - :param unique_columns: Columns in DataFrame or keys in dict uniquely identifying each document, - required for 'upsert' operations. + :param document_column: Column in DataFrame that identifying source document, required for 'replace' + operation. :param uuid_column: Column with pre-generated UUIDs. If not provided, UUIDs will be generated. :param vector_column: Column with embedding vectors for pre-embedded data. :param batch_config_params: Additional parameters for Weaviate batch configuration. @@ -788,10 +863,7 @@ def create_or_replace_objects( if isinstance(data, list): data = pd.json_normalize(data) - if isinstance(unique_columns, str): - unique_columns = [unique_columns] - elif unique_columns is None: - unique_columns = sorted(data.columns.to_list()) + unique_columns = sorted(data.columns.to_list()) self.log.info("Inserting %s objects.", data.shape[0]) @@ -806,41 +878,54 @@ def create_or_replace_objects( vector_column=vector_column, uuid_column=uuid_column, ) - # drop duplicate rows, using uuid_column and unique_columns - data = data.drop_duplicates( - subset=list({*unique_columns, uuid_column} - {vector_column, None}), keep="first" - ) - uuids_to_create = set() - existing_uuid, non_existing_uuid = self._check_existing_objects( - data=data, uuid_column=uuid_column, class_name=class_name, existing=existing + # drop duplicate rows, using uuid_column and unique_columns. Removed `None` as it can be added to + # set when `uuid_column` is None. + data = data.drop_duplicates(subset=[document_column, uuid_column], keep="first") + batch_delete_error: list = [] + existing_documents, non_existing_documents = self._check_existing_documents( + data=data, + document_column=document_column, + uuid_column=uuid_column, + class_name=class_name, ) - if existing == "error" and len(existing_uuid): - self.log.info("Found duplicate UUIDs %s", " ,".join(existing_uuid)) + if existing == "error" and len(existing_documents): raise ValueError( - f"Found {len(existing_uuid)} object with duplicate UUIDs. You can either ignore or replace" + f"Documents {', '.join(existing_documents)} already exists. You can either skip or replace" f" them by passing 'existing=skip' or 'existing=replace' respectively." ) - elif existing == "replace": - uuids_to_create = existing_uuid.union(non_existing_uuid) - self._delete_objects(existing_uuid, class_name=class_name) elif existing == "skip": - uuids_to_create = non_existing_uuid - data = data[data[uuid_column].isin(uuids_to_create)] - insertion_errors = [] + data = data[data[document_column].isin(non_existing_documents)] + elif existing == "replace": + batch_delete_error = self._delete_all_documents_objects( + document_keys=list(existing_documents), + document_column=document_column, + class_name=class_name, + batch_delete_error=batch_delete_error, + tenant=tenant, + batch_config_params=batch_config_params, + ) + data = data[data[document_column].isin(non_existing_documents.union(existing_documents))] + + insertion_errors: list = [] if data.shape[0]: self.log.info("Batch inserting %s objects.", data.shape[0]) insertion_errors = self.batch_data( class_name=class_name, data=data, + insertion_errors=insertion_errors, batch_config_params=batch_config_params, vector_col=vector_column, uuid_col=uuid_column, tenant=tenant, ) - if insertion_errors: - self.log.info("Failed to insert %s objects.", len(insertion_errors)) + if insertion_errors or batch_delete_error: + if insertion_errors: + self.log.info("Failed to insert %s objects.", len(insertion_errors)) + if batch_delete_error: + self.log.info("Failed to delete %s objects.", len(insertion_errors)) # Rollback object that were not created properly - self._delete_objects([item["uuid"] for item in insertion_errors], class_name=class_name) - + self._delete_objects( + [item["uuid"] for item in insertion_errors + batch_delete_error], class_name=class_name + ) return insertion_errors diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 0b415649de3c7..716e133f0cee8 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -18,7 +18,7 @@ from contextlib import ExitStack from unittest import mock -from unittest.mock import MagicMock, Mock +from unittest.mock import ANY, MagicMock, Mock import pandas as pd import pytest @@ -428,7 +428,7 @@ def test_batch_data(data, expected_length, weaviate_hook): test_class_name = "TestClass" # Test the batch_data method - weaviate_hook.batch_data(test_class_name, data) + weaviate_hook.batch_data(test_class_name, data, insertion_errors=[]) # Assert that the batch_data method was called with the correct arguments mock_client.batch.configure.assert_called_once() @@ -446,7 +446,7 @@ def test_batch_data_retry(get_conn, weaviate_hook): error.response = response side_effect = [None, error, None, error, None] get_conn.return_value.batch.__enter__.return_value.add_data_object.side_effect = side_effect - weaviate_hook.batch_data("TestClass", data) + weaviate_hook.batch_data("TestClass", data, insertion_errors=[]) assert get_conn.return_value.batch.__enter__.return_value.add_data_object.call_count == len(side_effect) @@ -697,61 +697,6 @@ def test___generate_uuids(generate_uuid5, weaviate_hook): ) -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.object_exists") -def test_error__check_existing_objects(object_exists, weaviate_hook): - df = pd.DataFrame.from_dict( - { - "id": ["1", "2", "3"], - "name": ["ross", "bob", "joy"], - "age": ["12", "22", "15"], - "gender": ["m", "m", "f"], - } - ) - object_exists.return_value = True - existing_uuid, non_existing_uuid = weaviate_hook._check_existing_objects( - data=df, uuid_column="id", class_name="test", existing="error" - ) - assert existing_uuid == {"1"} - - -def test_replace__check_existing_objects(weaviate_hook): - df = pd.DataFrame.from_dict( - { - "id": ["1", "2", "3"], - "name": ["ross", "bob", "joy"], - "age": ["12", "22", "15"], - "gender": ["m", "m", "f"], - } - ) - existing_uuid, non_existing_uuid = weaviate_hook._check_existing_objects( - data=df, uuid_column="id", class_name="test", existing="replace" - ) - assert existing_uuid == {"1", "2", "3"} - - -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.object_exists") -def test_skip__check_existing_objects(object_exists, weaviate_hook): - df = pd.DataFrame.from_dict( - { - "id": ["1", "2", "3"], - "name": ["ross", "bob", "joy"], - "age": ["12", "22", "15"], - "gender": ["m", "m", "f"], - } - ) - resp = requests.Response() - resp.status_code = 429 - requests.exceptions.HTTPError(response=resp) - http_exception = requests.exceptions.HTTPError(response=resp) - object_exists.side_effect = [True, http_exception, http_exception, True, False] - existing_uuid, non_existing_uuid = weaviate_hook._check_existing_objects( - data=df, uuid_column="id", class_name="test", existing="skip" - ) - assert existing_uuid == {"1", "2"} - assert non_existing_uuid == {"3"} - assert object_exists.call_count == 5 - - @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.delete_object") def test__delete_objects(delete_object, weaviate_hook): resp = requests.Response() @@ -770,30 +715,37 @@ def test__delete_objects(delete_object, weaviate_hook): assert delete_object.call_count == 5 -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_documents") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") -def test_error_option_of_create_or_replace_objects(_generate_uuids, _check_existing_objects, weaviate_hook): +def test_error_option_of_create_or_replace_document_objects( + _generate_uuids, _check_existing_documents, weaviate_hook +): df = pd.DataFrame.from_dict( { "id": ["1", "2", "3"], "name": ["ross", "bob", "joy"], "age": ["12", "22", "15"], "gender": ["m", "m", "f"], + "doc": ["abc.xml", "zyx.html", "zyx.html"], } ) - _check_existing_objects.return_value = ({"1"}, {"2", "3"}) + _check_existing_documents.return_value = ({"abc.xml"}, {"zyx.html"}) _generate_uuids.return_value = (df, "id") - with pytest.raises(ValueError, match=f"Found {len({'1'})} object with duplicate UUIDs. You can either"): - weaviate_hook.create_or_replace_objects(data=df, class_name="test", existing="error") + with pytest.raises( + ValueError, match="Documents abc.xml already exists. You can either" " skip or replace" + ): + weaviate_hook.create_or_replace_document_objects( + data=df, document_column="doc", class_name="test", existing="error" + ) @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._delete_objects") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.batch_data") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_documents") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") -def test_skip_option_of_create_or_replace_objects( - _generate_uuids, _check_existing_objects, batch_data, _delete_objects, weaviate_hook +def test_skip_option_of_create_or_replace_document_objects( + _generate_uuids, _check_existing_documents, batch_data, _delete_objects, weaviate_hook ): df = pd.DataFrame.from_dict( { @@ -801,32 +753,30 @@ def test_skip_option_of_create_or_replace_objects( "name": ["ross", "bob", "joy"], "age": ["12", "22", "15"], "gender": ["m", "m", "f"], + "doc": ["abc.xml", "zyx.html", "zyx.html"], } ) class_name = "test" - existing_uuid, non_existing_uuid = ({"1"}, {"2", "3"}) - batch_data_return_value = [{"uuid": i} for i in non_existing_uuid] - batch_data.return_value = batch_data_return_value - _check_existing_objects.return_value = (existing_uuid, non_existing_uuid) + existing_documents, non_existing_documents = ({"abc.xml"}, {"zyx.html"}) + _check_existing_documents.return_value = (existing_documents, non_existing_documents) _generate_uuids.return_value = (df, "id") - insertion_errors = weaviate_hook.create_or_replace_objects( - data=df, class_name=class_name, existing="skip" + weaviate_hook.create_or_replace_document_objects( + data=df, class_name=class_name, existing="skip", document_column="doc" ) - assert batch_data_return_value == insertion_errors pd.testing.assert_frame_equal( - batch_data.call_args_list[0].kwargs["data"], df[df["id"].isin(non_existing_uuid)] + batch_data.call_args_list[0].kwargs["data"], df[df["doc"].isin(non_existing_documents)] ) -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._delete_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._delete_all_documents_objects") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.batch_data") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_objects") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_documents") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") -def test_replace_option_of_create_or_replace_objects( - _generate_uuids, _check_existing_objects, batch_data, _delete_objects, weaviate_hook +def test_replace_option_of_create_or_replace_document_objects( + _generate_uuids, _check_existing_documents, batch_data, _delete_all_documents_objects, weaviate_hook ): df = pd.DataFrame.from_dict( { @@ -834,16 +784,27 @@ def test_replace_option_of_create_or_replace_objects( "name": ["ross", "bob", "joy"], "age": ["12", "22", "15"], "gender": ["m", "m", "f"], + "doc": ["abc.xml", "zyx.html", "zyx.html"], } ) class_name = "test" - existing_uuid, non_existing_uuid = ({"1"}, {"2", "3"}) + existing_documents, non_existing_documents = ({"abc.xml"}, {"zyx.html"}) batch_data.return_value = [] - _check_existing_objects.return_value = (existing_uuid, non_existing_uuid) + _check_existing_documents.return_value = (existing_documents, non_existing_documents) _generate_uuids.return_value = (df, "id") - weaviate_hook.create_or_replace_objects(data=df, class_name=class_name, existing="replace") - _delete_objects.assert_called_with(existing_uuid, class_name=class_name) + weaviate_hook.create_or_replace_document_objects( + data=df, class_name=class_name, existing="replace", document_column="doc" + ) + _delete_all_documents_objects.assert_called_with( + document_keys=list(existing_documents), + class_name=class_name, + document_column="doc", + batch_delete_error=ANY, + tenant=None, + batch_config_params=ANY, + ) pd.testing.assert_frame_equal( - batch_data.call_args_list[0].kwargs["data"], df[df["id"].isin(existing_uuid.union(non_existing_uuid))] + batch_data.call_args_list[0].kwargs["data"], + df[df["doc"].isin(existing_documents.union(non_existing_documents))], ) From 23beef43afa36ad36c3c8ce7437ee148af3414d1 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Wed, 20 Dec 2023 12:02:27 +0530 Subject: [PATCH 21/34] Update docstring for 'create_or_replace_document_objects' method --- airflow/providers/weaviate/hooks/weaviate.py | 23 +++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 6d58c804c1ede..af836971c38ec 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -446,7 +446,7 @@ def batch_data( uuid=uuid, tenant=tenant, ) - self.log.debug("Inserted object with uuid: %s", uuid) + self.log.debug("Inserted object with uuid: %s into batch", uuid) return insertion_errors def query_with_vector( @@ -837,18 +837,25 @@ def create_or_replace_document_objects( """ create or replace objects belonging to documents. + In real-world scenarios, information sources like Airflow docs, Stack Overflow, or other issues + are considered 'documents' here. It's crucial to keep the database objects in sync with these sources. + If any changes occur in these documents, this function aims to reflect those changes in the database. + + Note: This function assumes responsibility for identifying changes in documents, dropping relevant + database objects, and recreating them based on updated information. It's crucial to handle this + process with care, ensuring backups and validation are in place to prevent data loss or + inconsistencies. + Provides users with multiple ways of dealing with existing values. - 1. replace: replace the existing object with new object. This option requires to identify the rows - uniquely, which by default is done by using all columns(except `vector`) to create a uuid. - User can modify this behaviour by providing `unique_columns` params. - 2. skip: skip the existing objects. - 3. error: raise an error if an existing object is found. + 1. replace: replace the existing objects with new objects. This option requires to identify the + objects belonging to a document. which by default is done by using document_column field. + 2. skip: skip the existing objects and only add the missing objects of a document. + 3. error: raise an error if an object belonging to a existing document is tried to be created. :param data: A single pandas DataFrame or a list of dicts to be ingested. :param class_name: Name of the class in Weaviate schema where data is to be ingested. :param existing: Strategy for handling existing data: 'skip', or 'replace'. Default is 'skip'. - :param document_column: Column in DataFrame that identifying source document, required for 'replace' - operation. + :param document_column: Column in DataFrame that identifying source document. :param uuid_column: Column with pre-generated UUIDs. If not provided, UUIDs will be generated. :param vector_column: Column with embedding vectors for pre-embedded data. :param batch_config_params: Additional parameters for Weaviate batch configuration. From 6698529cb2d9c6853613cc69a462130e43aa4297 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 00:35:10 +0530 Subject: [PATCH 22/34] refactored code --- airflow/providers/weaviate/hooks/weaviate.py | 268 ++++++++++++++----- 1 file changed, 198 insertions(+), 70 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index af836971c38ec..84e986a245f30 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -20,7 +20,7 @@ import contextlib import json import warnings -from functools import cached_property, partial +from functools import cached_property from typing import TYPE_CHECKING, Any, Dict, List, cast import requests @@ -35,7 +35,7 @@ from airflow.hooks.base import BaseHook if TYPE_CHECKING: - from typing import Collection, Literal, Sequence + from typing import Callable, Collection, Literal, Sequence import pandas as pd from weaviate import ConsistencyLevel @@ -367,6 +367,7 @@ def check_subset_of_schema(self, classes_objects: list) -> bool: """ # When the class properties are not in same order or not the same length. We convert them to dicts # with property `name` as the key. This way we ensure, the subset is checked. + classes_objects = self._convert_properties_to_dict(classes_objects) exiting_classes_list = self._convert_properties_to_dict(self.get_schema()["classes"]) @@ -404,6 +405,44 @@ def batch_data( :param uuid_col: Name of the column containing the UUID. :param insertion_errors: list to hold errors while inserting. """ + data = self._convert_dataframe_to_list(data) + total_results = 0 + error_results = 0 + + def _process_batch_errors( + results: list, + verbose: bool = True, + ) -> None: + """ + Helper function to processes the results from insert or delete batch operation and collects any errors. + + :param results: Results from the batch operation. + :param verbose: Flag to enable verbose logging. + """ + nonlocal total_results + nonlocal error_results + total_batch_results = len(results) + error_batch_results = 0 + for item in results: + if "errors" in item["result"]: + error_batch_results = error_batch_results + 1 + item_error = {"uuid": item["id"], "errors": item["result"]["errors"]} + if verbose: + self.log.info( + f"Error occurred in batch process for {item['id']} with error {item['result']['errors']}" + ) + insertion_errors.append(item_error) + if verbose: + total_results = total_results + (total_batch_results - error_batch_results) + error_results = error_results + error_batch_results + + self.log.info( + "Total Objects %s / Objects %s successfully inserted and Objects %s had errors.", + len(data), + total_results, + error_results, + ) + client = self.conn if not batch_config_params: batch_config_params = {} @@ -411,7 +450,7 @@ def batch_data( # configuration for context manager for __exit__ method to callback on errors for weaviate # batch ingestion. if not batch_config_params.get("callback"): - batch_config_params.update({"callback": partial(self.process_batch_errors, insertion_errors)}) + batch_config_params.update({"callback": _process_batch_errors}) if not batch_config_params.get("timeout_retries"): batch_config_params.update({"timeout_retries": 5}) @@ -420,7 +459,6 @@ def batch_data( batch_config_params.update({"connection_error_retries": 5}) client.batch.configure(**batch_config_params) - data = self._convert_dataframe_to_list(data) with client.batch as batch: # Batch import all data for index, data_obj in enumerate(data): @@ -717,11 +755,10 @@ def _generate_uuids( return df, uuid_column - def _check_existing_documents( - self, data: pd.DataFrame, document_column: str, class_name: str, uuid_column: str - ) -> tuple[set, set]: - """ - Get all object uuids belonging to a document. + def _get_documents_to_uuid_map( + self, data: pd.DataFrame, document_column: str, uuid_column: str, class_name: str + ) -> dict[str, set]: + """Helper function to get the document to uuid map of existing objects in db. :param data: A single pandas DataFrame. :param document_column: The name of the property to query. @@ -731,9 +768,7 @@ def _check_existing_documents( offset = 0 limit = 2000 documents_to_uuid: dict = {} - existing_documents = set() document_keys = set(data[document_column]) - non_existing_documents = document_keys.copy() while True: data_objects = ( self.conn.query.get(properties=[document_column], class_name=class_name) @@ -753,79 +788,121 @@ def _check_existing_documents( ) if len(data_objects) == 0: break - for data_object in data_objects: - document_url = data_object[document_column] + offset = offset + limit + documents_to_uuid.update( + self._prepare_document_to_uuid_map( + data=data_objects, + group_key=document_column, + get_value=lambda x: x["_additional"][uuid_column], + ) + ) + return documents_to_uuid - if document_url not in documents_to_uuid: - documents_to_uuid[document_url] = set() - existing_documents.add(document_url) - non_existing_documents.remove(document_url) + @staticmethod + def _prepare_document_to_uuid_map( + data: list[dict], group_key: str, get_value: Callable[[dict], str] + ) -> dict[str, set]: + """Helper function to prepare the map of grouped_key to set.""" + grouped_key_to_set: dict = {} + for item in data: + document_url = item[group_key] - documents_to_uuid[document_url].add(data_object["_additional"][uuid_column]) - offset = offset + limit - return existing_documents, non_existing_documents + if document_url not in grouped_key_to_set: + grouped_key_to_set[document_url] = set() + + grouped_key_to_set[document_url].add(get_value(item)) + return grouped_key_to_set + + def _check_existing_documents( + self, data: pd.DataFrame, document_column: str, class_name: str, uuid_column: str + ) -> tuple[dict[str, set], set, set, set]: + """ + Get all object uuids belonging to a document. + + :param data: A single pandas DataFrame. + :param document_column: The name of the property to query. + :param class_name: The name of the class to query. + :param uuid_column: The name of the column containing the UUID. + """ + changed_documents = set() + unchanged_docs = set() + non_existing_documents = set() + documents_to_uuid = self._get_documents_to_uuid_map( + data=data, uuid_column=uuid_column, document_column=document_column, class_name=class_name + ) + + insert_documents_to_uuid = self._prepare_document_to_uuid_map( + data=data.to_dict("records"), + group_key=document_column, + get_value=lambda x: x[uuid_column], + ) + + for doc_url, doc_set in insert_documents_to_uuid.items(): + if doc_url in documents_to_uuid: + if documents_to_uuid[doc_url].difference(doc_set) or doc_set.difference( + documents_to_uuid[doc_url] + ): + changed_documents.add(doc_url) + else: + unchanged_docs.add(doc_url) + else: + non_existing_documents.add(doc_url) + + return documents_to_uuid, changed_documents, unchanged_docs, non_existing_documents def _delete_all_documents_objects( self, document_keys: list[str], document_column: str, class_name: str, + total_objects_count: int = 1, batch_delete_error: list | None = None, tenant: str | None = None, batch_config_params: dict[str, Any] | None = None, + verbose: bool = False, ): if not batch_config_params: batch_config_params = {} - # configuration for context manager for __exit__ method to callback on errors for weaviate - # batch ingestion. - if not batch_config_params.get("callback"): - batch_config_params.update({"callback": partial(self.process_batch_errors, batch_delete_error)}) + # This limit is imposed by Weavaite database + MAX_LIMIT_ON_TOTAL_DELETABLE_OBJECTS = 10000 self.conn.batch.configure(**batch_config_params) - with self.conn.batch as batch: batch.consistency_level = weaviate.data.replication.ConsistencyLevel.ALL - batch.delete_objects( - class_name=class_name, - # same where operator as in the GraphQL API - where={ - "operator": "Or", - "operands": [ - { - "path": [document_column], - "operator": "Equal", - "valueText": key, - } - for key in document_keys - ], - }, - output="verbose", - dry_run=False, - tenant=tenant, - ) - return batch_delete_error - - def process_batch_errors(self, batch_errors: list, results: list, verbose: bool = True) -> None: - """ - Processes the results from batch operation and collects any errors. - - :param batch_errors: list to populate in case of error - :param results: Results from the batch operation. - :param verbose: Flag to enable verbose logging. - """ - for item in results: - if "errors" in item["result"]: - item_error = {"uuid": item["id"], "errors": item["result"]["errors"]} + while total_objects_count > 0: + document_objects = batch.delete_objects( + class_name=class_name, + where={ + "operator": "Or", + "operands": [ + { + "path": [document_column], + "operator": "Equal", + "valueText": key, + } + for key in document_keys + ], + }, + output="verbose", + dry_run=False, + tenant=tenant, + ) + total_objects_count = total_objects_count - MAX_LIMIT_ON_TOTAL_DELETABLE_OBJECTS + matched_objects = document_objects["results"]["matches"] + batch_delete_error = [ + {"uuid": obj["id"]} + for obj in document_objects["results"]["objects"] + if "error" in obj["status"] + ] if verbose: - self.log.info( - f"Error occurred in batch process for {item['id']} with error {item['result']['errors']}" - ) - batch_errors.append(item_error) + self.log.info("Deleted %s Objects", matched_objects) + + return batch_delete_error def create_or_replace_document_objects( self, - data: pd.DataFrame | list[dict[str, Any]], + data: pd.DataFrame | list[dict[str, Any]] | list[pd.DataFrame], class_name: str, document_column: str, existing: str = "skip", @@ -833,6 +910,7 @@ def create_or_replace_document_objects( vector_column: str = "Vector", batch_config_params: dict | None = None, tenant: str | None = None, + verbose: bool = False, ): """ create or replace objects belonging to documents. @@ -860,6 +938,7 @@ def create_or_replace_document_objects( :param vector_column: Column with embedding vectors for pre-embedded data. :param batch_config_params: Additional parameters for Weaviate batch configuration. :param tenant: The tenant to which the object will be added. + :param verbose: Flag to enable verbose output during the ingestion process. :return: list of UUID which failed to create """ import pandas as pd @@ -867,12 +946,19 @@ def create_or_replace_document_objects( if existing not in ["skip", "replace", "error"]: raise ValueError("Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'error'.") - if isinstance(data, list): - data = pd.json_normalize(data) + # data = list(data) + + if isinstance(data, list) and len(data) and isinstance(data[0], dict): + data = cast(pd.DataFrame, pd.json_normalize(data)) + elif isinstance(data, list) and len(data) and isinstance(data[0], pd.DataFrame): + data = cast(pd.DataFrame, pd.concat(data, ignore_index=True)) + else: + data = cast(pd.DataFrame, data) unique_columns = sorted(data.columns.to_list()) - self.log.info("Inserting %s objects.", data.shape[0]) + if verbose: + self.log.info("%s objects came in for insertion.", data.shape[0]) if uuid_column is None or uuid_column not in data.columns: ( @@ -889,30 +975,65 @@ def create_or_replace_document_objects( # drop duplicate rows, using uuid_column and unique_columns. Removed `None` as it can be added to # set when `uuid_column` is None. data = data.drop_duplicates(subset=[document_column, uuid_column], keep="first") + if verbose: + self.log.info("%s objects remain after deduplication.", data.shape[0]) + batch_delete_error: list = [] - existing_documents, non_existing_documents = self._check_existing_documents( + ( + documents_to_uuid_map, + changed_documents, + unchanged_docs, + non_existing_documents, + ) = self._check_existing_documents( data=data, document_column=document_column, uuid_column=uuid_column, class_name=class_name, ) - if existing == "error" and len(existing_documents): + if verbose: + self.log.info( + "Found %s changed documents, %s unchanged documents and %s non-existing documents", + len(changed_documents), + len(unchanged_docs), + len(non_existing_documents), + ) + for document in changed_documents: + self.log.info( + "Changed document: %s has %s objects.", document, len(documents_to_uuid_map[document]) + ) + + self.log.info("Non-existing document: %s", ", ".join(non_existing_documents)) + + if existing == "error" and len(changed_documents): raise ValueError( - f"Documents {', '.join(existing_documents)} already exists. You can either skip or replace" + f"Documents {', '.join(changed_documents)} already exists. You can either skip or replace" f" them by passing 'existing=skip' or 'existing=replace' respectively." ) elif existing == "skip": data = data[data[document_column].isin(non_existing_documents)] + if verbose: + self.log.info( + "Since existing=skip, ingesting only non-existing document's object %s", data.shape[0] + ) elif existing == "replace": + total_objects_count = sum([len(documents_to_uuid_map[doc]) for doc in changed_documents]) + if verbose: + self.log.info( + "Since existing='replace', deleting %s objects belonging changed documents %s", + total_objects_count, + changed_documents, + ) batch_delete_error = self._delete_all_documents_objects( - document_keys=list(existing_documents), + document_keys=list(changed_documents), + total_objects_count=total_objects_count, document_column=document_column, class_name=class_name, batch_delete_error=batch_delete_error, tenant=tenant, batch_config_params=batch_config_params, + verbose=verbose, ) - data = data[data[document_column].isin(non_existing_documents.union(existing_documents))] + data = data[data[document_column].isin(non_existing_documents.union(changed_documents))] insertion_errors: list = [] if data.shape[0]: @@ -935,4 +1056,11 @@ def create_or_replace_document_objects( self._delete_objects( [item["uuid"] for item in insertion_errors + batch_delete_error], class_name=class_name ) - return insertion_errors + + if verbose: + self.log.info( + "Total objects in class %s : %s ", + class_name, + self.conn.query.aggregate(class_name).with_meta_count().do(), + ) + return insertion_errors, batch_delete_error From 14ceffdb873541665149649a4a2e1807695f3592 Mon Sep 17 00:00:00 2001 From: Utkarsh Sharma Date: Fri, 22 Dec 2023 00:40:05 +0530 Subject: [PATCH 23/34] Update airflow/providers/weaviate/hooks/weaviate.py Co-authored-by: Josh Fell <48934154+josh-fell@users.noreply.github.com> --- airflow/providers/weaviate/hooks/weaviate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 84e986a245f30..2826887096726 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -681,7 +681,7 @@ def _delete_objects(self, uuids: Collection, class_name: str, retry_attempts_per :param uuids: Collection of uuids. :param class_name: Name of the class in Weaviate schema where data is to be ingested. - :param retry_attempts_per_object: number of time to try in case of failure before giving up. + :param retry_attempts_per_object: number of times to try in case of failure before giving up. """ for uuid in uuids: for attempt in Retrying( From 1642f31de8b334855338e10f95010c6154b00377 Mon Sep 17 00:00:00 2001 From: Utkarsh Sharma Date: Fri, 22 Dec 2023 00:41:26 +0530 Subject: [PATCH 24/34] Update airflow/providers/weaviate/hooks/weaviate.py Co-authored-by: Josh Fell <48934154+josh-fell@users.noreply.github.com> --- airflow/providers/weaviate/hooks/weaviate.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 2826887096726..7d55ec5a873c2 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -919,10 +919,11 @@ def create_or_replace_document_objects( are considered 'documents' here. It's crucial to keep the database objects in sync with these sources. If any changes occur in these documents, this function aims to reflect those changes in the database. - Note: This function assumes responsibility for identifying changes in documents, dropping relevant - database objects, and recreating them based on updated information. It's crucial to handle this - process with care, ensuring backups and validation are in place to prevent data loss or - inconsistencies. + .. note:: + This function assumes responsibility for identifying changes in documents, dropping relevant + database objects, and recreating them based on updated information. It's crucial to handle this + process with care, ensuring backups and validation are in place to prevent data loss or + inconsistencies. Provides users with multiple ways of dealing with existing values. 1. replace: replace the existing objects with new objects. This option requires to identify the From 34b92ed67213b07c26522ba0c4a94eb14cbd56ae Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 00:48:33 +0530 Subject: [PATCH 25/34] Remove typo --- airflow/providers/weaviate/hooks/weaviate.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 7d55ec5a873c2..77c546693a97e 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -947,8 +947,6 @@ def create_or_replace_document_objects( if existing not in ["skip", "replace", "error"]: raise ValueError("Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'error'.") - # data = list(data) - if isinstance(data, list) and len(data) and isinstance(data[0], dict): data = cast(pd.DataFrame, pd.json_normalize(data)) elif isinstance(data, list) and len(data) and isinstance(data[0], pd.DataFrame): From 4f2e4a3a5a4e294578e1b5f094d1beaed140f7cc Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 15:13:55 +0530 Subject: [PATCH 26/34] Add better names --- airflow/providers/weaviate/hooks/weaviate.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 77c546693a97e..752422d35129f 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -827,28 +827,27 @@ def _check_existing_documents( changed_documents = set() unchanged_docs = set() non_existing_documents = set() - documents_to_uuid = self._get_documents_to_uuid_map( + existing_documents_to_uuid = self._get_documents_to_uuid_map( data=data, uuid_column=uuid_column, document_column=document_column, class_name=class_name ) - insert_documents_to_uuid = self._prepare_document_to_uuid_map( + input_documents_to_uuid = self._prepare_document_to_uuid_map( data=data.to_dict("records"), group_key=document_column, get_value=lambda x: x[uuid_column], ) - for doc_url, doc_set in insert_documents_to_uuid.items(): - if doc_url in documents_to_uuid: - if documents_to_uuid[doc_url].difference(doc_set) or doc_set.difference( - documents_to_uuid[doc_url] - ): + # segregate documents into changed, unchanged and non-existing documents. + for doc_url, doc_set in input_documents_to_uuid.items(): + if doc_url in existing_documents_to_uuid: + if existing_documents_to_uuid[doc_url] != doc_set: changed_documents.add(doc_url) else: unchanged_docs.add(doc_url) else: non_existing_documents.add(doc_url) - return documents_to_uuid, changed_documents, unchanged_docs, non_existing_documents + return existing_documents_to_uuid, changed_documents, unchanged_docs, non_existing_documents def _delete_all_documents_objects( self, From 5c464aaa3f9d68033ac68f68bd7d4fdc7d69e0fb Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 16:35:17 +0530 Subject: [PATCH 27/34] Add testcases --- airflow/providers/weaviate/hooks/weaviate.py | 2 +- .../providers/weaviate/hooks/test_weaviate.py | 41 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 752422d35129f..585e6353e28c4 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -1032,10 +1032,10 @@ def create_or_replace_document_objects( verbose=verbose, ) data = data[data[document_column].isin(non_existing_documents.union(changed_documents))] + self.log.info("Batch inserting %s objects for non-existing and changed documents.", data.shape[0]) insertion_errors: list = [] if data.shape[0]: - self.log.info("Batch inserting %s objects.", data.shape[0]) insertion_errors = self.batch_data( class_name=class_name, data=data, diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 716e133f0cee8..6788b277d461c 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -715,6 +715,47 @@ def test__delete_objects(delete_object, weaviate_hook): assert delete_object.call_count == 5 +def test__prepare_document_to_uuid_map(weaviate_hook): + input_data = [ + {"id": "1", "name": "ross", "age": "12", "gender": "m"}, + {"id": "2", "name": "bob", "age": "22", "gender": "m"}, + {"id": "3", "name": "joy", "age": "15", "gender": "f"}, + ] + grouped_data = weaviate_hook._prepare_document_to_uuid_map( + data=input_data, group_key="gender", get_value=lambda x: x["name"] + ) + assert grouped_data == {"m": {"ross", "bob"}, "f": {"joy"}} + + +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._prepare_document_to_uuid_map") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._get_documents_to_uuid_map") +def test___check_existing_documents(_get_documents_to_uuid_map, _prepare_document_to_uuid_map, weaviate_hook): + _get_documents_to_uuid_map.return_value = { + "abc.doc": {"uuid1", "uuid2", "uuid2"}, + "xyz.doc": {"uuid4", "uuid5"}, + "dfg.doc": {"uuid8", "uuid0", "uuid12"}, + } + _prepare_document_to_uuid_map.return_value = { + "abc.doc": {"uuid1", "uuid56", "uuid2"}, + "xyz.doc": {"uuid4", "uuid5"}, + "hjk.doc": {"uuid8", "uuid0", "uuid12"}, + } + ( + _, + changed_documents, + unchanged_docs, + non_existing_documents, + ) = weaviate_hook._check_existing_documents( + data=pd.DataFrame(), + document_column="doc_key", + uuid_column="id", + class_name="doc", + ) + assert changed_documents == {"abc.doc"} + assert unchanged_docs == {"xyz.doc"} + assert non_existing_documents == {"hjk.doc"} + + @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_documents") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") def test_error_option_of_create_or_replace_document_objects( From 4f506cca17388a9e13f6dd4f4dc1a449ba2a017f Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 18:16:48 +0530 Subject: [PATCH 28/34] Addressed PR comments --- airflow/providers/weaviate/hooks/weaviate.py | 47 +++++++++++++++----- 1 file changed, 36 insertions(+), 11 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 585e6353e28c4..4770bafa3be0a 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -28,6 +28,7 @@ from tenacity import Retrying, retry, retry_if_exception, retry_if_exception_type, stop_after_attempt from weaviate import Client as WeaviateClient from weaviate.auth import AuthApiKey, AuthBearerToken, AuthClientCredentials, AuthClientPassword +from weaviate.data.replication import ConsistencyLevel from weaviate.exceptions import ObjectAlreadyExistsException from weaviate.util import generate_uuid5 @@ -38,7 +39,6 @@ from typing import Callable, Collection, Literal, Sequence import pandas as pd - from weaviate import ConsistencyLevel from weaviate.types import UUID ExitingSchemaOptions = Literal["replace", "fail", "ignore"] @@ -397,13 +397,13 @@ def batch_data( :param class_name: The name of the class that objects belongs to. :param data: list or dataframe of objects we want to add. + :param insertion_errors: list to hold errors while inserting. :param batch_config_params: dict of batch configuration option. .. seealso:: `batch_config_params options `__ :param vector_col: name of the column containing the vector. + :param uuid_col: Name of the column containing the UUID. :param retry_attempts_per_object: number of time to try in case of failure before giving up. :param tenant: The tenant to which the object will be added. - :param uuid_col: Name of the column containing the UUID. - :param insertion_errors: list to hold errors while inserting. """ data = self._convert_dataframe_to_list(data) total_results = 0 @@ -429,7 +429,9 @@ def _process_batch_errors( item_error = {"uuid": item["id"], "errors": item["result"]["errors"]} if verbose: self.log.info( - f"Error occurred in batch process for {item['id']} with error {item['result']['errors']}" + "Error occurred in batch process for %s with error %s", + item["id"], + item["result"]["errors"], ) insertion_errors.append(item_error) if verbose: @@ -450,13 +452,13 @@ def _process_batch_errors( # configuration for context manager for __exit__ method to callback on errors for weaviate # batch ingestion. if not batch_config_params.get("callback"): - batch_config_params.update({"callback": _process_batch_errors}) + batch_config_params["callback"] = _process_batch_errors if not batch_config_params.get("timeout_retries"): - batch_config_params.update({"timeout_retries": 5}) + batch_config_params["timeout_retries"] = 5 if not batch_config_params.get("connection_error_retries"): - batch_config_params.update({"connection_error_retries": 5}) + batch_config_params["connection_error_retries"] = 5 client.batch.configure(**batch_config_params) with client.batch as batch: @@ -756,7 +758,13 @@ def _generate_uuids( return df, uuid_column def _get_documents_to_uuid_map( - self, data: pd.DataFrame, document_column: str, uuid_column: str, class_name: str + self, + data: pd.DataFrame, + document_column: str, + uuid_column: str, + class_name: str, + offset: int = 0, + limit: int = 2000, ) -> dict[str, set]: """Helper function to get the document to uuid map of existing objects in db. @@ -764,9 +772,9 @@ def _get_documents_to_uuid_map( :param document_column: The name of the property to query. :param class_name: The name of the class to query. :param uuid_column: The name of the column containing the UUID. + :param offset: pagination parameter to indicate the which object to start fetching data. + :param limit: pagination param to indicate the number of records to fetch from start object. """ - offset = 0 - limit = 2000 documents_to_uuid: dict = {} document_keys = set(data[document_column]) while True: @@ -860,6 +868,20 @@ def _delete_all_documents_objects( batch_config_params: dict[str, Any] | None = None, verbose: bool = False, ): + """Delete all object that belong to list of documents. + + :param document_keys: list of unique documents identifiers. + :param document_column: Column in DataFrame that identifying source document. + :param class_name: Name of the class in Weaviate schema where data is to be ingested. + :param total_objects_count: total number of objects to delete, needed as max limit on one delete + query is 10,000, if we have more objects to delete we need to run query multiple times. + :param batch_delete_error: list to hold errors while inserting. + :param tenant: The tenant to which the object will be added. + :param batch_config_params: Additional parameters for Weaviate batch configuration. + :param verbose: Flag to enable verbose output during the ingestion process. + """ + batch_delete_error = batch_delete_error or [] + if not batch_config_params: batch_config_params = {} @@ -868,7 +890,10 @@ def _delete_all_documents_objects( self.conn.batch.configure(**batch_config_params) with self.conn.batch as batch: - batch.consistency_level = weaviate.data.replication.ConsistencyLevel.ALL + # ConsistencyLevel.ALL is essential here to guarantee complete deletion of objects + # across all nodes. Maintaining this level ensures data integrity, preventing + # irrelevant objects from providing misleading context for LLM models. + batch.consistency_level = ConsistencyLevel.ALL while total_objects_count > 0: document_objects = batch.delete_objects( class_name=class_name, From 48dc2136ab397775d7ba97e6b5aeb0144ff6752c Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 18:59:55 +0530 Subject: [PATCH 29/34] Fix testcases --- airflow/providers/weaviate/hooks/weaviate.py | 30 ++++----- .../providers/weaviate/hooks/test_weaviate.py | 66 ++++++++++++------- 2 files changed, 59 insertions(+), 37 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 4770bafa3be0a..fa772ec0e556f 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -821,11 +821,11 @@ def _prepare_document_to_uuid_map( grouped_key_to_set[document_url].add(get_value(item)) return grouped_key_to_set - def _check_existing_documents( + def _get_segregated_documents( self, data: pd.DataFrame, document_column: str, class_name: str, uuid_column: str ) -> tuple[dict[str, set], set, set, set]: """ - Get all object uuids belonging to a document. + Segregate documents into changed, unchanged and new document, when compared to Weaviate db. :param data: A single pandas DataFrame. :param document_column: The name of the property to query. @@ -834,7 +834,7 @@ def _check_existing_documents( """ changed_documents = set() unchanged_docs = set() - non_existing_documents = set() + new_documents = set() existing_documents_to_uuid = self._get_documents_to_uuid_map( data=data, uuid_column=uuid_column, document_column=document_column, class_name=class_name ) @@ -853,9 +853,9 @@ def _check_existing_documents( else: unchanged_docs.add(doc_url) else: - non_existing_documents.add(doc_url) + new_documents.add(doc_url) - return existing_documents_to_uuid, changed_documents, unchanged_docs, non_existing_documents + return existing_documents_to_uuid, changed_documents, unchanged_docs, new_documents def _delete_all_documents_objects( self, @@ -972,9 +972,9 @@ def create_or_replace_document_objects( raise ValueError("Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'error'.") if isinstance(data, list) and len(data) and isinstance(data[0], dict): - data = cast(pd.DataFrame, pd.json_normalize(data)) + data = pd.json_normalize(data) elif isinstance(data, list) and len(data) and isinstance(data[0], pd.DataFrame): - data = cast(pd.DataFrame, pd.concat(data, ignore_index=True)) + data = pd.concat(data, ignore_index=True) else: data = cast(pd.DataFrame, data) @@ -1005,9 +1005,9 @@ def create_or_replace_document_objects( ( documents_to_uuid_map, changed_documents, - unchanged_docs, - non_existing_documents, - ) = self._check_existing_documents( + unchanged_documents, + new_documents, + ) = self._get_segregated_documents( data=data, document_column=document_column, uuid_column=uuid_column, @@ -1017,15 +1017,15 @@ def create_or_replace_document_objects( self.log.info( "Found %s changed documents, %s unchanged documents and %s non-existing documents", len(changed_documents), - len(unchanged_docs), - len(non_existing_documents), + len(unchanged_documents), + len(new_documents), ) for document in changed_documents: self.log.info( "Changed document: %s has %s objects.", document, len(documents_to_uuid_map[document]) ) - self.log.info("Non-existing document: %s", ", ".join(non_existing_documents)) + self.log.info("Non-existing document: %s", ", ".join(new_documents)) if existing == "error" and len(changed_documents): raise ValueError( @@ -1033,7 +1033,7 @@ def create_or_replace_document_objects( f" them by passing 'existing=skip' or 'existing=replace' respectively." ) elif existing == "skip": - data = data[data[document_column].isin(non_existing_documents)] + data = data[data[document_column].isin(new_documents)] if verbose: self.log.info( "Since existing=skip, ingesting only non-existing document's object %s", data.shape[0] @@ -1056,7 +1056,7 @@ def create_or_replace_document_objects( batch_config_params=batch_config_params, verbose=verbose, ) - data = data[data[document_column].isin(non_existing_documents.union(changed_documents))] + data = data[data[document_column].isin(new_documents.union(changed_documents))] self.log.info("Batch inserting %s objects for non-existing and changed documents.", data.shape[0]) insertion_errors: list = [] diff --git a/tests/providers/weaviate/hooks/test_weaviate.py b/tests/providers/weaviate/hooks/test_weaviate.py index 6788b277d461c..5b8d5ce65d55d 100644 --- a/tests/providers/weaviate/hooks/test_weaviate.py +++ b/tests/providers/weaviate/hooks/test_weaviate.py @@ -18,7 +18,7 @@ from contextlib import ExitStack from unittest import mock -from unittest.mock import ANY, MagicMock, Mock +from unittest.mock import MagicMock, Mock import pandas as pd import pytest @@ -729,7 +729,7 @@ def test__prepare_document_to_uuid_map(weaviate_hook): @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._prepare_document_to_uuid_map") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._get_documents_to_uuid_map") -def test___check_existing_documents(_get_documents_to_uuid_map, _prepare_document_to_uuid_map, weaviate_hook): +def test___get_segregated_documents(_get_documents_to_uuid_map, _prepare_document_to_uuid_map, weaviate_hook): _get_documents_to_uuid_map.return_value = { "abc.doc": {"uuid1", "uuid2", "uuid2"}, "xyz.doc": {"uuid4", "uuid5"}, @@ -744,8 +744,8 @@ def test___check_existing_documents(_get_documents_to_uuid_map, _prepare_documen _, changed_documents, unchanged_docs, - non_existing_documents, - ) = weaviate_hook._check_existing_documents( + new_documents, + ) = weaviate_hook._get_segregated_documents( data=pd.DataFrame(), document_column="doc_key", uuid_column="id", @@ -753,13 +753,13 @@ def test___check_existing_documents(_get_documents_to_uuid_map, _prepare_documen ) assert changed_documents == {"abc.doc"} assert unchanged_docs == {"xyz.doc"} - assert non_existing_documents == {"hjk.doc"} + assert new_documents == {"hjk.doc"} -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_documents") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._get_segregated_documents") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") def test_error_option_of_create_or_replace_document_objects( - _generate_uuids, _check_existing_documents, weaviate_hook + _generate_uuids, _get_segregated_documents, weaviate_hook ): df = pd.DataFrame.from_dict( { @@ -771,7 +771,7 @@ def test_error_option_of_create_or_replace_document_objects( } ) - _check_existing_documents.return_value = ({"abc.xml"}, {"zyx.html"}) + _get_segregated_documents.return_value = ({}, {"abc.xml"}, {}, {"zyx.html"}) _generate_uuids.return_value = (df, "id") with pytest.raises( ValueError, match="Documents abc.xml already exists. You can either" " skip or replace" @@ -783,10 +783,10 @@ def test_error_option_of_create_or_replace_document_objects( @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._delete_objects") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.batch_data") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_documents") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._get_segregated_documents") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") def test_skip_option_of_create_or_replace_document_objects( - _generate_uuids, _check_existing_documents, batch_data, _delete_objects, weaviate_hook + _generate_uuids, _get_segregated_documents, batch_data, _delete_objects, weaviate_hook ): df = pd.DataFrame.from_dict( { @@ -799,8 +799,18 @@ def test_skip_option_of_create_or_replace_document_objects( ) class_name = "test" - existing_documents, non_existing_documents = ({"abc.xml"}, {"zyx.html"}) - _check_existing_documents.return_value = (existing_documents, non_existing_documents) + documents_to_uuid_map, changed_documents, unchanged_documents, new_documents = ( + {}, + {"abc.xml"}, + {}, + {"zyx.html"}, + ) + _get_segregated_documents.return_value = ( + documents_to_uuid_map, + changed_documents, + unchanged_documents, + new_documents, + ) _generate_uuids.return_value = (df, "id") weaviate_hook.create_or_replace_document_objects( @@ -808,16 +818,16 @@ def test_skip_option_of_create_or_replace_document_objects( ) pd.testing.assert_frame_equal( - batch_data.call_args_list[0].kwargs["data"], df[df["doc"].isin(non_existing_documents)] + batch_data.call_args_list[0].kwargs["data"], df[df["doc"].isin(new_documents)] ) @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._delete_all_documents_objects") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook.batch_data") -@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._check_existing_documents") +@mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._get_segregated_documents") @mock.patch("airflow.providers.weaviate.hooks.weaviate.WeaviateHook._generate_uuids") def test_replace_option_of_create_or_replace_document_objects( - _generate_uuids, _check_existing_documents, batch_data, _delete_all_documents_objects, weaviate_hook + _generate_uuids, _get_segregated_documents, batch_data, _delete_all_documents_objects, weaviate_hook ): df = pd.DataFrame.from_dict( { @@ -830,22 +840,34 @@ def test_replace_option_of_create_or_replace_document_objects( ) class_name = "test" - existing_documents, non_existing_documents = ({"abc.xml"}, {"zyx.html"}) + documents_to_uuid_map, changed_documents, unchanged_documents, new_documents = ( + {"abc.xml": {"uuid"}}, + {"abc.xml"}, + {}, + {"zyx.html"}, + ) batch_data.return_value = [] - _check_existing_documents.return_value = (existing_documents, non_existing_documents) + _get_segregated_documents.return_value = ( + documents_to_uuid_map, + changed_documents, + unchanged_documents, + new_documents, + ) _generate_uuids.return_value = (df, "id") weaviate_hook.create_or_replace_document_objects( data=df, class_name=class_name, existing="replace", document_column="doc" ) _delete_all_documents_objects.assert_called_with( - document_keys=list(existing_documents), - class_name=class_name, + document_keys=list(changed_documents), + total_objects_count=1, document_column="doc", - batch_delete_error=ANY, + class_name="test", + batch_delete_error=[], tenant=None, - batch_config_params=ANY, + batch_config_params=None, + verbose=False, ) pd.testing.assert_frame_equal( batch_data.call_args_list[0].kwargs["data"], - df[df["doc"].isin(existing_documents.union(non_existing_documents))], + df[df["doc"].isin(changed_documents.union(new_documents))], ) From c208b72c0896d5131762076c861da1391c15c982 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 20:22:51 +0530 Subject: [PATCH 30/34] Fix docstring --- airflow/providers/weaviate/hooks/weaviate.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index fa772ec0e556f..981b78c35b5ce 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -569,7 +569,7 @@ def get_or_create_object( :param class_name: Class name associated with the object given. This is required to create a new object. :param vector: Vector associated with the object given. This argument is only used when creating object. :param consistency_level: Consistency level to be used. Applies to both create and get operations. - :tenant: Tenant to be used. Applies to both create and get operations. + :param tenant: Tenant to be used. Applies to both create and get operations. :param kwargs: Additional parameters to be passed to weaviate_client.data_object.create() and weaviate_client.data_object.get() """ @@ -944,16 +944,17 @@ def create_or_replace_document_objects( If any changes occur in these documents, this function aims to reflect those changes in the database. .. note:: + This function assumes responsibility for identifying changes in documents, dropping relevant database objects, and recreating them based on updated information. It's crucial to handle this process with care, ensuring backups and validation are in place to prevent data loss or inconsistencies. Provides users with multiple ways of dealing with existing values. - 1. replace: replace the existing objects with new objects. This option requires to identify the + replace: replace the existing objects with new objects. This option requires to identify the objects belonging to a document. which by default is done by using document_column field. - 2. skip: skip the existing objects and only add the missing objects of a document. - 3. error: raise an error if an object belonging to a existing document is tried to be created. + skip: skip the existing objects and only add the missing objects of a document. + error: raise an error if an object belonging to a existing document is tried to be created. :param data: A single pandas DataFrame or a list of dicts to be ingested. :param class_name: Name of the class in Weaviate schema where data is to be ingested. From 4eaf7d6f9704c291469f236b2f13884e3d392c9d Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 21:11:49 +0530 Subject: [PATCH 31/34] Fix docstring --- airflow/providers/weaviate/hooks/weaviate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 981b78c35b5ce..7a76fbbda8923 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -63,7 +63,7 @@ def check_http_error_is_retryable(exc: BaseException): class WeaviateHook(BaseHook): """ - Interact with Weaviate database to store vectors. This hook uses the `conn_id`. + Interact with Weaviate database to store vectors. This hook uses the 'conn_id'. :param conn_id: The connection id to use when connecting to Weaviate. """ @@ -952,7 +952,7 @@ def create_or_replace_document_objects( Provides users with multiple ways of dealing with existing values. replace: replace the existing objects with new objects. This option requires to identify the - objects belonging to a document. which by default is done by using document_column field. + objects belonging to a document. which by default is done by using document_column field. skip: skip the existing objects and only add the missing objects of a document. error: raise an error if an object belonging to a existing document is tried to be created. From 383180068fa35b7808e4e0304266525107290fc3 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Fri, 22 Dec 2023 21:41:49 +0530 Subject: [PATCH 32/34] Address static code issue --- airflow/providers/weaviate/hooks/weaviate.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 7a76fbbda8923..0a131e792010d 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -972,9 +972,12 @@ def create_or_replace_document_objects( if existing not in ["skip", "replace", "error"]: raise ValueError("Invalid parameter for 'existing'. Choices are 'skip', 'replace', 'error'.") - if isinstance(data, list) and len(data) and isinstance(data[0], dict): + if len(data) == 0: + return [] + + if isinstance(data, list) and isinstance(data[0], dict): data = pd.json_normalize(data) - elif isinstance(data, list) and len(data) and isinstance(data[0], pd.DataFrame): + elif isinstance(data, list) and isinstance(data[0], pd.DataFrame): data = pd.concat(data, ignore_index=True) else: data = cast(pd.DataFrame, data) From d7a5356b11b35771bf3bb9f13c618d95960c9872 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Sat, 23 Dec 2023 10:44:03 +0530 Subject: [PATCH 33/34] Fix typing issue --- airflow/providers/weaviate/hooks/weaviate.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 0a131e792010d..80a96a3827938 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -976,9 +976,13 @@ def create_or_replace_document_objects( return [] if isinstance(data, list) and isinstance(data[0], dict): - data = pd.json_normalize(data) + # This is done to narrow the type to list[dict[str, Any]. + dict_list: list[dict[str, Any]] = data + data = pd.json_normalize(dict_list) elif isinstance(data, list) and isinstance(data[0], pd.DataFrame): - data = pd.concat(data, ignore_index=True) + # This is done to narrow the type to list[pd.DataFrame]. + df_list: list[pd.DataFrame] = data + data = pd.concat(df_list, ignore_index=True) else: data = cast(pd.DataFrame, data) From 15a087bd2344735ccfffb1a67f1cc9bfe581fc37 Mon Sep 17 00:00:00 2001 From: utkarsh sharma Date: Sat, 23 Dec 2023 11:11:59 +0530 Subject: [PATCH 34/34] Fix typing issue --- airflow/providers/weaviate/hooks/weaviate.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/airflow/providers/weaviate/hooks/weaviate.py b/airflow/providers/weaviate/hooks/weaviate.py index 80a96a3827938..d0b8db37cb51a 100644 --- a/airflow/providers/weaviate/hooks/weaviate.py +++ b/airflow/providers/weaviate/hooks/weaviate.py @@ -976,13 +976,11 @@ def create_or_replace_document_objects( return [] if isinstance(data, list) and isinstance(data[0], dict): - # This is done to narrow the type to list[dict[str, Any]. - dict_list: list[dict[str, Any]] = data - data = pd.json_normalize(dict_list) + # This is done to narrow the type to List[Dict[str, Any]. + data = pd.json_normalize(cast(List[Dict[str, Any]], data)) elif isinstance(data, list) and isinstance(data[0], pd.DataFrame): - # This is done to narrow the type to list[pd.DataFrame]. - df_list: list[pd.DataFrame] = data - data = pd.concat(df_list, ignore_index=True) + # This is done to narrow the type to List[pd.DataFrame]. + data = pd.concat(cast(List[pd.DataFrame], data), ignore_index=True) else: data = cast(pd.DataFrame, data)