React hooks for Mosaic

[willium]

It'd be great if we can support creating clients as idiomatically correct react hooks.

I started on a useMosaicClient hook here that essentially monkey-patches an instance of the class. It's very unpolished, but seems to work. I thought I'd share here for anyone who may want a starting point, and to begin a conversation about how we might support this in a more first-party way.

"use client";

import type * as Arrow from "apache-arrow";
import { useEffect, useMemo, useState } from "react";
import {
	coordinator,
	type Field,
	type FieldInfo,
	MosaicClient,
	type Selection as MosaicSelection,
} from "@uwdata/mosaic-core";
import * as SQL from "@uwdata/mosaic-sql";
import { useMosaic } from "./mosaic-coordinator-provider";

type ReactMosaicClientProps<T extends Arrow.TypeMap> = {
	table: string;
	columns?: string[];

	selection?: MosaicSelection;

	// if query is undefined, uses this.defaultQuery
	query?: SQL.Query | null;
	orderby?: unknown[];
	offset?: number;
	limit?: number;

	// callbacks
	onSchemaChange?: (schema: FieldInfo[]) => void;
	onRequestQuery?: (query: SQL.Query) => void;
	onPrefetchQuery?: (nextQuery: SQL.Query) => void;
	onQueryPending?: () => void;
	onQueryResult?: (data: Arrow.Table<T>) => void;
	onQueryError?: (error: unknown) => void;
	onUpdate?: () => void;
};

export class ReactMosaicClient<T extends Arrow.TypeMap> extends MosaicClient {
	private _data: Arrow.Table<T> | null = null;
	private _selection: MosaicSelection | undefined;

	private _table: string;
	private _columns: string[] = ["*"];

	private _query?: SQL.Query | null;
	private _orderby: Set<unknown>;
	private _offset: number;
	private _limit: number;

	private _schema: FieldInfo[] | null;

	private _onSchemaChange?: (schema: FieldInfo[]) => void;
	private _onRequestQuery?: (query: SQL.Query) => void;
	private _onPrefetchQuery?: (nextQuery: SQL.Query) => void;
	private _onQueryPending?: () => void;
	private _onQueryResult?: (data: Arrow.Table<T>) => void;
	private _onQueryError?: (error: unknown) => void;
	private _onUpdate?: () => void;

	constructor({
		table,
		columns = ["*"],

		selection,

		query,
		orderby = [],
		offset = 0,
		limit = -1, // -1 means no limit

		onSchemaChange,
		onRequestQuery,
		onPrefetchQuery,
		onQueryPending,
		onQueryResult,
		onQueryError,
		onUpdate,
	}: ReactMosaicClientProps<T>) {
		super(selection); // Call the parent constructor with the filterBy selection

		// Initialize client properties from the options
		this._table = table;
		this._columns = columns;

		this._selection = selection;

		this._query = query;
		this._orderby = new Set(orderby);
		this._offset = offset;
		this._limit = limit;

		this._onSchemaChange = onSchemaChange;
		this._onRequestQuery = onRequestQuery;
		this._onPrefetchQuery = onPrefetchQuery;
		this._onQueryPending = onQueryPending;
		this._onQueryResult = onQueryResult;
		this._onQueryError = onQueryError;
		this._onUpdate = onUpdate;

		// Setup internal states
		this._schema = null;
	}

	/**
	 * Get the data of the client
	 * @returns The data of the client
	 */
	get data(): Arrow.Table<T> | null {
		return this._data;
	}

	/**
	 * Get the selection of the client
	 * @returns The selection of the client
	 */
	get selection(): MosaicSelection | undefined {
		return this._selection;
	}

	/**
	 * Get the table of the client
	 * @returns The table of the client
	 */
	get table(): string {
		return this._table;
	}

	/**
	 * Get the columns of the client
	 * @returns The columns of the client
	 */
	get columns(): string[] {
		return this._columns;
	}

	/**
	 * Return an array of fields queried by this client
	 * @returns An array of field objects or null if no fields are needed
	 */
	fields(): Field[] | null {
		// derive fields from columns
		return this._columns.map((column) => {
			// HACK: https://github.com/uwdata/mosaic/issues/366
			return SQL.column(this._table, column) as Field;
		});
	}

	/**
	 * Process the field information provided by the coordinator
	 * @param schema - The field information object
	 * @returns The client instance (this)
	 */
	fieldInfo(schema: FieldInfo[]): this {
		this._onSchemaChange?.(schema);
		this._schema = schema;
		return this;
	}

	/**
	 * Get the default query for the client
	 * @returns The default query
	 */
	get defaultQuery(): SQL.Query {
		return SQL.Query.from(this._table).select(
			this._schema?.map(({ column }) => column) ?? this._columns,
		);
	}

