// SPDX-License-Identifier: Apache-2.0

// SPDX-FileCopyrightText: Copyright the Vortex contributors

#include "multi_file_reader.hpp"

#include "error.hpp"

#include "table_function.h"

#include "vortex_duckdb.h"

#include "vortex.h"

#include "duckdb/execution/operator/scan/physical_table_scan.hpp"

unique_ptr<FunctionData> VortexBindData::Copy() const {

auto result = make_uniq<VortexBindData>();

if (ffi_bind_data) {

const duckdb_vx_data copy = duckdb_table_function_bind_data_clone(ffi_bind_data->DataPtr());

result->ffi_bind_data = unique_ptr<CData>(reinterpret_cast<CData *>(copy));

}

return result;

}

bool VortexBindData::Equals(const FunctionData &other_base) const {

const VortexBindData &other = other_base.Cast<VortexBindData>();

return ffi_bind_data.get() == other.ffi_bind_data.get();

}

bool SkipMultiFileReaderColumn(size_t i,

optional_idx filename_idx,

const std::vector<HivePartitioningIndex> &hive_indices) {

if (filename_idx.IsValid() && i == filename_idx.GetIndex()) {

return true;

}

for (const auto &idx : hive_indices) {

if (idx.index == i) {

return true;

}

}

return false;

}

ReaderInitializeType

VortexMultiFileReader::InitializeReader(MultiFileReaderData &reader_data,

const MultiFileBindData &bind_data,

const vector<MultiFileColumnDefinition> &global_columns,

const vector<ColumnIndex> &global_column_ids,

optional_ptr<TableFilterSet> table_filters,

ClientContext &context,

MultiFileGlobalState &gstate) {

D_ASSERT(reader_data.reader != nullptr);

D_ASSERT(gstate.global_state != nullptr);

VortexBaseReader &reader = reader_data.reader->Cast<VortexBaseReader>();

const VortexBindData &bind = bind_data.bind_data->Cast<VortexBindData>();

const void *const ffi_bind = bind.ffi_bind_data->DataPtr();

reader.columns = global_columns;

const optional_idx filename_idx = bind_data.reader_bind.filename_idx;

const auto &hive_indices = bind_data.reader_bind.hive_partitioning_indexes;

// Aggregate scan columns are aggregate results so base column mapping is

// wrong.

if (!duckdb_reader_is_aggregate(ffi_bind)) {

// Projection expression pushdown changes types of columns

for (size_t i = 0; i < global_columns.size(); i++) {

if (SkipMultiFileReaderColumn(i, filename_idx, hive_indices)) {

continue;

}

const duckdb_logical_type type = duckdb_reader_bind_column_type(ffi_bind, i);

reader.columns[i].type = *reinterpret_cast<const LogicalType *>(type);

}

// this prunes files for virtual columns like "file_index"

const ReaderInitializeType base_skip = MultiFileReader::InitializeReader(reader_data,

bind_data,

reader.columns,

global_column_ids,

table_filters,

context,

gstate);

if (base_skip == ReaderInitializeType::SKIP_READING_FILE) {

return base_skip;

}

}

const VortexGlobalState &global = gstate.global_state->Cast<VortexGlobalState>();

duckdb_vx_error error = nullptr;

const void *const ffi_global = global.ffi_global_state->DataPtr();

void *const ffi_file = reader.ffi_file->DataPtr();

const bool skip = duckdb_reader_initialize(ffi_global, ffi_file, &error);

if (error) {

throw InvalidInputException(IntoErrString(error));

}

return skip ? ReaderInitializeType::SKIP_READING_FILE : ReaderInitializeType::INITIALIZED;

}

void VortexReaderInterface::BindReader(ClientContext &context,

vector<LogicalType> &types,

vector<string> &names,

MultiFileBindData &bind_data) {

BaseFileReaderOptions options;

VortexBindResult result = {types, names};

VortexBindData &bind = bind_data.bind_data->Cast<VortexBindData>();

const OpenFileInfo first_file = bind_data.file_list->GetFirstFile();

bind_data.initial_reader = CreateReader(context, first_file, options, bind_data.file_options);

VortexBaseReader &initial_reader = bind_data.initial_reader->Cast<VortexBaseReader>();

duckdb_vx_error error = nullptr;

const void *const ffi_file = initial_reader.ffi_file->DataPtr();

duckdb_bind_result ffi_result = reinterpret_cast<duckdb_bind_result>(&result);

duckdb_vx_data ffi_bind_data = duckdb_reader_bind(ffi_file, ffi_result, &error);

if (error) {

throw BinderException(IntoErrString(error));

}

bind.ffi_bind_data = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_bind_data));

initial_reader.ffi_bind = bind.ffi_bind_data->DataPtr();

// Fills bind_data.file_options which are used for hive partitioning and

// "filename" column.

bind_data.multi_file_reader->BindOptions(bind_data.file_options,

*bind_data.file_list,

types,

names,

bind_data.reader_bind);

}

unique_ptr<GlobalTableFunctionState>

VortexReaderInterface::InitializeGlobalState(ClientContext &context,

MultiFileBindData &bind_data,

MultiFileGlobalState &input) {

const VortexBindData &bind = bind_data.bind_data->Cast<VortexBindData>();

const optional_idx filename_idx = bind_data.reader_bind.filename_idx;

const auto &hive_indices = bind_data.reader_bind.hive_partitioning_indexes;

vector<idx_t> column_ids(input.column_indexes.size());

for (size_t i = 0; i < input.column_indexes.size(); ++i) {

idx_t storage_index = input.column_indexes[i].GetPrimaryIndex();

if (SkipMultiFileReaderColumn(storage_index, filename_idx, hive_indices)) {

// replace with a virtual column which Vortex filters in Rust

storage_index = COLUMN_IDENTIFIER_EMPTY;

}

column_ids[i] = storage_index;

}

// MultiFileGlobalState projection_ids are filled only when this call

// returns. Take these from a physical operator.

const idx_t *projection_ids = nullptr;

size_t projection_ids_count = 0;

if (input.op && input.op->type == PhysicalOperatorType::TABLE_SCAN) {

const PhysicalTableScan &scan = input.op->Cast<PhysicalTableScan>();

if (!scan.projection_ids.empty() && scan.projection_ids.size() != column_ids.size()) {

projection_ids = scan.projection_ids.data();

projection_ids_count = scan.projection_ids.size();

}

}

void *const ffi_bind = bind.ffi_bind_data->DataPtr();

duckdb_vx_tfunc_init_input ffi_input = {

.bind_data = ffi_bind,

.column_ids = column_ids.data(),

.column_ids_count = column_ids.size(),

.projection_ids = projection_ids,

.projection_ids_count = projection_ids_count,

.filters = reinterpret_cast<duckdb_vx_table_filter_set>(input.filters.get()),

.client_context = reinterpret_cast<duckdb_client_context>(&context),

};

duckdb_vx_error error_out = nullptr;

duckdb_vx_data ffi_global_state = duckdb_table_function_init_global(&ffi_input, &error_out);

if (error_out) {

throw BinderException(IntoErrString(error_out));

}

auto result = make_uniq<VortexGlobalState>();

result->ffi_bind_data = ffi_bind;

result->ffi_global_state = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_global_state));

return result;

}

unique_ptr<LocalTableFunctionState>

VortexReaderInterface::InitializeLocalState(ExecutionContext &, GlobalTableFunctionState &global_state) {

auto &global = global_state.Cast<VortexGlobalState>();

const void *const ffi_global = global.ffi_global_state->DataPtr();

duckdb_vx_data ffi_local_state = duckdb_table_function_init_local(global.ffi_bind_data, ffi_global);

auto result = make_uniq<VortexLocalState>();

result->ffi_local_state = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_local_state));

return result;

}

static shared_ptr<BaseFileReader> OpenReader(const OpenFileInfo &file) {

duckdb_vx_error error = nullptr;

const char *const ffi_file_path = file.path.c_str();

const size_t ffi_file_size = file.path.size();

duckdb_vx_data ffi_file = duckdb_reader_open(ffi_file_path, ffi_file_size, &error);

if (error) {

throw IOException(IntoErrString(error));

}

auto cdata = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_file));

return make_shared_ptr<VortexBaseReader>(file, std::move(cdata));

}

shared_ptr<BaseFileReader> VortexReaderInterface::CreateReader(ClientContext &,

GlobalTableFunctionState &,

const OpenFileInfo &file,

idx_t,

const MultiFileBindData &) {

return OpenReader(file);

}

shared_ptr<BaseFileReader> VortexReaderInterface::CreateReader(ClientContext &,

const OpenFileInfo &file,

BaseFileReaderOptions &,

const MultiFileOptions &) {

return OpenReader(file);

}

unique_ptr<NodeStatistics> VortexReaderInterface::GetCardinality(const MultiFileBindData &data,

idx_t file_count) {

const VortexBindData &bind_data = data.bind_data->Cast<VortexBindData>();

const void *const ffi_bind = bind_data.ffi_bind_data->DataPtr();

duckdb_vx_node_statistics stats = {};

duckdb_table_function_cardinality(ffi_bind, file_count, &stats);

auto out = make_uniq<NodeStatistics>();

out->has_estimated_cardinality = stats.has_estimated_cardinality;

out->estimated_cardinality = stats.estimated_cardinality;

out->has_max_cardinality = stats.has_max_cardinality;

out->max_cardinality = stats.max_cardinality;

return out;

}

bool VortexBaseReader::TryInitializeScan(ClientContext &,

GlobalTableFunctionState &,

LocalTableFunctionState &local_state) {

VortexLocalState &local = local_state.Cast<VortexLocalState>();

void *const ffi_local = local.ffi_local_state->DataPtr();

void *const ffi_file_ptr = ffi_file->DataPtr();

return duckdb_reader_try_initialize_scan(ffi_local, ffi_file_ptr);

}

AsyncResult VortexBaseReader::Scan(ClientContext &,

GlobalTableFunctionState &global_state,

LocalTableFunctionState &local_state,

DataChunk &chunk) {

VortexGlobalState &global = global_state.Cast<VortexGlobalState>();

VortexLocalState &local = local_state.Cast<VortexLocalState>();

duckdb_vx_error error = nullptr;

duckdb_data_chunk ffi_chunk = reinterpret_cast<duckdb_data_chunk>(&chunk);

const void *const ffi_global = global.ffi_global_state->DataPtr();

void *const ffi_local = local.ffi_local_state->DataPtr();

const void *const ffi_file_ptr = ffi_file->DataPtr();

const bool has_more_data = duckdb_reader_scan(ffi_file_ptr, ffi_global, ffi_local, ffi_chunk, &error);

if (error) {

throw InvalidInputException(IntoErrString(error));

}

return has_more_data ? SourceResultType::HAVE_MORE_OUTPUT : SourceResultType::FINISHED;

}

void VortexReaderInterface::GetVirtualColumns(ClientContext &,

MultiFileBindData &,

virtual_column_map_t &result) {

// "filename", "file_index" and "empty" come from MultiFileReader

result.insert(

{MultiFileReader::COLUMN_IDENTIFIER_FILE_ROW_NUMBER, {"file_row_number", LogicalType::UBIGINT}});

}

bool VortexReaderInterface::FinalizeScan(ClientContext &,

GlobalTableFunctionState &global_state,

DataChunk &output) {

const VortexGlobalState &global = global_state.Cast<VortexGlobalState>();

duckdb_vx_error error = nullptr;

duckdb_data_chunk ffi_chunk = reinterpret_cast<duckdb_data_chunk>(&output);

const void *const ffi_global = global.ffi_global_state->DataPtr();

const bool filled = duckdb_reader_finalize_scan(ffi_global, ffi_chunk, &error);

if (error) {

throw InvalidInputException(IntoErrString(error));

}

return filled;

}

