feat(bigframes): support offset-based column access via iloc#17367
feat(bigframes): support offset-based column access via iloc#17367tswast wants to merge 16 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for iloc column-based assignment (__setitem__) in DataFrames and refactors iloc indexing to dynamically enforce ordering only when row subsets or specific orderings are requested. It also updates JSON column serialization in _get_display_df to use the new iloc assignment, bypassing limitations with duplicate or non-string column names. Feedback on these changes includes a recommendation to use _assign_multi_items instead of assign in __setitem__ to prevent potential TypeErrors with non-string column labels, removing a useless expression (df._block.apply_analytic) in _get_display_df, and reverting a hardcoded development project ID in the generative AI notebook to a placeholder.
tswast
changed the title
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
| if ( | ||
| isinstance(key, slice) | ||
| and (key.start is None or key.start == 0) | ||
| and (key.step is None or key.step == 1) | ||
| and key.stop is None | ||
| ): | ||
| requires_ordering = False |
There was a problem hiding this comment.
Maybe we extract this out to a helper? The idea being we want iloc[:] to work without ordering?
There was a problem hiding this comment.
Yes, helper makes sense to me. Might be able to keep it from getting out of sync that way.
Yes, that's the idea. If they select all rows, we don't care about the ordering.
| _assign_multi_items_by_offset = _assign_multi_items_by_offsets | ||
| _assign_multi_items_by_label = _assign_multi_items | ||
| _assign_multi_items_by_labels = _assign_multi_items |
There was a problem hiding this comment.
why do we need these aliases?
There was a problem hiding this comment.
Not strictly necessary, but kinda nice to be able to be clear about the usage. Alternatively, I could remove _assign_multi_items and just rename it to _assign_multi_items_by_labels. What do you think?
| col_indexer_list = list(col_indexer) | ||
|
|
||
| if len(col_indexer_list) > 0 and all( | ||
| isinstance(x, bool) for x in col_indexer_list |
There was a problem hiding this comment.
Might we need something more general than isintance to handle eg numpy bool?
There was a problem hiding this comment.
Also similar concern with other isinstance checks elsewhere
There was a problem hiding this comment.
Good point. I might be able to use the typing.abc classes instead. I'll check on this.
There was a problem hiding this comment.
Fixed in the latest commits. Added support for pyarrow scalars too.
| if isinstance(col_indexer, int): | ||
| col_offset = col_indexer | ||
| if col_offset < 0: | ||
| col_offset += n_cols | ||
| if col_offset < 0 or col_offset >= n_cols: | ||
| raise IndexError("single positional indexer is out-of-bounds") | ||
|
|
||
| df = self._dataframe._assign_multi_items_by_offsets([col_offset], value) | ||
| self._dataframe._set_block(df._get_block()) | ||
|
|
||
| elif isinstance(col_indexer, slice): | ||
| col_offsets = list(range(*col_indexer.indices(n_cols))) | ||
| if not col_offsets: | ||
| return | ||
| df = self._dataframe._assign_multi_items_by_offsets(col_offsets, value) | ||
| self._dataframe._set_block(df._get_block()) | ||
|
|
||
| elif pd.api.types.is_list_like(col_indexer): | ||
| col_indexer_list = list(col_indexer) | ||
|
|
||
| if len(col_indexer_list) > 0 and all( | ||
| isinstance(x, bool) for x in col_indexer_list | ||
| ): | ||
| if len(col_indexer_list) != n_cols: | ||
| raise ValueError( | ||
| f"Boolean index has wrong length: {len(col_indexer_list)} instead of {n_cols}" | ||
| ) | ||
| col_offsets = [i for i, val in enumerate(col_indexer_list) if val] | ||
| else: | ||
| col_offsets = [] | ||
| for idx in col_indexer_list: | ||
| if isinstance(idx, bool): | ||
| raise TypeError("pos list must contain only integers") | ||
| if not isinstance(idx, int): | ||
| raise TypeError("pos list must contain only integers") | ||
| if idx < 0: | ||
| idx += n_cols | ||
| if idx < 0 or idx >= n_cols: | ||
| raise IndexError("positional indexer is out-of-bounds") | ||
| col_offsets.append(idx) | ||
|
|
||
| if not col_offsets: | ||
| return |
There was a problem hiding this comment.
Seems all of this is about essentially converting a key to an offset list. and then doing the same _assign_multi_items_by_offsets and then _set_block sequence given that.
There was a problem hiding this comment.
That's a good callout. I added _iloc_col_indexer_to_offsets and _iloc_clip_to_offset helper functions along these lines. I also realized that getitem is sharing quite a bit of logic here too and updated it to use the same helpers.
|
|
||
| def _iloc_col_indexer_to_offsets( | ||
| df: bigframes.dataframe.DataFrame, col_indexer: Any | ||
| ) -> Tuple[Sequence[int], str]: |
There was a problem hiding this comment.
Probably better to put the return type logic in the getitem implementation even if it means a little bit of duplicate code to determine if col_indexer is a scalar integer.
There was a problem hiding this comment.
Done in latest commit.
| return False | ||
|
|
||
|
|
||
| def pyarrow_array_from_sequence(data: Any, pa_type: pa.DataType) -> pa.Array: |
There was a problem hiding this comment.
q: can we not just use the ManagedArrowTable.from_pyarrow to ingest this data? AFAICT we are using raw string datatype for json in our representation, is this not working to dodge the extension type issue?
2 participants
🦕