diff --git a/velox/connectors/hive/HiveDataSink.cpp b/velox/connectors/hive/HiveDataSink.cpp index 36ec297220d..f6f3c8d8c65 100644 --- a/velox/connectors/hive/HiveDataSink.cpp +++ b/velox/connectors/hive/HiveDataSink.cpp @@ -1012,9 +1012,6 @@ HiveWriterParameters::UpdateMode HiveDataSink::getUpdateMode() const { insertBehavior)); } } else { - if (insertTableHandle_->isBucketed()) { - VELOX_USER_FAIL("Cannot insert into bucketed unpartitioned Hive table"); - } if (hiveConfig_->immutablePartitions()) { VELOX_USER_FAIL("Unpartitioned Hive tables are immutable."); } diff --git a/velox/exec/tests/TableWriterTest.cpp b/velox/exec/tests/TableWriterTest.cpp index 461d3a124dc..db3a0664624 100644 --- a/velox/exec/tests/TableWriterTest.cpp +++ b/velox/exec/tests/TableWriterTest.cpp @@ -315,6 +315,48 @@ class UnpartitionedTableWriterTest } }; +class BucketedUnpartitionedTableWriterTest + : public TableWriterTestBase, + public testing::WithParamInterface { + public: + BucketedUnpartitionedTableWriterTest() : TableWriterTestBase(GetParam()) {} + + static std::vector getTestParams() { + std::vector testParams; + const std::vector multiDriverOptions = {false, true}; + std::vector fileFormats = {FileFormat::DWRF}; + if (hasWriterFactory(FileFormat::PARQUET)) { + fileFormats.push_back(FileFormat::PARQUET); + } + const std::vector bucketModes = {TestMode::kOnlyBucketed}; + for (bool multiDrivers : multiDriverOptions) { + for (FileFormat fileFormat : fileFormats) { + testParams.push_back(TestParam{ + fileFormat, + TestMode::kOnlyBucketed, + CommitStrategy::kNoCommit, + HiveBucketProperty::Kind::kHiveCompatible, + true, + multiDrivers, + facebook::velox::common::CompressionKind_ZSTD, + /*scaleWriter=*/false} + .value); + testParams.push_back(TestParam{ + fileFormat, + TestMode::kOnlyBucketed, + CommitStrategy::kTaskCommit, + HiveBucketProperty::Kind::kHiveCompatible, + true, + multiDrivers, + facebook::velox::common::CompressionKind_NONE, + /*scaleWriter=*/false} + .value); + } + } + return testParams; + } +}; + class BucketedTableOnlyWriteTest : public TableWriterTestBase, public testing::WithParamInterface { @@ -1553,6 +1595,37 @@ TEST_P(UnpartitionedTableWriterTest, immutableSettings) { } } +TEST_P(BucketedUnpartitionedTableWriterTest, bucketNonPartitioned) { + SCOPED_TRACE(testParam_.toString()); + auto input = makeVectors(1, 100); + createDuckDbTable(input); + + auto outputDirectory = TempDirectoryPath::create(); + setBucketProperty( + bucketProperty_->kind(), + bucketProperty_->bucketCount(), + bucketProperty_->bucketedBy(), + bucketProperty_->bucketedTypes(), + bucketProperty_->sortedBy()); + auto plan = createInsertPlan( + PlanBuilder().values({input}), + rowType_, + outputDirectory->getPath(), + {}, + bucketProperty_, + compressionKind_, + getNumWriters(), + connector::hive::LocationHandle::TableType::kExisting, + commitStrategy_); + assertQueryWithWriterConfigs(plan, "SELECT count(*) FROM tmp"); + + assertQuery( + PlanBuilder().tableScan(rowType_).planNode(), + makeHiveConnectorSplits(outputDirectory), + "SELECT * FROM tmp"); + verifyTableWriterOutput(outputDirectory->getPath(), rowType_); +} + TEST_P(BucketedTableOnlyWriteTest, bucketCountLimit) { SCOPED_TRACE(testParam_.toString()); auto input = makeVectors(1, 100); @@ -2603,6 +2676,11 @@ VELOX_INSTANTIATE_TEST_SUITE_P( UnpartitionedTableWriterTest, testing::ValuesIn(UnpartitionedTableWriterTest::getTestParams())); +VELOX_INSTANTIATE_TEST_SUITE_P( + TableWriterTest, + BucketedUnpartitionedTableWriterTest, + testing::ValuesIn(BucketedUnpartitionedTableWriterTest::getTestParams())); + VELOX_INSTANTIATE_TEST_SUITE_P( TableWriterTest, PartitionedTableWriterTest,