Skip to content

fix: Address slicing for pyarrow array in data_color #741

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

fix: Address slicing for pyarrow array in data_color #741

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ff4e4fd
fix: Address slicing pyarrow array in data_color
FBruzzesi Jul 30, 2025
c0af9e6
rename function
FBruzzesi Jul 30, 2025
1f09f61
collapse pyarrow into one func:
FBruzzesi Aug 2, 2025
e810e7a
Merge branch 'main' into fix/data-color-pyarrow
FBruzzesi Aug 9, 2025
c1c0b3c
Merge branch 'main' into fix/data-color-pyarrow
FBruzzesi Aug 18, 2025
657eacf
Merge branch 'main' into fix/data-color-pyarrow
FBruzzesi Sep 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions great_tables/_data_color/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,14 @@
from typing_extensions import TypeAlias

from great_tables._locations import RowSelectExpr, resolve_cols_c, resolve_rows_i
from great_tables._tbl_data import DataFrameLike, SelectExpr, get_column_names, is_na
from great_tables._tbl_data import (
DataFrameLike,
SelectExpr,
get_column_names,
get_rows,
is_na,
to_list,
)
from great_tables.loc import body
from great_tables.style import fill, text

Expand Down Expand Up @@ -227,7 +234,7 @@ def data_color(
# For each column targeted, get the data values as a new list object
for col in columns_resolved:
# This line handles both pandas and polars dataframes
column_vals = data_table[col][row_pos].to_list()
column_vals = to_list(get_rows(data_table[col], indexes=row_pos))

# Filter out NA values from `column_vals`
filtered_column_vals = [x for x in column_vals if not is_na(data_table, x)]
Expand Down
30 changes: 26 additions & 4 deletions great_tables/_tbl_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,10 @@
PlSelectExpr = _selector_proxy_
PlExpr = pl.Expr

PdSeries = pd.Series
PdSeries = pd.Series[Any]
PlSeries = pl.Series
PyArrowArray = pa.Array
PyArrowChunkedArray = pa.ChunkedArray
PyArrowArray = pa.Array[Any]
PyArrowChunkedArray = pa.ChunkedArray[Any]

PdNA = pd.NA
PlNull = pl.Null
Expand Down Expand Up @@ -763,7 +763,7 @@ def _(df: PyArrowTable, x: Any) -> bool:
import pyarrow as pa

arr = pa.array([x])
return arr.is_null().to_pylist()[0] or arr.is_nan().to_pylist()[0]
return arr.is_null(nan_is_null=True).to_pylist()[0]


@singledispatch
Expand Down Expand Up @@ -874,3 +874,25 @@ def _(ser: PyArrowChunkedArray, name: Optional[str] = None) -> PyArrowTable:
import pyarrow as pa

return pa.table({name: ser})


@singledispatch
def get_rows(ser: SeriesLike, indexes: list[int]) -> SeriesLike:
"""Returns values of the series at `indexes` position.`"""
raise NotImplementedError(f"Unsupported type: {type(ser)}")


@get_rows.register
def _(ser: PdSeries, indexes: list[int]) -> PdSeries:
return ser.iloc[indexes]
Copy link
Contributor Author

@FBruzzesi FBruzzesi Jul 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it was intended to use index location in the first place but the issue never arise?



@get_rows.register
def _(ser: PlSeries, indexes: list[int]) -> PlSeries:
return ser[indexes]


@get_rows.register(PyArrowArray)
@get_rows.register(PyArrowChunkedArray)
def _(ser: Any, indexes: list[int]) -> PyArrowArray | PyArrowChunkedArray:
return ser.take(indexes)
182 changes: 182 additions & 0 deletions tests/data_color/__snapshots__/test_data_color.ambr
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,32 @@
</tbody>
'''
# ---
# name: test_data_color_autocolor_text_false[pyarrow]
'''
<tbody class="gt_table_body">
<tr>
<td class="gt_row gt_right">0.1111</td>
<td class="gt_row gt_left">apricot</td>
<td style="background-color: #ff0000;" class="gt_row gt_right">49.95</td>
</tr>
<tr>
<td class="gt_row gt_right">2.222</td>
<td class="gt_row gt_left">banana</td>
<td style="background-color: #5c5200;" class="gt_row gt_right">17.95</td>
</tr>
<tr>
<td class="gt_row gt_right">33.33</td>
<td class="gt_row gt_left">coconut</td>
<td style="background-color: #077c00;" class="gt_row gt_right">1.39</td>
</tr>
<tr>
<td class="gt_row gt_right">444.4</td>
<td class="gt_row gt_left">durian</td>
<td style="background-color: #0000FF;" class="gt_row gt_right">65100</td>
</tr>
</tbody>
'''
# ---
# name: test_data_color_colorbrewer_snap
'''
<tbody class="gt_table_body">
Expand Down Expand Up @@ -206,6 +232,32 @@
</tbody>
'''
# ---
# name: test_data_color_domain_na_color_reverse_snap[pyarrow]
'''
<tbody class="gt_table_body">
<tr>
<td class="gt_row gt_right">0.1111</td>
<td class="gt_row gt_left">apricot</td>
<td style="color: #000000; background-color: #ff0000;" class="gt_row gt_right">49.95</td>
</tr>
<tr>
<td class="gt_row gt_right">2.222</td>
<td class="gt_row gt_left">banana</td>
<td style="color: #FFFFFF; background-color: #5c5200;" class="gt_row gt_right">17.95</td>
</tr>
<tr>
<td class="gt_row gt_right">33.33</td>
<td class="gt_row gt_left">coconut</td>
<td style="color: #FFFFFF; background-color: #077c00;" class="gt_row gt_right">1.39</td>
</tr>
<tr>
<td class="gt_row gt_right">444.4</td>
<td class="gt_row gt_left">durian</td>
<td style="color: #FFFFFF; background-color: #0000FF;" class="gt_row gt_right">65100</td>
</tr>
</tbody>
'''
# ---
# name: test_data_color_domain_na_color_snap[pandas]
'''
<tbody class="gt_table_body">
Expand Down Expand Up @@ -258,6 +310,32 @@
</tbody>
'''
# ---
# name: test_data_color_domain_na_color_snap[pyarrow]
'''
<tbody class="gt_table_body">
<tr>
<td class="gt_row gt_right">0.1111</td>
<td class="gt_row gt_left">apricot</td>
<td style="color: #FFFFFF; background-color: #008000;" class="gt_row gt_right">49.95</td>
</tr>
<tr>
<td class="gt_row gt_right">2.222</td>
<td class="gt_row gt_left">banana</td>
<td style="color: #FFFFFF; background-color: #a32e00;" class="gt_row gt_right">17.95</td>
</tr>
<tr>
<td class="gt_row gt_right">33.33</td>
<td class="gt_row gt_left">coconut</td>
<td style="color: #000000; background-color: #f80400;" class="gt_row gt_right">1.39</td>
</tr>
<tr>
<td class="gt_row gt_right">444.4</td>
<td class="gt_row gt_left">durian</td>
<td style="color: #FFFFFF; background-color: #0000FF;" class="gt_row gt_right">65100</td>
</tr>
</tbody>
'''
# ---
# name: test_data_color_overlapping_domain[pandas]
'''
<tbody class="gt_table_body">
Expand Down Expand Up @@ -310,6 +388,32 @@
</tbody>
'''
# ---
# name: test_data_color_overlapping_domain[pyarrow]
'''
<tbody class="gt_table_body">
<tr>
<td class="gt_row gt_right">0.1111</td>
<td class="gt_row gt_left">apricot</td>
<td style="color: #000000; background-color: #FF0000;" class="gt_row gt_right">49.95</td>
</tr>
<tr>
<td class="gt_row gt_right">2.222</td>
<td class="gt_row gt_left">banana</td>
<td style="color: #000000; background-color: #FF0000;" class="gt_row gt_right">17.95</td>
</tr>
<tr>
<td class="gt_row gt_right">33.33</td>
<td class="gt_row gt_left">coconut</td>
<td style="color: #000000; background-color: #FF0000;" class="gt_row gt_right">1.39</td>
</tr>
<tr>
<td class="gt_row gt_right">444.4</td>
<td class="gt_row gt_left">durian</td>
<td style="color: #FFFFFF; background-color: #673498;" class="gt_row gt_right">65100</td>
</tr>
</tbody>
'''
# ---
# name: test_data_color_palette_snap[pandas]
'''
<tbody class="gt_table_body">
Expand Down Expand Up @@ -362,6 +466,32 @@
</tbody>
'''
# ---
# name: test_data_color_palette_snap[pyarrow]
'''
<tbody class="gt_table_body">
<tr>
<td style="color: #000000; background-color: #ff0000;" class="gt_row gt_right">0.1111</td>
<td class="gt_row gt_left">apricot</td>
<td style="color: #000000; background-color: #ff0000;" class="gt_row gt_right">49.95</td>
</tr>
<tr>
<td style="color: #000000; background-color: #fe0100;" class="gt_row gt_right">2.222</td>
<td class="gt_row gt_left">banana</td>
<td style="color: #000000; background-color: #ff0000;" class="gt_row gt_right">17.95</td>
</tr>
<tr>
<td style="color: #000000; background-color: #ec0a00;" class="gt_row gt_right">33.33</td>
<td class="gt_row gt_left">coconut</td>
<td style="color: #000000; background-color: #ff0000;" class="gt_row gt_right">1.39</td>
</tr>
<tr>
<td style="color: #FFFFFF; background-color: #008000;" class="gt_row gt_right">444.4</td>
<td class="gt_row gt_left">durian</td>
<td style="color: #FFFFFF; background-color: #008000;" class="gt_row gt_right">65100</td>
</tr>
</tbody>
'''
# ---
# name: test_data_color_pd_cols_rows_snap
'''
<tbody class="gt_table_body">
Expand Down Expand Up @@ -495,6 +625,32 @@
</tbody>
'''
# ---
# name: test_data_color_simple_exibble_snap[pyarrow]
'''
<tbody class="gt_table_body">
<tr>
<td style="color: #FFFFFF; background-color: #000000;" class="gt_row gt_right">0.1111</td>
<td style="color: #FFFFFF; background-color: #000000;" class="gt_row gt_left">apricot</td>
<td style="color: #FFFFFF; background-color: #010001;" class="gt_row gt_right">49.95</td>
</tr>
<tr>
<td style="color: #FFFFFF; background-color: #070304;" class="gt_row gt_right">2.222</td>
<td style="color: #000000; background-color: #4cbd81;" class="gt_row gt_left">banana</td>
<td style="color: #FFFFFF; background-color: #000000;" class="gt_row gt_right">17.95</td>
</tr>
<tr>
<td style="color: #FFFFFF; background-color: #752b38;" class="gt_row gt_right">33.33</td>
<td style="color: #FFFFFF; background-color: #9653ca;" class="gt_row gt_left">coconut</td>
<td style="color: #FFFFFF; background-color: #000000;" class="gt_row gt_right">1.39</td>
</tr>
<tr>
<td style="color: #000000; background-color: #9e9e9e;" class="gt_row gt_right">444.4</td>
<td style="color: #000000; background-color: #9e9e9e;" class="gt_row gt_left">durian</td>
<td style="color: #000000; background-color: #9e9e9e;" class="gt_row gt_right">65100</td>
</tr>
</tbody>
'''
# ---
# name: test_data_color_subset_domain[pandas]
'''
<tbody class="gt_table_body">
Expand Down Expand Up @@ -547,6 +703,32 @@
</tbody>
'''
# ---
# name: test_data_color_subset_domain[pyarrow]
'''
<tbody class="gt_table_body">
<tr>
<td class="gt_row gt_right">0.1111</td>
<td class="gt_row gt_left">apricot</td>
<td style="color: #000000; background-color: #FF0000;" class="gt_row gt_right">49.95</td>
</tr>
<tr>
<td class="gt_row gt_right">2.222</td>
<td class="gt_row gt_left">banana</td>
<td style="color: #000000; background-color: #FF0000;" class="gt_row gt_right">17.95</td>
</tr>
<tr>
<td class="gt_row gt_right">33.33</td>
<td class="gt_row gt_left">coconut</td>
<td style="color: #000000; background-color: #FF0000;" class="gt_row gt_right">1.39</td>
</tr>
<tr>
<td class="gt_row gt_right">444.4</td>
<td class="gt_row gt_left">durian</td>
<td style="color: #000000; background-color: #FF0000;" class="gt_row gt_right">65100</td>
</tr>
</tbody>
'''
# ---
# name: test_data_color_viridis_snap
'''
<tbody class="gt_table_body">
Expand Down
2 changes: 2 additions & 0 deletions tests/data_color/test_data_color.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import numpy as np
import pandas as pd
import polars as pl
import pyarrow as pa
import pytest

from great_tables import GT, style
Expand All @@ -16,6 +17,7 @@
params_frames = [
pytest.param(pd.DataFrame, id="pandas"),
pytest.param(pl.DataFrame, id="polars"),
pytest.param(pa.table, id="pyarrow"),
]


Expand Down
Loading