fix: Snowflake Destination schema example and vector type embeddings (#502)

1. Fixes table schema example
2. Fixes issue with embeddings column when using vector type
3. Fixes issue with dropping/adding case insensitive column names

Author: mpolomdeepsense

CHANGELOG.md

@@ -1,3 +1,11 @@
+## 1.0.27
+
+### Fixes
+
+* **Fix table schema example for Snowflake Destination connector**
+* **Fix Snowflake Destination issue with dropping/removing case insensitive column names when populating the table**
+* **Fix Snowflake Destination issue with `embeddings` column when using `VECTOR` type**
+
 ## 1.0.26
 
 * **Fix Notion connector error with FileIcons**

test/integration/connectors/env_setup/sql/snowflake/destination/snowflake-schema.sql

@@ -1,43 +1,46 @@
 CREATE TABLE elements (
-    id UUID PRIMARY KEY,
+    id VARCHAR(36) PRIMARY KEY NOT NULL DEFAULT UUID_STRING(),
     record_id VARCHAR,
     element_id VARCHAR,
-    text TEXT,
-    embeddings ARRAY(FLOAT),
+    text VARCHAR,
+    embeddings VECTOR(FLOAT, 1536),
     type VARCHAR,
     system VARCHAR,
     layout_width DECIMAL,
     layout_height DECIMAL,
-    points TEXT,
-    url TEXT,
+    points VARCHAR,
+    url VARCHAR,
     version VARCHAR,
     date_created TIMESTAMP_TZ,
     date_modified TIMESTAMP_TZ,
     date_processed TIMESTAMP_TZ,
-    permissions_data TEXT,
-    record_locator TEXT,
+    permissions_data VARCHAR,
+    record_locator VARCHAR,
     category_depth INTEGER,
     parent_id VARCHAR,
     attached_filename VARCHAR,
     filetype VARCHAR,
     last_modified TIMESTAMP_TZ,
     file_directory VARCHAR,
     filename VARCHAR,
-    languages ARRAY(VARCHAR),
+    languages ARRAY,
     page_number VARCHAR,
-    links TEXT,
+    links VARCHAR,
     page_name VARCHAR,
-    link_urls ARRAY(VARCHAR),
-    link_texts ARRAY(VARCHAR),
-    sent_from ARRAY(VARCHAR),
-    sent_to ARRAY(VARCHAR),
+    link_urls ARRAY,
+    link_texts ARRAY,
+    sent_from ARRAY,
+    sent_to ARRAY,
     subject VARCHAR,
     section VARCHAR,
     header_footer_type VARCHAR,
-    emphasized_text_contents ARRAY(VARCHAR),
-    emphasized_text_tags ARRAY(VARCHAR),
-    text_as_html TEXT,
-    regex_metadata TEXT,
-    detection_class_prob DECIMAL
+    emphasized_text_contents ARRAY,
+    emphasized_text_tags ARRAY,
+    text_as_html VARCHAR,
+    regex_metadata VARCHAR,
+    detection_class_prob DECIMAL,
+    image_base64 VARCHAR,
+    image_mime_type VARCHAR,
+    orig_elements VARCHAR,
+    is_continuation BOOLEAN
 );
-

test/unit/connectors/sql/test_snowflake.py

@@ -0,0 +1,149 @@
+from unittest.mock import MagicMock
+
+import pytest
+from pydantic import Secret
+from pytest_mock import MockerFixture
+
+from unstructured_ingest.processes.connectors.sql.snowflake import (
+    SnowflakeAccessConfig,
+    SnowflakeConnectionConfig,
+    SnowflakeUploader,
+)
+
+
+@pytest.fixture
+def snowflake_access_config():
+    return SnowflakeAccessConfig(
+        password="password",
+    )
+
+
+@pytest.fixture
+def snowflake_connection_config(snowflake_access_config: SnowflakeAccessConfig):
+    return SnowflakeConnectionConfig(
+        account="account",
+        user="user",
+        host="host",
+        port=1234,
+        database="database",
+        schema="schema",
+        role="role",
+        access_config=Secret(snowflake_access_config),
+    )
+
+
+@pytest.fixture
+def snowflake_uploader(snowflake_connection_config: SnowflakeConnectionConfig):
+    uploader = SnowflakeUploader(
+        connection_config=snowflake_connection_config,
+        upload_config=MagicMock(table_name="test_table"),
+    )
+    uploader._embeddings_dimension = None
+    return uploader
+
+
+@pytest.fixture
+def mock_cursor(mocker: MockerFixture):
+    return mocker.MagicMock()
+
+
+@pytest.fixture
+def mock_get_cursor(mocker: MockerFixture, mock_cursor: MagicMock):
+    mock = mocker.patch(
+        "unstructured_ingest.processes.connectors.sql.snowflake.SnowflakeConnectionConfig.get_cursor",
+        autospec=True,
+    )
+    mock.return_value.__enter__.return_value = mock_cursor
+    return mock
+
+
+def test_embeddings_dimension_vector_column(
+    mock_cursor: MagicMock, snowflake_uploader: SnowflakeUploader, mock_get_cursor: MagicMock
+):
+    mock_cursor.execute.return_value.fetchone.return_value = (
+        "",
+        "",
+        "",
+        '{"type": "VECTOR", "dimension": 128}',
+    )
+
+    dimension = snowflake_uploader.embeddings_dimension
+
+    assert dimension == 128
+    mock_cursor.execute.assert_called_once_with("SHOW COLUMNS LIKE 'embeddings' IN test_table")
+
+
+def test_embeddings_dimension_vector_column_dict(
+    mock_cursor: MagicMock, snowflake_uploader: SnowflakeUploader, mock_get_cursor: MagicMock
+):
+    mock_cursor.execute.return_value.fetchone.return_value = {
+        "col_1": "test_1",
+        "col_2": "test_2",
+        "data_type": '{"type": "VECTOR", "dimension": 64}',
+    }
+
+    dimension = snowflake_uploader.embeddings_dimension
+
+    assert dimension == 64
+    mock_cursor.execute.assert_called_once_with("SHOW COLUMNS LIKE 'embeddings' IN test_table")
+
+
+def test_embeddings_dimension_no_column(
+    mock_cursor: MagicMock, snowflake_uploader: SnowflakeUploader, mock_get_cursor: MagicMock
+):
+    mock_cursor.execute.return_value.fetchone.return_value = None
+
+    dimension = snowflake_uploader.embeddings_dimension
+
+    assert dimension == 0
+    mock_cursor.execute.assert_called_once_with("SHOW COLUMNS LIKE 'embeddings' IN test_table")
+
+
+def test_embeddings_dimension_non_vector_column(
+    mock_cursor: MagicMock, snowflake_uploader: SnowflakeUploader, mock_get_cursor: MagicMock
+):
+    mock_cursor.execute.return_value.fetchone.return_value = ("", "", "", '{"type": "STRING"}')
+
+    dimension = snowflake_uploader.embeddings_dimension
+
+    assert dimension == 0
+    mock_cursor.execute.assert_called_once_with("SHOW COLUMNS LIKE 'embeddings' IN test_table")
+
+
+def test_parse_values_with_vector_columns(snowflake_uploader: SnowflakeUploader):
+    snowflake_uploader._embeddings_dimension = 128
+    columns = ["embeddings", "languages", "other_column"]
+
+    result = snowflake_uploader._parse_values(columns)
+
+    expected_result = "PARSE_JSON(?)::VECTOR(FLOAT,128),PARSE_JSON(?),?"
+    assert result == expected_result
+
+
+def test_parse_values_with_vector_columns_embedding_dimension_zero(
+    snowflake_uploader: SnowflakeUploader,
+):
+    snowflake_uploader._embeddings_dimension = 0
+    columns = ["embeddings", "languages", "other_column"]
+
+    result = snowflake_uploader._parse_values(columns)
+
+    expected_result = "PARSE_JSON(?),PARSE_JSON(?),?"
+    assert result == expected_result
+
+
+def test_parse_values_with_vector_columns_embedding_dimension_none(
+    mocker: MockerFixture, snowflake_uploader: SnowflakeUploader
+):
+    mock_embeddings_dimension = mocker.PropertyMock
+    mock_embeddings_dimension.return_value = None
+    mocker.patch(
+        "unstructured_ingest.processes.connectors.sql.snowflake.SnowflakeUploader.embeddings_dimension",
+        new_callable=mock_embeddings_dimension,
+    )
+    columns = ["embeddings", "languages", "other_column"]
+
+    result = snowflake_uploader._parse_values(columns)
+
+    expected_result = "PARSE_JSON(?),PARSE_JSON(?),?"
+    assert result == expected_result

test/unit/connectors/sql/test_sql.py

@@ -150,3 +150,39 @@ def test_fit_to_schema_no_add_missing_columns(mocker: MockerFixture, mock_upload
 
     assert "col2" not in result.columns
     assert "col1" in result.columns
+
+
+def test_fit_to_schema_case_sensitive(mocker: MockerFixture, mock_uploader: SQLUploader):
+    df = pd.DataFrame(
+        {
+            "col1": [1, 2],
+            "col2": [3, 4],
+            "col3": [5, 6],
+        }
+    )
+    mocker.patch.object(mock_uploader, "get_table_columns", return_value=["COL1", "COL2", "col3"])
+
+    result = mock_uploader._fit_to_schema(df)
+
+    assert "col1" not in result.columns
+    assert "col2" not in result.columns
+    assert "col3" in result.columns
+    assert "COL1" in result.columns
+    assert "COL2" in result.columns
+
+
+def test_fit_to_schema_not_case_sensitive(mocker: MockerFixture, mock_uploader: SQLUploader):
+    df = pd.DataFrame(
+        {
+            "col1": [1, 2],
+            "col2": [3, 4],
+            "col3": [5, 6],
+        }
+    )
+    mocker.patch.object(mock_uploader, "get_table_columns", return_value=["COL1", "COL2"])
+
+    result = mock_uploader._fit_to_schema(df, add_missing_columns=False, case_sensitive=False)
+
+    assert "col3" not in result.columns
+    assert "col1" in result.columns
+    assert "col2" in result.columns

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "1.0.26"  # pragma: no cover
+__version__ = "1.0.27"  # pragma: no cover

unstructured_ingest/processes/connectors/sql/snowflake.py

@@ -37,8 +37,8 @@
 
 CONNECTOR_TYPE = "snowflake"
 
+EMBEDDINGS_COLUMN = "embeddings"
 _ARRAY_COLUMNS = (
-    "embeddings",
     "languages",
     "link_urls",
     "link_texts",
@@ -47,6 +47,7 @@
     "emphasized_text_contents",
     "emphasized_text_tags",
 )
+_VECTOR_COLUMNS = (EMBEDDINGS_COLUMN,)
 
 
 class SnowflakeAccessConfig(SQLAccessConfig):
@@ -174,6 +175,33 @@ class SnowflakeUploader(SQLUploader):
     connector_type: str = CONNECTOR_TYPE
     values_delimiter: str = "?"
 
+    _embeddings_dimension: Optional[int] = None
+
+    @property
+    def embeddings_dimension(self) -> Optional[int]:
+        """
+        Get the dimension of the embeddings column in the Snowflake table.
+        If the column is not present or is not of type VECTOR, returns None.
+        """
+        if self._embeddings_dimension is None:
+            with self.connection_config.get_cursor() as cursor:
+                embeddings_column = cursor.execute(
+                    f"SHOW COLUMNS LIKE '{EMBEDDINGS_COLUMN}' IN {self.upload_config.table_name}"
+                ).fetchone()
+            if embeddings_column:
+                data_type = {}
+                if isinstance(embeddings_column, dict):
+                    data_type = json.loads(embeddings_column.get("data_type", "{}"))
+                elif isinstance(embeddings_column, tuple):
+                    data_type = json.loads(embeddings_column[3] or "{}")
+                if isinstance(data_type, dict) and data_type.get("type") == "VECTOR":
+                    self._embeddings_dimension = data_type.get("dimension")
+        # If the _embeddings_dimension is still None, it means the column
+        # is not present or not a VECTOR type
+        if self._embeddings_dimension is None:
+            self._embeddings_dimension = 0
+        return self._embeddings_dimension
+
     @requires_dependencies(["pandas"], extras="snowflake")
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         super().run(path=path, file_data=file_data, **kwargs)
@@ -193,7 +221,7 @@ def prepare_data(
                         parsed.append(None)
                     else:
                         parsed.append(parse_date_string(value))
-                elif column_name in _ARRAY_COLUMNS:
+                elif column_name in _ARRAY_COLUMNS or column_name in _VECTOR_COLUMNS:
                     if not isinstance(value, list) and (
                         value is None or pd.isna(value)
                     ):  # pandas is nan
@@ -206,16 +234,18 @@ def prepare_data(
         return output
 
     def _parse_values(self, columns: list[str]) -> str:
-        return ",".join(
-            [
-                (
-                    f"PARSE_JSON({self.values_delimiter})"
-                    if col in _ARRAY_COLUMNS
-                    else self.values_delimiter
+        embeddings_dimension = self.embeddings_dimension
+        parsed_values = []
+        for col in columns:
+            if col in _VECTOR_COLUMNS and embeddings_dimension:
+                parsed_values.append(
+                    f"PARSE_JSON({self.values_delimiter})::VECTOR(FLOAT,{embeddings_dimension})"
                 )
-                for col in columns
-            ]
-        )
+            elif col in _ARRAY_COLUMNS or col in _VECTOR_COLUMNS:
+                parsed_values.append(f"PARSE_JSON({self.values_delimiter})")
+            else:
+                parsed_values.append(self.values_delimiter)
+        return ",".join(parsed_values)
 
     def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
         import numpy as np
@@ -228,8 +258,8 @@ def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
                 f"record id column "
                 f"{self.upload_config.record_id_key}, skipping delete"
             )
+        df = self._fit_to_schema(df=df, add_missing_columns=True, case_sensitive=False)
         df.replace({np.nan: None}, inplace=True)
-        self._fit_to_schema(df=df)
 
         columns = list(df.columns)
         stmt = "INSERT INTO {table_name} ({columns}) SELECT {values}".format(

unstructured_ingest/processes/connectors/sql/sql.py

@@ -339,12 +339,16 @@ def prepare_data(
             output.append(tuple(parsed))
         return output
 
-    def _fit_to_schema(self, df: "DataFrame", add_missing_columns: bool = True) -> "DataFrame":
+    def _fit_to_schema(
+        self, df: "DataFrame", add_missing_columns: bool = True, case_sensitive: bool = True
+    ) -> "DataFrame":
         import pandas as pd
 
         table_columns = self.get_table_columns()
-        columns = set(df.columns)
-        schema_fields = set(table_columns)
+        columns = set(df.columns if case_sensitive else df.columns.str.lower())
+        schema_fields = set(
+            table_columns if case_sensitive else {col.lower() for col in table_columns}
+        )
         columns_to_drop = columns - schema_fields
         missing_columns = schema_fields - columns