@nshiab/simple-data-analysis-core 2.0.5 → 2.0.7

package/script/methods/updateWithJS.js

@@ -4,10 +4,15 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
 };
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.default = updateWithJS;
+const prepareGeometry_js_1 = __importDefault(require("../helpers/prepareGeometry.js"));
+const geometryFromJSON_js_1 = __importDefault(require("../helpers/geometryFromJSON.js"));
 const quoteIdentifier_js_1 = __importDefault(require("../helpers/quoteIdentifier.js"));
 const tableRegistry_js_1 = require("../helpers/tableRegistry.js");
 const queueOp_js_1 = __importDefault(require("../helpers/queueOp.js"));
 const loadArray_js_1 = require("./loadArray.js");
+const node_api_1 = require("@duckdb/node-api");
+const readMutationRows_js_1 = __importDefault(require("../helpers/readMutationRows.js"));
+const parseDuckDBType_js_1 = __importDefault(require("../helpers/parseDuckDBType.js"));
 function updateWithJS(simpleTable, dataModifier, options = {}) {
     if (options.batchSize !== undefined &&
         (!Number.isInteger(options.batchSize) || options.batchSize < 1)) {
@@ -23,39 +28,35 @@ function updateWithJS(simpleTable, dataModifier, options = {}) {
 }
 async function executeUpdateWithJS(simpleTable, dataModifier, options) {
     const types = await simpleTable.getTypes();
-    // Geometry columns are typed as GEOMETRY('<crs>'), not a bare "GEOMETRY".
-    if (Object.values(types).some((d) => d.toUpperCase().startsWith("GEOMETRY"))) {
-        throw new Error("updateWithJS doesn't work with tables containing geometries.");
-    }
-    if (options.batchSize === undefined) {
-        const oldData = await simpleTable.getData();
-        if (!oldData) {
-            throw new Error("No data from getData.");
+    const geometryColumns = Object.keys(types).filter((key) => types[key].toUpperCase().startsWith("GEOMETRY"));
+    for (const column of geometryColumns) {
+        if (types[column].toUpperCase() !== "GEOMETRY('EPSG:4326')") {
+            throw new Error(`Column ${JSON.stringify(column)} has CRS ${types[column]}. Call .reproject("EPSG:4326", { column: ${JSON.stringify(column)} }) before updateWithJS().`);
         }
-        const newData = await dataModifier(oldData);
-        if (newData.length === 0) {
-            if (oldData.length === 0) {
-                // Empty table in, empty data out: nothing to update.
-                return;
-            }
-            throw new Error("The dataModifier returned no rows. updateWithJS can't infer the table schema from zero rows.");
+        const name = (0, quoteIdentifier_js_1.default)(column);
+        // Compare the actual binary representation, including dimensions and double
+        // precision, rather than assuming GeoJSON serialization is lossless.
+        const result = await simpleTable.connection.runAndReadAll(`SELECT count(*) AS invalid FROM ${(0, quoteIdentifier_js_1.default)(simpleTable.name)}
+       WHERE ${name} IS NOT NULL AND
+       ST_AsWKB(${name}) != ST_AsWKB(${(0, geometryFromJSON_js_1.default)(`ST_AsGeoJSON(${name})`)})`);
+        if (result.getRows()[0][0] !== 0n) {
+            throw new Error(`Column ${JSON.stringify(column)} contains geometry that cannot round-trip through GeoJSON without losing precision or dimensions; unsupported representations cannot be updated with updateWithJS().`);
         }
-        await (0, loadArray_js_1.executePreparedArray)(simpleTable, (0, loadArray_js_1.prepareArray)(newData));
-        return;
     }
-    const batchSize = options.batchSize;
-    if (Object.keys(types).includes("__sda_rowid")) {
-        throw new Error('The table has a column named "__sda_rowid", which conflicts with the internal column used by the batchSize option. Rename it or run updateWithJS without batchSize.');
+    for (const key of Object.keys(options.columnTypes ?? {})) {
+        if (Object.keys(types).some((name) => name.toLowerCase() === key.toLowerCase())) {
+            throw new Error(`updateWithJS() columnTypes may only declare new columns; ${JSON.stringify(key)} already exists.`);
+        }
     }
-    const rowCount = await simpleTable.getRowCount();
-    if (rowCount === 0) {
-        // Same behavior as the non-batched path on an empty table.
-        const newData = await dataModifier([]);
-        if (newData.length === 0) {
-            return;
+    const batchSize = options.batchSize;
+    const geometryBatch = batchSize !== undefined && geometryColumns.length > 0;
+    if (batchSize !== undefined) {
+        // DuckDB resolves identifiers case-insensitively, and a real rowid column
+        // shadows its hidden row identifier used for pagination.
+        const conflictingColumn = Object.keys(types).find((column) => ["rowid", "__sda_rowid"].includes(column.toLowerCase()));
+        if (conflictingColumn !== undefined) {
+            throw new Error(`The table has a column named "${conflictingColumn}", which conflicts with the internal column used by the batchSize option. Rename it or run updateWithJS without batchSize.`);
         }
-        await (0, loadArray_js_1.executePreparedArray)(simpleTable, (0, loadArray_js_1.prepareArray)(newData));
-        return;
     }
     // Rows are pulled in batches by rowid, passed through the modifier, and
     // accumulated in a temporary table, so only one batch of rows is
@@ -66,27 +67,112 @@ async function executeUpdateWithJS(simpleTable, dataModifier, options) {
     try {
         let first = true;
         let lastRowid = null;
+        let sawRows = false;
+        let outputOffset = 0;
+        let inputOffset = 0;
         while (true) {
-            const batch = (await simpleTable.sdb.customQuery(`SELECT *, rowid AS __sda_rowid FROM ${(0, quoteIdentifier_js_1.default)(simpleTable.name)}${lastRowid === null ? "" : ` WHERE rowid > ${lastRowid}`} ORDER BY rowid LIMIT ${batchSize}`, { returnData: true }));
-            if (batch.length === 0) {
+            const page = `${(0, quoteIdentifier_js_1.default)(simpleTable.name)}${lastRowid === null ? "" : ` WHERE rowid > ${lastRowid}`}${batchSize === undefined ? "" : ` ORDER BY rowid LIMIT ${batchSize}`}`;
+            // Select the batch before serializing geometry. Otherwise DuckDB can
+            // evaluate ST_AsGeoJSON on remaining rows before applying the LIMIT,
+            // repeatedly converting polygons that belong to later batches.
+            const from = geometryBatch
+                ? `(SELECT *, rowid AS __sda_rowid FROM ${page}) AS __sda_batch ORDER BY __sda_rowid`
+                : page;
+            const source = await (0, readMutationRows_js_1.default)(simpleTable.connection, `SELECT ${Object.keys(types).map((key) => geometryColumns.includes(key)
+                ? `ST_AsGeoJSON(${(0, quoteIdentifier_js_1.default)(key)})::VARCHAR AS ${(0, quoteIdentifier_js_1.default)(key)}`
+                : (0, quoteIdentifier_js_1.default)(key)).join(", ")}${batchSize === undefined
+                ? ""
+                : geometryBatch
+                    ? ", __sda_rowid"
+                    : ", rowid AS __sda_rowid"} FROM ${from}`, batchSize === undefined ? undefined : "__sda_rowid");
+            const batch = source.rows;
+            const inputCount = batch.length;
+            if (inputCount === 0 && sawRows)
                 break;
+            sawRows ||= inputCount > 0;
+            if (batchSize !== undefined && batch.length > 0) {
+                // Keep the cursor separate: adding and deleting it on every row also
+                // makes ordinary callback operations such as object spread slower.
+                lastRowid = source.lastCursorValue;
             }
-            lastRowid = batch[batch.length - 1].__sda_rowid;
-            for (const row of batch) {
-                delete row.__sda_rowid;
+            for (const [index, row] of batch.entries()) {
+                for (const key of geometryColumns) {
+                    row[key] = row[key] === null ? null : JSON.parse(row[key]);
+                    (0, prepareGeometry_js_1.default)(row[key], key, inputOffset + index + 1);
+                }
             }
+            inputOffset += inputCount;
             const modified = await dataModifier(batch);
             if (modified.length === 0) {
+                if (!sawRows)
+                    return;
+                if (batchSize === undefined)
+                    break;
                 continue;
             }
-            await (0, loadArray_js_1.executePreparedArray)(scratch, (0, loadArray_js_1.prepareArray)(modified));
+            const keys = Object.keys(modified[0]);
+            for (const key of Object.keys(options.columnTypes ?? {})) {
+                if (!keys.includes(key)) {
+                    throw new Error(`Unknown column ${JSON.stringify(key)} in updateWithJS columnTypes.`);
+                }
+            }
+            const newKeys = keys.filter((key) => !source.types.has(key));
+            const added = newKeys.length === 0 ? undefined : (0, loadArray_js_1.prepareArray)(modified.map((row) => Object.fromEntries(newKeys.map((key) => [key, row[key]]))), options.columnTypes);
+            const outputGeometry = keys.filter((key) => geometryColumns.includes(key) ||
+                (added && added.types[newKeys.indexOf(key)] === "GEOMETRY('EPSG:4326')"));
+            const outputJSON = keys.filter((key) => types[key] === "JSON" ||
+                (added && added.types[newKeys.indexOf(key)] === "JSON"));
+            if (outputGeometry.length && !simpleTable.sdb.spatialLoaded) {
+                await simpleTable.sdb.customQuery("INSTALL spatial; LOAD spatial; SET geometry_always_xy = true;");
+                simpleTable.sdb.spatialLoaded = true;
+            }
+            const geometryData = new Map(outputGeometry.map((key) => [
+                key,
+                geometryColumns.includes(key)
+                    ? modified.map((row, i) => (0, prepareGeometry_js_1.default)(row[key], key, outputOffset + i + 1))
+                    : added.columnsData[newKeys.indexOf(key)],
+            ]));
+            outputOffset += modified.length;
+            const columnTypes = keys.map((key) => outputGeometry.includes(key) || outputJSON.includes(key)
+                ? (0, parseDuckDBType_js_1.default)("VARCHAR")
+                : source.types.get(key) ??
+                    (0, parseDuckDBType_js_1.default)(added.types[newKeys.indexOf(key)]));
+            await simpleTable.sdb.customQuery(`CREATE OR REPLACE TABLE ${(0, quoteIdentifier_js_1.default)(scratch.name)} (${keys.map((key, i) => `${(0, quoteIdentifier_js_1.default)(key)} ${columnTypes[i]}`)
+                .join(", ")})`);
+            const appender = await simpleTable.connection.createAppender(scratch.name);
+            try {
+                for (let start = 0; start < modified.length; start += 2000) {
+                    const end = Math.min(start + 2000, modified.length);
+                    const chunk = node_api_1.DuckDBDataChunk.create(columnTypes, end - start);
+                    for (let i = 0; i < keys.length; i++) {
+                        const key = keys[i];
+                        chunk.setColumnValues(i, outputGeometry.includes(key)
+                            ? geometryData.get(key).slice(start, end)
+                            : source.types.has(key)
+                                ? modified.slice(start, end).map((row) => source.toNative(row[key], columnTypes[i]))
+                                : added.columnsData[newKeys.indexOf(key)].slice(start, end));
+                    }
+                    appender.appendDataChunk(chunk);
+                }
+                appender.flushSync();
+            }
+            finally {
+                appender.closeSync();
+            }
+            const projection = keys.map((key) => outputGeometry.includes(key)
+                ? `${(0, geometryFromJSON_js_1.default)((0, quoteIdentifier_js_1.default)(key))} AS ${(0, quoteIdentifier_js_1.default)(key)}`
+                : outputJSON.includes(key)
+                    ? `${(0, quoteIdentifier_js_1.default)(key)}::JSON AS ${(0, quoteIdentifier_js_1.default)(key)}`
+                    : (0, quoteIdentifier_js_1.default)(key)).join(", ");
             if (first) {
-                await simpleTable.sdb.customQuery(`CREATE OR REPLACE TABLE ${(0, quoteIdentifier_js_1.default)(accumulator)} AS SELECT * FROM ${(0, quoteIdentifier_js_1.default)(scratch.name)}`);
+                await simpleTable.sdb.customQuery(`CREATE OR REPLACE TABLE ${(0, quoteIdentifier_js_1.default)(accumulator)} AS SELECT ${projection} FROM ${(0, quoteIdentifier_js_1.default)(scratch.name)}`);
                 first = false;
             }
             else {
-                await simpleTable.sdb.customQuery(`INSERT INTO ${(0, quoteIdentifier_js_1.default)(accumulator)} BY NAME SELECT * FROM ${(0, quoteIdentifier_js_1.default)(scratch.name)}`);
+                await simpleTable.sdb.customQuery(`INSERT INTO ${(0, quoteIdentifier_js_1.default)(accumulator)} BY NAME SELECT ${projection} FROM ${(0, quoteIdentifier_js_1.default)(scratch.name)}`);
             }
+            if (batchSize === undefined || inputCount === 0)
+                break;
         }
         if (first) {
             throw new Error("The dataModifier returned no rows. updateWithJS can't infer the table schema from zero rows.");