Merge pull request ClickHouse#78775 from ClickHouse/refactor-code-around-data-lakes

Small refactoring around data lakes

Author: kssenii

src/Storages/ObjectStorage/DataLakes/DataLakeConfiguration.h

@@ -64,8 +64,6 @@ class DataLakeConfiguration : public BaseStorageConfiguration, public std::enabl
                     ErrorCodes::FORMAT_VERSION_TOO_OLD,
                     "Metadata is not consinsent with the one which was used to infer table schema. Please, retry the query.");
             }
-            if (!supportsFileIterator())
-                BaseStorageConfiguration::setPaths(current_metadata->getDataFiles());
         }
     }
 
@@ -81,14 +79,6 @@ class DataLakeConfiguration : public BaseStorageConfiguration, public std::enabl
         return std::nullopt;
     }
 
-    void implementPartitionPruning(const ActionsDAG & filter_dag) override
-    {
-        if (!current_metadata || !current_metadata->supportsPartitionPruning())
-            return;
-        BaseStorageConfiguration::setPaths(current_metadata->makePartitionPruning(filter_dag));
-    }
-
-
     std::optional<size_t> totalRows() override
     {
         if (!current_metadata)
@@ -123,20 +113,11 @@ class DataLakeConfiguration : public BaseStorageConfiguration, public std::enabl
         ContextPtr context) override
     {
         BaseStorageConfiguration::update(object_storage, context);
-        if (updateMetadataObjectIfNeeded(object_storage, context))
-        {
-            if (!supportsFileIterator())
-                BaseStorageConfiguration::setPaths(current_metadata->getDataFiles());
-        }
-
+        updateMetadataObjectIfNeeded(object_storage, context);
         return ColumnsDescription{current_metadata->getTableSchema()};
     }
 
-    bool supportsFileIterator() const override
-    {
-        chassert(current_metadata);
-        return current_metadata->supportsFileIterator();
-    }
+    bool supportsFileIterator() const override { return true; }
 
     ObjectIterator iterate(
         const ActionsDAG * filter_dag,

src/Storages/ObjectStorage/DataLakes/DeltaLakeMetadata.cpp

@@ -600,6 +600,7 @@ DeltaLakeMetadata::DeltaLakeMetadata(ObjectStoragePtr object_storage_, Configura
     data_files = result.data_files;
     schema = result.schema;
     partition_columns = result.partition_columns;
+    object_storage = object_storage_;
 
     LOG_TRACE(impl.log, "Found {} data files, {} partition files, schema: {}",
              data_files.size(), partition_columns.size(), schema.toString());
@@ -712,6 +713,14 @@ Field DeltaLakeMetadata::getFieldValue(const String & value, DataTypePtr data_ty
     throw Exception(ErrorCodes::BAD_ARGUMENTS, "Unsupported DeltaLake type for {}", check_type->getColumnType());
 }
 
+ObjectIterator DeltaLakeMetadata::iterate(
+    const ActionsDAG * filter_dag,
+    FileProgressCallback callback,
+    size_t /* list_batch_size */) const
+{
+    return createKeysIterator(getDataFiles(filter_dag), object_storage, callback);
+}
+
 }
 
 #endif

src/Storages/ObjectStorage/DataLakes/DeltaLakeMetadata.h

@@ -41,8 +41,6 @@ class DeltaLakeMetadata final : public IDataLakeMetadata
 
     DeltaLakeMetadata(ObjectStoragePtr object_storage_, ConfigurationObserverPtr configuration_, ContextPtr context_);
 
-    Strings getDataFiles() const override { return data_files; }
-
     NamesAndTypesList getTableSchema() const override { return schema; }
 
     DeltaLakePartitionColumns getPartitionColumns() const { return partition_columns; }
@@ -80,10 +78,19 @@ class DeltaLakeMetadata final : public IDataLakeMetadata
     static DataTypePtr getFieldValue(const Poco::JSON::Object::Ptr & field, const String & type_key, bool is_nullable);
     static Field getFieldValue(const String & value, DataTypePtr data_type);
 
+protected:
+    ObjectIterator iterate(
+        const ActionsDAG * filter_dag,
+        FileProgressCallback callback,
+        size_t list_batch_size) const override;
+
 private:
     mutable Strings data_files;
     NamesAndTypesList schema;
     DeltaLakePartitionColumns partition_columns;
+    ObjectStoragePtr object_storage;
+
+    Strings getDataFiles(const ActionsDAG *) const { return data_files; }
 };
 
 }

src/Storages/ObjectStorage/DataLakes/DeltaLakeMetadataDeltaKernel.cpp

@@ -32,11 +32,6 @@ bool DeltaLakeMetadataDeltaKernel::update(const ContextPtr &)
     return table_snapshot->update();
 }
 
