Skip to content

Commit

Permalink
reuse test plan
Browse files Browse the repository at this point in the history
  • Loading branch information
majetideepak committed Oct 25, 2024
1 parent 8c657d8 commit d864ae7
Show file tree
Hide file tree
Showing 5 changed files with 99 additions and 154 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,16 +17,16 @@
#include <folly/init/Init.h>
#include <gtest/gtest.h>

#include "velox/connectors/hive/storage_adapters/test_plans/InsertTestPlan.h"
#include "velox/connectors/hive/storage_adapters/gcs/RegisterGCSFileSystem.h"
#include "velox/connectors/hive/storage_adapters/gcs/tests/GcsTestbench.h"
#include "velox/connectors/hive/storage_adapters/test_common/InsertTest.h"

using namespace facebook::velox::exec::test;

namespace facebook::velox::filesystems {
namespace {

class GCSInsertTest : public testing::Test, public test::InsertTestPlan {
class GCSInsertTest : public testing::Test, public test::InsertTest {
protected:
static void SetUpTestSuite() {
registerGCSFileSystem();
Expand Down Expand Up @@ -78,7 +78,7 @@ std::shared_ptr<GcsTestbench> GCSInsertTest::testbench_ = nullptr;
TEST_F(GCSInsertTest, gcsInsertTest) {
const int64_t kExpectedRows = 1'000;
const auto gcsBucket = gcsURI(testbench_->preexistingBucketName());
runInsertPlan(gcsBucket, kExpectedRows, pool());
runInsertTest(gcsBucket, kExpectedRows, pool());
}
} // namespace facebook::velox::filesystems

Expand Down
65 changes: 3 additions & 62 deletions velox/connectors/hive/storage_adapters/s3fs/tests/S3InsertTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -17,21 +17,14 @@
#include <folly/init/Init.h>
#include <gtest/gtest.h>

#include "velox/common/memory/Memory.h"
#include "velox/connectors/hive/storage_adapters/s3fs/RegisterS3FileSystem.h"
#include "velox/connectors/hive/storage_adapters/s3fs/tests/S3Test.h"
#include "velox/dwio/parquet/RegisterParquetReader.h"
#include "velox/dwio/parquet/RegisterParquetWriter.h"
#include "velox/exec/TableWriter.h"
#include "velox/exec/tests/utils/AssertQueryBuilder.h"
#include "velox/exec/tests/utils/PlanBuilder.h"

using namespace facebook::velox::exec::test;
#include "velox/connectors/hive/storage_adapters/test_common/InsertTest.h"

namespace facebook::velox {
namespace {

class S3InsertTest : public S3Test {
class S3InsertTest : public S3Test, public test::InsertTest {
protected:
static void SetUpTestCase() {
memory::MemoryManager::testingSetInstance({});
Expand Down Expand Up @@ -69,61 +62,9 @@ class S3InsertTest : public S3Test {
TEST_F(S3InsertTest, s3InsertTest) {
const int64_t kExpectedRows = 1'000;
const std::string_view kOutputDirectory{"s3://writedata/"};

auto rowType = ROW(
{"c0", "c1", "c2", "c3"}, {BIGINT(), INTEGER(), SMALLINT(), DOUBLE()});

auto input = makeRowVector(
{makeFlatVector<int64_t>(kExpectedRows, [](auto row) { return row; }),
makeFlatVector<int32_t>(kExpectedRows, [](auto row) { return row; }),
makeFlatVector<int16_t>(kExpectedRows, [](auto row) { return row; }),
makeFlatVector<double>(kExpectedRows, [](auto row) { return row; })});

minioServer_->addBucket("writedata");

// Insert into s3 with one writer.
auto plan =
PlanBuilder()
.values({input})
.tableWrite(
kOutputDirectory.data(), dwio::common::FileFormat::PARQUET)
.planNode();

// Execute the write plan.
auto results = AssertQueryBuilder(plan).copyResults(pool());

// First column has number of rows written in the first row and nulls in other
// rows.
auto rowCount = results->childAt(exec::TableWriteTraits::kRowCountChannel)
->as<FlatVector<int64_t>>();
ASSERT_FALSE(rowCount->isNullAt(0));
ASSERT_EQ(kExpectedRows, rowCount->valueAt(0));
ASSERT_TRUE(rowCount->isNullAt(1));

// Second column contains details about written files.
auto details = results->childAt(exec::TableWriteTraits::kFragmentChannel)
->as<FlatVector<StringView>>();
ASSERT_TRUE(details->isNullAt(0));
ASSERT_FALSE(details->isNullAt(1));
folly::dynamic obj = folly::parseJson(details->valueAt(1));

ASSERT_EQ(kExpectedRows, obj["rowCount"].asInt());
auto fileWriteInfos = obj["fileWriteInfos"];
ASSERT_EQ(1, fileWriteInfos.size());

auto writeFileName = fileWriteInfos[0]["writeFileName"].asString();

// Read from 'writeFileName' and verify the data matches the original.
plan = PlanBuilder().tableScan(rowType).planNode();

auto filePath = fmt::format("{}{}", kOutputDirectory, writeFileName);
const int64_t fileSize = fileWriteInfos[0]["fileSize"].asInt();
auto split = HiveConnectorSplitBuilder(filePath)
.fileFormat(dwio::common::FileFormat::PARQUET)
.length(fileSize)
.build();
auto copy = AssertQueryBuilder(plan).split(split).copyResults(pool());
assertEqualResults({input}, {copy});
runInsertTest(kOutputDirectory, kExpectedRows, pool());
}
} // namespace facebook::velox

Expand Down
2 changes: 1 addition & 1 deletion velox/connectors/hive/storage_adapters/s3fs/tests/S3Test.h
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ constexpr int kOneMB = 1 << 20;

static constexpr std::string_view kDummyPath = "s3://dummy/foo.txt";

class S3Test : public testing::Test, public ::test::VectorTestBase {
class S3Test : public testing::Test {
protected:
void SetUp() override {
minioServer_ = std::make_unique<MinioServer>();
Expand Down
92 changes: 92 additions & 0 deletions velox/connectors/hive/storage_adapters/test_common/InsertTest.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
/*
* Copyright (c) Facebook, Inc. and its affiliates.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

#include <folly/init/Init.h>
#include <gtest/gtest.h>

#include "velox/common/memory/Memory.h"
#include "velox/dwio/parquet/RegisterParquetReader.h"
#include "velox/dwio/parquet/RegisterParquetWriter.h"
#include "velox/exec/TableWriter.h"
#include "velox/exec/tests/utils/AssertQueryBuilder.h"
#include "velox/exec/tests/utils/HiveConnectorTestBase.h"
#include "velox/exec/tests/utils/PlanBuilder.h"
#include "velox/vector/tests/utils/VectorTestBase.h"

namespace facebook::velox::test {

class InsertTest : public test::VectorTestBase {
public:
void runInsertTest(
std::string_view outputDirectory,
int numRows,
memory::MemoryPool* pool) {
auto rowType = ROW(
{"c0", "c1", "c2", "c3"}, {BIGINT(), INTEGER(), SMALLINT(), DOUBLE()});

auto input = makeRowVector(
{makeFlatVector<int64_t>(numRows, [](auto row) { return row; }),
makeFlatVector<int32_t>(numRows, [](auto row) { return row; }),
makeFlatVector<int16_t>(numRows, [](auto row) { return row; }),
makeFlatVector<double>(numRows, [](auto row) { return row; })});

// Insert with one writer.
auto plan =
exec::test::PlanBuilder()
.values({input})
.tableWrite(
outputDirectory.data(), dwio::common::FileFormat::PARQUET)
.planNode();

// Execute the write plan.
auto results = exec::test::AssertQueryBuilder(plan).copyResults(pool);

// First column has number of rows written in the first row and nulls in
// other rows.
auto rowCount = results->childAt(exec::TableWriteTraits::kRowCountChannel)
->as<FlatVector<int64_t>>();
ASSERT_FALSE(rowCount->isNullAt(0));
ASSERT_EQ(numRows, rowCount->valueAt(0));
ASSERT_TRUE(rowCount->isNullAt(1));

// Second column contains details about written files.
auto details = results->childAt(exec::TableWriteTraits::kFragmentChannel)
->as<FlatVector<StringView>>();
ASSERT_TRUE(details->isNullAt(0));
ASSERT_FALSE(details->isNullAt(1));
folly::dynamic obj = folly::parseJson(details->valueAt(1));

ASSERT_EQ(numRows, obj["rowCount"].asInt());
auto fileWriteInfos = obj["fileWriteInfos"];
ASSERT_EQ(1, fileWriteInfos.size());

auto writeFileName = fileWriteInfos[0]["writeFileName"].asString();

// Read from 'writeFileName' and verify the data matches the original.
plan = exec::test::PlanBuilder().tableScan(rowType).planNode();

auto filePath = fmt::format("{}{}", outputDirectory, writeFileName);
const int64_t fileSize = fileWriteInfos[0]["fileSize"].asInt();
auto split = exec::test::HiveConnectorSplitBuilder(filePath)
.fileFormat(dwio::common::FileFormat::PARQUET)
.length(fileSize)
.build();
auto copy =
exec::test::AssertQueryBuilder(plan).split(split).copyResults(pool);
exec::test::assertEqualResults({input}, {copy});
}
};
} // namespace facebook::velox::test
88 changes: 0 additions & 88 deletions velox/connectors/hive/storage_adapters/test_plans/InsertTestPlan.h

This file was deleted.

0 comments on commit d864ae7

Please sign in to comment.