Repository: duckdb/duckdb
File: src/parser/transform/helpers/transform_sample.cpp (C++)

Code before:
	if (!options) {
		return nullptr;
	}
	auto result = make_uniq<SampleOptions>();
	auto &sample_options = PGCast<duckdb_libpgquery::PGSampleOptions>(*options);
	auto &sample_size = *PGPointerCast<duckdb_libpgquery::PGSampleSize>(sample_options.sample_size);
	auto sample_expression = TransformExpression(sample_size.sample_size);
	if (sample_expression->GetExpressionType() != ExpressionType::VALUE_CONSTANT) {
		throw ParserException(sample_expression->GetQueryLocation(),
		                      "Only constants are supported in sample clause currently");
	}
	auto &const_expr = sample_expression->Cast<ConstantExpression>();
	auto &sample_value = const_expr.value;
	result->is_percentage = sample_size.is_percentage;
	if (sample_size.is_percentage) {
		// sample size is given in sample_size: use system sampling
		auto percentage = sample_value.GetValue<double>();
		if (percentage < 0 || percentage > 100) {
			throw ParserException("Sample sample_size %llf out of range, must be between 0 and 100", percentage);
		}
		result->sample_size = Value::DOUBLE(percentage);
		result->method = SampleMethod::SYSTEM_SAMPLE;
	} else {
		// sample size is given in rows: use reservoir sampling
		auto rows = sample_value.GetValue<int64_t>();
		if (rows < 0 || sample_value.GetValue<uint64_t>() >= Allocator::MAXIMUM_ALLOC_SIZE) {
			throw ParserException("Sample rows %lld out of range, must be between 0 and %lld", rows,
			                      Allocator::MAXIMUM_ALLOC_SIZE);
		}
		result->sample_size = Value::BIGINT(rows);
		result->method = SampleMethod::RESERVOIR_SAMPLE;
	}
	if (sample_options.method) {
		result->method = GetSampleMethod(sample_options.method);
	}
	if (sample_options.has_seed && sample_options.seed >= 0) {
		result->seed = static_cast<idx_t>(sample_options.seed);
		result->repeatable = true;
	}
	return result;
}

} // namespace duckdb


Review comment (suggestion):
The sample size is in rows but the allocator max size is in bytes - I suspect that sampling for `281474976710656-1` rows will still throw the allocation error. Should we reduce the limit further / add a test for this?

Alternatively, we might want to not allocate all the space for the sample up-front if the sample size is very large - but that's obviously a larger change.

Code after:
	if (!options) {
		return nullptr;
	}
	auto result = make_uniq<SampleOptions>();
	auto &sample_options = PGCast<duckdb_libpgquery::PGSampleOptions>(*options);
	auto &sample_size = *PGPointerCast<duckdb_libpgquery::PGSampleSize>(sample_options.sample_size);
	auto sample_expression = TransformExpression(sample_size.sample_size);
	if (sample_expression->GetExpressionType() != ExpressionType::VALUE_CONSTANT) {
		throw ParserException(sample_expression->GetQueryLocation(),
		                      "Only constants are supported in sample clause currently");
	}
	auto &const_expr = sample_expression->Cast<ConstantExpression>();
	auto &sample_value = const_expr.value;
	result->is_percentage = sample_size.is_percentage;
	if (sample_size.is_percentage) {
		// sample size is given in sample_size: use system sampling
		auto percentage = sample_value.GetValue<double>();
		if (percentage < 0 || percentage > 100) {
			throw ParserException("Sample sample_size %llf out of range, must be between 0 and 100", percentage);
		}
		result->sample_size = Value::DOUBLE(percentage);
		result->method = SampleMethod::SYSTEM_SAMPLE;
	} else {
		// sample size is given in rows: use reservoir sampling
		auto rows = sample_value.GetValue<int64_t>();
		if (rows < 0 || sample_value.GetValue<uint64_t>() > SampleOptions::MAX_SAMPLE_ROWS) {
			throw ParserException("Sample rows %lld out of range, must be between 0 and %lld", rows,
			                      SampleOptions::MAX_SAMPLE_ROWS);
		}
		result->sample_size = Value::BIGINT(rows);
		result->method = SampleMethod::RESERVOIR_SAMPLE;
	}
	if (sample_options.method) {
		result->method = GetSampleMethod(sample_options.method);
	}
	if (sample_options.has_seed && sample_options.seed >= 0) {
		result->seed = static_cast<idx_t>(sample_options.seed);
		result->repeatable = true;
	}
	return result;
}

} // namespace duckdb