Spaces:
Running
Running
File size: 2,375 Bytes
5534d69 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import { AsyncDuckDB, AsyncDuckDBConnection } from "@duckdb/duckdb-wasm";
import { useCallback, useEffect, useState } from "react";
// https://duckdb.org/docs/api/wasm/query#arrow-table-to-json
// TODO: import the arrow lib and use the correct type
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const arrowResultToJson = (arrowResult: any) => {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return arrowResult.toArray().map((row: any) => row.toJSON());
};
export const useParquetTable = (db: AsyncDuckDB | null, datasetUrl: string) => {
const [loading, setLoading] = useState(false);
const [dataset, setDataset] = useState<null>(null);
const [error, setError] = useState<string | null>(null);
const clearDataset = useCallback(() => {
setDataset(null);
setError(null);
}, []);
const loadDataset = useCallback(() => {
if (db && datasetUrl) {
setLoading(true);
setDataset(null);
setError(null);
console.log("Loading", datasetUrl);
const status: { conn: AsyncDuckDBConnection | null; killed: boolean } = {
conn: null,
killed: false,
};
db.connect().then((conn) => {
if (status.killed) {
conn.close();
return;
}
status.conn = conn;
conn
// we get httpfs for free with cors restrictions
// we get parquet downloaded as we load it
// https://duckdb.org/docs/api/wasm/extensions
// https://duckdb.org/docs/api/wasm/data_ingestion#parquet
.query(
`LOAD parquet;LOAD httpfs;SELECT * FROM '${datasetUrl}' LIMIT 10`
)
.then((result) => setDataset(arrowResultToJson(result)))
.catch((err) => {
console.error(err);
setError(err.message);
setDataset(null);
})
.finally(() => {
conn.close();
setLoading(false);
});
});
return () => {
status.killed = true;
if (status.conn) {
console.log("Closing connection");
status.conn.close();
}
};
} else if (db) {
console.log("Resetting db");
db.reset();
}
}, [db, datasetUrl]);
useEffect(() => {
loadDataset();
}, [loadDataset]);
return { loading, dataset, error, clearDataset, loadDataset };
};
|