Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions vortex-duckdb/cpp/include/multi_file_reader.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -154,16 +154,14 @@ struct VortexBaseReader final : BaseFileReader {
}

unique_ptr<CData> ffi_file;
vector<column_t> 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
* file after BindReader().
*/
const void *ffi_bind = nullptr;

inline void AddVirtualColumn(column_t id) override {
virtual_ids.push_back(id);
inline void AddVirtualColumn(column_t) override {
}

/*
Expand Down
45 changes: 42 additions & 3 deletions vortex-duckdb/cpp/multi_file_reader.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<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,
Expand All @@ -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<const LogicalType *>(type);
}
Expand Down Expand Up @@ -79,12 +100,12 @@ void VortexReaderInterface::BindReader(ClientContext &context,
vector<string> &names,
MultiFileBindData &bind_data) {
BaseFileReaderOptions options;
MultiFileOptions file_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, file_options);
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;
Expand All @@ -98,6 +119,14 @@ void VortexReaderInterface::BindReader(ClientContext &context,

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>
Expand All @@ -106,9 +135,19 @@ VortexReaderInterface::InitializeGlobalState(ClientContext &context,
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) {
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();
Expand Down
5 changes: 3 additions & 2 deletions vortex-duckdb/cpp/table_function.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -133,8 +133,9 @@ unique_ptr<MultiFileReader> get_multi_file_reader(const TableFunction &) {
duckdb_state register_table_function(DatabaseInstance &db, LogicalType parameter, const std::string &name) {
MultiFileFunction<VortexReaderInterface> 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;
Expand Down
13 changes: 12 additions & 1 deletion vortex-duckdb/src/convert/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()?;
Expand Down
143 changes: 143 additions & 0 deletions vortex-sqllogictest/slt/duckdb/filename.slt
Original file line number Diff line number Diff line change
@@ -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
Loading