Extract sharding key from conditions

DifferentialOrange · DifferentialOrange · commit 5cac02b7dae3 · 2021-11-30T16:33:10.000+03:00
PR #181 introduced support of DDL sharding keys. But if sharding key hasn't got a separate index in schema, select with equal conditions for all required sharding key fields still led to map-reduce instead of a single storage call. This patch introduces impoved support of sharding keys extraction and fixes the issue. Closes #213
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -20,6 +20,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 * Use tuple-merger backed select implementation on tarantool 2.10+ (it gives
   less pressure on Lua GC).
+* DDL sharding key now can be extracted from select conditions even if
+  there are no separate index.
 
 ## [0.9.0] - 20-10-21
 
diff --git a/README.md b/README.md
@@ -101,8 +101,6 @@ Current limitations for using custom sharding key:
   updated on storages, see
   [#212](https://github.com/tarantool/crud/issues/212). However it is possible
   to do it manually with `require('crud.sharding_key').update_cache()`.
-- CRUD select may lead map reduce in some cases, see
-  [#213](https://github.com/tarantool/crud/issues/213).
 - No support of JSON path for sharding key, see
   [#219](https://github.com/tarantool/crud/issues/219).
 - `primary_index_fieldno_map` is not cached, see
diff --git a/crud/select/plan.lua b/crud/select/plan.lua
@@ -12,7 +12,9 @@ if has_keydef then
     keydef_lib = compat.require('tuple.keydef', 'key_def')
 end
 
-local select_plan = {}
+local select_plan = {
+    _internal = {},
+}
 
 local IndexTypeError = errors.new_class('IndexTypeError', {capture_stack = false})
 local FilterFieldsError = errors.new_class('FilterFieldsError', {capture_stack = false})
@@ -48,49 +50,92 @@ local function get_index_for_condition(space_indexes, space_format, condition)
     end
 end
 
-local function extract_sharding_key_from_scan_value(scan_value, scan_index, sharding_index)
-    if #scan_value < #sharding_index.parts then
-        return nil
-    end
+function select_plan._internal.extract_sharding_key_from_conditions(conditions, sharding_index,
+                                                                    space_indexes, fieldno_map)
+    dev_checks('table', 'table', 'table', 'table')
 
-    if scan_index.id == sharding_index.id then
-        return scan_value
-    end
+    -- If name is both valid index name and field name,
+    -- it is interpreted as index name.
+    local filled_fields = {}
+    for _, condition in ipairs(conditions) do
+        if condition.operator ~= compare_conditions.operators.EQ then
+            goto continue
+        end
 
-    local scan_value_fields_values = {}
-    for i, scan_index_part in ipairs(scan_index.parts) do
-        scan_value_fields_values[scan_index_part.fieldno] = scan_value[i]
-    end
+        local index = space_indexes[condition.operand]
+        if index ~= nil then
+            for i, part in ipairs(index.parts) do
+            -- Consider the following case:
+            --     index_0: {'foo', 'bar'},
+            --     index_1: {'baz', 'bar'},
+            --     conditions: {{'==', 'index_0', {1, 2}}, {'==', 'index_1', {3, nil}}}.
+            -- To check that nil parts will not overwrite already filled_fields,
+            -- we verify that filled_fields[part.fieldno] is empty. If there are
+            -- more than one non-null different value in conditions with equal operator,
+            -- request is already in conflict and it doesn't matter what sharding key we
+            -- will return.
+                if filled_fields[part.fieldno] == nil then
+                    filled_fields[part.fieldno] = condition.values[i]
+                end
+            end
 
-    -- check that sharding key is included in the scan index fields
-    local sharding_key = {}
-    for _, sharding_key_part in ipairs(sharding_index.parts) do
-        local fieldno = sharding_key_part.fieldno
+            goto continue
+        end
 
-        -- sharding key isn't included in scan key
-        if scan_value_fields_values[fieldno] == nil then
-            return nil
+        local fieldno = fieldno_map[condition.operand]
+        if fieldno == nil then
+            goto continue
         end
+        filled_fields[fieldno] = condition.values[1]
 
-        local field_value = scan_value_fields_values[fieldno]
+        ::continue::
+    end
 
-        -- sharding key contains nil values
-        if field_value == nil then
+    local sharding_key = {}
+    for i, v in ipairs(sharding_index.parts) do
+        if filled_fields[v.fieldno] == nil then
             return nil
         end
 
-        table.insert(sharding_key, field_value)
+        sharding_key[i] = filled_fields[v.fieldno]
     end
 
     return sharding_key
 end
 
+function select_plan._internal.get_sharding_key_from_scan_value(scan_value, scan_index, scan_iter, sharding_index)
+    dev_checks('?', 'table', 'number', 'table')
+
+    if scan_value == nil then
+        return nil
+    end
+
+    if scan_iter ~= box.index.EQ and scan_iter ~= box.index.REQ then
+        return nil
+    end
+
+    if scan_index.id == sharding_index.id then
+        if type(scan_value) ~= 'table' then
+            return scan_value
+        end
+
+        for i, _ in ipairs(sharding_index.parts) do
+            if scan_value[i] == nil then return nil end
+        end
+        return scan_value
+    end
+
+    return nil
+end
+
 -- We need to construct after_tuple by field_names
 -- because if `fields` option is specified we have after_tuple with partial fields
 -- and these fields are ordered by field_names + primary key + scan key
 -- this order can be differ from order in space format
 -- so we need to cast after_tuple to space format for scrolling tuples on storage
-local function construct_after_tuple_by_fields(space_format, field_names, tuple)
+local function construct_after_tuple_by_fields(fieldno_map, field_names, tuple)
+    dev_checks('?table', '?table', '?table|cdata')
+
     if tuple == nil then
         return nil
     end
@@ -99,15 +144,10 @@ local function construct_after_tuple_by_fields(space_format, field_names, tuple)
         return tuple
     end
 
-    local positions = {}
     local transformed_tuple = {}
 
-    for i, field in ipairs(space_format) do
-        positions[field.name] = i
-    end
-
     for i, field_name in ipairs(field_names) do
-        local fieldno = positions[field_name]
+        local fieldno = fieldno_map[field_name]
         if fieldno == nil then
             return nil, FilterFieldsError:new(
                     'Space format doesn\'t contain field named %q', field_name
@@ -145,6 +185,8 @@ function select_plan.new(space, conditions, opts)
     local scan_value
     local scan_condition_num
 
+    local fieldno_map = utils.get_format_fieldno_map(space_format)
+
     -- search index to iterate over
     for i, condition in ipairs(conditions) do
         scan_index = get_index_for_condition(space_indexes, space_format, condition)
@@ -176,9 +218,7 @@ function select_plan.new(space, conditions, opts)
 
     -- handle opts.first
     local total_tuples_count
-    local scan_after_tuple, err = construct_after_tuple_by_fields(
-            space_format, field_names, opts.after_tuple
-    )
+    local scan_after_tuple, err = construct_after_tuple_by_fields(fieldno_map, field_names, opts.after_tuple)
     if err ~= nil then
         return nil, err
     end
@@ -230,9 +270,12 @@ function select_plan.new(space, conditions, opts)
     local sharding_index = opts.sharding_key_as_index_obj or primary_index
 
     -- get sharding key value
-    local sharding_key
-    if scan_value ~= nil and (scan_iter == box.index.EQ or scan_iter == box.index.REQ) then
-        sharding_key = extract_sharding_key_from_scan_value(scan_value, scan_index, sharding_index)
+    local sharding_key = select_plan._internal.get_sharding_key_from_scan_value(scan_value, scan_index,
+                                                                                scan_iter, sharding_index)
+
+    if sharding_key == nil then
+        sharding_key = select_plan._internal.extract_sharding_key_from_conditions(conditions, sharding_index,
+                                                                                  space_indexes, fieldno_map)
     end
 
     local plan = {
diff --git a/test/entrypoint/srv_ddl.lua b/test/entrypoint/srv_ddl.lua
@@ -61,6 +61,14 @@ package.preload['customers-storage'] = function()
                     {path = 'name', is_nullable = false, type = 'string'},
                 },
             }
+            local age_index = {
+                name = 'age',
+                type = 'TREE',
+                unique = false,
+                parts = {
+                    {path = 'age', is_nullable = false, type = 'number'},
+                },
+            }
             local secondary_index = {
                 name = 'secondary',
                 type = 'TREE',
@@ -71,6 +79,17 @@ package.preload['customers-storage'] = function()
                 },
             }
 
+            local three_fields_index = {
+                name = 'three_fields',
+                type = 'TREE',
+                unique = false,
+                parts = {
+                    {path = 'age', is_nullable = false, type = 'number'},
+                    {path = 'name', is_nullable = false, type = 'string'},
+                    {path = 'id', is_nullable = false, type = 'unsigned'},
+                },
+            }
+
             local customers_name_key_schema = table.deepcopy(customers_schema)
             customers_name_key_schema.sharding_key = {'name'}
             table.insert(customers_name_key_schema.indexes, primary_index)
@@ -100,13 +119,27 @@ package.preload['customers-storage'] = function()
             table.insert(customers_age_key_schema.indexes, primary_index)
             table.insert(customers_age_key_schema.indexes, bucket_id_index)
 
+            local customers_name_age_key_different_indexes_schema = table.deepcopy(customers_schema)
+            customers_name_age_key_different_indexes_schema.sharding_key = {'name', 'age'}
+            table.insert(customers_name_age_key_different_indexes_schema.indexes, primary_index)
+            table.insert(customers_name_age_key_different_indexes_schema.indexes, bucket_id_index)
+            table.insert(customers_name_age_key_different_indexes_schema.indexes, age_index)
+
+            local customers_name_age_key_three_fields_index_schema = table.deepcopy(customers_schema)
+            customers_name_age_key_three_fields_index_schema.sharding_key = {'name', 'age'}
+            table.insert(customers_name_age_key_three_fields_index_schema.indexes, primary_index_id)
+            table.insert(customers_name_age_key_three_fields_index_schema.indexes, bucket_id_index)
+            table.insert(customers_name_age_key_three_fields_index_schema.indexes, three_fields_index)
+
             local schema = {
                 spaces = {
                     customers_name_key = customers_name_key_schema,
                     customers_name_key_uniq_index = customers_name_key_uniq_index_schema,
                     customers_name_key_non_uniq_index = customers_name_key_non_uniq_index_schema,
                     customers_secondary_idx_name_key = customers_secondary_idx_name_key_schema,
                     customers_age_key = customers_age_key_schema,
+                    customers_name_age_key_different_indexes = customers_name_age_key_different_indexes_schema,
+                    customers_name_age_key_three_fields_index = customers_name_age_key_three_fields_index_schema,
                 }
             }
 
diff --git a/test/helpers/storage_stat.lua b/test/helpers/storage_stat.lua
@@ -95,4 +95,16 @@ function storage_stat.diff(a, b)
     return diff
 end
 
+-- Accepts collect (or diff) return value and returns
+-- total number of select requests across all storages.
+function storage_stat.total(stats)
+    local total = 0
+
+    for _, stat in pairs(stats) do
+        total = total + (stat.select_requests or 0)
+    end
+
+    return total
+end
+
 return storage_stat
diff --git a/test/integration/ddl_sharding_key_test.lua b/test/integration/ddl_sharding_key_test.lua
diff --git a/test/unit/select_plan_test.lua b/test/unit/select_plan_test.lua
