diff --git a/src/iceberg/catalog/rest/rest_table_scan.cc b/src/iceberg/catalog/rest/rest_table_scan.cc index e7dae0f49..fa2bc6737 100644 --- a/src/iceberg/catalog/rest/rest_table_scan.cc +++ b/src/iceberg/catalog/rest/rest_table_scan.cc @@ -200,9 +200,20 @@ Result FetchPlanningResult( ICEBERG_ASSIGN_OR_RAISE(auto path, ctx.paths->Plan(ctx.identifier, plan_id)); auto delay_ms = kMinSleepMs; - auto start = std::chrono::steady_clock::now(); + const auto start = std::chrono::steady_clock::now(); + const auto remaining_wait = [&]() -> Result { + auto elapsed_ms = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count(); + if (elapsed_ms >= poll_config.max_wait_ms) { + return IOError("Scan planning timed out after {}ms waiting for plan_id={}", + elapsed_ms, plan_id); + } + return std::chrono::milliseconds(poll_config.max_wait_ms - elapsed_ms); + }; for (int32_t retry = 0; retry <= poll_config.max_retries; ++retry) { + ICEBERG_RETURN_UNEXPECTED(remaining_wait()); ICEBERG_ASSIGN_OR_RAISE(auto response, ctx.client->Get(path, /*params=*/{}, /*headers=*/{}, *PlanErrorHandler::Instance(), *ctx.session)); @@ -218,15 +229,10 @@ Result FetchPlanningResult( return result; } case PlanStatus::kSubmitted: { - auto elapsed_ms = std::chrono::duration_cast( - std::chrono::steady_clock::now() - start) - .count(); - if (elapsed_ms >= poll_config.max_wait_ms) { - return IOError("Scan planning timed out after {}ms waiting for plan_id={}", - elapsed_ms, plan_id); - } + ICEBERG_ASSIGN_OR_RAISE(auto remaining, remaining_wait()); if (retry == poll_config.max_retries) break; - std::this_thread::sleep_for(std::chrono::milliseconds(delay_ms)); + std::this_thread::sleep_for( + std::min(std::chrono::milliseconds(delay_ms), remaining)); delay_ms = std::min(delay_ms * 2, kMaxSleepMs); continue; } @@ -278,9 +284,11 @@ class RestFileScanTaskStream final : public FileScanTaskStream { auto task = std::move(buffer_[buffer_pos_++]); if (ignore_residuals_ || !task->residual_filter()) { // The server may leave row filtering to the client. - return std::make_shared( - task->data_file(), task->delete_files(), - ignore_residuals_ ? True::Instance() : filter_); + ICEBERG_ASSIGN_OR_RAISE( + task, + FileScanTask::MakeSplit(task->data_file(), task->start(), task->length(), + task->delete_files(), + ignore_residuals_ ? True::Instance() : filter_)); } return task; } diff --git a/src/iceberg/test/rest_table_scan_test.cc b/src/iceberg/test/rest_table_scan_test.cc index a1b8e6e86..7610f8249 100644 --- a/src/iceberg/test/rest_table_scan_test.cc +++ b/src/iceberg/test/rest_table_scan_test.cc @@ -20,9 +20,11 @@ #include "iceberg/catalog/rest/rest_table_scan.h" #include +#include #include #include #include +#include #include #include #include @@ -935,7 +937,7 @@ class RestScanResidualTest : public RestIncrementalAppendScanTest, public ::testing::WithParamInterface> {}; -TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) { +TEST_P(RestScanResidualTest, PreservesSplitRangesFilteringAndDeleteFiles) { const auto [incremental, submitted, fetch_token, server_residual, ignore_residuals] = GetParam(); auto filter = Expressions::GreaterThan("id", Literal::Int(10)); @@ -944,6 +946,8 @@ TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) { ICEBERG_UNWRAP_OR_FAIL(auto residual_json, iceberg::ToJson(*residual)); auto task = FileTaskJson("s3://bucket/data.parquet"); + task["start"] = 25; + task["length"] = 50; task["delete-file-references"] = {0}; if (server_residual) task["residual-filter"] = residual_json; auto delete_file = FileTaskJson("s3://bucket/delete.parquet")["data-file"]; @@ -985,6 +989,9 @@ TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) { })); ASSERT_EQ(tasks.size(), 1); EXPECT_EQ(tasks[0]->data_file()->file_path, "s3://bucket/data.parquet"); + EXPECT_TRUE(tasks[0]->is_split()); + EXPECT_EQ(tasks[0]->start(), 25); + EXPECT_EQ(tasks[0]->length(), 50); ASSERT_EQ(tasks[0]->delete_files().size(), 1); EXPECT_EQ(tasks[0]->delete_files()[0]->file_path, "s3://bucket/delete.parquet"); ASSERT_NE(tasks[0]->residual_filter(), nullptr); @@ -1039,6 +1046,115 @@ TEST_P(RestScanRegressionTest, AllColumnStatsIncludeUnprojectedColumns) { EXPECT_TRUE(tasks.empty()); } +TEST_P(RestScanRegressionTest, PollingSleepRespectsTimeout) { + constexpr std::string_view kSubmitted = + R"({"status":"submitted","plan-id":"timeout-plan"})"; + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kSubmitted)))); + int poll_count = 0; + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillRepeatedly([&](auto&&...) -> Result { + ++poll_count; + return HttpResponse( + 200, poll_count == 1 ? std::string(kSubmitted) : R"({"status":"completed"})"); + }); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, "{}"))); + + auto ctx = MakeContext(); + ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}}; + auto start = std::chrono::steady_clock::now(); + EXPECT_THAT( + PlanScan(GetParam(), std::move(ctx), [](auto&) {}), + ::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("timed out"))); + auto elapsed = std::chrono::steady_clock::now() - start; + EXPECT_EQ(poll_count, 1); + EXPECT_GE(elapsed, std::chrono::milliseconds(10)); + EXPECT_LT(elapsed, std::chrono::milliseconds(500)); +} + +TEST_P(RestScanRegressionTest, PollingCompletionAfterTimeoutIsPreserved) { + nlohmann::json completed = { + {"status", "completed"}, + {"file-scan-tasks", + nlohmann::json::array({FileTaskJson("s3://bucket/completed.parquet")})}}; + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce( + Return(HttpResponse(200, R"({"status":"submitted","plan-id":"slow-plan"})"))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillOnce([&](auto&&...) -> Result { + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + return HttpResponse(200, completed.dump()); + }); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + + auto ctx = MakeContext(); + ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}}; + ICEBERG_UNWRAP_OR_FAIL(auto tasks, PlanScan(GetParam(), std::move(ctx), [](auto&) {})); + ASSERT_EQ(tasks.size(), 1); + EXPECT_EQ(tasks[0]->data_file()->file_path, "s3://bucket/completed.parquet"); +} + +TEST_P(RestScanRegressionTest, PollingFailureAfterTimeoutPreservesServerError) { + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce( + Return(HttpResponse(200, R"({"status":"submitted","plan-id":"slow-plan"})"))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillOnce([](auto&&...) -> Result { + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + return HttpResponse( + 200, + R"({"status":"failed","error":{"message":"server error","type":"ServerError","code":500}})"); + }); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, "{}"))); + + auto ctx = MakeContext(); + ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}}; + EXPECT_THAT( + PlanScan(GetParam(), std::move(ctx), [](auto&) {}), + ::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("server error"))); +} + +TEST_P(RestScanRegressionTest, PollingCancellationAfterTimeoutDoesNotCancelAgain) { + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce( + Return(HttpResponse(200, R"({"status":"submitted","plan-id":"slow-plan"})"))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillOnce([](auto&&...) -> Result { + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + return HttpResponse(200, R"({"status":"cancelled"})"); + }); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); + + auto ctx = MakeContext(); + ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}}; + EXPECT_THAT( + PlanScan(GetParam(), std::move(ctx), [](auto&) {}), + ::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("was cancelled"))); +} + +TEST_P(RestScanRegressionTest, PollingSubmittedAfterTimeoutReturnsTimeout) { + constexpr std::string_view kSubmitted = + R"({"status":"submitted","plan-id":"slow-plan"})"; + EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, std::string(kSubmitted)))); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + .WillOnce([&](auto&&...) -> Result { + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + return HttpResponse(200, std::string(kSubmitted)); + }); + EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) + .WillOnce(Return(HttpResponse(200, "{}"))); + + auto ctx = MakeContext(); + ctx.catalog_config = {{"rest-scan-planning.poll-timeout-ms", "10"}, + {"rest-scan-planning.poll-num-retries", "0"}}; + EXPECT_THAT( + PlanScan(GetParam(), std::move(ctx), [](auto&) {}), + ::testing::AllOf(IsError(ErrorKind::kIOError), HasErrorMessage("timed out"))); +} + TEST_P(RestScanRegressionTest, InvalidPollingConfigFailsBeforePlanning) { struct PropertyCase { std::string_view key;