resolve comments

snowflakedb · Aug 16, 2024 · 2c2de22 · 2c2de22
1 parent 24dad97
commit 2c2de22
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 82 deletions.
diff --git a/src/snowflake/snowpark/modin/plugin/_internal/concat_utils.py b/src/snowflake/snowpark/modin/plugin/_internal/concat_utils.py
@@ -228,8 +228,8 @@ def union_all(
         data_column_pandas_index_names=frame1.data_column_pandas_index_names,
         index_column_pandas_labels=frame1.index_column_pandas_labels,
         index_column_snowflake_quoted_identifiers=frame1.index_column_snowflake_quoted_identifiers,
-        data_column_types=frame1.cached_data_column_snowpark_pandas_types,
-        index_column_types=frame1.cached_index_column_snowpark_pandas_types,
+        data_column_types=None,
+        index_column_types=None,
     )
 
 

diff --git a/src/snowflake/snowpark/modin/plugin/_internal/frame.py b/src/snowflake/snowpark/modin/plugin/_internal/frame.py
@@ -1098,12 +1098,12 @@ def get_updated_identifiers(identifiers: list[str]) -> list[str]:
                 self.data_column_snowflake_quoted_identifiers
             ),
             data_column_pandas_index_names=self.data_column_pandas_index_names,
-            data_column_types=self.cached_data_column_snowpark_pandas_types,
             index_column_pandas_labels=self.index_column_pandas_labels,
             index_column_snowflake_quoted_identifiers=get_updated_identifiers(
                 self.index_column_snowflake_quoted_identifiers
             ),
-            index_column_types=self.cached_index_column_snowpark_pandas_types,
+            data_column_types=None,
+            index_column_types=None,
         )
 
     def update_snowflake_quoted_identifiers_with_expressions(

diff --git a/src/snowflake/snowpark/modin/plugin/_internal/indexing_utils.py b/src/snowflake/snowpark/modin/plugin/_internal/indexing_utils.py
@@ -3163,8 +3163,16 @@ def get_row_position_index_from_bool_indexer(index: InternalFrame) -> InternalFr
         index_column_snowflake_quoted_identifiers=[
             index_column_snowflake_quoted_identifier
         ],
-        data_column_types=index.cached_data_column_snowpark_pandas_types,
-        index_column_types=index.cached_index_column_snowpark_pandas_types,
+        data_column_types=[
+            index.snowflake_quoted_identifier_to_snowpark_pandas_type.get(
+                data_column_snowflake_quoted_identifier, None
+            )
+        ],
+        index_column_types=[
+            index.snowflake_quoted_identifier_to_snowpark_pandas_type.get(
+                index_column_snowflake_quoted_identifier, None
+            )
+        ],
     )
     return index