Send optimal batch size

2018-07-10 19:06:49 +02:00
parent 286a75fa8e
commit 531ad28158
10 changed files with 153 additions and 90 deletions
--- a/client/cdb_dataservices_client--0.25.0.sql
+++ b/client/cdb_dataservices_client--0.25.0.sql
@@ -112,7 +112,8 @@ CREATE TYPE cdb_dataservices_client.service_quota_info AS (
    monthly_quota NUMERIC,
    used_quota NUMERIC,
    soft_limit BOOLEAN,
-    provider TEXT
+    provider TEXT,
+    max_batch_size NUMERIC
 );
 --
 -- Public dataservices API function
@@ -1987,25 +1988,36 @@ CREATE OR REPLACE FUNCTION cdb_dataservices_client._DST_DisconnectUserTable(
    TARGET cdb_dataservices_server._DST_DisconnectUserTable;
 $$ LANGUAGE plproxy VOLATILE PARALLEL UNSAFE;
 CREATE OR REPLACE FUNCTION cdb_dataservices_client.cdb_bulk_geocode_street_point (query text,
-    street_column text, city_column text default null, state_column text default null, country_column text default null, batch_size integer DEFAULT 50)
+    street_column text, city_column text default null, state_column text default null, country_column text default null, batch_size integer DEFAULT NULL)
 RETURNS SETOF cdb_dataservices_client.geocoding AS $$
 DECLARE
  query_row_count integer;
  enough_quota boolean;
  remaining_quota integer;
+  max_batch_size integer;

  cartodb_id_batch integer;
  batches_n integer;
  DEFAULT_BATCH_SIZE CONSTANT numeric := 100;
-  MAX_BATCH_SIZE CONSTANT numeric := 10000;
+  MAX_SAFE_BATCH_SIZE CONSTANT numeric := 5000;
  current_row_count integer ;

  temp_table_name text;
 BEGIN
+  SELECT csqi.monthly_quota - csqi.used_quota AS remaining_quota, csqi.max_batch_size
+  INTO remaining_quota, max_batch_size
+  FROM cdb_dataservices_client.cdb_service_quota_info() csqi
+  WHERE service = 'hires_geocoder';
+  RAISE DEBUG 'remaining_quota: %; max_batch_size: %', remaining_quota, max_batch_size;
+
  IF batch_size IS NULL THEN
-    RAISE EXCEPTION 'batch_size can''t be null';
-  ELSIF batch_size > MAX_BATCH_SIZE THEN
-    RAISE EXCEPTION 'batch_size must be lower than %', MAX_BATCH_SIZE + 1;
+    batch_size := max_batch_size;
+  ELSIF batch_size > max_batch_size THEN
+    RAISE EXCEPTION 'batch_size must be lower than %', max_batch_size + 1;
+  END IF;
+
+  IF batch_size > MAX_SAFE_BATCH_SIZE THEN
+    batch_size := MAX_SAFE_BATCH_SIZE;
  END IF;

  EXECUTE format('SELECT COUNT(1) from (%s) _x', query) INTO query_row_count;
@@ -2013,11 +2025,7 @@ BEGIN
  RAISE DEBUG 'cdb_bulk_geocode_street_point --> query_row_count: %; query: %; country: %; state: %; city: %; street: %',
      query_row_count, query, country_column, state_column, city_column, street_column;
  SELECT cdb_dataservices_client.cdb_enough_quota('hires_geocoder', query_row_count) INTO enough_quota;
-  IF enough_quota IS NOT NULL AND NOT enough_quota THEN
-    SELECT csqi.monthly_quota - csqi.used_quota AS remaining_quota
-    INTO remaining_quota
-    FROM cdb_dataservices_client.cdb_service_quota_info() csqi
-    WHERE service = 'hires_geocoder';
+  IF remaining_quota < query_row_count THEN
    RAISE EXCEPTION 'Remaining quota: %. Estimated cost: %', remaining_quota, query_row_count;
  END IF;

@@ -2036,25 +2044,27 @@ BEGIN
    coalesce(state_column, ''''''), coalesce(country_column, '''''')
  into street_column, city_column, state_column, country_column;

-  FOR cartodb_id_batch in 0..(batches_n - 1)
-  LOOP
+  IF batches_n > 0 THEN
+    FOR cartodb_id_batch in 0..(batches_n - 1)
+    LOOP

