table_scan.h

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#pragma once
 
#include <functional>
#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
 
#include "iceberg/arrow_c_data.h"
#include "iceberg/iceberg_export.h"
#include "iceberg/result.h"
#include "iceberg/table_metadata.h"
#include "iceberg/type_fwd.h"
#include "iceberg/util/error_collector.h"
 
namespace iceberg {
 
class ICEBERG_EXPORT ScanTask {
 public:
  enum class Kind : uint8_t {
    kFileScanTask,
    kChangelogScanTask,
  };
 
  virtual Kind kind() const = 0;
 
  virtual int64_t size_bytes() const = 0;
 
  virtual int32_t files_count() const = 0;
 
  virtual int64_t estimated_row_count() const = 0;
 
  virtual ~ScanTask();
};
 
class ICEBERG_EXPORT FileScanTask : public ScanTask {
 public:
  explicit FileScanTask(std::shared_ptr<DataFile> data_file,
                        std::vector<std::shared_ptr<DataFile>> delete_files = {},
                        std::shared_ptr<Expression> filter = nullptr);
 
  const std::shared_ptr<DataFile>& data_file() const { return data_file_; }
 
  const std::vector<std::shared_ptr<DataFile>>& delete_files() const {
    return delete_files_;
  }
 
  const std::shared_ptr<Expression>& residual_filter() const { return residual_filter_; }
 
  Kind kind() const override { return Kind::kFileScanTask; }
  int64_t size_bytes() const override;
  int32_t files_count() const override;
  int64_t estimated_row_count() const override;
 
  Result<ArrowArrayStream> ToArrow(const std::shared_ptr<FileIO>& io,
                                   std::shared_ptr<Schema> projected_schema) const;
 
 private:
  std::shared_ptr<DataFile> data_file_;
  std::vector<std::shared_ptr<DataFile>> delete_files_;
  std::shared_ptr<Expression> residual_filter_;
};
 
enum class ChangelogOperation : uint8_t {
  kInsert,
  kDelete,
  kUpdateBefore,
  kUpdateAfter,
};
 
class ICEBERG_EXPORT ChangelogScanTask : public ScanTask {
 public:
  ChangelogScanTask(int32_t change_ordinal, int64_t commit_snapshot_id,
                    std::shared_ptr<DataFile> data_file,
                    std::vector<std::shared_ptr<DataFile>> delete_files = {},
                    std::shared_ptr<Expression> residual_filter = nullptr)
      : change_ordinal_(change_ordinal),
        commit_snapshot_id_(commit_snapshot_id),
        data_file_(std::move(data_file)),
        delete_files_(std::move(delete_files)),
        residual_filter_(std::move(residual_filter)) {}
 
  Kind kind() const override { return Kind::kChangelogScanTask; }
 
  int64_t size_bytes() const override;
  int32_t files_count() const override;
  int64_t estimated_row_count() const override;
 
  virtual ChangelogOperation operation() const = 0;
 
  int32_t change_ordinal() const { return change_ordinal_; }
 
  int64_t commit_snapshot_id() const { return commit_snapshot_id_; }
 
  const std::shared_ptr<Expression>& residual_filter() const { return residual_filter_; }
 
 protected:
  int32_t change_ordinal_;
  int64_t commit_snapshot_id_;
  std::shared_ptr<DataFile> data_file_;
  std::vector<std::shared_ptr<DataFile>> delete_files_;
  std::shared_ptr<Expression> residual_filter_;
};
 
class ICEBERG_EXPORT AddedRowsScanTask : public ChangelogScanTask {
 public:
  using ChangelogScanTask::ChangelogScanTask;
 
  ChangelogOperation operation() const override { return ChangelogOperation::kInsert; }
 
  const std::shared_ptr<DataFile>& data_file() const { return data_file_; }
 
  const std::vector<std::shared_ptr<DataFile>>& delete_files() const {
    return delete_files_;
  }
};
 
class ICEBERG_EXPORT DeletedDataFileScanTask : public ChangelogScanTask {
 public:
  using ChangelogScanTask::ChangelogScanTask;
 
  ChangelogOperation operation() const override { return ChangelogOperation::kDelete; }
 
  const std::shared_ptr<DataFile>& data_file() const { return data_file_; }
 
  const std::vector<std::shared_ptr<DataFile>>& existing_deletes() const {
    return delete_files_;
  }
};
 
namespace internal {
 
// Internal table scan context used by different scan implementations.
struct TableScanContext {
  std::optional<int64_t> snapshot_id;
  std::shared_ptr<Expression> filter;
  bool ignore_residuals{false};
  bool case_sensitive{true};
  bool return_column_stats{false};
  std::unordered_set<int32_t> columns_to_keep_stats;
  std::vector<std::string> selected_columns;
  std::shared_ptr<Schema> projected_schema;
  std::unordered_map<std::string, std::string> options;
  bool from_snapshot_id_inclusive{false};
  std::optional<int64_t> from_snapshot_id;
  std::optional<int64_t> to_snapshot_id;
  std::string branch{};
  std::optional<int64_t> min_rows_requested;
 
  // Validate the context parameters to see if they have conflicts.
  [[nodiscard]] Status Validate() const;
};
 
}  // namespace internal
 
// Concept to check if a type is an incremental scan
template <typename T>
concept IsIncrementalScan = std::is_base_of_v<IncrementalScan<FileScanTask>, T> ||
                            std::is_base_of_v<IncrementalScan<ChangelogScanTask>, T>;
 
template <typename ScanType = DataTableScan>
class ICEBERG_TEMPLATE_CLASS_EXPORT TableScanBuilder : public ErrorCollector {
 public:
  static Result<std::unique_ptr<TableScanBuilder<ScanType>>> Make(
      std::shared_ptr<TableMetadata> metadata, std::shared_ptr<FileIO> io);
 
  TableScanBuilder& Option(std::string key, std::string value);
 
  TableScanBuilder& Project(std::shared_ptr<Schema> schema);
 
  TableScanBuilder& CaseSensitive(bool case_sensitive);
 
  TableScanBuilder& IncludeColumnStats();
 
  TableScanBuilder& IncludeColumnStats(const std::vector<std::string>& requested_columns);
 
  TableScanBuilder& Select(const std::vector<std::string>& column_names);
 
  TableScanBuilder& Filter(std::shared_ptr<Expression> filter);
 
  TableScanBuilder& IgnoreResiduals();
 
  TableScanBuilder& MinRowsRequested(int64_t num_rows);
 
  TableScanBuilder& UseSnapshot(int64_t snapshot_id);
 
  TableScanBuilder& UseRef(const std::string& ref);
 
  TableScanBuilder& AsOfTime(int64_t timestamp_millis);
 
  TableScanBuilder& FromSnapshot(int64_t from_snapshot_id, bool inclusive = false)
    requires IsIncrementalScan<ScanType>;
 
  TableScanBuilder& FromSnapshot(const std::string& ref, bool inclusive = false)
    requires IsIncrementalScan<ScanType>;
 
  TableScanBuilder& ToSnapshot(int64_t to_snapshot_id)
    requires IsIncrementalScan<ScanType>;
 
  TableScanBuilder& ToSnapshot(const std::string& ref)
    requires IsIncrementalScan<ScanType>;
 
  TableScanBuilder& UseBranch(const std::string& branch)
    requires IsIncrementalScan<ScanType>;
 
  Result<std::unique_ptr<ScanType>> Build();
 
 protected:
  TableScanBuilder(std::shared_ptr<TableMetadata> metadata, std::shared_ptr<FileIO> io);
 
  // Return the schema bound to the specified snapshot.
  Result<std::reference_wrapper<const std::shared_ptr<Schema>>> ResolveSnapshotSchema();
 
  std::shared_ptr<TableMetadata> metadata_;
  std::shared_ptr<FileIO> io_;
  internal::TableScanContext context_;
  std::shared_ptr<Schema> snapshot_schema_;
};
 
class ICEBERG_EXPORT TableScan {
 public:
  virtual ~TableScan();
 
  const std::shared_ptr<TableMetadata>& metadata() const;
 
  Result<std::shared_ptr<Snapshot>> snapshot() const;
 
  Result<std::shared_ptr<Schema>> schema() const;
 
  const internal::TableScanContext& context() const;
 
