Pass clustering_row_ranges to mutation readers.

This will allow readers to reduce the amount of data read.

Signed-off-by: Piotr Jastrzebski <piotr@scylladb.com>

database.cc

@@ -147,8 +147,11 @@ column_family::make_partition_presence_checker(lw_shared_ptr<sstable_list> old_s
 
 mutation_source
 column_family::sstables_as_mutation_source() {
-    return mutation_source([this] (schema_ptr s, const query::partition_range& r, const io_priority_class& pc) {
-        return make_sstable_reader(std::move(s), r, pc);
+    return mutation_source([this] (schema_ptr s,
+                                   const query::partition_range& r,
+                                   query::clustering_key_filtering_context ck_filtering,
+                                   const io_priority_class& pc) {
+        return make_sstable_reader(std::move(s), r, ck_filtering, pc);
     });
 }
 
@@ -180,17 +183,23 @@ class range_sstable_reader final : public mutation_reader::impl {
     // Use a pointer instead of copying, so we don't need to regenerate the reader if
     // the priority changes.
     const io_priority_class& _pc;
+    query::clustering_key_filtering_context _ck_filtering;
 public:
-    range_sstable_reader(schema_ptr s, lw_shared_ptr<sstable_list> sstables, const query::partition_range& pr, const io_priority_class& pc)
+    range_sstable_reader(schema_ptr s,
+                         lw_shared_ptr<sstable_list> sstables,
+                         const query::partition_range& pr,
+                         query::clustering_key_filtering_context ck_filtering,
+                         const io_priority_class& pc)
         : _pr(pr)
         , _sstables(std::move(sstables))
         , _pc(pc)
+        , _ck_filtering(ck_filtering)
     {
         std::vector<mutation_reader> readers;
         for (const lw_shared_ptr<sstables::sstable>& sst : *_sstables | boost::adaptors::map_values) {
             // FIXME: make sstable::read_range_rows() return ::mutation_reader so that we can drop this wrapper.
             mutation_reader reader =
-                make_mutation_reader<sstable_range_wrapping_reader>(sst, s, pr, query::no_clustering_key_filtering, pc);
+                make_mutation_reader<sstable_range_wrapping_reader>(sst, s, pr, _ck_filtering, _pc);
             if (sst->is_shared()) {
                 reader = make_filtering_reader(std::move(reader), belongs_to_current_shard);
             }
@@ -215,23 +224,30 @@ class single_key_sstable_reader final : public mutation_reader::impl {
     // Use a pointer instead of copying, so we don't need to regenerate the reader if
     // the priority changes.
     const io_priority_class& _pc;
+    query::clustering_key_filtering_context _ck_filtering;
 public:
-    single_key_sstable_reader(schema_ptr schema, lw_shared_ptr<sstable_list> sstables, const partition_key& key, const io_priority_class& pc)
+    single_key_sstable_reader(schema_ptr schema,
+                              lw_shared_ptr<sstable_list> sstables,
+                              const partition_key& key,
+                              query::clustering_key_filtering_context ck_filtering,
+                              const io_priority_class& pc)
         : _schema(std::move(schema))
         , _key(sstables::key::from_partition_key(*_schema, key))
         , _sstables(std::move(sstables))
         , _pc(pc)
+        , _ck_filtering(ck_filtering)
     { }
 
     virtual future<mutation_opt> operator()() override {
         if (_done) {
             return make_ready_future<mutation_opt>();
         }
-        return parallel_for_each(*_sstables | boost::adaptors::map_values, [this](const lw_shared_ptr<sstables::sstable>& sstable) {
-            return sstable->read_row(_schema, _key, query::no_clustering_key_filtering, _pc)
-                .then([this](mutation_opt mo) {
-                    apply(_m, std::move(mo));
-            });
+        return parallel_for_each(*_sstables | boost::adaptors::map_values,
+            [this](const lw_shared_ptr<sstables::sstable>& sstable) {
+                return sstable->read_row(_schema, _key, _ck_filtering, _pc)
+                    .then([this](mutation_opt mo) {
+                        apply(_m, std::move(mo));
+                });
         }).then([this] {
             _done = true;
             return std::move(_m);
@@ -240,16 +256,19 @@ public:
 };
 
 mutation_reader
-column_family::make_sstable_reader(schema_ptr s, const query::partition_range& pr, const io_priority_class& pc) const {
+column_family::make_sstable_reader(schema_ptr s,
+                                   const query::partition_range& pr,
+                                   query::clustering_key_filtering_context ck_filtering,
+                                   const io_priority_class& pc) const {
     if (pr.is_singular() && pr.start()->value().has_key()) {
         const dht::ring_position& pos = pr.start()->value();
         if (dht::shard_of(pos.token()) != engine().cpu_id()) {
             return make_empty_reader(); // range doesn't belong to this shard
         }
-        return make_mutation_reader<single_key_sstable_reader>(std::move(s), _sstables, *pos.key(), pc);
+        return make_mutation_reader<single_key_sstable_reader>(std::move(s), _sstables, *pos.key(), ck_filtering, pc);
     } else {
         // range_sstable_reader is not movable so we need to wrap it
-        return make_mutation_reader<range_sstable_reader>(std::move(s), _sstables, pr, pc);
+        return make_mutation_reader<range_sstable_reader>(std::move(s), _sstables, pr, ck_filtering, pc);
     }
 }
 
@@ -308,7 +327,10 @@ column_family::find_row(schema_ptr s, const dht::decorated_key& partition_key, c
 }
 
 mutation_reader
-column_family::make_reader(schema_ptr s, const query::partition_range& range, const io_priority_class& pc) const {
+column_family::make_reader(schema_ptr s,
+                           const query::partition_range& range,
+                           const query::clustering_key_filtering_context& ck_filtering,
+                           const io_priority_class& pc) const {
     if (query::is_wrap_around(range, *s)) {
         // make_combined_reader() can't handle streams that wrap around yet.
         fail(unimplemented::cause::WRAP_AROUND);
@@ -338,13 +360,13 @@ column_family::make_reader(schema_ptr s, const query::partition_range& range, co
     // https://github.com/scylladb/scylla/issues/185
 
     for (auto&& mt : *_memtables) {
-        readers.emplace_back(mt->make_reader(s, range, query::no_clustering_key_filtering, pc));
+        readers.emplace_back(mt->make_reader(s, range, ck_filtering, pc));
     }
 
     if (_config.enable_cache) {
-        readers.emplace_back(_cache.make_reader(s, range, pc));
+        readers.emplace_back(_cache.make_reader(s, range, ck_filtering, pc));
     } else {
-        readers.emplace_back(make_sstable_reader(s, range, pc));
+        readers.emplace_back(make_sstable_reader(s, range, ck_filtering, pc));
     }
 
     return make_combined_reader(std::move(readers));
@@ -1826,8 +1848,11 @@ column_family::query(schema_ptr s, const query::read_command& cmd, query::result
 
 mutation_source
 column_family::as_mutation_source() const {
-    return mutation_source([this] (schema_ptr s, const query::partition_range& range, const io_priority_class& pc) {
-        return this->make_reader(std::move(s), range, pc);
+    return mutation_source([this] (schema_ptr s,
+                                   const query::partition_range& range,
+                                   query::clustering_key_filtering_context ck_filtering,
+                                   const io_priority_class& pc) {
+        return this->make_reader(std::move(s), range, ck_filtering, pc);
     });
 }
 

database.hh

@@ -367,7 +367,10 @@ private:
     // Caller needs to ensure that column_family remains live (FIXME: relax this).
     // The 'range' parameter must be live as long as the reader is used.
     // Mutations returned by the reader will all have given schema.
-    mutation_reader make_sstable_reader(schema_ptr schema, const query::partition_range& range, const io_priority_class& pc) const;
+    mutation_reader make_sstable_reader(schema_ptr schema,
+                                        const query::partition_range& range,
+                                        query::clustering_key_filtering_context ck_filtering,
+                                        const io_priority_class& pc) const;
 
     mutation_source sstables_as_mutation_source();
     key_source sstables_as_key_source() const;
@@ -405,6 +408,7 @@ public:
     // will be scheduled under the priority class given by pc.
     mutation_reader make_reader(schema_ptr schema,
             const query::partition_range& range = query::full_partition_range,
+            const query::clustering_key_filtering_context& ck_filtering = query::no_clustering_key_filtering,
             const io_priority_class& pc = default_priority_class()) const;
 
     mutation_source as_mutation_source() const;

memtable.cc

@@ -115,7 +115,7 @@ class scanning_reader final : public mutation_reader::impl {
     stdx::optional<query::partition_range> _delegate_range;
     mutation_reader _delegate;
     const io_priority_class& _pc;
-    const query::clustering_key_filtering_context& _ck_filtering;
+    query::clustering_key_filtering_context _ck_filtering;
 private:
     memtable::partitions_type::iterator lookup_end() {
         auto cmp = partition_entry::compare(_memtable->_schema);

mutation_query.cc

@@ -81,7 +81,8 @@ querying_reader::querying_reader(schema_ptr s,
 { }
 
 future<> querying_reader::read() {
-    _reader = _source(_schema, _range, service::get_local_sstable_query_read_priority());
+    _reader = _source(_schema, _range, query::clustering_key_filtering_context::create(_schema, _slice),
+            service::get_local_sstable_query_read_priority());
     return consume(*_reader, [this](mutation&& m) {
         // FIXME: Make data sources respect row_ranges so that we don't have to filter them out here.
         auto is_distinct = _slice.options.contains(query::partition_slice::option::distinct);

mutation_reader.hh

@@ -153,19 +153,30 @@ future<> consume(mutation_reader& reader, Consumer consumer) {
 // The reader returns mutations having all the same schema, the one passed
 // when invoking the source.
 class mutation_source {
-    std::function<mutation_reader(schema_ptr, const query::partition_range& range, const io_priority_class& pc)> _fn;
+    using partition_range = const query::partition_range&;
+    using clustering_filter = query::clustering_key_filtering_context;
+    using io_priority = const io_priority_class&;
+    std::function<mutation_reader(schema_ptr, partition_range, clustering_filter, io_priority)> _fn;
 public:
-    mutation_source(std::function<mutation_reader(schema_ptr, const query::partition_range& range, const io_priority_class& pc)> fn) : _fn(std::move(fn)) {}
-    mutation_source(std::function<mutation_reader(schema_ptr, const query::partition_range& range)> fn)
-        : _fn([fn = std::move(fn)] (schema_ptr s, const query::partition_range& range, const io_priority_class& pc) {
+    mutation_source(std::function<mutation_reader(schema_ptr, partition_range, clustering_filter, io_priority)> fn)
+        : _fn(std::move(fn)) {}
+    mutation_source(std::function<mutation_reader(schema_ptr, partition_range, clustering_filter)> fn)
+        : _fn([fn = std::move(fn)] (schema_ptr s, partition_range range, clustering_filter ck_filtering, io_priority) {
+            return fn(s, range, ck_filtering);
+        }) {}
+    mutation_source(std::function<mutation_reader(schema_ptr, partition_range range)> fn)
+        : _fn([fn = std::move(fn)] (schema_ptr s, partition_range range, clustering_filter, io_priority) {
             return fn(s, range);
         }) {}
 
-    mutation_reader operator()(schema_ptr s, const query::partition_range& range, const io_priority_class& pc) const {
-        return _fn(std::move(s), range, pc);
+    mutation_reader operator()(schema_ptr s, partition_range range, clustering_filter ck_filtering, io_priority pc) const {
+        return _fn(std::move(s), range, ck_filtering, pc);
     }
-    mutation_reader operator()(schema_ptr s, const query::partition_range& range) const {
-        return _fn(std::move(s), range, default_priority_class());
+    mutation_reader operator()(schema_ptr s, partition_range range, clustering_filter ck_filtering) const {
+        return _fn(std::move(s), range, ck_filtering, default_priority_class());
+    }
+    mutation_reader operator()(schema_ptr s, partition_range range) const {
+        return _fn(std::move(s), range, query::no_clustering_key_filtering, default_priority_class());
     }
 };
 

repair/repair.cc

@@ -327,7 +327,11 @@ static future<partition_checksum> checksum_range_shard(database &db,
         const ::range<dht::token>& range) {
     auto& cf = db.find_column_family(keyspace_name, cf_name);
     return do_with(query::to_partition_range(range), [&cf] (const auto& partition_range) {
-        return do_with(cf.make_reader(cf.schema(), partition_range, service::get_local_streaming_read_priority()), partition_checksum(),
+        auto reader = cf.make_reader(cf.schema(),
+                                     partition_range,
+                                     query::no_clustering_key_filtering,
+                                     service::get_local_streaming_read_priority());
+        return do_with(std::move(reader), partition_checksum(),
             [] (auto& reader, auto& checksum) {
             return repeat([&reader, &checksum] () {
                 return reader().then([&checksum] (auto mopt) {

row_cache.cc

@@ -279,7 +279,10 @@ class scanning_and_populating_reader final : public mutation_reader::impl {
     dht::decorated_key_opt _last_secondary_key;
     const io_priority_class _pc;
 public:
-    scanning_and_populating_reader(schema_ptr s, row_cache& cache, const query::partition_range& range, const io_priority_class& pc)
+    scanning_and_populating_reader(schema_ptr s,
+                                   row_cache& cache,
+                                   const query::partition_range& range,
+                                   const io_priority_class& pc)
         : _cache(cache), _schema(s),
           _primary(make_mutation_reader<just_cache_scanning_reader>(s, cache, range)),
           _underlying(cache._underlying), _original_range(range), _underlying_keys(cache._underlying_keys),
@@ -319,7 +322,7 @@ public:
                 _range = query::partition_range(query::partition_range::bound { std::move(*dk), true }, std::move(end));
                 _last_secondary_key = {};
                 _secondary_phase = _cache._populate_phaser.phase();
-                _secondary = _underlying(_cache._schema, _range, _pc);
+                _secondary = _underlying(_cache._schema, _range, query::no_clustering_key_filtering, _pc);
                 _secondary_only = true;
                 return next_secondary();
             });
@@ -332,7 +335,7 @@ private:
             auto cmp = dht::ring_position_comparator(*_schema);
             _range = _range.split_after(*_last_secondary_key, cmp);
             _secondary_phase = _cache._populate_phaser.phase();
-            _secondary = _underlying(_cache._schema, _range, _pc);
+            _secondary = _underlying(_cache._schema, _range, query::no_clustering_key_filtering, _pc);
         }
         return _secondary().then([this, op = _cache._populate_phaser.start()] (mutation_opt&& mo) {
             if (!mo && _next_primary) {
@@ -361,7 +364,9 @@ private:
 };
 
 mutation_reader
-row_cache::make_scanning_reader(schema_ptr s, const query::partition_range& range, const io_priority_class& pc) {
+row_cache::make_scanning_reader(schema_ptr s,
+                                const query::partition_range& range,
+                                const io_priority_class& pc) {
     if (range.is_wrap_around(dht::ring_position_comparator(*s))) {
         warn(unimplemented::cause::WRAP_AROUND);
         throw std::runtime_error("row_cache doesn't support wrap-around ranges");
@@ -369,13 +374,49 @@ row_cache::make_scanning_reader(schema_ptr s, const query::partition_range& rang
     return make_mutation_reader<scanning_and_populating_reader>(std::move(s), *this, range, pc);
 }
 
+class slicing_reader : public mutation_reader::impl {
+private:
+    mutation_reader _underlying;
+    query::clustering_key_filtering_context _ck_filtering;
+
+    future<mutation_opt> filter(mutation_opt&& mut) {
+        while (mut && !mut->partition().empty()) {
+            const query::clustering_row_ranges& ck_ranges = _ck_filtering.get_ranges(mut->key());
+            mutation_partition filtered_partition = mutation_partition(mut->partition(), *(mut->schema()), ck_ranges);
+
+            if (!filtered_partition.empty()) {
+                mut->partition() = std::move(filtered_partition);
+                return make_ready_future<mutation_opt>(std::move(mut));
+            }
+
+            future<mutation_opt> next = _underlying();
+            if (!next.available()) {
+                return next.then([this] (mutation_opt&& mut) { return filter(std::move(mut)); });
+            }
+            mut = std::move(next.get0());
+        }
+        return make_ready_future<mutation_opt>(std::move(mut));
+    }
+
+public:
+    slicing_reader(mutation_reader&& reader, query::clustering_key_filtering_context ck_filtering)
+        : _underlying(std::move(reader)), _ck_filtering(std::move(ck_filtering)) {}
+
+    virtual future<mutation_opt> operator()() override {
+        return _underlying().then([this] (mutation_opt&& mut) { return filter(std::move(mut)); });
+    }
+};
+
 mutation_reader
-row_cache::make_reader(schema_ptr s, const query::partition_range& range, const io_priority_class& pc) {
+row_cache::make_reader(schema_ptr s,
+                       const query::partition_range& range,
+                       query::clustering_key_filtering_context ck_filtering,
+                       const io_priority_class& pc) {
     if (range.is_singular()) {
         const query::ring_position& pos = range.start()->value();
 
         if (!pos.has_key()) {
-            return make_scanning_reader(std::move(s), range, pc);
+            return make_mutation_reader<slicing_reader>(make_scanning_reader(std::move(s), range, pc), ck_filtering);
         }
 
         return _read_section(_tracker.region(), [&] {
@@ -387,16 +428,18 @@ row_cache::make_reader(schema_ptr s, const query::partition_range& range, const
                 _tracker.touch(e);
                 on_hit();
                 upgrade_entry(e);
-                return make_reader_returning(e.read(s));
+                return make_reader_returning(e.read(s, ck_filtering));
             } else {
                 on_miss();
-                return make_mutation_reader<populating_reader>(s, *this, _underlying(_schema, range, pc));
+                return make_mutation_reader<slicing_reader>(
+                    make_mutation_reader<populating_reader>(s, *this, _underlying(_schema, range, query::no_clustering_key_filtering, pc)),
+                    ck_filtering);
             }
           });
         });
     }
 
-    return make_scanning_reader(std::move(s), range, pc);
+    return make_mutation_reader<slicing_reader>(make_scanning_reader(std::move(s), range, pc), ck_filtering);
 }
 
 row_cache::~row_cache() {
@@ -622,6 +665,16 @@ mutation cache_entry::read(const schema_ptr& s) {
     return m;
 }
 
+mutation cache_entry::read(const schema_ptr& s, query::clustering_key_filtering_context ck_filtering) {
+    const query::clustering_row_ranges& ck_ranges = ck_filtering.get_ranges(_key.key());
+    mutation_partition filtered_partition = mutation_partition(_p, *_schema, ck_ranges);
+    auto m = mutation(_schema, _key, std::move(filtered_partition));
+    if (_schema != s) {
+        m.upgrade(s);
+    }
+    return m;
+}
+
 const schema_ptr& row_cache::schema() const {
     return _schema;
 }

row_cache.hh

@@ -84,6 +84,7 @@ public:
     const schema_ptr& schema() const { return _schema; }
     schema_ptr& schema() { return _schema; }
     mutation read(const schema_ptr&);
+    mutation read(const schema_ptr&, query::clustering_key_filtering_context);
 
     struct compare {
         dht::decorated_key::less_comparator _c;
@@ -195,7 +196,9 @@ private:
     logalloc::allocating_section _update_section;
     logalloc::allocating_section _populate_section;
     logalloc::allocating_section _read_section;
-    mutation_reader make_scanning_reader(schema_ptr, const query::partition_range&, const io_priority_class& pc);
+    mutation_reader make_scanning_reader(schema_ptr,
+                                         const query::partition_range&,
+                                         const io_priority_class& pc);
     void on_hit();
     void on_miss();
     void upgrade_entry(cache_entry&);
@@ -212,7 +215,10 @@ public:
     // User needs to ensure that the row_cache object stays alive
     // as long as the reader is used.
     // The range must not wrap around.
-    mutation_reader make_reader(schema_ptr, const query::partition_range& = query::full_partition_range, const io_priority_class& = default_priority_class());
+    mutation_reader make_reader(schema_ptr,
+                                const query::partition_range& = query::full_partition_range,
+                                query::clustering_key_filtering_context = query::no_clustering_key_filtering,
+                                const io_priority_class& = default_priority_class());
 
     const stats& stats() const { return _stats; }
 public:

service/storage_proxy.cc

@@ -3159,7 +3159,8 @@ private:
         return _db.invoke_on(_shard, [this] (database& db) {
             schema_ptr s = _schema;
             column_family& cf = db.find_column_family(s->id());
-            return make_foreign(std::make_unique<remote_state>(remote_state{cf.make_reader(std::move(s), _range, *_pc)}));
+            return make_foreign(std::make_unique<remote_state>(
+                remote_state{cf.make_reader(std::move(s), _range, query::no_clustering_key_filtering, *_pc)}));
         }).then([this] (auto&& ptr) {
             _remote = std::move(ptr);
         });

streaming/stream_transfer_task.cc

@@ -110,7 +110,7 @@ future<stop_iteration> do_send_mutations(auto si, auto fm) {
 future<> send_mutations(auto si) {
     auto& cf = si->db.find_column_family(si->cf_id);
     auto& priority = service::get_local_streaming_read_priority();
-    return do_with(cf.make_reader(cf.schema(), si->pr, priority), [si] (auto& reader) {
+    return do_with(cf.make_reader(cf.schema(), si->pr, query::no_clustering_key_filtering, priority), [si] (auto& reader) {
         return repeat([si, &reader] () {
             return reader().then([si] (auto mopt) {
                 if (mopt && si->db.column_family_exists(si->cf_id)) {