-    EXECUTE format(
-      'WITH geocoding_data as (' ||
-      '   SELECT ' ||
-      '      json_build_object(''id'', cartodb_id, ''address'', %s, ''city'', %s, ''state'', %s, ''country'', %s) as data , ' ||
-      '      floor((cartodb_id-1)::float/$1) as batch' ||
-      '   FROM (%s) _x' ||
-      ') ' ||
-      'INSERT INTO %s SELECT (cdb_dataservices_client._cdb_bulk_geocode_street_point(jsonb_agg(data))).* ' ||
-      'FROM geocoding_data ' ||
-      'WHERE batch = $2', street_column, city_column, state_column, country_column, query, temp_table_name)
-    USING batch_size, cartodb_id_batch;
+      EXECUTE format(
+        'WITH geocoding_data as (' ||
+        '   SELECT ' ||
+        '      json_build_object(''id'', cartodb_id, ''address'', %s, ''city'', %s, ''state'', %s, ''country'', %s) as data , ' ||
+        '      floor((cartodb_id-1)::float/$1) as batch' ||
+        '   FROM (%s) _x' ||
+        ') ' ||
+        'INSERT INTO %s SELECT (cdb_dataservices_client._cdb_bulk_geocode_street_point(jsonb_agg(data))).* ' ||
+        'FROM geocoding_data ' ||
+        'WHERE batch = $2', street_column, city_column, state_column, country_column, query, temp_table_name)
+      USING batch_size, cartodb_id_batch;

-    GET DIAGNOSTICS current_row_count = ROW_COUNT;
-    RAISE DEBUG 'Batch % --> %', cartodb_id_batch, current_row_count;
+      GET DIAGNOSTICS current_row_count = ROW_COUNT;
+      RAISE DEBUG 'Batch % --> %', cartodb_id_batch, current_row_count;

-  END LOOP;
+    END LOOP;
+  END IF;

  RETURN QUERY EXECUTE 'SELECT * FROM ' || quote_ident(temp_table_name);
 END;
--- a/client/sql/16_custom_types.sql
+++ b/client/sql/16_custom_types.sql
@@ -39,5 +39,6 @@ CREATE TYPE cdb_dataservices_client.service_quota_info AS (
    monthly_quota NUMERIC,
    used_quota NUMERIC,
    soft_limit BOOLEAN,
-    provider TEXT
+    provider TEXT,
+    max_batch_size NUMERIC
 );
--- a/client/sql/21_bulk_geocoding_functions.sql
+++ b/client/sql/21_bulk_geocoding_functions.sql
@@ -1,23 +1,34 @@
 CREATE OR REPLACE FUNCTION cdb_dataservices_client.cdb_bulk_geocode_street_point (query text,
-    street_column text, city_column text default null, state_column text default null, country_column text default null, batch_size integer DEFAULT 50)
+    street_column text, city_column text default null, state_column text default null, country_column text default null, batch_size integer DEFAULT NULL)
 RETURNS SETOF cdb_dataservices_client.geocoding AS $$
 DECLARE
  query_row_count integer;
  enough_quota boolean;
  remaining_quota integer;
+  max_batch_size integer;

  cartodb_id_batch integer;
  batches_n integer;
  DEFAULT_BATCH_SIZE CONSTANT numeric := 100;
-  MAX_BATCH_SIZE CONSTANT numeric := 10000;
+  MAX_SAFE_BATCH_SIZE CONSTANT numeric := 5000;
  current_row_count integer ;

  temp_table_name text;
 BEGIN
+  SELECT csqi.monthly_quota - csqi.used_quota AS remaining_quota, csqi.max_batch_size
+  INTO remaining_quota, max_batch_size
+  FROM cdb_dataservices_client.cdb_service_quota_info() csqi
+  WHERE service = 'hires_geocoder';
+  RAISE DEBUG 'remaining_quota: %; max_batch_size: %', remaining_quota, max_batch_size;
+
  IF batch_size IS NULL THEN
-    RAISE EXCEPTION 'batch_size can''t be null';
-  ELSIF batch_size > MAX_BATCH_SIZE THEN
-    RAISE EXCEPTION 'batch_size must be lower than %', MAX_BATCH_SIZE + 1;
+    batch_size := max_batch_size;
+  ELSIF batch_size > max_batch_size THEN
+    RAISE EXCEPTION 'batch_size must be lower than %', max_batch_size + 1;
+  END IF;
+
+  IF batch_size > MAX_SAFE_BATCH_SIZE THEN
+    batch_size := MAX_SAFE_BATCH_SIZE;
  END IF;

  EXECUTE format('SELECT COUNT(1) from (%s) _x', query) INTO query_row_count;
@@ -25,11 +36,7 @@ BEGIN
  RAISE DEBUG 'cdb_bulk_geocode_street_point --> query_row_count: %; query: %; country: %; state: %; city: %; street: %',
      query_row_count, query, country_column, state_column, city_column, street_column;
  SELECT cdb_dataservices_client.cdb_enough_quota('hires_geocoder', query_row_count) INTO enough_quota;