  const std::shared_ptr<FileIO>& io() const;
 
  const std::shared_ptr<Expression>& filter() const;
 
  bool is_case_sensitive() const;
 
 protected:
  TableScan(std::shared_ptr<TableMetadata> metadata, std::shared_ptr<Schema> schema,
            std::shared_ptr<FileIO> io, internal::TableScanContext context);
 
  Result<std::reference_wrapper<const std::shared_ptr<Schema>>> ResolveProjectedSchema()
      const;
 
  virtual const std::vector<std::string>& ScanColumns() const;
 
  const std::shared_ptr<TableMetadata> metadata_;
  const std::shared_ptr<Schema> schema_;
  const std::shared_ptr<FileIO> io_;
  const internal::TableScanContext context_;
  mutable std::shared_ptr<Schema> projected_schema_;
};
 
class ICEBERG_EXPORT DataTableScan : public TableScan {
 public:
  ~DataTableScan() override = default;
 
  static Result<std::unique_ptr<DataTableScan>> Make(
      std::shared_ptr<TableMetadata> metadata, std::shared_ptr<Schema> schema,
      std::shared_ptr<FileIO> io, internal::TableScanContext context);
 
  Result<std::vector<std::shared_ptr<FileScanTask>>> PlanFiles() const;
 
 protected:
  using TableScan::TableScan;
};
 
template <typename ScanTaskType>
class IncrementalScan : public TableScan {
 public:
  ~IncrementalScan() override = default;
 
  virtual Result<std::vector<std::shared_ptr<ScanTaskType>>> PlanFiles() const = 0;
 
 protected:
  virtual Result<std::vector<std::shared_ptr<ScanTaskType>>> PlanFiles(
      std::optional<int64_t> from_snapshot_id_exclusive,
      int64_t to_snapshot_id_inclusive) const = 0;
 
  using TableScan::TableScan;
 
  // Allow the free function ResolvePlanFiles to access protected members.
  template <typename T>
  friend Result<std::vector<std::shared_ptr<T>>> ResolvePlanFiles(
      const IncrementalScan<T>& scan);
};
 
class ICEBERG_EXPORT IncrementalAppendScan : public IncrementalScan<FileScanTask> {
 public:
  static Result<std::unique_ptr<IncrementalAppendScan>> Make(
      std::shared_ptr<TableMetadata> metadata, std::shared_ptr<Schema> schema,
      std::shared_ptr<FileIO> io, internal::TableScanContext context);
 
  ~IncrementalAppendScan() override = default;
 
  Result<std::vector<std::shared_ptr<FileScanTask>>> PlanFiles() const override;
 
 protected:
  Result<std::vector<std::shared_ptr<FileScanTask>>> PlanFiles(
      std::optional<int64_t> from_snapshot_id_exclusive,
      int64_t to_snapshot_id_inclusive) const override;
 
  using IncrementalScan::IncrementalScan;
};
 
class ICEBERG_EXPORT IncrementalChangelogScan
    : public IncrementalScan<ChangelogScanTask> {
 public:
  static Result<std::unique_ptr<IncrementalChangelogScan>> Make(
      std::shared_ptr<TableMetadata> metadata, std::shared_ptr<Schema> schema,
      std::shared_ptr<FileIO> io, internal::TableScanContext context);
 
  ~IncrementalChangelogScan() override = default;
 
  Result<std::vector<std::shared_ptr<ChangelogScanTask>>> PlanFiles() const override;
 
 protected:
  Result<std::vector<std::shared_ptr<ChangelogScanTask>>> PlanFiles(
      std::optional<int64_t> from_snapshot_id_exclusive,
      int64_t to_snapshot_id_inclusive) const override;
 
  using IncrementalScan::IncrementalScan;
};
 
extern template class ICEBERG_EXTERN_TEMPLATE_CLASS_EXPORT
    TableScanBuilder<DataTableScan>;
extern template class ICEBERG_EXTERN_TEMPLATE_CLASS_EXPORT
    TableScanBuilder<IncrementalAppendScan>;
extern template class ICEBERG_EXTERN_TEMPLATE_CLASS_EXPORT
    TableScanBuilder<IncrementalChangelogScan>;
 
}  // namespace iceberg