@@ -1858,17 +1858,20 @@ def _check_dataframe_all_leaves(df: nw.DataFrame) -> None:
1858
1858
null_mask = nw .any_horizontal (nw .all ())
1859
1859
).get_column ("null_mask" )
1860
1860
1861
- for row_idx , row in zip (
1862
- null_indices_mask , null_mask .filter (null_indices_mask ).iter_rows ()
1863
- ):
1864
-
1865
- i = row .index (True )
1866
-
1867
- if not all (row [i :]):
1868
- raise ValueError (
1869
- "None entries cannot have not-None children" ,
1870
- df_sorted .row (row_idx ),
1861
+ null_mask_filtered = null_mask .filter (null_indices_mask )
1862
+ if not null_mask_filtered .is_empty ():
1863
+ for col_idx in range (1 , null_mask_filtered .shape [1 ]):
1864
+ # For each row, if a True value is encountered, then check that
1865
+ # all values in subsequent columns are also True
1866
+ null_entries_with_non_null_children = (
1867
+ ~ null_mask_filtered [:, col_idx ] & null_mask_filtered [:, col_idx - 1 ]
1871
1868
)
1869
+ if nw .to_py_scalar (null_entries_with_non_null_children .any ()):
1870
+ row_idx = null_entries_with_non_null_children .to_list ().index (True )
1871
+ raise ValueError (
1872
+ "None entries cannot have not-None children" ,
1873
+ df_sorted .row (row_idx ),
1874
+ )
1872
1875
1873
1876
fill_series = nw .new_series (
1874
1877
name = "fill_value" ,
@@ -1893,13 +1896,13 @@ def _check_dataframe_all_leaves(df: nw.DataFrame) -> None:
1893
1896
)
1894
1897
1895
1898
null_indices = set (null_indices_mask .arg_true ().to_list ())
1896
- for i , (current_row , next_row ) in enumerate (
1899
+ for col_idx , (current_row , next_row ) in enumerate (
1897
1900
zip (row_strings [:- 1 ], row_strings [1 :]), start = 1
1898
1901
):
1899
- if (next_row in current_row ) and (i in null_indices ):
1902
+ if (next_row in current_row ) and (col_idx in null_indices ):
1900
1903
raise ValueError (
1901
1904
"Non-leaves rows are not permitted in the dataframe \n " ,
1902
- df_sorted .row (i ),
1905
+ df_sorted .row (col_idx ),
1903
1906
"is not a leaf." ,
1904
1907
)
1905
1908
0 commit comments