SNOW-1051741: df.apply(axis=1) should preserve the original index #3955
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
sfc-gh-jkew
added
the
NO-PANDAS-CHANGEDOC-UPDATES
sfc-gh-helmeleegy
approved these changes
Oct 24, 2025
Contributor
sfc-gh-helmeleegy
left a comment
sfc-gh-helmeleegy
left a comment
There was a problem hiding this comment.
LGTM, just had one question.
sfc-gh-jkew
commented
Oct 24, 2025
src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
Outdated
Show resolved
Hide resolved
…b/snowpark-python into jkew/apply.axis.1.row.index.0
Comment on lines
469
to
489
| if num_index_columns > 0: | ||
| # Columns after row position are index columns, then data columns | ||
| index_cols = df.iloc[:, 1 : 1 + num_index_columns] | ||
| data_cols = df.iloc[:, 1 + num_index_columns :] | ||
|
|
||
| # Set the index using the index columns | ||
| if num_index_columns == 1: | ||
| index = index_cols.iloc[:, 0] | ||
| if index_column_pandas_labels: | ||
| index.name = index_column_pandas_labels[0] | ||
| else: | ||
| # Multi-index case | ||
| index = native_pd.MultiIndex.from_arrays( | ||
| [index_cols.iloc[:, i] for i in range(num_index_columns)], | ||
| names=index_column_pandas_labels | ||
| if index_column_pandas_labels | ||
| else None, | ||
| ) | ||
| data_cols.index = index | ||
| df = data_cols | ||
| else: |
Contributor
There was a problem hiding this comment.
can't we use set_index() in both cases?
Contributor
There was a problem hiding this comment.
I meant that you can replace most of the code here with set_index(). See #3979.
src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
Outdated
Show resolved
Hide resolved
src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
Outdated
Show resolved
Hide resolved
…b/snowpark-python into jkew/apply.axis.1.row.index.0
…b/snowpark-python into jkew/apply.axis.1.row.index.0
Comment on lines
447
to
448
| input_types: Snowpark column types of the input data columns (including index columns). | ||
| index_column_pandas_labels: The pandas labels for the index columns, if any. |
Contributor
There was a problem hiding this comment.
Suggested change
| input_types: Snowpark column types of the input data columns (including index columns). | |
| index_column_pandas_labels: The pandas labels for the index columns, if any. | |
| input_types: Snowpark column types of the input data columns (including index columns). |
|
|
||
|
|
||
| @sql_count_checker(query_count=5, join_count=2, udtf_count=1) | ||
| def test_apply_axis_1_multiindex_preservation(): |
Contributor
There was a problem hiding this comment.
Could we also test
funcwith return type annotations. We'll use vectorized UDFs instead of UDTFs.funcreturning a series- apply() on series (with func typed, untyped, or returning a series)
Comment on lines
469
to
489
| if num_index_columns > 0: | ||
| # Columns after row position are index columns, then data columns | ||
| index_cols = df.iloc[:, 1 : 1 + num_index_columns] | ||
| data_cols = df.iloc[:, 1 + num_index_columns :] | ||
|
|
||
| # Set the index using the index columns | ||
| if num_index_columns == 1: | ||
| index = index_cols.iloc[:, 0] | ||
| if index_column_pandas_labels: | ||
| index.name = index_column_pandas_labels[0] | ||
| else: | ||
| # Multi-index case | ||
| index = native_pd.MultiIndex.from_arrays( | ||
| [index_cols.iloc[:, i] for i in range(num_index_columns)], | ||
| names=index_column_pandas_labels | ||
| if index_column_pandas_labels | ||
| else None, | ||
| ) | ||
| data_cols.index = index | ||
| df = data_cols | ||
| else: |
Contributor
There was a problem hiding this comment.
I meant that you can replace most of the code here with set_index(). See #3979.
Comment on lines
+9812
to
+9813
| # Determine if we should pass index columns to the UDTF | ||
| # We pass index columns when the index is not the row position itself |
Contributor
There was a problem hiding this comment.
We always pass the index column names here. We can keep doing that, but we should update the comment and make the parameter required, since there don't seem to be any other invocations of that function.
| column_index: native_pd.Index, | ||
| input_types: list[DataType], | ||
| session: Session, | ||
| index_column_labels: list[Hashable] | None = None, |
Contributor
There was a problem hiding this comment.
It turns out that just passing the number of index columns is enough:
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
df.apply(axis=1)should preserve the original index. Previously we would return a RangeIndex regardless of the original index. This approach passes the index data into the underlying UDTF.Mostly AI written approach, but with original tests for verification.
Fixes SNOW-1051741
Fill out the following pre-review checklist: