fix conflicts

Author: arthurpassos

src/Storages/ObjectStorage/HDFS/Configuration.cpp

@@ -59,18 +59,6 @@ ObjectStoragePtr StorageHDFSConfiguration::createObjectStorage( /// NOLINT
         url, std::move(hdfs_settings), context->getConfigRef(), /* lazy_initialize */true);
 }
 
-std::string StorageHDFSConfiguration::getPathWithoutGlobs() const
-{
-    /// Unlike s3 and azure, which are object storages,
-    /// hdfs is a filesystem, so it cannot list files by partual prefix,
-    /// only by directory.
-    auto first_glob_pos = path.find_first_of("*?{");
-    auto end_of_path_without_globs = path.substr(0, first_glob_pos).rfind('/');
-    if (end_of_path_without_globs == std::string::npos || end_of_path_without_globs == 0)
-        return "/";
-    return path.substr(0, end_of_path_without_globs);
-}
-
 StorageObjectStorage::QuerySettings StorageHDFSConfiguration::getQuerySettings(const ContextPtr & context) const
 {
     const auto & settings = context->getSettingsRef();

src/Storages/ObjectStorage/StorageObjectStorage.cpp

@@ -280,6 +280,20 @@ bool StorageObjectStorage::Configuration::update( ///NOLINT
     return true;
 }
 
+void StorageObjectStorage::Configuration::initPartitionStrategy(ASTPtr partition_by, const ColumnsDescription & columns, ContextPtr context)
+{
+
+    partition_strategy = PartitionStrategyFactory::get(
+        partition_strategy_type,
+        partition_by,
+        columns.getOrdinary(),
+        context,
+        format,
+        getRawPath().withGlobs(),
+        getRawPath().withPartitionWildcard(),
+        partition_columns_in_data_file);
+}
+
 IDataLakeMetadata * StorageObjectStorage::getExternalMetadata(ContextPtr query_context)
 {
     configuration->update(
@@ -454,7 +468,7 @@ class ReadFromObjectStorageStep : public SourceStepWithFilter
         auto context = getContext();
         iterator_wrapper = StorageObjectStorageSource::createFileIterator(
             configuration, configuration->getQuerySettings(context), object_storage, distributed_processing,
-            context, predicate, filter_actions_dag, virtual_columns, nullptr, context->getFileProgressCallback());
+            context, predicate, filter_actions_dag, virtual_columns, info.hive_partition_columns_to_read_from_file_path, nullptr, context->getFileProgressCallback());
     }
 };
 }
@@ -464,9 +478,10 @@ ReadFromFormatInfo StorageObjectStorage::Configuration::prepareReadingFromFormat
     const Strings & requested_columns,
     const StorageSnapshotPtr & storage_snapshot,
     bool supports_subset_of_columns,
-    ContextPtr local_context)
+    ContextPtr local_context,
+    const PrepareReadingFromFormatHiveParams & hive_params)
 {
-    return DB::prepareReadingFromFormat(requested_columns, storage_snapshot, local_context, supports_subset_of_columns);
+    return DB::prepareReadingFromFormat(requested_columns, storage_snapshot, local_context, supports_subset_of_columns, hive_params);
 }
 
 std::optional<ColumnsDescription> StorageObjectStorage::Configuration::tryGetTableStructureFromMetadata() const
@@ -752,7 +767,7 @@ void StorageObjectStorage::Configuration::initialize(
         {
             configuration_to_initialize.format
                 = FormatFactory::instance()
-                      .tryGetFormatFromFileName(configuration_to_initialize.isArchive() ? configuration_to_initialize.getPathInArchive() : configuration_to_initialize.getPath())
+                      .tryGetFormatFromFileName(configuration_to_initialize.isArchive() ? configuration_to_initialize.getPathInArchive() : configuration_to_initialize.getRawPath().path)
                       .value_or("auto");
         }
     }
@@ -767,33 +782,65 @@ void StorageObjectStorage::Configuration::check(ContextPtr) const
     FormatFactory::instance().checkFormatName(format);
 }
 
-bool StorageObjectStorage::Configuration::withPartitionWildcard() const
+StorageObjectStorage::Configuration::Path StorageObjectStorage::Configuration::getPathForRead() const
+{
+    auto raw_path = getRawPath();
+
+    if (!partition_strategy)
+    {
+        return raw_path;
+    }
+
+    return Path {partition_strategy->getPathForRead(raw_path.path)};
+}
+
+StorageObjectStorage::Configuration::Path StorageObjectStorage::Configuration::getPathForWrite(const std::string & partition_id) const
+{
+    auto raw_path = getRawPath();
+
+    if (!partition_strategy)
+    {
+        return raw_path;
+    }
+
+    return Path {partition_strategy->getPathForWrite(raw_path.path, partition_id)};
+}
+
+bool StorageObjectStorage::Configuration::Path::withPartitionWildcard() const
 {
     static const String PARTITION_ID_WILDCARD = "{_partition_id}";
-    return getPath().find(PARTITION_ID_WILDCARD) != String::npos
-        || getNamespace().find(PARTITION_ID_WILDCARD) != String::npos;
+    return path.find(PARTITION_ID_WILDCARD) != String::npos;
 }
 
