feat(UNTRACKED): snowflake uploader executemany (#599)

Author: mpolomdeepsense

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 1.2.17
+
+* **Enhancement: Use a single `executemany` instead of per element `execute` in Snowflake Uploader**
+
 ## 1.2.16
 
 * **Fix: Catch databricks client auth errors that were being missed**

test/unit/connectors/sql/test_snowflake.py

@@ -110,29 +110,29 @@ def test_embeddings_dimension_non_vector_column(
     mock_cursor.execute.assert_called_once_with("SHOW COLUMNS LIKE 'embeddings' IN test_table")
 
 
-def test_parse_values_with_vector_columns(snowflake_uploader: SnowflakeUploader):
+def test_parse_select_with_vector_columns(snowflake_uploader: SnowflakeUploader):
     snowflake_uploader._embeddings_dimension = 128
     columns = ["embeddings", "languages", "other_column"]
 
-    result = snowflake_uploader._parse_values(columns)
+    result = snowflake_uploader._parse_select(columns)
 
-    expected_result = "PARSE_JSON(?)::VECTOR(FLOAT,128),PARSE_JSON(?),?"
+    expected_result = "PARSE_JSON($1)::VECTOR(FLOAT,128),PARSE_JSON($2),$3"
     assert result == expected_result
 
 
-def test_parse_values_with_vector_columns_embedding_dimension_zero(
+def test_parse_select_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)
+    result = snowflake_uploader._parse_select(columns)
 
-    expected_result = "PARSE_JSON(?),PARSE_JSON(?),?"
+    expected_result = "PARSE_JSON($1),PARSE_JSON($2),$3"
     assert result == expected_result
 
 
-def test_parse_values_with_vector_columns_embedding_dimension_none(
+def test_parse_select_with_vector_columns_embedding_dimension_none(
     mocker: MockerFixture, snowflake_uploader: SnowflakeUploader
 ):
     mock_embeddings_dimension = mocker.PropertyMock
@@ -143,7 +143,7 @@ def test_parse_values_with_vector_columns_embedding_dimension_none(
     )
     columns = ["embeddings", "languages", "other_column"]
 
-    result = snowflake_uploader._parse_values(columns)
+    result = snowflake_uploader._parse_select(columns)
 
-    expected_result = "PARSE_JSON(?),PARSE_JSON(?),?"
+    expected_result = "PARSE_JSON($1),PARSE_JSON($2),$3"
     assert result == expected_result

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "1.2.16"  # pragma: no cover
+__version__ = "1.2.17"  # pragma: no cover

unstructured_ingest/processes/connectors/sql/snowflake.py

@@ -233,18 +233,19 @@ def prepare_data(
             output.append(tuple(parsed))
         return output
 
-    def _parse_values(self, columns: list[str]) -> str:
+    def _parse_select(self, columns: list[str]) -> str:
         embeddings_dimension = self.embeddings_dimension
         parsed_values = []
-        for col in columns:
+        for i, col in enumerate(columns):
+            argument_selector = f"${i + 1}"
             if col in _VECTOR_COLUMNS and embeddings_dimension:
                 parsed_values.append(
-                    f"PARSE_JSON({self.values_delimiter})::VECTOR(FLOAT,{embeddings_dimension})"
+                    f"PARSE_JSON({argument_selector})::VECTOR(FLOAT,{embeddings_dimension})"
                 )
             elif col in _ARRAY_COLUMNS or col in _VECTOR_COLUMNS:
-                parsed_values.append(f"PARSE_JSON({self.values_delimiter})")
+                parsed_values.append(f"PARSE_JSON({argument_selector})")
             else:
-                parsed_values.append(self.values_delimiter)
+                parsed_values.append(argument_selector)
         return ",".join(parsed_values)
 
     def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
@@ -262,10 +263,11 @@ def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
         df.replace({np.nan: None}, inplace=True)
 
         columns = list(df.columns)
-        stmt = "INSERT INTO {table_name} ({columns}) SELECT {values}".format(
+        stmt = "INSERT INTO {table_name} ({columns}) SELECT {select} FROM VALUES ({values})".format(
             table_name=self.upload_config.table_name,
             columns=",".join(columns),
-            values=self._parse_values(columns),
+            select=self._parse_select(columns),
+            values=",".join([self.values_delimiter for _ in columns]),
         )
         logger.info(
             f"writing a total of {len(df)} elements via"
@@ -276,10 +278,7 @@ def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
         for rows in split_dataframe(df=df, chunk_size=self.upload_config.batch_size):
             with self.connection_config.get_cursor() as cursor:
                 values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
-                # TODO: executemany break on 'Binding data in type (list) is not supported'
-                for val in values:
-                    logger.debug(f"running query: {stmt}\nwith values: {val}")
-                    cursor.execute(stmt, val)
+                cursor.executemany(stmt, values)
 
 
 snowflake_source_entry = SourceRegistryEntry(