-  IF enough_quota IS NOT NULL AND NOT enough_quota THEN
-    SELECT csqi.monthly_quota - csqi.used_quota AS remaining_quota
-    INTO remaining_quota
-    FROM cdb_dataservices_client.cdb_service_quota_info() csqi
-    WHERE service = 'hires_geocoder';
+  IF remaining_quota < query_row_count THEN
    RAISE EXCEPTION 'Remaining quota: %. Estimated cost: %', remaining_quota, query_row_count;
  END IF;

@@ -48,25 +55,27 @@ BEGIN
    coalesce(state_column, ''''''), coalesce(country_column, '''''')
  into street_column, city_column, state_column, country_column;

-  FOR cartodb_id_batch in 0..(batches_n - 1)
-  LOOP
+  IF batches_n > 0 THEN
+    FOR cartodb_id_batch in 0..(batches_n - 1)
+    LOOP

-    EXECUTE format(
-      'WITH geocoding_data as (' ||
-      '   SELECT ' ||
-      '      json_build_object(''id'', cartodb_id, ''address'', %s, ''city'', %s, ''state'', %s, ''country'', %s) as data , ' ||
-      '      floor((cartodb_id-1)::float/$1) as batch' ||
-      '   FROM (%s) _x' ||
-      ') ' ||
-      'INSERT INTO %s SELECT (cdb_dataservices_client._cdb_bulk_geocode_street_point(jsonb_agg(data))).* ' ||
-      'FROM geocoding_data ' ||
-      'WHERE batch = $2', street_column, city_column, state_column, country_column, query, temp_table_name)
-    USING batch_size, cartodb_id_batch;
+      EXECUTE format(
+        'WITH geocoding_data as (' ||
+        '   SELECT ' ||
+        '      json_build_object(''id'', cartodb_id, ''address'', %s, ''city'', %s, ''state'', %s, ''country'', %s) as data , ' ||
+        '      floor((cartodb_id-1)::float/$1) as batch' ||
+        '   FROM (%s) _x' ||
+        ') ' ||
+        'INSERT INTO %s SELECT (cdb_dataservices_client._cdb_bulk_geocode_street_point(jsonb_agg(data))).* ' ||
+        'FROM geocoding_data ' ||
+        'WHERE batch = $2', street_column, city_column, state_column, country_column, query, temp_table_name)
+      USING batch_size, cartodb_id_batch;

-    GET DIAGNOSTICS current_row_count = ROW_COUNT;
-    RAISE DEBUG 'Batch % --> %', cartodb_id_batch, current_row_count;
+      GET DIAGNOSTICS current_row_count = ROW_COUNT;
+      RAISE DEBUG 'Batch % --> %', cartodb_id_batch, current_row_count;

-  END LOOP;
+    END LOOP;
+  END IF;

  RETURN QUERY EXECUTE 'SELECT * FROM ' || quote_ident(temp_table_name);
 END;
--- a/client/test/expected/21_bulk_geocoding_functions_test.out
+++ b/client/test/expected/21_bulk_geocoding_functions_test.out
@@ -1,21 +1,21 @@
 \set VERBOSITY terse
-- Test bulk size mandatory
-SELECT cdb_dataservices_client.cdb_bulk_geocode_street_point('select 1 as cartodb_id', '''Valladolid, Spain''', null, null, null, null);
-ERROR:  batch_size can't be null
-- Test quota check by mocking quota 0
+ALTER FUNCTION cdb_dataservices_client.cdb_service_quota_info() RENAME TO cdb_service_quota_info_mocked;
+CREATE FUNCTION cdb_dataservices_client.cdb_service_quota_info ()
+RETURNS SETOF cdb_dataservices_client.service_quota_info AS $$
+  SELECT 'hires_geocoder'::cdb_dataservices_client.service_type AS service, 0::NUMERIC AS monthly_quota, 0::NUMERIC AS used_quota, FALSE AS soft_limit, 'google' AS provider, 1::NUMERIC AS max_batch_size;
+$$ LANGUAGE SQL;
 ALTER FUNCTION cdb_dataservices_client.cdb_enough_quota (service TEXT ,input_size NUMERIC) RENAME TO cdb_enough_quota_mocked;
 CREATE FUNCTION cdb_dataservices_client.cdb_enough_quota (service TEXT ,input_size NUMERIC)
 RETURNS BOOLEAN as $$
  SELECT FALSE;
 $$ LANGUAGE SQL;