void VortexReaderInterface::FinishReading(ClientContext &,

GlobalTableFunctionState &global_state,

LocalTableFunctionState &local_state) {

const VortexGlobalState &global = global_state.Cast<VortexGlobalState>();

const VortexLocalState &local = local_state.Cast<VortexLocalState>();

const void *const ffi_global = global.ffi_global_state->DataPtr();

void *const ffi_local = local.ffi_local_state->DataPtr();

duckdb_reader_finish_reading(ffi_global, ffi_local);

}

static Value &UnwrapValue(duckdb_value value) {

return *(reinterpret_cast<Value *>(value));

}

static unique_ptr<BaseStatistics> numeric_stats(duckdb_column_statistics &stats, LogicalType type) {

BaseStatistics out = NumericStats::CreateUnknown(type);

if (stats.min) {

NumericStats::SetMin(out, UnwrapValue(stats.min));

duckdb_destroy_value(&stats.min);

}

if (stats.max) {

NumericStats::SetMax(out, UnwrapValue(stats.max));

duckdb_destroy_value(&stats.max);

}

if (!stats.has_null) {

out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES);

}

return out.ToUnique();

}

static unique_ptr<BaseStatistics> string_stats(duckdb_column_statistics &stats, LogicalType type) {

BaseStatistics out = StringStats::CreateUnknown(type);

if (stats.min) {

StringStats::SetMin(out, StringValue::Get(UnwrapValue(stats.min)));

duckdb_destroy_value(&stats.min);

}

if (stats.max) {

StringStats::SetMax(out, StringValue::Get(UnwrapValue(stats.max)));

duckdb_destroy_value(&stats.max);

}

if (stats.max_string_length >> 63) {

StringStats::SetMaxStringLength(out, uint32_t(stats.max_string_length));

}

if (!stats.has_null) {

out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES);

}

return out.ToUnique();

}

static unique_ptr<BaseStatistics> base_stats(duckdb_column_statistics &stats, LogicalType type) {

BaseStatistics out = BaseStatistics::CreateUnknown(type);

if (stats.min) {

duckdb_destroy_value(&stats.min);

}

if (stats.max) {

duckdb_destroy_value(&stats.max);

}

if (!stats.has_null) {

out.Set(StatsInfo::CANNOT_HAVE_NULL_VALUES);

}

return out.ToUnique();

}

unique_ptr<BaseStatistics> to_duckdb_statistics(duckdb_column_statistics &statistics) {

using enum LogicalTypeId;

const unique_ptr<LogicalType> type(reinterpret_cast<LogicalType *>(statistics.type));

switch (type->id()) {

case BOOLEAN:

case TINYINT:

case SMALLINT:

case INTEGER:

case BIGINT:

case FLOAT:

case DOUBLE:

case UTINYINT:

case USMALLINT:

case UINTEGER:

case UBIGINT:

case UHUGEINT:

case HUGEINT:

case DATE:

case TIME:

case TIME_TZ:

case TIMESTAMP_SEC:

case TIMESTAMP_MS:

case TIMESTAMP:

case TIMESTAMP_NS:

case TIMESTAMP_TZ: {

return numeric_stats(statistics, *type);

}

case VARCHAR:

case BLOB: {

return string_stats(statistics, *type);

}

case STRUCT: {

// TODO(myrrc)

// Duckdb's has_null has a different semantics for structs.

// If we propagate our has_null, this breaks Duckdb optimizer.

// You can reproduce it in struct.slt test in vortex-sqllogictests:

if (statistics.min) {

duckdb_destroy_value(&statistics.min);

}

if (statistics.max) {

duckdb_destroy_value(&statistics.max);

}

return {};

}

default:

return base_stats(statistics, *type);

}

}

unique_ptr<BaseStatistics> VortexBaseReader::GetStatistics(ClientContext &, const string &name) {

D_ASSERT(ffi_bind);

duckdb_column_statistics statistics = {};

if (!duckdb_reader_get_statistics(ffi_file->DataPtr(),

ffi_bind,

name.c_str(),

name.size(),

&statistics)) {

return {};

}

return to_duckdb_statistics(statistics);

}

double VortexBaseReader::GetProgressInFile(ClientContext &) {

return duckdb_reader_get_progress_in_file(ffi_file->DataPtr());

}