fix max_page

app.py

@@ -74,24 +74,25 @@ def index(dataset: str, config: str, split: str) -> Tuple[np.ndarray, RowGroupRe
             raise AppError(f"Invalid config {config}. Available configs are: {', '.join(get_parquet_configs(dataset))}.")
         else:
             raise AppError(f"Invalid split {split}. Available splits are: {', '.join(get_parquet_splits(dataset, config))}.")
-    all_pf: List[pq.ParquetFile] = thread_map(partial(pq.ParquetFile, filesystem=fs), sources)
+    desc = f"{dataset}/{config}/{split}"
+    all_pf: List[pq.ParquetFile] = thread_map(partial(pq.ParquetFile, filesystem=fs), sources, desc=desc, unit="pq")
     features = Features.from_arrow_schema(all_pf[0].schema.to_arrow_schema())
     columns = [col for col in features if all(bad_type not in str(features[col]) for bad_type in ["Image(", "Audio(", "'binary'"])]
     info = "" if len(columns) == len(features) else f"Some columns are not supported yet: {sorted(set(features) - set(columns))}"
     rg_offsets = np.cumsum([pf.metadata.row_group(i).num_rows for pf in all_pf for i in range(pf.metadata.num_row_groups)])
     rg_readers = [partial(pf.read_row_group, i, columns=columns) for pf in all_pf for i in range(pf.metadata.num_row_groups)]
-    max_page = rg_offsets[-1] // PAGE_SIZE
+    max_page = 1 + (rg_offsets[-1] - 1) // PAGE_SIZE
     return rg_offsets, rg_readers, max_page, info
 
 
 def query(page: int, page_size: int, rg_offsets: np.ndarray, rg_readers: RowGroupReaders) -> pd.DataFrame:
-    start_row, end_row = (page - 1) * page_size, page * page_size
-    start_rg, end_rg = np.searchsorted(rg_offsets, [start_row, end_row], side="right")
-    if page < 1 or end_rg >= len(rg_readers):
-        raise AppError(f"Page {page} does not exist")
+    start_row, end_row = (page - 1) * page_size, min(page * page_size, rg_offsets[-1] - 1) # both included
+    # rg_offsets[start_rg - 1] <= start_row < rg_offsets[start_rg]
+    # rg_offsets[end_rg - 1]   <=  end_row  < rg_offsets[end_rg]
+    start_rg, end_rg = np.searchsorted(rg_offsets, [start_row, end_row], side="right") # both included
     pa_table = pa.concat_tables([rg_readers[i]() for i in range(start_rg, end_rg + 1)])
-    offset = start_row - rg_offsets[start_rg - 1] if start_rg else start_row
-    pa_table = pa_table.slice(offset, end_row - start_row)
+    offset = start_row - (rg_offsets[start_rg - 1] if start_rg > 0 else 0)
+    pa_table = pa_table.slice(offset, page_size)
     return pa_table.to_pandas()
 
 
@@ -99,6 +100,8 @@ def query(page: int, page_size: int, rg_offsets: np.ndarray, rg_readers: RowGrou
 def get_page(dataset: str, config: str, split: str, page: str) -> Tuple[str, int, str]:
     dataset, config, split, page = sanitize_inputs(dataset, config, split, page)
     rg_offsets, rg_readers, max_page, info = index(dataset, config, split)
+    if page > max_page:
+        raise AppError(f"Page {page} does not exist")
     df = query(page, PAGE_SIZE, rg_offsets=rg_offsets, rg_readers=rg_readers)
     buf = StringIO()
     df.to_json(buf, lines=True, orient="records")