-bool StorageObjectStorage::Configuration::withGlobsIgnorePartitionWildcard() const
+bool StorageObjectStorage::Configuration::Path::withGlobsIgnorePartitionWildcard() const
 {
     if (!withPartitionWildcard())
         return withGlobs();
-    return PartitionedSink::replaceWildcards(getPath(), "").find_first_of("*?{") != std::string::npos;
+    return PartitionedSink::replaceWildcards(path, "").find_first_of("*?{") != std::string::npos;
 }
 
-bool StorageObjectStorage::Configuration::isPathWithGlobs() const
+bool StorageObjectStorage::Configuration::Path::withGlobs() const
 {
-    return getPath().find_first_of("*?{") != std::string::npos;
+    return path.find_first_of("*?{") != std::string::npos;
 }
 
-bool StorageObjectStorage::Configuration::isNamespaceWithGlobs() const
+std::string StorageObjectStorage::Configuration::Path::getWithoutGlobs() const
 {
-    return getNamespace().find_first_of("*?{") != std::string::npos;
+    if (allow_partial_prefix)
+    {
+        return path.substr(0, path.find_first_of("*?{"));
+    }
+
+    auto first_glob_pos = path.find_first_of("*?{");
+    auto end_of_path_without_globs = path.substr(0, first_glob_pos).rfind('/');
+    if (end_of_path_without_globs == std::string::npos || end_of_path_without_globs == 0)
+        return "/";
+    return path.substr(0, end_of_path_without_globs);
 }
 
-std::string StorageObjectStorage::Configuration::getPathWithoutGlobs() const
+bool StorageObjectStorage::Configuration::isNamespaceWithGlobs() const
 {
-    return getPath().substr(0, getPath().find_first_of("*?{"));
+    return getNamespace().find_first_of("*?{") != std::string::npos;
 }
 
 bool StorageObjectStorage::Configuration::isPathInArchiveWithGlobs() const
@@ -803,7 +850,7 @@ bool StorageObjectStorage::Configuration::isPathInArchiveWithGlobs() const
 
 std::string StorageObjectStorage::Configuration::getPathInArchive() const
 {
-    throw Exception(ErrorCodes::LOGICAL_ERROR, "Path {} is not archive", getPath());
+    throw Exception(ErrorCodes::LOGICAL_ERROR, "Path {} is not archive", getRawPath().path);
 }
 
 void StorageObjectStorage::Configuration::assertInitialized() const

src/Storages/ObjectStorage/StorageObjectStorage.h

@@ -187,7 +187,24 @@ class StorageObjectStorage::Configuration
     Configuration() = default;
     virtual ~Configuration() = default;
 
-    using Path = std::string;
+    struct Path
+    {
+        Path() = default;
+        /// A partial prefix is a prefix that does not represent an actual object (directory or file), usually strings that do not end with a slash character.
+        /// Example: `table_root/year=20`. AWS S3 supports partial prefixes, but HDFS does not.
+        Path(const std::string & path_, bool allow_partial_prefix_ = true) : path(path_), allow_partial_prefix(allow_partial_prefix_) {} /// NOLINT(google-explicit-constructor)
+
+        std::string path;
+
+        bool withPartitionWildcard() const;
+        bool withGlobsIgnorePartitionWildcard() const;
+        bool withGlobs() const;
+        std::string getWithoutGlobs() const;
+
+    private:
+        bool allow_partial_prefix;
+    };
+
     using Paths = std::vector<Path>;
 
     /// Initialize configuration from either AST or NamedCollection.
@@ -206,10 +223,20 @@ class StorageObjectStorage::Configuration
     /// buckets in S3. If object storage doesn't have any namepaces return empty string.
     virtual std::string getNamespaceType() const { return "namespace"; }
 
-    virtual Path getFullPath() const { return ""; }
-    virtual Path getPath() const = 0;
-    virtual void setPath(const Path & path) = 0;
+    // Path provided by the user in the query
+    virtual Path getRawPath() const = 0;
+    // Path used for reading, it is usually a globbed path like `'table_root/**.parquet'
+    Path getPathForRead() const;
+    // Path used for writing, it should not be globbed and might contain a partition key
+    Path getPathForWrite(const std::string & partition_id = "") const;
+
+    virtual void setRawPath(const Path & path) = 0;
 