	/**
	 * Sort the data by the given column
	 * @param column - The column to sort by
	 */
	sort(column: string) {
		if (this._orderby.has(column)) {
			this._orderby.delete(column);
		} else {
			this._orderby.add(column);
		}
		this.runQuery();
	}

	/**
	 * Return a query specifying the data needed by this client
	 * @param filter - The filtering criteria to apply in the query (optional)
	 * @returns The client query
	 */
	query(filter: unknown = []): SQL.Query | null {
		if (this._query === null) return null;

		let query = this._query?.clone() ?? this.defaultQuery.clone();
		query = query.where(filter);
		query = query.offset(this._offset);
		if (this._limit > 0) query = query.limit(this._limit);
		query = query.orderby(Array.from(this._orderby).map(SQL.desc));

		return query;
	}

	/**
	 * Run the query with the provided offset and limit
	 * @param offset - The offset for the query
	 * @param limit - The limit for the query
	 */
	async runQuery(opts?: {
		offset?: number;
		limit?: number;
	}): Promise<Arrow.Table<T> | null> {
		this._offset = opts?.offset ?? this._offset;
		this._limit = opts?.limit ?? this._limit;

		const query = this.query(this._selection?.predicate(this));
		if (!query) return null;

		// request next data batch
		this._onRequestQuery?.(query);
		await this.requestQuery(query);

		if (this._limit > 0) {
			// prefetch subsequent data batch
			const nextQuery = query.clone().offset(this._offset + this._limit);
			this._onPrefetchQuery?.(nextQuery);
			this.coordinator?.prefetch(nextQuery);
		}

		return this.data;
	}

	/**
	 * Called by the coordinator to inform the client that a query is pending.
	 */
	queryPending() {
		this._onQueryPending?.();
		return this;
	}

	/**
	 * Handle the query result returned by the coordinator
	 * @param data - The query result data
	 * @returns The client instance (this)
	 */
	queryResult(data: Arrow.Table): this {
		this._onQueryResult?.(data);
		this._data = data; // store previous result?
		return this;
	}

	/**
	 * Handle query errors reported by the coordinator
	 * @param error - The error object
	 * @returns The client instance (this)
	 */
	queryError(error: unknown): this {
		this._onQueryError?.(error);
		return super.queryError(error); // Call the parent method to log the error
	}

	/**
	 * Update the client when new data is available
	 * @returns The client instance (this)
	 */
	update(): this {
		this._onUpdate?.();
		return this;
	}
}

export function useMosaicClient<T extends Arrow.TypeMap>(
	props: ReactMosaicClientProps<T>,
	deps: unknown[] = [],
) {
	const [schema, setSchema] = useState<FieldInfo[] | null>(null);
	const [error, setError] = useState<unknown | null>(null);
	const [data, setData] = useState<Arrow.Table<T> | null>(null);

	const client = useMemo(() => {
		function onSchemaChange(schema: FieldInfo[]) {
			props.onSchemaChange?.(schema);
			setSchema(schema);
		}

		function onQueryError(error: unknown) {
			props.onQueryError?.(error);
			setError(error); // TODO: convert error to String
		}

		function onQueryResult(data: Arrow.Table<T>) {
			props.onQueryResult?.(data);
			setData(data);
		}

		return new ReactMosaicClient<T>({
			...props,
			onSchemaChange,
			onQueryError,
			onQueryResult,
		});
		// eslint-disable-next-line react-hooks/exhaustive-deps
	}, deps);

	const hasLoadedData = useMosaic();
	const [isConnected, setIsConnected] = useState(false);
	useEffect(() => {
		if (hasLoadedData) {
			coordinator()
				.connect(client)
				.then(() => setIsConnected(true));
		}
		return () => {
			coordinator().disconnect(client);
			setIsConnected(false);
		};
	}, [client, hasLoadedData]);

	return { isConnected, client, data, schema, error };
}

Here is a macro hook, just for getting the count of a table (+ selection)

export function useMosaicCountClient<T extends Arrow.TypeMap>(
	props: Omit<ReactMosaicClientProps<T>, "query">,
	deps: unknown[] = [],
) {
	const query = useMemo(
		() => SQL.Query.from(props.table).select({ count: SQL.count() }),
		[props.table],
	);

	const { isConnected, data } = useMosaicClient({ ...props, query }, deps);

	const count = useMemo(() => {
		if (!isConnected) return 0;
		if (!data) return 0;
		return data.toArray()[0].count;
	}, [data, isConnected]);

	return count;
}

[jheer]

This is great, thanks for sharing! However, I don't see an active issue for maintainers to address so I'm moving this over to discussions.