fix single-threading failsafe when number of files cannot be estimated

zvonand · zvonand · commit c9a31599c08f · 2024-04-24T22:14:47.000Z
diff --git a/src/Storages/StorageS3.cpp b/src/Storages/StorageS3.cpp
@@ -229,6 +229,14 @@ class StorageS3Source::DisclosedGlobIterator::Impl : WithContext
         return buffer.size();
     }
 
+    bool hasMore()
+    {
+        if (!buffer.size())
+            return !(expanded_keys_iter == expanded_keys.end() && is_finished_for_key);
+        else
+            return true;
+    }
+
     ~Impl()
     {
         list_objects_pool.wait();
@@ -481,6 +489,11 @@ size_t StorageS3Source::DisclosedGlobIterator::estimatedKeysCount()
     return pimpl->objectsCount();
 }
 
+bool StorageS3Source::DisclosedGlobIterator::hasMore()
+{
+    return pimpl->hasMore();
+}
+
 class StorageS3Source::KeysIterator::Impl
 {
 public:
@@ -1243,8 +1256,16 @@ void ReadFromStorageS3Step::initializePipeline(QueryPipelineBuilder & pipeline,
     if (estimated_keys_count > 1)
         num_streams = std::min(num_streams, estimated_keys_count);
     else
-        /// Disclosed glob iterator can underestimate the amount of keys in some cases. We will keep one stream for this particular case.
-        num_streams = 1;
+    {
+        const auto glob_iter = std::dynamic_pointer_cast<StorageS3Source::DisclosedGlobIterator>(iterator_wrapper);
+        if (!(glob_iter && glob_iter->hasMore()))
+        {
+            /// Disclosed glob iterator can underestimate the amount of keys in some cases. We will keep one stream for this particular case.
+            num_streams = 1;
+        }
+        /// Otherwise, 1000 files were already listed, but none of them is actually what we are looking for.
+        /// We cannot estimate _how many_ there are left, but if there are more files to list, it's faster to do it in many streams.
+    }
 
     const size_t max_threads = context->getSettingsRef().max_threads;
     const size_t max_parsing_threads = num_streams >= max_threads ? 1 : (max_threads / std::max(num_streams, 1ul));
diff --git a/src/Storages/StorageS3.h b/src/Storages/StorageS3.h
@@ -83,6 +83,7 @@ class StorageS3Source : public SourceWithKeyCondition, WithContext
 
         KeyWithInfoPtr next(size_t idx = 0) override; /// NOLINT
         size_t estimatedKeysCount() override;
+        bool hasMore();
 
     private:
         class Impl;