-ALTER FUNCTION cdb_dataservices_client.cdb_service_quota_info() RENAME TO cdb_service_quota_info_mocked;
-CREATE FUNCTION cdb_dataservices_client.cdb_service_quota_info ()
-RETURNS SETOF cdb_dataservices_client.service_quota_info AS $$
-  SELECT 'hires_geocoder'::cdb_dataservices_client.service_type AS service, 0::NUMERIC AS monthly_quota, 0::NUMERIC AS used_quota, FALSE AS soft_limit, 'google' AS provider;
-$$ LANGUAGE SQL;
+-- Test bulk size not mandatory (it will get the optimal)
+SELECT cdb_dataservices_client.cdb_bulk_geocode_street_point('select 1 as cartodb_id', '''Valladolid, Spain''', null, null, null, null);
+ERROR:  Remaining quota: 0. Estimated cost: 1
+-- Test quota check by mocking quota 0
 SELECT cdb_dataservices_client.cdb_bulk_geocode_street_point('select 1 as cartodb_id', '''Valladolid, Spain''');
 ERROR:  Remaining quota: 0. Estimated cost: 1
 DROP FUNCTION cdb_dataservices_client.cdb_service_quota_info;
 DROP FUNCTION cdb_dataservices_client.cdb_enough_quota;
-ALTER FUNCTION cdb_dataservices_client.cdb_service_quota_info_mocked() RENAME TO cdb_service_quota_info;
 ALTER FUNCTION cdb_dataservices_client.cdb_enough_quota_mocked (service TEXT ,input_size NUMERIC) RENAME TO cdb_enough_quota;
+ALTER FUNCTION cdb_dataservices_client.cdb_service_quota_info_mocked() RENAME TO cdb_service_quota_info;
--- a/client/test/sql/21_bulk_geocoding_functions_test.sql
+++ b/client/test/sql/21_bulk_geocoding_functions_test.sql
@@ -1,26 +1,26 @@
 \set VERBOSITY terse

-- Test bulk size mandatory
-SELECT cdb_dataservices_client.cdb_bulk_geocode_street_point('select 1 as cartodb_id', '''Valladolid, Spain''', null, null, null, null);
+ALTER FUNCTION cdb_dataservices_client.cdb_service_quota_info() RENAME TO cdb_service_quota_info_mocked;
+CREATE FUNCTION cdb_dataservices_client.cdb_service_quota_info ()
+RETURNS SETOF cdb_dataservices_client.service_quota_info AS $$
+  SELECT 'hires_geocoder'::cdb_dataservices_client.service_type AS service, 0::NUMERIC AS monthly_quota, 0::NUMERIC AS used_quota, FALSE AS soft_limit, 'google' AS provider, 1::NUMERIC AS max_batch_size;
+$$ LANGUAGE SQL;

-- Test quota check by mocking quota 0
 ALTER FUNCTION cdb_dataservices_client.cdb_enough_quota (service TEXT ,input_size NUMERIC) RENAME TO cdb_enough_quota_mocked;
 CREATE FUNCTION cdb_dataservices_client.cdb_enough_quota (service TEXT ,input_size NUMERIC)
 RETURNS BOOLEAN as $$
  SELECT FALSE;
 $$ LANGUAGE SQL;

-ALTER FUNCTION cdb_dataservices_client.cdb_service_quota_info() RENAME TO cdb_service_quota_info_mocked;
-CREATE FUNCTION cdb_dataservices_client.cdb_service_quota_info ()
-RETURNS SETOF cdb_dataservices_client.service_quota_info AS $$
-  SELECT 'hires_geocoder'::cdb_dataservices_client.service_type AS service, 0::NUMERIC AS monthly_quota, 0::NUMERIC AS used_quota, FALSE AS soft_limit, 'google' AS provider;
-$$ LANGUAGE SQL;
+-- Test bulk size not mandatory (it will get the optimal)
+SELECT cdb_dataservices_client.cdb_bulk_geocode_street_point('select 1 as cartodb_id', '''Valladolid, Spain''', null, null, null, null);

+-- Test quota check by mocking quota 0
 SELECT cdb_dataservices_client.cdb_bulk_geocode_street_point('select 1 as cartodb_id', '''Valladolid, Spain''');

 DROP FUNCTION cdb_dataservices_client.cdb_service_quota_info;
 DROP FUNCTION cdb_dataservices_client.cdb_enough_quota;

-ALTER FUNCTION cdb_dataservices_client.cdb_service_quota_info_mocked() RENAME TO cdb_service_quota_info;
 ALTER FUNCTION cdb_dataservices_client.cdb_enough_quota_mocked (service TEXT ,input_size NUMERIC) RENAME TO cdb_enough_quota;
+ALTER FUNCTION cdb_dataservices_client.cdb_service_quota_info_mocked() RENAME TO cdb_service_quota_info;