-Strings DeltaLakeMetadataDeltaKernel::getDataFiles() const
-{
-    throwNotImplemented("getDataFiles()");
-}
-
 ObjectIterator DeltaLakeMetadataDeltaKernel::iterate(
     const ActionsDAG * filter_dag,
     FileProgressCallback callback,

src/Storages/ObjectStorage/DataLakes/DeltaLakeMetadataDeltaKernel.h

@@ -40,8 +40,6 @@ class DeltaLakeMetadataDeltaKernel final : public IDataLakeMetadata
 
     bool update(const ContextPtr & context) override;
 
-    Strings getDataFiles() const override;
-
     NamesAndTypesList getTableSchema() const override;
 
     NamesAndTypesList getReadSchema() const override;
@@ -61,8 +59,6 @@ class DeltaLakeMetadataDeltaKernel final : public IDataLakeMetadata
             settings_ref[StorageObjectStorageSetting::delta_lake_read_schema_same_as_table_schema]);
     }
 
-    bool supportsFileIterator() const override { return true; }
-
     ObjectIterator iterate(
         const ActionsDAG * filter_dag,
         FileProgressCallback callback,

src/Storages/ObjectStorage/DataLakes/HudiMetadata.cpp

@@ -91,11 +91,19 @@ HudiMetadata::HudiMetadata(ObjectStoragePtr object_storage_, ConfigurationObserv
 {
 }
 
-Strings HudiMetadata::getDataFiles() const
+Strings HudiMetadata::getDataFiles(const ActionsDAG *) const
 {
     if (data_files.empty())
         data_files = getDataFilesImpl();
     return data_files;
 }
 
+ObjectIterator HudiMetadata::iterate(
+    const ActionsDAG * filter_dag,
+    FileProgressCallback callback,
+    size_t /* list_batch_size */) const
+{
+    return createKeysIterator(getDataFiles(filter_dag), object_storage, callback);
+}
+
 }

src/Storages/ObjectStorage/DataLakes/HudiMetadata.h

@@ -19,8 +19,6 @@ class HudiMetadata final : public IDataLakeMetadata, private WithContext
 
     HudiMetadata(ObjectStoragePtr object_storage_, ConfigurationObserverPtr configuration_, ContextPtr context_);
 
-    Strings getDataFiles() const override;
-
     NamesAndTypesList getTableSchema() const override { return {}; }
 
     bool operator ==(const IDataLakeMetadata & other) const override
@@ -39,12 +37,19 @@ class HudiMetadata final : public IDataLakeMetadata, private WithContext
         return std::make_unique<HudiMetadata>(object_storage, configuration, local_context);
     }
 
+protected:
+    ObjectIterator iterate(
+        const ActionsDAG * filter_dag,
+        FileProgressCallback callback,
+        size_t list_batch_size) const override;
+
 private:
     const ObjectStoragePtr object_storage;
     const ConfigurationObserverPtr configuration;
     mutable Strings data_files;
 
     Strings getDataFilesImpl() const;
+    Strings getDataFiles(const ActionsDAG * filter_dag) const;
 };
 
 }

src/Storages/ObjectStorage/DataLakes/IDataLakeMetadata.cpp

@@ -0,0 +1,63 @@
+#include "IDataLakeMetadata.h"
+#include <Storages/ObjectStorage/StorageObjectStorageSource.h>
+
+namespace DB
+{
+
+namespace
+{
+
+class KeysIterator : public IObjectIterator
+{
+public:
+    KeysIterator(
+        Strings && data_files_,
+        ObjectStoragePtr object_storage_,
+        IDataLakeMetadata::FileProgressCallback callback_)
+        : data_files(data_files_)
+        , object_storage(object_storage_)
+        , callback(callback_)
+    {
+    }
+
+    size_t estimatedKeysCount() override
+    {
+        return data_files.size();
+    }
+
+    ObjectInfoPtr next(size_t) override
+    {
+        while (true)
+        {
+            size_t current_index = index.fetch_add(1, std::memory_order_relaxed);
+            if (current_index >= data_files.size())
+                return nullptr;
+
+            auto key = data_files[current_index];
+            auto object_metadata = object_storage->getObjectMetadata(key);
+
+            if (callback)
+                callback(FileProgress(0, object_metadata.size_bytes));
+
+            return std::make_shared<ObjectInfo>(key, std::move(object_metadata));
+        }
+    }
+
+private:
+    Strings data_files;
+    ObjectStoragePtr object_storage;
+    std::atomic<size_t> index = 0;
+    IDataLakeMetadata::FileProgressCallback callback;
+};
+
+}
+
+ObjectIterator IDataLakeMetadata::createKeysIterator(
+    Strings && data_files_,
+    ObjectStoragePtr object_storage_,
+    IDataLakeMetadata::FileProgressCallback callback_) const
+{
+    return std::make_shared<KeysIterator>(std::move(data_files_), object_storage_, callback_);
+}
+
+}

src/Storages/ObjectStorage/DataLakes/IDataLakeMetadata.h

@@ -20,17 +20,12 @@ class IDataLakeMetadata : boost::noncopyable
 
     virtual bool operator==(const IDataLakeMetadata & other) const = 0;
 
