diff --git a/vortex-duckdb/cpp/include/multi_file_reader.hpp b/vortex-duckdb/cpp/include/multi_file_reader.hpp index 442e39bd4b4..c0dce58c5f6 100644 --- a/vortex-duckdb/cpp/include/multi_file_reader.hpp +++ b/vortex-duckdb/cpp/include/multi_file_reader.hpp @@ -154,7 +154,6 @@ struct VortexBaseReader final : BaseFileReader { } unique_ptr ffi_file; - vector virtual_ids; /* * Populated only for first file reader in scan when BindReader() is * called on it. Used in GetStatistics() which is called only for first @@ -162,8 +161,7 @@ struct VortexBaseReader final : BaseFileReader { */ const void *ffi_bind = nullptr; - inline void AddVirtualColumn(column_t id) override { - virtual_ids.push_back(id); + inline void AddVirtualColumn(column_t) override { } /* diff --git a/vortex-duckdb/cpp/multi_file_reader.cpp b/vortex-duckdb/cpp/multi_file_reader.cpp index ae8ef61d980..9184bb01239 100644 --- a/vortex-duckdb/cpp/multi_file_reader.cpp +++ b/vortex-duckdb/cpp/multi_file_reader.cpp @@ -21,6 +21,20 @@ bool VortexBindData::Equals(const FunctionData &other_base) const { return ffi_bind_data.get() == other.ffi_bind_data.get(); } +bool SkipMultiFileReaderColumn(size_t i, + optional_idx filename_idx, + const std::vector &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, @@ -39,11 +53,18 @@ VortexMultiFileReader::InitializeReader(MultiFileReaderData &reader_data, 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(type); } @@ -79,12 +100,12 @@ void VortexReaderInterface::BindReader(ClientContext &context, vector &names, MultiFileBindData &bind_data) { BaseFileReaderOptions options; - MultiFileOptions file_options; + VortexBindResult result = {types, names}; VortexBindData &bind = bind_data.bind_data->Cast(); const OpenFileInfo first_file = bind_data.file_list->GetFirstFile(); - bind_data.initial_reader = CreateReader(context, first_file, options, file_options); + bind_data.initial_reader = CreateReader(context, first_file, options, bind_data.file_options); VortexBaseReader &initial_reader = bind_data.initial_reader->Cast(); duckdb_vx_error error = nullptr; @@ -98,6 +119,14 @@ void VortexReaderInterface::BindReader(ClientContext &context, bind.ffi_bind_data = unique_ptr(reinterpret_cast(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 @@ -106,9 +135,19 @@ VortexReaderInterface::InitializeGlobalState(ClientContext &context, MultiFileGlobalState &input) { const VortexBindData &bind = bind_data.bind_data->Cast(); + const optional_idx filename_idx = bind_data.reader_bind.filename_idx; + const auto &hive_indices = bind_data.reader_bind.hive_partitioning_indexes; + vector column_ids(input.column_indexes.size()); for (size_t i = 0; i < input.column_indexes.size(); ++i) { - column_ids[i] = input.column_indexes[i].GetPrimaryIndex(); + 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; } void *const ffi_bind = bind.ffi_bind_data->DataPtr(); diff --git a/vortex-duckdb/cpp/table_function.cpp b/vortex-duckdb/cpp/table_function.cpp index 3a7bb302d2e..7b6e118d7c6 100644 --- a/vortex-duckdb/cpp/table_function.cpp +++ b/vortex-duckdb/cpp/table_function.cpp @@ -133,8 +133,9 @@ unique_ptr get_multi_file_reader(const TableFunction &) { duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter, const std::string &name) { MultiFileFunction fn(name); fn.arguments[0] = parameter; - // We neither support UNION BY NAME nor hive partitioning as for now - fn.named_parameters = {}; + fn.named_parameters = {{"filename", LogicalType::ANY}, + {"allow_empty", LogicalType::BOOLEAN}, + {"hive_partitioning", LogicalType::BOOLEAN}}; fn.filter_pushdown = true; fn.filter_prune = true; diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index e7b4a31c6a5..66bb7aa6df9 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -588,7 +588,18 @@ fn try_from_expression_inner( }; col.clone() } - BoundColumnRef(col_ref) => col(col_ref.name.as_ref()), + BoundColumnRef(col_ref) => { + let name = col_ref.name.as_ref(); + // Duckdb generates some columns (e.g. hive partitions) after we + // load file data, so filters on these columns can't be evaluated + if ctx + .fields + .is_some_and(|fields| !fields.iter().any(|field| field.name == name)) + { + return Ok(None); + } + col(name) + } BoundConstant(const_) => lit(Scalar::try_from(const_.value)?), BoundComparison(compare) => { let operator: Operator = compare.op.try_into()?; diff --git a/vortex-sqllogictest/slt/duckdb/filename.slt b/vortex-sqllogictest/slt/duckdb/filename.slt new file mode 100644 index 00000000000..c36d5cd38c0 --- /dev/null +++ b/vortex-sqllogictest/slt/duckdb/filename.slt @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +include ../setup.slt.no + +statement ok +CREATE TABLE partitioned AS SELECT range % 3 AS i, range AS j FROM range(4); + +statement ok +COPY partitioned TO '${WORK_DIR}/filename' (FORMAT vortex, PARTITION_BY (i)); + +# hive auto detection is turned on, so we get output j, [filename], [i = hive index] +query ITI +SELECT * FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +ORDER BY j; +---- +0 ${WORK_DIR}/filename/i=0/data_0.vortex 0 +1 ${WORK_DIR}/filename/i=1/data_0.vortex 1 +2 ${WORK_DIR}/filename/i=2/data_0.vortex 2 +3 ${WORK_DIR}/filename/i=0/data_0.vortex 0 + +query IT +SELECT * FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true, hive_partitioning=false) +ORDER BY j; +---- +0 ${WORK_DIR}/filename/i=0/data_0.vortex +1 ${WORK_DIR}/filename/i=1/data_0.vortex +2 ${WORK_DIR}/filename/i=2/data_0.vortex +3 ${WORK_DIR}/filename/i=0/data_0.vortex + +query II +SELECT * FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=false) +ORDER BY j; +---- +0 0 +1 1 +2 2 +3 0 + +query TI +SELECT filename, j FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +ORDER BY j; +---- +${WORK_DIR}/filename/i=0/data_0.vortex 0 +${WORK_DIR}/filename/i=1/data_0.vortex 1 +${WORK_DIR}/filename/i=2/data_0.vortex 2 +${WORK_DIR}/filename/i=0/data_0.vortex 3 + +query ITI +SELECT j, fname, i FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename='fname') +ORDER BY j; +---- +0 ${WORK_DIR}/filename/i=0/data_0.vortex 0 +1 ${WORK_DIR}/filename/i=1/data_0.vortex 1 +2 ${WORK_DIR}/filename/i=2/data_0.vortex 2 +3 ${WORK_DIR}/filename/i=0/data_0.vortex 0 + +query ITI +SELECT * FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +WHERE filename NOT LIKE '%i=0%' +ORDER BY j; +---- +1 ${WORK_DIR}/filename/i=1/data_0.vortex 1 +2 ${WORK_DIR}/filename/i=2/data_0.vortex 2 + +query ITI +SELECT * FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +WHERE j != 0 +ORDER BY j; +---- +1 ${WORK_DIR}/filename/i=1/data_0.vortex 1 +2 ${WORK_DIR}/filename/i=2/data_0.vortex 2 +3 ${WORK_DIR}/filename/i=0/data_0.vortex 0 + +query I +SELECT j FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +WHERE filename = '${WORK_DIR}/filename/i=1/data_0.vortex'; +---- +1 + +query I +SELECT j FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +WHERE filename = 'no such file'; +---- + +query I +SELECT j FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +WHERE i = 0 AND j > 1; +---- +3 + +query I +SELECT j FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +WHERE filename LIKE '%i=0%' AND j >= 0 +ORDER BY j; +---- +0 +3 + +query I +SELECT count(*) FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true); +---- +4 + +query TI +SELECT filename, count(*) FROM read_vortex('${WORK_DIR}/filename/**/*.vortex', filename=true) +GROUP BY filename +ORDER BY filename; +---- +${WORK_DIR}/filename/i=0/data_0.vortex 2 +${WORK_DIR}/filename/i=1/data_0.vortex 1 +${WORK_DIR}/filename/i=2/data_0.vortex 1 + +query TI +SELECT filename, j FROM '${WORK_DIR}/filename/**/*.vortex' +ORDER BY j; +---- +${WORK_DIR}/filename/i=0/data_0.vortex 0 +${WORK_DIR}/filename/i=1/data_0.vortex 1 +${WORK_DIR}/filename/i=2/data_0.vortex 2 +${WORK_DIR}/filename/i=0/data_0.vortex 3 + +query I +SELECT j FROM '${WORK_DIR}/filename/**/*.vortex' +WHERE filename LIKE '%i=1%'; +---- +1 + +statement ok +COPY (SELECT i AS filename, j FROM partitioned) +TO '${WORK_DIR}/filename-collision.vortex' (FORMAT vortex); + +statement error +SELECT * FROM read_vortex('${WORK_DIR}/filename-collision.vortex', filename=true); + +query IIT +SELECT * FROM read_vortex('${WORK_DIR}/filename-collision.vortex', filename='other_filename') +ORDER BY j; +---- +0 0 ${WORK_DIR}/filename-collision.vortex +1 1 ${WORK_DIR}/filename-collision.vortex +2 2 ${WORK_DIR}/filename-collision.vortex +0 3 ${WORK_DIR}/filename-collision.vortex