duckdb 1.3.2 → 1.3.3-dev4.0

package/src/duckdb/src/planner/binder/statement/bind_copy.cpp

@@ -19,8 +19,8 @@
 #include "duckdb/planner/operator/logical_get.hpp"
 #include "duckdb/planner/operator/logical_insert.hpp"
 #include "duckdb/planner/operator/logical_projection.hpp"
-
-#include <algorithm>
+#include "duckdb/planner/expression_binder/table_function_binder.hpp"
+#include "duckdb/common/algorithm.hpp"
 
 #include "duckdb/main/extension_entries.hpp"
 
@@ -41,26 +41,54 @@ void IsFormatExtensionKnown(const string &format) {
 	}
 }
 
-BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type) {
-	// Let's first bind our format
-	auto on_entry_do =
-	    stmt.info->is_format_auto_detected ? OnEntryNotFound::RETURN_NULL : OnEntryNotFound::THROW_EXCEPTION;
-	CatalogEntryRetriever entry_retriever {context};
-	auto &catalog = Catalog::GetSystemCatalog(context);
-	auto entry = catalog.GetEntry(entry_retriever, DEFAULT_SCHEMA,
-	                              {CatalogType::COPY_FUNCTION_ENTRY, stmt.info->format}, on_entry_do);
+case_insensitive_map_t<CopyOption> Binder::GetFullCopyOptionsList(const CopyFunction &function, CopyOptionMode mode) {
+	case_insensitive_map_t<CopyOption> copy_options;
+	CopyOptionsInput input(copy_options);
+	function.copy_options(context, input);
+
+	// first erase all options that don't match this type
+	if (mode != CopyOptionMode::READ_WRITE) {
+		vector<string> erased_options;
+		for (auto &entry : copy_options) {
+			if (entry.second.mode == CopyOptionMode::READ_WRITE) {
+				// used for both
+				continue;
+			}
+			if (entry.second.mode != mode) {
+				erased_options.push_back(entry.first);
+			}
+		}
+		for (auto &erased : erased_options) {
+			copy_options.erase(erased);
+		}
+	}
 
-	if (!entry) {
-		IsFormatExtensionKnown(stmt.info->format);
-		// If we did not find an entry, we default to a CSV
-		entry = catalog.GetEntry(entry_retriever, DEFAULT_SCHEMA, {CatalogType::COPY_FUNCTION_ENTRY, "csv"},
-		                         OnEntryNotFound::THROW_EXCEPTION);
+	// now we have a list of all options for this copy type
+	// add generic options
+	if (mode != CopyOptionMode::READ_ONLY) {
+		copy_options["use_tmp_file"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["overwrite_or_ignore"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["overwrite"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["append"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["filename_pattern"] = CopyOption(LogicalType::VARCHAR, CopyOptionMode::WRITE_ONLY);
+		copy_options["file_extension"] = CopyOption(LogicalType::VARCHAR, CopyOptionMode::WRITE_ONLY);
+		copy_options["per_thread_output"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["file_size_bytes"] = CopyOption(LogicalType::ANY, CopyOptionMode::WRITE_ONLY);
+		copy_options["partition_by"] = CopyOption(LogicalType::ANY, CopyOptionMode::WRITE_ONLY);
+		copy_options["return_files"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["preserve_order"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["return_stats"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["write_partition_columns"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["write_empty_file"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
+		copy_options["hive_file_pattern"] = CopyOption(LogicalType::BOOLEAN, CopyOptionMode::WRITE_ONLY);
 	}
-	// lookup the format in the catalog
-	auto &copy_function = entry->Cast<CopyFunctionCatalogEntry>();
-	if (copy_function.function.plan) {
+	return copy_options;
+}
+
+BoundStatement Binder::BindCopyTo(CopyStatement &stmt, const CopyFunction &function, CopyToType copy_to_type) {
+	if (function.plan) {
 		// plan rewrite COPY TO
-		return copy_function.function.plan(*this, stmt);
+		return function.plan(*this, stmt);
 	}
 
 	auto &copy_info = *stmt.info;
@@ -68,7 +96,7 @@ BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type)
 	auto node_copy = copy_info.select_statement->Copy();
 	auto select_node = Bind(*node_copy);
 
-	if (!copy_function.function.copy_to_bind) {
+	if (!function.copy_to_bind) {
 		throw NotImplementedException("COPY TO is not supported for FORMAT \"%s\"", stmt.info->format);
 	}
 
@@ -89,7 +117,7 @@ BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type)
 
 	CopyFunctionBindInput bind_input(*stmt.info);
 
-	bind_input.file_extension = copy_function.function.extension;
+	bind_input.file_extension = function.extension;
 
 	auto original_options = stmt.info->options;
 	stmt.info->options.clear();
@@ -135,7 +163,7 @@ BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type)
 			if (option.second.empty()) {
 				throw BinderException("FILE_SIZE_BYTES cannot be empty");
 			}
-			if (!copy_function.function.rotate_files) {
+			if (!function.rotate_files) {
 				throw NotImplementedException("FILE_SIZE_BYTES not implemented for FORMAT \"%s\"", stmt.info->format);
 			}
 			if (option.second[0].GetTypeMutable().id() == LogicalTypeId::VARCHAR) {
@@ -207,7 +235,7 @@ BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type)
 	}
 
 	// Allow the copy function to intercept the select list and types and push a new projection on top of the plan
-	if (copy_function.function.copy_to_select) {
+	if (function.copy_to_select) {
 		auto bindings = select_node.plan->GetColumnBindings();
 
 		CopyToSelectInput input = {context, stmt.info->options, {}, copy_to_type};
@@ -221,7 +249,7 @@ BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type)
 			input.select_list.push_back(make_uniq<BoundColumnRefExpression>(name, type, binding));
 		}
 
-		auto new_select_list = copy_function.function.copy_to_select(input);
+		auto new_select_list = function.copy_to_select(input);
 		if (!new_select_list.empty()) {
 
 			// We have a new select list, create a projection on top of the current plan
@@ -248,13 +276,12 @@ BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type)
 	    LogicalCopyToFile::GetNamesWithoutPartitions(unique_column_names, partition_cols, write_partition_columns);
 	auto types_to_write =
 	    LogicalCopyToFile::GetTypesWithoutPartitions(select_node.types, partition_cols, write_partition_columns);
-	auto function_data = copy_function.function.copy_to_bind(context, bind_input, names_to_write, types_to_write);
+	auto function_data = function.copy_to_bind(context, bind_input, names_to_write, types_to_write);
 
-	const auto rotate =
-	    copy_function.function.rotate_files && copy_function.function.rotate_files(*function_data, file_size_bytes);
+	const auto rotate = function.rotate_files && function.rotate_files(*function_data, file_size_bytes);
 	if (rotate) {
-		if (!copy_function.function.rotate_next_file) {
-			throw InternalException("rotate_next_file not implemented for \"%s\"", copy_function.function.extension);
+		if (!function.rotate_next_file) {
+			throw InternalException("rotate_next_file not implemented for \"%s\"", function.extension);
 		}
 		if (user_set_use_tmp_file) {
 			throw NotImplementedException(
@@ -273,13 +300,12 @@ BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type)
 			throw NotImplementedException("Can't combine WRITE_EMPTY_FILE false with PARTITION_BY");
 		}
 	}
-	if (return_type == CopyFunctionReturnType::WRITTEN_FILE_STATISTICS &&
-	    !copy_function.function.copy_to_get_written_statistics) {
+	if (return_type == CopyFunctionReturnType::WRITTEN_FILE_STATISTICS && !function.copy_to_get_written_statistics) {
 		throw NotImplementedException("RETURN_STATS is not supported for the \"%s\" copy format", stmt.info->format);
 	}
 
 	// now create the copy information
-	auto copy = make_uniq<LogicalCopyToFile>(copy_function.function, std::move(function_data), std::move(stmt.info));
+	auto copy = make_uniq<LogicalCopyToFile>(function, std::move(function_data), std::move(stmt.info));
 	copy->file_path = file_path;
 	copy->use_tmp_file = use_tmp_file;
 	copy->overwrite_mode = overwrite_mode;
@@ -324,7 +350,7 @@ BoundStatement Binder::BindCopyTo(CopyStatement &stmt, CopyToType copy_to_type)
 	return result;
 }
 
-BoundStatement Binder::BindCopyFrom(CopyStatement &stmt) {
+BoundStatement Binder::BindCopyFrom(CopyStatement &stmt, const CopyFunction &function) {
 	BoundStatement result;
 	result.types = {LogicalType::BIGINT};
 	result.names = {"Count"};
@@ -332,6 +358,9 @@ BoundStatement Binder::BindCopyFrom(CopyStatement &stmt) {
 	if (stmt.info->table.empty()) {
 		throw ParserException("COPY FROM requires a table name to be specified");
 	}
+	if (!function.copy_from_bind) {
+		throw NotImplementedException("COPY FROM is not supported for FORMAT \"%s\"", stmt.info->format);
+	}
 	// COPY FROM a file
 	// generate an insert statement for the to-be-inserted table
 	InsertStatement insert;
@@ -346,24 +375,6 @@ BoundStatement Binder::BindCopyFrom(CopyStatement &stmt) {
 
 	auto &bound_insert = insert_statement.plan->Cast<LogicalInsert>();
 
-	// lookup the format in the catalog
-	auto &catalog = Catalog::GetSystemCatalog(context);
-	auto on_entry_do =
-	    stmt.info->is_format_auto_detected ? OnEntryNotFound::RETURN_NULL : OnEntryNotFound::THROW_EXCEPTION;
-	CatalogEntryRetriever entry_retriever {context};
-	auto entry = catalog.GetEntry(entry_retriever, DEFAULT_SCHEMA,
-	                              {CatalogType::COPY_FUNCTION_ENTRY, stmt.info->format}, on_entry_do);
-	if (!entry) {
-		IsFormatExtensionKnown(stmt.info->format);
-		// If we did not find an entry, we default to a CSV
-		entry = catalog.GetEntry(entry_retriever, DEFAULT_SCHEMA, {CatalogType::COPY_FUNCTION_ENTRY, "csv"},
-		                         OnEntryNotFound::THROW_EXCEPTION);
-	}
-	// lookup the format in the catalog
-	auto &copy_function = entry->Cast<CopyFunctionCatalogEntry>();
-	if (!copy_function.function.copy_from_bind) {
-		throw NotImplementedException("COPY FROM is not supported for FORMAT \"%s\"", stmt.info->format);
-	}
 	// lookup the table to copy into
 	BindSchemaOrCatalog(stmt.info->catalog, stmt.info->schema);
 	auto &table =
@@ -383,11 +394,11 @@ BoundStatement Binder::BindCopyFrom(CopyStatement &stmt) {
 			expected_names.push_back(col.Name());
 		}
 	}
-
-	auto function_data =
-	    copy_function.function.copy_from_bind(context, *stmt.info, expected_names, bound_insert.expected_types);
-	auto get = make_uniq<LogicalGet>(GenerateTableIndex(), copy_function.function.copy_from_function,
-	                                 std::move(function_data), bound_insert.expected_types, expected_names);
+	auto copy_from_function = function.copy_from_function;
+	CopyFromFunctionBindInput input(*stmt.info, copy_from_function);
+	auto function_data = function.copy_from_bind(context, input, expected_names, bound_insert.expected_types);
+	auto get = make_uniq<LogicalGet>(GenerateTableIndex(), std::move(copy_from_function), std::move(function_data),
+	                                 bound_insert.expected_types, expected_names);
 	for (idx_t i = 0; i < bound_insert.expected_types.size(); i++) {
 		get->AddColumnId(i);
 	}
@@ -396,7 +407,89 @@ BoundStatement Binder::BindCopyFrom(CopyStatement &stmt) {
 	return result;
 }
 
+vector<Value> BindCopyOption(ClientContext &context, TableFunctionBinder &option_binder, const string &name,
+                             unique_ptr<ParsedExpression> &expr) {
+	vector<Value> result;
+	if (!expr) {
+		return result;
+	}
+	if (expr->type == ExpressionType::STAR) {
+		auto &star = expr->Cast<StarExpression>();
+		// for compatibility with previous copy implementation - turn a raw * into a * string literal
+		if (star.relation_name.empty() && star.exclude_list.empty() && star.replace_list.empty() &&
+		    star.rename_list.empty() && !star.expr && !star.columns) {
+			result.push_back("*");
+			return result;
+		}
+	}
+	auto bound_expr = option_binder.Bind(expr);
+	auto val = ExpressionExecutor::EvaluateScalar(context, *bound_expr);
+	if (val.IsNull()) {
+		throw BinderException("NULL is not supported as a valid option for COPY option \"" + name + "\"");
+	}
+	if (val.type().id() == LogicalTypeId::STRUCT && StructType::IsUnnamed(val.type())) {
+		// unpack unnamed structs into a list of options
+		return StructValue::GetChildren(val);
+	}
+	result.push_back(std::move(val));
+	return result;
+}
+
+string ExtractFormat(const string &file_path) {
+	auto format = StringUtil::Lower(file_path);
+	// We first remove extension suffixes
+	if (StringUtil::EndsWith(format, CompressionExtensionFromType(FileCompressionType::GZIP))) {
+		format = format.substr(0, format.size() - 3);
+	} else if (StringUtil::EndsWith(format, CompressionExtensionFromType(FileCompressionType::ZSTD))) {
+		format = format.substr(0, format.size() - 4);
+	}
+	// Now lets check for the last .
+	size_t dot_pos = format.rfind('.');
+	if (dot_pos == std::string::npos || dot_pos == format.length() - 1) {
+		// No format found
+		return "";
+	}
+	// We found something
+	return format.substr(dot_pos + 1);
+}
+
+void Binder::BindCopyOptions(CopyInfo &info) {
+	TableFunctionBinder option_binder(*this, context, "Copy", "Copy options");
+	if (info.file_path_expression) {
+		auto inputs = BindCopyOption(context, option_binder, "filename", info.file_path_expression);
+		if (inputs.size() != 1 || inputs[0].type().id() != LogicalTypeId::VARCHAR) {
+			throw InternalException("Unsupported parameter type for filename: expected e.g. TARGET 'file.parquet'");
+		}
+		if (!info.file_path.empty()) {
+			throw InternalException("Both a file path and a file path expression were provided for COPY - only one of "
+			                        "the two can be provided");
+		}
+		info.file_path = inputs[0].ToString();
+		info.file_path_expression.reset();
+	}
+	for (auto &entry : info.parsed_options) {
+		auto inputs = BindCopyOption(context, option_binder, entry.first, entry.second);
+		if (StringUtil::CIEquals(entry.first, "format")) {
+			// format specifier: interpret this option
+			if (inputs.size() != 1 || inputs[0].type().id() != LogicalTypeId::VARCHAR) {
+				throw ParserException("Unsupported parameter type for FORMAT: expected e.g. FORMAT 'csv', 'parquet'");
+			}
+			info.format = StringUtil::Lower(inputs[0].ToString());
+			info.is_format_auto_detected = false;
+			continue;
+		}
+		info.options[entry.first] = std::move(inputs);
+	}
+	if (info.is_format_auto_detected && info.format.empty()) {
+		info.format = ExtractFormat(info.file_path);
+	}
+	info.parsed_options.clear();
+}
+
 BoundStatement Binder::Bind(CopyStatement &stmt, CopyToType copy_to_type) {
+	// bind the copy options
+	BindCopyOptions(*stmt.info);
+
 	if (!stmt.info->is_from && !stmt.info->select_statement) {
 		// copy table into file without a query
 		// generate SELECT * FROM table;
@@ -417,13 +510,101 @@ BoundStatement Binder::Bind(CopyStatement &stmt, CopyToType copy_to_type) {
 		stmt.info->select_statement = std::move(statement);
 	}
 
+	// Let's first bind our format
+	// lookup the format in the catalog
+	auto on_entry_do =
+	    stmt.info->is_format_auto_detected ? OnEntryNotFound::RETURN_NULL : OnEntryNotFound::THROW_EXCEPTION;
+	CatalogEntryRetriever entry_retriever {context};
+	auto &catalog = Catalog::GetSystemCatalog(context);
+	auto entry = catalog.GetEntry(entry_retriever, DEFAULT_SCHEMA,
+	                              {CatalogType::COPY_FUNCTION_ENTRY, stmt.info->format}, on_entry_do);
+
+	if (!entry) {
+		IsFormatExtensionKnown(stmt.info->format);
+		// If we did not find an entry, we default to a CSV
+		entry = catalog.GetEntry(entry_retriever, DEFAULT_SCHEMA, {CatalogType::COPY_FUNCTION_ENTRY, "csv"},
+		                         OnEntryNotFound::THROW_EXCEPTION);
+	}
+	auto &copy_function = entry->Cast<CopyFunctionCatalogEntry>();
+	auto &function = copy_function.function;
+
+	if (function.copy_options) {
+		// list all copy options - then bind them and offer alternatives
+		auto copy_mode = stmt.info->is_from ? CopyOptionMode::READ_ONLY : CopyOptionMode::WRITE_ONLY;
+		auto copy_options = GetFullCopyOptionsList(function, CopyOptionMode::READ_WRITE);
+		for (auto &provided_entry : stmt.info->options) {
+			auto &provided_option = provided_entry.first;
+			auto option_entry = copy_options.find(provided_option);
+			if (option_entry == copy_options.end()) {
+				// option not found - offer an alternative suggestion
+				vector<string> candidates;
+				for (auto &copy_entry : copy_options) {
+					candidates.push_back(copy_entry.first);
+				}
+				string candidate_str = StringUtil::CandidatesMessage(
+				    StringUtil::TopNJaroWinkler(candidates, provided_option), "Candidate options");
+
+				throw NotImplementedException("Unrecognized option \"%s\" for %s\n%s", provided_option,
+				                              stmt.info->format, candidate_str);
+			}
+			auto &copy_option = option_entry->second;
+			// check if this matches the mode
+			if (copy_option.mode != CopyOptionMode::READ_WRITE && copy_option.mode != copy_mode) {
+				throw InvalidInputException("Option \"%s\" is not supported for %s - only for %s", provided_option,
+				                            stmt.info->is_from ? "reading" : "writing",
+				                            stmt.info->is_from ? "writing" : "reading");
+			}
+			if (copy_option.type.id() != LogicalTypeId::ANY) {
+				if (provided_entry.second.empty()) {
+					if (copy_option.type.id() == LogicalTypeId::BOOLEAN) {
+						// boolean can be empty (e.g. "HEADER")
+						continue;
+					}
+					throw InvalidInputException("Copy option \"%s\" requires an argument of type %s", provided_option,
+					                            copy_option.type.ToString());
+				}
+				if (provided_entry.second.size() > 1) {
+					throw InvalidInputException("Copy option \"%s\" did not expect a list as argument",
+					                            provided_option);
+				}
+				auto &original_value = provided_entry.second[0];
+				if (copy_option.type == original_value.type()) {
+					// types match
+					continue;
+				}
+				bool can_cast =
+				    CastFunctionSet::ImplicitCastCost(context, original_value.type(), copy_option.type) >= 0;
+				if (!can_cast) {
+					// for backwards compatibility - we are more lax on casting rules for copy options
+					if (copy_option.type.IsNumeric()) {
+						can_cast = original_value.type().IsNumeric();
+					} else if (copy_option.type.id() == LogicalTypeId::BOOLEAN) {
+						can_cast = original_value.type().IsIntegral();
+					}
+					if (original_value.type().id() == LogicalTypeId::VARCHAR) {
+						can_cast = true;
+					}
+				}
+
+				Value new_value;
+				if (!can_cast || !original_value.TryCastAs(context, copy_option.type, new_value, nullptr)) {
+					throw InvalidInputException("Copy option \"%s\" expected an argument of type %s - the argument "
+					                            "\"%s\" of type %s could not be cast as this type",
+					                            provided_option, copy_option.type, original_value.ToString(),
+					                            original_value.type());
+				}
+				original_value = std::move(new_value);
+			}
+		}
+	}
+
 	auto &properties = GetStatementProperties();
 	properties.allow_stream_result = false;
 	properties.return_type = StatementReturnType::CHANGED_ROWS;
 	if (stmt.info->is_from) {
-		return BindCopyFrom(stmt);
+		return BindCopyFrom(stmt, function);
 	} else {
-		return BindCopyTo(stmt, copy_to_type);
+		return BindCopyTo(stmt, function, copy_to_type);
 	}
 }