-    /// List all data files.
-    /// For better parallelization, iterate() method should be used.
-    virtual Strings getDataFiles() const = 0;
-    /// Whether `iterate()` method is supported for the data lake.
-    virtual bool supportsFileIterator() const { return false; }
     /// Return iterator to `data files`.
     using FileProgressCallback = std::function<void(FileProgress)>;
     virtual ObjectIterator iterate(
         const ActionsDAG * /* filter_dag */,
         FileProgressCallback /* callback */,
-        size_t /* list_batch_size */) const { throwNotImplemented("iterate()"); }
+        size_t /* list_batch_size */) const = 0;
 
     /// Table schema from data lake metadata.
     virtual NamesAndTypesList getTableSchema() const = 0;
@@ -39,9 +34,6 @@ class IDataLakeMetadata : boost::noncopyable
     /// Return nothing if read schema is the same as table schema.
     virtual NamesAndTypesList getReadSchema() const { return {}; }
 
-    virtual bool supportsPartitionPruning() { return false; }
-    virtual Strings makePartitionPruning(const ActionsDAG &) { throwNotImplemented("makePartitionPrunning()"); }
-
     virtual std::shared_ptr<NamesAndTypesList> getInitialSchemaByPath(const String &) const { return {}; }
     virtual std::shared_ptr<const ActionsDAG> getSchemaTransformer(const String &) const { return {}; }
 
@@ -56,7 +48,13 @@ class IDataLakeMetadata : boost::noncopyable
 
     virtual std::optional<size_t> totalRows() const { return {}; }
     virtual std::optional<size_t> totalBytes() const { return {}; }
+
 protected:
+    ObjectIterator createKeysIterator(
+        Strings && data_files_,
+        ObjectStoragePtr object_storage_,
+        IDataLakeMetadata::FileProgressCallback callback_) const;
+
     [[noreturn]] void throwNotImplemented(std::string_view method) const
     {
         throw Exception(ErrorCodes::UNSUPPORTED_METHOD, "Method `{}` is not implemented", method);

src/Storages/ObjectStorage/DataLakes/Iceberg/IcebergMetadata.cpp

@@ -49,6 +49,7 @@ namespace Setting
 {
 extern const SettingsInt64 iceberg_timestamp_ms;
 extern const SettingsInt64 iceberg_snapshot_id;
+extern const SettingsBool use_iceberg_partition_pruning;
 }
 
 
@@ -571,18 +572,23 @@ ManifestListPtr IcebergMetadata::getManifestList(const String & filename) const
     return manifest_list_ptr;
 }
 
-Strings IcebergMetadata::getDataFilesImpl(const ActionsDAG * filter_dag) const
+Strings IcebergMetadata::getDataFiles(const ActionsDAG * filter_dag) const
 {
     if (!relevant_snapshot)
         return {};
 
-    if (!filter_dag && cached_unprunned_files_for_last_processed_snapshot.has_value())
+    bool use_partition_pruning = filter_dag && getContext()->getSettingsRef()[Setting::use_iceberg_partition_pruning];
+
+    if (!use_partition_pruning && cached_unprunned_files_for_last_processed_snapshot.has_value())
         return cached_unprunned_files_for_last_processed_snapshot.value();
 
     Strings data_files;
     for (const auto & manifest_file_ptr : *(relevant_snapshot->manifest_list))
     {
-        ManifestFilesPruner pruner(schema_processor, relevant_snapshot_schema_id, filter_dag, *manifest_file_ptr, getContext());
+        ManifestFilesPruner pruner(
+            schema_processor, relevant_snapshot_schema_id,
+            use_partition_pruning ? filter_dag : nullptr,
+            *manifest_file_ptr, getContext());
         const auto & data_files_in_manifest = manifest_file_ptr->getFiles();
         for (const auto & manifest_file_entry : data_files_in_manifest)
         {
@@ -597,7 +603,7 @@ Strings IcebergMetadata::getDataFilesImpl(const ActionsDAG * filter_dag) const
         }
     }
 
-    if (!filter_dag)
+    if (!use_partition_pruning)
     {
         cached_unprunned_files_for_last_processed_snapshot = data_files;
         return cached_unprunned_files_for_last_processed_snapshot.value();
@@ -606,16 +612,6 @@ Strings IcebergMetadata::getDataFilesImpl(const ActionsDAG * filter_dag) const
     return data_files;
 }
 
-Strings IcebergMetadata::makePartitionPruning(const ActionsDAG & filter_dag)
-{
-    auto configuration_ptr = configuration.lock();
-    if (!configuration_ptr)
-    {
-        throw Exception(ErrorCodes::LOGICAL_ERROR, "Configuration is expired");
-    }
-    return getDataFilesImpl(&filter_dag);
-}
-
 std::optional<size_t> IcebergMetadata::totalRows() const
 {
     auto configuration_ptr = configuration.lock();
@@ -678,6 +674,14 @@ std::optional<size_t> IcebergMetadata::totalBytes() const
     return result;
 }
 
+ObjectIterator IcebergMetadata::iterate(
+    const ActionsDAG * filter_dag,
+    FileProgressCallback callback,
+    size_t /* list_batch_size */) const
+{
+    return createKeysIterator(getDataFiles(filter_dag), object_storage, callback);
+}
+
 }
 
 #endif