From 4bb8914d9a02488ddf37bbc726eff74ed30ef4fd Mon Sep 17 00:00:00 2001
From: Javier Goizueta <jgoizueta@gmail.com>
Date: Mon, 28 May 2018 16:08:31 +0200
Subject: [PATCH] Add parameters to select metadata sample columns

---
 .../layer-stats/mapnik-layer-stats.js         | 35 ++++++++++++++++---
 lib/cartodb/utils/query-utils.js              | 20 ++++++++---
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/lib/cartodb/backends/layer-stats/mapnik-layer-stats.js b/lib/cartodb/backends/layer-stats/mapnik-layer-stats.js
index 3a7a3358..0caf2d0e 100644
--- a/lib/cartodb/backends/layer-stats/mapnik-layer-stats.js
+++ b/lib/cartodb/backends/layer-stats/mapnik-layer-stats.js
@@ -146,13 +146,36 @@ function mergeColumns(results) {
     }
 }
 
-function _sample(ctx, numRows) {
+const SAMPLE_SEED = 0.5;
+const DEFAULT_SAMPLE_ROWS = 100;
+
+function exclude(items, excludedItems) {
+    if (excludedItems) {
+        return items.filter(item => !excludedItems.includes(item));
+    }
+    return items;
+}
+
+function _sample(ctx, numRows, availableColumns = null) {
     if (ctx.metaOptions.sample) {
-        const sampleProb = Math.min(ctx.metaOptions.sample / numRows, 1);
+        const sampleProb = Math.min(ctx.metaOptions.sample.num_rows / numRows, 1);
         // We'll use a safety limit just in case numRows is a bad estimate
-        const limit = Math.ceil(ctx.metaOptions.sample * 1.5);
-        return queryPromise(ctx.dbConnection, _getSQL(ctx, sql => queryUtils.getQuerySample(sql, sampleProb, limit)))
-            .then(res => ({ sample: res.rows }));
+        const requestedRows = ctx.metaOptions.sample.num_rows || DEFAULT_SAMPLE_ROWS;
+        const limit = Math.ceil(requestedRows * 1.5);
+        let columns = ctx.metaOptions.sample.include_columns;
+        if (columns) {
+            columns = exclude(columns, ctx.metaOptions.sample.exclude_columns);
+        }
+        else if (ctx.metaOptions.sample.exclude_columns) {
+            if (availableColumns === null) {
+                return Promise.reject(new Error('Column stats are needed to use the sample exclude columns options'));
+            }
+            columns = exclude(availableColumns, ctx.metaOptions.sample.exclude_columns);
+        }
+        return queryPromise(ctx.dbConnection, _getSQL(
+            ctx,
+            sql => queryUtils.getQuerySample(sql, sampleProb, limit, SAMPLE_SEED, columns)
+        )).then(res => ({ sample: res.rows }));
     }
     return Promise.resolve();
 }
@@ -265,6 +288,8 @@ function (layer, dbConnection, callback) {
     // (if metaOptions.geometryType) from it.
 
     // TODO: compute _sample with _featureCount when available
+    // TODO: add support for sample.exclude option by, in that case, forcing the columns query and
+    // passing the results to the sample query function.
 
     Promise.all([
         _estimatedFeatureCount(ctx).then(
diff --git a/lib/cartodb/utils/query-utils.js b/lib/cartodb/utils/query-utils.js
index 22d99eef..ae1a2e2d 100644
--- a/lib/cartodb/utils/query-utils.js
+++ b/lib/cartodb/utils/query-utils.js
@@ -88,17 +88,27 @@ module.exports.getQueryTopCategories = function(query, column, topN, includeNull
     `;
 };
 
-module.exports.getQuerySample = function(query, sampleProb, limit = null, randomSeed = 0.5) {
+function columnSelector(columns) {
+    if (!columns) {
+        return '*';
+    }
+    if (typeof(columns) === 'string') {
+        return columns;
+    }
+    return columns.map(name => `"${name}"`).join(', ');
+}
+
+module.exports.getQuerySample = function(query, sampleProb, limit = null, randomSeed = 0.5, columns = null) {
     const singleTable = simpleQueryTable(query);
     if (singleTable) {
-        return getTableSample(singleTable.table, singleTable.columns, sampleProb, limit, randomSeed);
+        return getTableSample(singleTable.table, columns || singleTable.columns, sampleProb, limit, randomSeed);
     }
     const limitClause = limit ? `LIMIT ${limit}` : '';
     return `
         WITH __cdb_rndseed AS (
             SELECT setseed(${randomSeed})
         )
-        SELECT *
+        SELECT ${columnSelector(columns)}
             FROM (${query}) AS __cdb_query
             WHERE random() < ${sampleProb}
             ${limitClause}
@@ -110,7 +120,9 @@ function getTableSample(table, columns, sampleProb, limit = null, randomSeed = 0
     sampleProb *= 100;
     randomSeed *= Math.pow(2, 31) -1;
     return `
-        SELECT ${columns} FROM ${table} TABLESAMPLE BERNOULLI (${sampleProb}) REPEATABLE (${randomSeed}) ${limitClause}
+        SELECT ${columnSelector(columns)}
+        FROM ${table}
+        TABLESAMPLE BERNOULLI (${sampleProb}) REPEATABLE (${randomSeed}) ${limitClause}
     `;
 }