Optimize large S3 dataset listings

CHANGELOG.md

@@ -23,6 +23,22 @@ tag. Releases before 2.0.0 are recorded in the
   where it used to be a `404`. The page carries the colors, the type and the card layout of the
   documentation site, and it loads nothing from the network, so it also renders on a server with
   no route out. Swagger keeps its own path, so a bookmark to `/swagger` is unaffected.
+- **`browse_datasets` reads one directory level.** A folder view needs the files and the
+  sub-directories of one directory, and got them by listing the whole store and grouping the paths
+  in the browser. `browse_datasets([prefix[, offset[, limit]]])` reads a single level with one
+  delimiter request. On a bucket of 2 853 217 objects that measured 14 milliseconds, against 79.9
+  seconds to enumerate the same bucket. It shares its provider with `list_datasets`, so the columns
+  are the same, and a new `is_directory` column marks a sub-directory row. `GET
+  /api/browse-datasets?prefix=…` runs the same SQL and splits the rows into `folders` and
+  `datasets`, and the admin UI folder view calls it. A directory-shaped dataset (Zarr, Atlas) is a
+  folder at its own level; descend into it to see the dataset. Neither listing filters by grant
+  rules, as before. See
+  [Introspection](docs/docs/2.0.0-rc5/sql/table-functions-utility.md#browse-datasets) and
+  [the REST API](docs/docs/2.0.0-rc5/api/exploring-data.md#browse-one-directory-level).
+- **`list_datasets` takes a pattern, an offset and a limit.** All three were already accepted and
+  none were documented. `list_datasets([pattern[, offset[, limit]]])` globs (default `**/*`), skips
+  and caps, and the full row shape — `can_inspect`, `can_partial_explore`, `size`, `last_modified`
+  — is now written down beside it.
 - **`BEACON_TYPE_WIDENING_ON_CONFLICT` settles a column that no type holds.** A collection can
   type one column as a number in one file and as a string in another. No type holds both, so the
   schema merge refused the whole table and the table answered no query: `Incompatible types for
@@ -169,6 +185,36 @@ tag. Releases before 2.0.0 are recorded in the
 
 ### Changed
 
+- **A large S3 dataset listing is sharded, streamed and paged.** Beacon listed the whole datasets
+  store for every browse. A recursive listing is a chain of pages, and each page needs the
+  continuation token of the one before it, so a bucket of 2 853 217 objects cost 2854 strictly
+  sequential requests and 79.9 seconds. Three changes. Beacon asks for 5000 keys per page instead
+  of the server default of 1000, which removes 80% of the round trips. It then reads three
+  directory levels with a delimiter listing and walks each leaf prefix as its own page chain, 16 at
+  a time. The same walk now takes 19.4 seconds. And the listing streams: a plan node emits rows as
+  pages arrive, so `LIMIT 50` over a bucket of millions reads one page rather than all of it, and
+  the walk no longer holds every object in memory first. Two consequences. Shards interleave, so
+  objects no longer arrive sorted — `ObjectStore::list` never guaranteed an order, and a query that
+  needs one must say `ORDER BY`. And a directory-shaped dataset (Zarr, Atlas) reports no `size` or
+  `last_modified` in a recursive listing, because those are an aggregate over the directory and a
+  streamed row cannot wait for it. Raise `AWS_TIMEOUT` before raising the page size.
+- **The listing ran during query planning, on a worker thread.** `list_datasets` did its walk
+  inside `TableFunctionImpl::call`, a synchronous method, so it reached the store through
+  `block_in_place` and held a worker thread for the length of the walk — before anything decided to
+  read the result. The walk now belongs to the table provider, and it starts when the plan is
+  executed. `EXPLAIN SELECT … FROM list_datasets()` touches the store not at all, and names the
+  node and what it will list.
+- **A discovered Zarr store is a `*.zarr` directory.** Zarr v3 gives every group *and every array*
+  a `zarr.json`, so the marker alone does not say which one it belongs to. Beacon compared each
+  marker against the others to find the shallowest, which needs the whole listing in memory and
+  cannot be done as it streams. Discovery now reads the directory that holds the marker: only a
+  store carries the `.zarr` suffix. A store in a directory without the suffix stays out of
+  `list_datasets`. `read_zarr` is unaffected — the path already names the store, so that directory
+  needs no suffix. See [Zarr](docs/docs/2.0.0-rc5/formats/zarr.md#discovery-in-the-datasets-store).
+- **A dropped object-store error no longer shortens a dataset listing.** `list_datasets` discarded
+  a failure part-way through the walk, so a timeout on object 2 000 000 of 2 850 000 returned the
+  first two million as a complete listing. It now reports the error and names how far it got.
+
 - **A file statistics pass drains the queue, and one pass runs at a time.** A pass used to stop
   after one batch of `BEACON_FILE_STATS_BATCH_FILES` files, 10 000 by default. A fresh archive of
   a million files therefore needed 100 ticks, which is over 24 hours at the default interval of 900

beacon-clients/beacon-ts/src/client.ts

@@ -323,6 +323,19 @@ export class BeaconClient {
     return this.http.fetchJson<T>("GET", "/api/dataset-schema", { query: { file } });
   }
 
+  /**
+   * Reads one directory level of the datasets store (`GET /api/browse-datasets`).
+   *
+   * Prefer this over {@link datasets} for a folder view. `datasets` globs, so it
+   * enumerates every object under the pattern before it answers; this reads a
+   * single level and its cost does not grow with the store below it.
+   */
+  browseDatasets<T = unknown>(prefix?: string): Promise<T> {
+    return this.http.fetchJson<T>("GET", "/api/browse-datasets", {
+      query: { prefix },
+    });
+  }
+
   /** Counts the total number of datasets (`GET /api/total-datasets`). */
   totalDatasets(): Promise<number> {
     return this.http.fetchJson<number>("GET", "/api/total-datasets");

beacon-clients/beacon-ts/test/client.test.ts

@@ -175,6 +175,21 @@ describe("BeaconClient metadata paths", () => {
     expect(calls[0]!.url).toBe("http://beacon.test/api/table-schema?table_name=my+table");
   });
 
+  it("browses one directory level, and omits an absent prefix", async () => {
+    const { fn, calls } = stubFetch(() =>
+      jsonResponse({ prefix: "argo", folders: ["2024"], datasets: [] }),
+    );
+    const client = new BeaconClient({ url: "http://beacon.test", fetch: fn });
+
+    await client.browseDatasets("argo/floats");
+    expect(calls[0]!.url).toBe("http://beacon.test/api/browse-datasets?prefix=argo%2Ffloats");
+
+    // No prefix means the root, so the parameter is left off rather than sent
+    // empty.
+    await client.browseDatasets();
+    expect(calls[1]!.url).toBe("http://beacon.test/api/browse-datasets");
+  });
+
   it("respects basePath", async () => {
     const { fn, calls } = stubFetch(() => jsonResponse([]));
     const client = new BeaconClient({ url: "http://beacon.test", basePath: "/beacon", fetch: fn });

beacon-clients/beacon-web/src/pages/datasets.tsx

@@ -119,38 +119,11 @@ function normalize(raw: unknown): DatasetItem[] {
   });
 }
 
-interface FolderEntry {
-  name: string;
-  count: number;
-}
-
-/** Splits the datasets under `prefix` into immediate sub-folders and files. */
-function browse(items: DatasetItem[], prefix: string, filterText: string) {
-  const folderCounts = new Map<string, number>();
-  const files: DatasetItem[] = [];
-  for (const it of items) {
-    if (prefix && !it.path.startsWith(prefix)) continue;
-    const rest = it.path.slice(prefix.length);
-    if (!rest) continue;
-    const slash = rest.indexOf("/");
-    if (slash >= 0) {
-      const folder = rest.slice(0, slash);
-      folderCounts.set(folder, (folderCounts.get(folder) ?? 0) + 1);
-    } else {
-      files.push(it);
-    }
-  }
-  const f = filterText.trim().toLowerCase();
-  let folders: FolderEntry[] = [...folderCounts.entries()]
-    .map(([name, count]) => ({ name, count }))
-    .sort((a, b) => a.name.localeCompare(b.name));
-  files.sort((a, b) => a.path.localeCompare(b.path));
-  let fileList = files;
-  if (f) {
-    folders = folders.filter((x) => x.name.toLowerCase().includes(f));
-    fileList = files.filter((x) => baseName(x.path).toLowerCase().includes(f));
-  }
-  return { folders, files: fileList };
+/** `GET /api/browse-datasets` — one directory level. */
+interface BrowseResponse {
+  prefix: string;
+  folders: string[];
+  datasets: unknown[];
 }
 
 const baseName = (p: string) => p.slice(p.lastIndexOf("/") + 1);
@@ -165,7 +138,7 @@ function parentPrefix(p: string): string {
 /** One rendered row: the up (..) entry, a sub-folder, or a file. */
 type ListRow =
   | { type: "up" }
-  | { type: "folder"; name: string; count: number }
+  | { type: "folder"; name: string; count?: number }
   | { type: "file"; item: DatasetItem; label: string };
 
 /** Fixed row height (px) and shared column grid — must match between header and rows. */
@@ -267,8 +240,8 @@ export function DatasetsPage() {
   const dragDepth = React.useRef(0);
 
   const refreshDatasets = React.useCallback(() => {
+    qc.invalidateQueries({ queryKey: ["datasets-browse"] });
     qc.invalidateQueries({ queryKey: ["datasets-all"] });
-    qc.invalidateQueries({ queryKey: ["total-datasets"] });
     // An upload or a delete moves the used space, so the disk values are stale.
     qc.invalidateQueries({ queryKey: ["dataset-storage"] });
   }, [qc]);
@@ -355,14 +328,31 @@ export function DatasetsPage() {
     );
   }
 
-  const totalQuery = useQuery({ queryKey: ["total-datasets"], queryFn: () => beacon.totalDatasets() });
-  const datasetsQuery = useQuery({
+  const searching = search.trim().length > 0;
+
+  // One directory level. This is the normal path, and its cost does not grow
+  // with the store: a delimiter listing of a 2.85-million-object bucket
+  // measured 14 ms, where enumerating the same bucket took 79.9 s.
+  const browseQuery = useQuery({
+    queryKey: ["datasets-browse", path],
+    queryFn: () => beacon.browseDatasets<BrowseResponse>(path || undefined),
+    enabled: !searching,
+  });
+
+  // Search is the one view that genuinely needs every path, so it is the one
+  // view that pays for them — and only once the user has typed something.
+  const searchQuery = useQuery({
     queryKey: ["datasets-all"],
     queryFn: async () => normalize(await beacon.datasets({ limit: 100_000 })),
+    enabled: searching,
   });
 
-  const items = datasetsQuery.data ?? [];
-  const searching = search.trim().length > 0;
+  // `/api/total-datasets` enumerates the whole store to produce one number, so
+  // it is not run on mount. The badge shows what the current view holds.
+  const datasetsQuery = searching ? searchQuery : browseQuery;
+  const items = searching
+    ? searchQuery.data ?? []
+    : normalize(browseQuery.data?.datasets ?? []);
   const crumbs = path ? path.replace(/\/$/, "").split("/") : [];
 
   // Unified, virtualizable row list. A non-empty search switches to a flat list
@@ -376,15 +366,22 @@ export function DatasetsPage() {
         .sort((a, b) => compareFiles(a, b, sort))
         .map((item) => ({ type: "file", item, label: item.path }));
     }
-    const { folders, files } = browse(items, path, "");
     const list: ListRow[] = [];
     if (path) list.push({ type: "up" });
     // Folders stay grouped on top (name-ordered); files follow the active sort.
-    for (const f of folders) list.push({ type: "folder", name: f.name, count: f.count });
+    // They come from the browse response, so no client-side grouping is needed
+    // and no count is known without descending into each one.
+    const folders = (browseQuery.data?.folders ?? []).map((name) => ({ name }));
+    const files = items;
+    for (const f of folders) list.push({ type: "folder", name: f.name });
     for (const file of [...files].sort((a, b) => compareFiles(a, b, sort)))
       list.push({ type: "file", item: file, label: baseName(file.path) });
     return list;
-  }, [items, path, search, searching, sort]);
+  }, [items, path, search, searching, sort, browseQuery.data]);
+
+  // What the badge counts: the folders and files of this view. The `..` entry is
+  // navigation, not content, so it is left out.
+  const shown = rows.filter((row) => row.type !== "up").length;
 
   function toggleSort(key: SortKey) {
     setSort((s) => (s.key === key ? { key, dir: s.dir === "asc" ? "desc" : "asc" } : { key, dir: "asc" }));
@@ -443,7 +440,7 @@ export function DatasetsPage() {
             <span className="truncate">{row.name}</span>
           </span>
           <span className="text-muted-foreground">
-            {row.count} item{row.count === 1 ? "" : "s"}
+            {row.count === undefined ? "" : `${row.count} item${row.count === 1 ? "" : "s"}`}
           </span>
           <span />
           <span />
@@ -530,9 +527,7 @@ export function DatasetsPage() {
       description="Browse the datasets store. Preview rows or inspect a file's schema."
       actions={
         <div className="flex items-center gap-2">
-          {typeof totalQuery.data === "number" && (
-            <Badge variant="secondary">{totalQuery.data} total</Badge>
-          )}
+          {shown > 0 && <Badge variant="secondary">{shown} here</Badge>}
           <Button size="sm" className="gap-1.5" onClick={() => openUpload([])}>
             <Upload className="h-4 w-4" />
             Upload

beacon-db/beacon-core/src/crawler/engine.rs

@@ -1,7 +1,7 @@
 //! The crawl engine: turn a [`CrawlerDefinition`] into registered external tables.
 //!
 //! Reuses Beacon's existing primitives end-to-end:
-//! - [`beacon_functions::file_formats::list_datasets::list_datasets`] for scan +
+//! - [`beacon_functions::listing::list_datasets`] for scan +
 //!   per-format classification,
 //! - [`ExternalTableDefinition::build_provider`] for schema inference + partition
 //!   validation (the same code path used when loading persisted tables),
@@ -19,7 +19,7 @@ use beacon_datafusion_ext::table_ext::{ExternalTable, ExternalTableDefinition, T
 use datafusion::prelude::SessionContext;
 use serde::{Deserialize, Serialize};
 
-use beacon_functions::file_formats::list_datasets::list_datasets;
+use beacon_functions::listing::list_datasets;
 
 use crate::statement_plan::{upgrade_session, SessionCell};
 

beacon-db/beacon-core/src/file_stats.rs

@@ -72,7 +72,7 @@ use datafusion::datasource::listing::ListingTableUrl;
 use datafusion::execution::object_store::ObjectStoreUrl;
 use datafusion::prelude::SessionContext;
 use futures::StreamExt;
-use object_store::{path::Path, ObjectMeta, ObjectStore};
+use object_store::{path::Path, ObjectMeta};
 
 use crate::statement_plan::{upgrade_session, SessionCell};
 
@@ -450,7 +450,6 @@ pub struct FileStatsService {
     store: Arc<FileStatsStore>,
     collector: StatsCollector,
     session: SessionCell,
-    datasets_url: ObjectStoreUrl,
     config: FileStatsConfig,
     /// The timer, and the startup collection when one was asked for. Both are
     /// aborted on drop.
@@ -471,11 +470,13 @@ pub struct FileStatsService {
 }
 
 impl FileStatsService {
+    /// Takes no datasets store URL. Discovery resolves the scan prefix through
+    /// the listing factory on the session, so the store it reads is the store a
+    /// query would read.
     pub fn new(
         store: Arc<FileStatsStore>,
         analyzer: Arc<dyn FileAnalyzer>,
         session: SessionCell,
-        datasets_url: ObjectStoreUrl,
         config: FileStatsConfig,
     ) -> Arc<Self> {
         let collector = StatsCollector::new(
@@ -494,7 +495,6 @@ impl FileStatsService {
             store,
             collector,
             session,
-            datasets_url,
             config,
             tasks: parking_lot::Mutex::new(Vec::new()),
             pass_lock: Arc::new(tokio::sync::Mutex::new(())),
@@ -691,16 +691,23 @@ impl FileStatsService {
 
     /// The same, restricted to a prefix. `None` uses the configured scan prefix.
     async fn discover_under(&self, prefix: Option<&str>) -> anyhow::Result<usize> {
-        let session = self.session()?;
-        let store = session
-            .state()
-            .runtime_env()
-            .object_store(&self.datasets_url)
-            .map_err(|e| anyhow::anyhow!("datasets store unavailable: {e}"))?;
+        use beacon_datafusion_ext::listing_factory::ListingFactory;
 
+        let session = self.session()?;
+        let state = session.state();
         let scan_prefix = prefix.unwrap_or(self.config.scan_prefix.as_str());
-        let prefix = (!scan_prefix.is_empty()).then(|| Path::from(scan_prefix));
-        let mut listing = store.list(prefix.as_ref());
+
+        // Through the listing factory rather than the store directly, so the
+        // scan prefix resolves by the same rules a query would use: the
+        // configured datasets store, and a glob if one is given.
+        let factory = state
+            .config()
+            .get_extension::<ListingFactory>()
+            .ok_or_else(|| anyhow::anyhow!("the listing factory is not registered"))?;
+        let mut listing = factory
+            .listing(&state, scan_prefix)
+            .map_err(|e| anyhow::anyhow!("cannot resolve the scan prefix `{scan_prefix}`: {e}"))?
+            .stream();
 
         let mut batch: Vec<ObservedFile> = Vec::with_capacity(self.config.discovery_chunk);
         let mut total = 0usize;

beacon-db/beacon-core/src/runtime_builder.rs

@@ -449,13 +449,12 @@ async fn init_file_stats(
 
     let analyzer = Arc::new(crate::file_stats::FormatFileAnalyzer::new(
         session_cell.clone(),
-        datasets.clone(),
+        datasets,
     ));
     let service = crate::file_stats::FileStatsService::new(
         store,
         analyzer,
         session_cell,
-        datasets,
         builder.file_stats.clone(),
     );
     service.start();