+    /*
+     * When using `s3_create_new_file_on_insert`, each new file path generated will be appended to the path list.
+     * This list is used to determine the next file name and the set of files that shall be read from remote storage.
+     * This is not ideal, there are much better ways to implement reads and writes. It should be eventually removed
+     */
     virtual const Paths & getPaths() const = 0;
     virtual void setPaths(const Paths & paths) = 0;
 
@@ -222,12 +249,7 @@ class StorageObjectStorage::Configuration
     virtual void addStructureAndFormatToArgsIfNeeded(
         ASTs & args, const String & structure_, const String & format_, ContextPtr context, bool with_structure) = 0;
 
-    bool withPartitionWildcard() const;
-    bool withGlobs() const { return isPathWithGlobs() || isNamespaceWithGlobs(); }
-    bool withGlobsIgnorePartitionWildcard() const;
-    bool isPathWithGlobs() const;
     bool isNamespaceWithGlobs() const;
-    virtual std::string getPathWithoutGlobs() const;
 
     virtual bool isArchive() const { return false; }
     bool isPathInArchiveWithGlobs() const;
@@ -259,7 +281,10 @@ class StorageObjectStorage::Configuration
         const Strings & requested_columns,
         const StorageSnapshotPtr & storage_snapshot,
         bool supports_subset_of_columns,
-        ContextPtr local_context);
+        ContextPtr local_context,
+        const PrepareReadingFromFormatHiveParams & hive_params);
+
+    void initPartitionStrategy(ASTPtr partition_by, const ColumnsDescription & columns, ContextPtr context);
 
     virtual std::optional<ColumnsDescription> tryGetTableStructureFromMetadata() const;
 
@@ -290,6 +315,11 @@ class StorageObjectStorage::Configuration
     String format = "auto";
     String compression_method = "auto";
     String structure = "auto";
+    PartitionStrategyFactory::StrategyType partition_strategy_type = PartitionStrategyFactory::StrategyType::WILDCARD;
+    /// Whether partition column values are contained in the actual data.
+    /// And alternative is with hive partitioning, when they are contained in file path.
+    bool partition_columns_in_data_file = true;
+    std::shared_ptr<IPartitionStrategy> partition_strategy;
 
 protected:
     virtual void fromNamedCollection(const NamedCollection & collection, ContextPtr context) = 0;

src/Storages/ObjectStorage/StorageObjectStorageSink.cpp

@@ -38,7 +38,7 @@ namespace
         PartitionedSink::validatePartitionKey(str, true);
     }
 
-    void validateNamespace(const String & str, StorageObjectStorageConfigurationPtr configuration)
+    void validateNamespace(const String & str, StorageObjectStorage::ConfigurationPtr configuration)
     {
         configuration->validateNamespace(str);
 
@@ -167,7 +167,7 @@ SinkPtr PartitionedStorageObjectStorageSink::createSinkForPartition(const String
         object_storage,
         configuration,
         format_settings,
-        std::make_shared<Block>(partition_strategy->getFormatHeader()),
+        partition_strategy->getFormatHeader(),
         context
     );
 }

src/Storages/PartitionedSink.cpp

@@ -31,7 +31,7 @@ PartitionedSink::PartitionedSink(
     : SinkToStorage(sample_block_)
     , partition_strategy(partition_strategy_)
     , context(context_)
-    , source_header(source_header_)
+    , sample_block(sample_block_)
 {
 }
 

src/Storages/PartitionedSink.h

@@ -19,14 +19,10 @@ class PartitionedSink : public SinkToStorage
 public:
     static constexpr auto PARTITION_ID_WILDCARD = "{_partition_id}";
 
-<<<<<<< HEAD
-    PartitionedSink(const ASTPtr & partition_by, ContextPtr context_, const Block & sample_block_);
-=======
     PartitionedSink(
         std::shared_ptr<IPartitionStrategy> partition_strategy_,
         ContextPtr context_,
-        SharedHeader source_header_);
->>>>>>> 790e1be4c1e (Add Hive-style S3 partitioned reads/writes)
+        const Block & sample_block_);
 
     ~PartitionedSink() override;
 
@@ -49,15 +45,7 @@ class PartitionedSink : public SinkToStorage
 
 private:
     ContextPtr context;
-<<<<<<< HEAD
     Block sample_block;
-
-    ExpressionActionsPtr partition_by_expr;
-    String partition_by_column_name;
-=======
-    SharedHeader source_header;
->>>>>>> 790e1be4c1e (Add Hive-style S3 partitioned reads/writes)
-
     absl::flat_hash_map<StringRef, SinkPtr> partition_id_to_sink;
     HashMapWithSavedHash<StringRef, size_t> partition_id_to_chunk_index;
     IColumn::Selector chunk_row_index_to_partition_index;