--- /dev/null
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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 <cstdint>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "arrow/array/array_base.h"
+#include "arrow/array/data.h"
+#include "arrow/array/util.h"
+#include "arrow/chunked_array.h"
+#include "arrow/record_batch.h"
+#include "arrow/status.h"
+#include "arrow/table.h"
+#include "arrow/testing/gtest_common.h"
+#include "arrow/testing/gtest_util.h"
+#include "arrow/type.h"
+#include "arrow/util/key_value_metadata.h"
+
+namespace arrow {
+
+class TestTable : public TestBase {
+ public:
+ void MakeExample1(int length) {
+ auto f0 = field("f0", int32());
+ auto f1 = field("f1", uint8());
+ auto f2 = field("f2", int16());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1, f2};
+ schema_ = std::make_shared<Schema>(fields);
+
+ arrays_ = {MakeRandomArray<Int32Array>(length), MakeRandomArray<UInt8Array>(length),
+ MakeRandomArray<Int16Array>(length)};
+
+ columns_ = {std::make_shared<ChunkedArray>(arrays_[0]),
+ std::make_shared<ChunkedArray>(arrays_[1]),
+ std::make_shared<ChunkedArray>(arrays_[2])};
+ }
+
+ protected:
+ std::shared_ptr<Table> table_;
+ std::shared_ptr<Schema> schema_;
+
+ std::vector<std::shared_ptr<Array>> arrays_;
+ std::vector<std::shared_ptr<ChunkedArray>> columns_;
+};
+
+TEST_F(TestTable, EmptySchema) {
+ auto empty_schema = ::arrow::schema({});
+ table_ = Table::Make(empty_schema, columns_);
+ ASSERT_OK(table_->ValidateFull());
+ ASSERT_EQ(0, table_->num_rows());
+ ASSERT_EQ(0, table_->num_columns());
+}
+
+TEST_F(TestTable, Ctors) {
+ const int length = 100;
+ MakeExample1(length);
+
+ table_ = Table::Make(schema_, columns_);
+ ASSERT_OK(table_->ValidateFull());
+ ASSERT_EQ(length, table_->num_rows());
+ ASSERT_EQ(3, table_->num_columns());
+
+ auto array_ctor = Table::Make(schema_, arrays_);
+ ASSERT_TRUE(table_->Equals(*array_ctor));
+
+ table_ = Table::Make(schema_, columns_, length);
+ ASSERT_OK(table_->ValidateFull());
+ ASSERT_EQ(length, table_->num_rows());
+
+ table_ = Table::Make(schema_, arrays_);
+ ASSERT_OK(table_->ValidateFull());
+ ASSERT_EQ(length, table_->num_rows());
+ ASSERT_EQ(3, table_->num_columns());
+}
+
+TEST_F(TestTable, Metadata) {
+ const int length = 100;
+ MakeExample1(length);
+
+ table_ = Table::Make(schema_, columns_);
+
+ ASSERT_TRUE(table_->schema()->Equals(*schema_));
+
+ auto col = table_->column(0);
+ ASSERT_EQ(schema_->field(0)->type(), col->type());
+}
+
+TEST_F(TestTable, InvalidColumns) {
+ // Check that columns are all the same length
+ const int length = 100;
+ MakeExample1(length);
+
+ table_ = Table::Make(schema_, columns_, length - 1);
+ ASSERT_RAISES(Invalid, table_->ValidateFull());
+
+ columns_.clear();
+
+ // Wrong number of columns
+ table_ = Table::Make(schema_, columns_, length);
+ ASSERT_RAISES(Invalid, table_->ValidateFull());
+
+ columns_ = {std::make_shared<ChunkedArray>(MakeRandomArray<Int32Array>(length)),
+ std::make_shared<ChunkedArray>(MakeRandomArray<UInt8Array>(length)),
+ std::make_shared<ChunkedArray>(MakeRandomArray<Int16Array>(length - 1))};
+
+ table_ = Table::Make(schema_, columns_, length);
+ ASSERT_RAISES(Invalid, table_->ValidateFull());
+}
+
+TEST_F(TestTable, AllColumnsAndFields) {
+ const int length = 100;
+ MakeExample1(length);
+ table_ = Table::Make(schema_, columns_);
+
+ auto columns = table_->columns();
+ auto fields = table_->fields();
+
+ for (int i = 0; i < table_->num_columns(); ++i) {
+ AssertChunkedEqual(*table_->column(i), *columns[i]);
+ AssertFieldEqual(*table_->field(i), *fields[i]);
+ }
+
+ // Zero length
+ std::vector<std::shared_ptr<Array>> t2_columns;
+ auto t2 = Table::Make(::arrow::schema({}), t2_columns);
+ columns = t2->columns();
+ fields = t2->fields();
+
+ ASSERT_EQ(0, columns.size());
+ ASSERT_EQ(0, fields.size());
+}
+
+TEST_F(TestTable, Equals) {
+ const int length = 100;
+ MakeExample1(length);
+
+ table_ = Table::Make(schema_, columns_);
+
+ ASSERT_TRUE(table_->Equals(*table_));
+ // Differing schema
+ auto f0 = field("f3", int32());
+ auto f1 = field("f4", uint8());
+ auto f2 = field("f5", int16());
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1, f2};
+ auto other_schema = std::make_shared<Schema>(fields);
+ auto other = Table::Make(other_schema, columns_);
+ ASSERT_FALSE(table_->Equals(*other));
+ // Differing columns
+ std::vector<std::shared_ptr<ChunkedArray>> other_columns = {
+ std::make_shared<ChunkedArray>(MakeRandomArray<Int32Array>(length, 10)),
+ std::make_shared<ChunkedArray>(MakeRandomArray<UInt8Array>(length, 10)),
+ std::make_shared<ChunkedArray>(MakeRandomArray<Int16Array>(length, 10))};
+
+ other = Table::Make(schema_, other_columns);
+ ASSERT_FALSE(table_->Equals(*other));
+
+ // Differring schema metadata
+ other_schema = schema_->WithMetadata(::arrow::key_value_metadata({"key"}, {"value"}));
+ other = Table::Make(other_schema, columns_);
+ ASSERT_TRUE(table_->Equals(*other));
+ ASSERT_FALSE(table_->Equals(*other, /*check_metadata=*/true));
+}
+
+TEST_F(TestTable, FromRecordBatches) {
+ const int64_t length = 10;
+ MakeExample1(length);
+
+ auto batch1 = RecordBatch::Make(schema_, length, arrays_);
+
+ ASSERT_OK_AND_ASSIGN(auto result, Table::FromRecordBatches({batch1}));
+
+ auto expected = Table::Make(schema_, columns_);
+ ASSERT_TRUE(result->Equals(*expected));
+
+ std::vector<std::shared_ptr<ChunkedArray>> other_columns;
+ for (int i = 0; i < schema_->num_fields(); ++i) {
+ std::vector<std::shared_ptr<Array>> col_arrays = {arrays_[i], arrays_[i]};
+ other_columns.push_back(std::make_shared<ChunkedArray>(col_arrays));
+ }
+
+ ASSERT_OK_AND_ASSIGN(result, Table::FromRecordBatches({batch1, batch1}));
+ expected = Table::Make(schema_, other_columns);
+ ASSERT_TRUE(result->Equals(*expected));
+
+ // Error states
+ std::vector<std::shared_ptr<RecordBatch>> empty_batches;
+ ASSERT_RAISES(Invalid, Table::FromRecordBatches(empty_batches));
+
+ auto other_schema = ::arrow::schema({schema_->field(0), schema_->field(1)});
+
+ std::vector<std::shared_ptr<Array>> other_arrays = {arrays_[0], arrays_[1]};
+ auto batch2 = RecordBatch::Make(other_schema, length, other_arrays);
+ ASSERT_RAISES(Invalid, Table::FromRecordBatches({batch1, batch2}));
+}
+
+TEST_F(TestTable, FromRecordBatchesZeroLength) {
+ // ARROW-2307
+ MakeExample1(10);
+
+ ASSERT_OK_AND_ASSIGN(auto result, Table::FromRecordBatches(schema_, {}));
+
+ ASSERT_EQ(0, result->num_rows());
+ ASSERT_TRUE(result->schema()->Equals(*schema_));
+}
+
+TEST_F(TestTable, CombineChunksZeroColumn) {
+ // ARROW-11232
+ auto record_batch = RecordBatch::Make(schema({}), /*num_rows=*/10,
+ std::vector<std::shared_ptr<Array>>{});
+
+ ASSERT_OK_AND_ASSIGN(
+ auto table,
+ Table::FromRecordBatches(record_batch->schema(), {record_batch, record_batch}));
+ ASSERT_EQ(20, table->num_rows());
+
+ ASSERT_OK_AND_ASSIGN(auto combined, table->CombineChunks());
+
+ EXPECT_EQ(20, combined->num_rows());
+ EXPECT_TRUE(combined->Equals(*table));
+}
+
+TEST_F(TestTable, CombineChunksZeroRow) {
+ MakeExample1(10);
+
+ ASSERT_OK_AND_ASSIGN(auto table, Table::FromRecordBatches(schema_, {}));
+ ASSERT_EQ(0, table->num_rows());
+
+ ASSERT_OK_AND_ASSIGN(auto compacted, table->CombineChunks());
+
+ EXPECT_TRUE(compacted->Equals(*table));
+}
+
+TEST_F(TestTable, CombineChunks) {
+ MakeExample1(10);
+ auto batch1 = RecordBatch::Make(schema_, 10, arrays_);
+
+ MakeExample1(15);
+ auto batch2 = RecordBatch::Make(schema_, 15, arrays_);
+
+ ASSERT_OK_AND_ASSIGN(auto table, Table::FromRecordBatches({batch1, batch2}));
+ for (int i = 0; i < table->num_columns(); ++i) {
+ ASSERT_EQ(2, table->column(i)->num_chunks());
+ }
+
+ ASSERT_OK_AND_ASSIGN(auto compacted, table->CombineChunks());
+
+ EXPECT_TRUE(compacted->Equals(*table));
+ for (int i = 0; i < compacted->num_columns(); ++i) {
+ EXPECT_EQ(1, compacted->column(i)->num_chunks());
+ }
+}
+
+TEST_F(TestTable, LARGE_MEMORY_TEST(CombineChunksStringColumn)) {
+ schema_ = schema({field("str", utf8())});
+ arrays_ = {nullptr};
+
+ std::string value(1 << 16, '-');
+
+ auto num_rows = kBinaryMemoryLimit / static_cast<int64_t>(value.size());
+ StringBuilder builder;
+ ASSERT_OK(builder.Resize(num_rows));
+ ASSERT_OK(builder.ReserveData(value.size() * num_rows));
+ for (int i = 0; i < num_rows; ++i) builder.UnsafeAppend(value);
+ ASSERT_OK(builder.Finish(&arrays_[0]));
+
+ auto batch = RecordBatch::Make(schema_, num_rows, arrays_);
+
+ ASSERT_OK_AND_ASSIGN(auto table, Table::FromRecordBatches({batch, batch}));
+ ASSERT_EQ(table->column(0)->num_chunks(), 2);
+
+ ASSERT_OK_AND_ASSIGN(auto compacted, table->CombineChunks());
+ EXPECT_TRUE(compacted->Equals(*table));
+
+ // can't compact these columns any further; they contain too much character data
+ ASSERT_EQ(compacted->column(0)->num_chunks(), 2);
+}
+
+TEST_F(TestTable, ConcatenateTables) {
+ const int64_t length = 10;
+
+ MakeExample1(length);
+ auto batch1 = RecordBatch::Make(schema_, length, arrays_);
+
+ // generate different data
+ MakeExample1(length);
+ auto batch2 = RecordBatch::Make(schema_, length, arrays_);
+
+ ASSERT_OK_AND_ASSIGN(auto t1, Table::FromRecordBatches({batch1}));
+ ASSERT_OK_AND_ASSIGN(auto t2, Table::FromRecordBatches({batch2}));
+
+ ASSERT_OK_AND_ASSIGN(auto result, ConcatenateTables({t1, t2}));
+ ASSERT_OK_AND_ASSIGN(auto expected, Table::FromRecordBatches({batch1, batch2}));
+ AssertTablesEqual(*expected, *result);
+
+ // Error states
+ std::vector<std::shared_ptr<Table>> empty_tables;
+ ASSERT_RAISES(Invalid, ConcatenateTables(empty_tables));
+
+ auto other_schema = ::arrow::schema({schema_->field(0), schema_->field(1)});
+
+ std::vector<std::shared_ptr<Array>> other_arrays = {arrays_[0], arrays_[1]};
+ auto batch3 = RecordBatch::Make(other_schema, length, other_arrays);
+ ASSERT_OK_AND_ASSIGN(auto t3, Table::FromRecordBatches({batch3}));
+
+ ASSERT_RAISES(Invalid, ConcatenateTables({t1, t3}));
+}
+
+std::shared_ptr<Table> MakeTableWithOneNullFilledColumn(
+ const std::string& column_name, const std::shared_ptr<DataType>& data_type,
+ const int length) {
+ auto array_of_nulls = *MakeArrayOfNull(data_type, length);
+ return Table::Make(schema({field(column_name, data_type)}), {array_of_nulls});
+}
+
+using TestPromoteTableToSchema = TestTable;
+
+TEST_F(TestPromoteTableToSchema, IdenticalSchema) {
+ const int length = 10;
+ auto metadata =
+ std::shared_ptr<KeyValueMetadata>(new KeyValueMetadata({"foo"}, {"bar"}));
+ MakeExample1(length);
+ std::shared_ptr<Table> table = Table::Make(schema_, arrays_);
+
+ ASSERT_OK_AND_ASSIGN(auto result,
+ PromoteTableToSchema(table, schema_->WithMetadata(metadata)));
+
+ std::shared_ptr<Table> expected = table->ReplaceSchemaMetadata(metadata);
+
+ ASSERT_TRUE(result->Equals(*expected));
+}
+
+// The promoted table's fields are ordered the same as the promote-to schema.
+TEST_F(TestPromoteTableToSchema, FieldsReorderedAfterPromotion) {
+ const int length = 10;
+ MakeExample1(length);
+
+ std::vector<std::shared_ptr<Field>> reversed_fields(schema_->fields().crbegin(),
+ schema_->fields().crend());
+ std::vector<std::shared_ptr<Array>> reversed_arrays(arrays_.crbegin(), arrays_.crend());
+
+ std::shared_ptr<Table> table = Table::Make(schema(reversed_fields), reversed_arrays);
+
+ ASSERT_OK_AND_ASSIGN(auto result, PromoteTableToSchema(table, schema_));
+
+ ASSERT_TRUE(result->schema()->Equals(*schema_));
+}
+
+TEST_F(TestPromoteTableToSchema, PromoteNullTypeField) {
+ const int length = 10;
+ auto metadata =
+ std::shared_ptr<KeyValueMetadata>(new KeyValueMetadata({"foo"}, {"bar"}));
+ auto table_with_null_column = MakeTableWithOneNullFilledColumn("field", null(), length)
+ ->ReplaceSchemaMetadata(metadata);
+ auto promoted_schema = schema({field("field", int32())});
+
+ ASSERT_OK_AND_ASSIGN(auto result,
+ PromoteTableToSchema(table_with_null_column, promoted_schema));
+
+ ASSERT_TRUE(
+ result->Equals(*MakeTableWithOneNullFilledColumn("field", int32(), length)));
+}
+
+TEST_F(TestPromoteTableToSchema, AddMissingField) {
+ const int length = 10;
+ auto f0 = field("f0", int32());
+ auto table = Table::Make(schema({}), std::vector<std::shared_ptr<Array>>(), length);
+ auto promoted_schema = schema({field("field", int32())});
+
+ ASSERT_OK_AND_ASSIGN(auto result, PromoteTableToSchema(table, promoted_schema));
+
+ ASSERT_TRUE(
+ result->Equals(*MakeTableWithOneNullFilledColumn("field", int32(), length)));
+}
+
+TEST_F(TestPromoteTableToSchema, IncompatibleTypes) {
+ const int length = 10;
+ auto table = MakeTableWithOneNullFilledColumn("field", int32(), length);
+
+ // Invalid promotion: int32 to null.
+ ASSERT_RAISES(Invalid, PromoteTableToSchema(table, schema({field("field", null())})));
+
+ // Invalid promotion: int32 to uint32.
+ ASSERT_RAISES(Invalid, PromoteTableToSchema(table, schema({field("field", uint32())})));
+}
+
+TEST_F(TestPromoteTableToSchema, IncompatibleNullity) {
+ const int length = 10;
+ auto table = MakeTableWithOneNullFilledColumn("field", int32(), length);
+ ASSERT_RAISES(Invalid,
+ PromoteTableToSchema(
+ table, schema({field("field", uint32())->WithNullable(false)})));
+}
+
+TEST_F(TestPromoteTableToSchema, DuplicateFieldNames) {
+ const int length = 10;
+
+ auto table = Table::Make(
+ schema({field("field", int32()), field("field", null())}),
+ {MakeRandomArray<Int32Array>(length), MakeRandomArray<NullArray>(length)});
+
+ ASSERT_RAISES(Invalid, PromoteTableToSchema(table, schema({field("field", int32())})));
+}
+
+TEST_F(TestPromoteTableToSchema, TableFieldAbsentFromSchema) {
+ const int length = 10;
+
+ auto table =
+ Table::Make(schema({field("f0", int32())}), {MakeRandomArray<Int32Array>(length)});
+
+ std::shared_ptr<Table> result;
+ ASSERT_RAISES(Invalid, PromoteTableToSchema(table, schema({field("f1", int32())})));
+}
+
+class ConcatenateTablesWithPromotionTest : public TestTable {
+ protected:
+ ConcatenateTablesOptions GetOptions() {
+ ConcatenateTablesOptions options;
+ options.unify_schemas = true;
+ return options;
+ }
+
+ void MakeExample2(int length) {
+ auto f0 = field("f0", int32());
+ auto f1 = field("f1", null());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1};
+ schema_ = std::make_shared<Schema>(fields);
+
+ arrays_ = {MakeRandomArray<Int32Array>(length), MakeRandomArray<NullArray>(length)};
+
+ columns_ = {std::make_shared<ChunkedArray>(arrays_[0]),
+ std::make_shared<ChunkedArray>(arrays_[1])};
+ }
+
+ void AssertTablesEqualUnorderedFields(const Table& lhs, const Table& rhs) {
+ ASSERT_EQ(lhs.schema()->num_fields(), rhs.schema()->num_fields());
+ if (lhs.schema()->metadata()) {
+ ASSERT_NE(nullptr, rhs.schema()->metadata());
+ ASSERT_TRUE(lhs.schema()->metadata()->Equals(*rhs.schema()->metadata()));
+ } else {
+ ASSERT_EQ(nullptr, rhs.schema()->metadata());
+ }
+ for (int i = 0; i < lhs.schema()->num_fields(); ++i) {
+ const auto& lhs_field = lhs.schema()->field(i);
+ const auto& rhs_field = rhs.schema()->GetFieldByName(lhs_field->name());
+ ASSERT_NE(nullptr, rhs_field);
+ ASSERT_TRUE(lhs_field->Equals(rhs_field, true));
+ const auto& lhs_column = lhs.column(i);
+ const auto& rhs_column = rhs.GetColumnByName(lhs_field->name());
+ AssertChunkedEqual(*lhs_column, *rhs_column);
+ }
+ }
+};
+
+TEST_F(ConcatenateTablesWithPromotionTest, Simple) {
+ const int64_t length = 10;
+
+ MakeExample1(length);
+ auto batch1 = RecordBatch::Make(schema_, length, arrays_);
+
+ ASSERT_OK_AND_ASSIGN(auto f1_nulls, MakeArrayOfNull(schema_->field(1)->type(), length));
+ ASSERT_OK_AND_ASSIGN(auto f2_nulls, MakeArrayOfNull(schema_->field(2)->type(), length));
+
+ MakeExample2(length);
+ auto batch2 = RecordBatch::Make(schema_, length, arrays_);
+
+ auto batch2_null_filled =
+ RecordBatch::Make(batch1->schema(), length, {arrays_[0], f1_nulls, f2_nulls});
+
+ ASSERT_OK_AND_ASSIGN(auto t1, Table::FromRecordBatches({batch1}));
+ ASSERT_OK_AND_ASSIGN(auto t2, Table::FromRecordBatches({batch2}));
+ ASSERT_OK_AND_ASSIGN(auto t3, Table::FromRecordBatches({batch2_null_filled}));
+
+ ASSERT_OK_AND_ASSIGN(auto result, ConcatenateTables({t1, t2}, GetOptions()));
+ ASSERT_OK_AND_ASSIGN(auto expected, ConcatenateTables({t1, t3}));
+ AssertTablesEqualUnorderedFields(*expected, *result);
+
+ ASSERT_OK_AND_ASSIGN(result, ConcatenateTables({t2, t1}, GetOptions()));
+ ASSERT_OK_AND_ASSIGN(expected, ConcatenateTables({t3, t1}));
+ AssertTablesEqualUnorderedFields(*expected, *result);
+}
+
+TEST_F(TestTable, Slice) {
+ const int64_t length = 10;
+
+ MakeExample1(length);
+ auto batch = RecordBatch::Make(schema_, length, arrays_);
+
+ ASSERT_OK_AND_ASSIGN(auto half, Table::FromRecordBatches({batch}));
+ ASSERT_OK_AND_ASSIGN(auto whole, Table::FromRecordBatches({batch, batch}));
+ ASSERT_OK_AND_ASSIGN(auto three, Table::FromRecordBatches({batch, batch, batch}));
+
+ AssertTablesEqual(*whole->Slice(0, length), *half);
+ AssertTablesEqual(*whole->Slice(length), *half);
+ AssertTablesEqual(*whole->Slice(length / 3, 2 * (length - length / 3)),
+ *three->Slice(length + length / 3, 2 * (length - length / 3)));
+}
+
+TEST_F(TestTable, RemoveColumn) {
+ const int64_t length = 10;
+ MakeExample1(length);
+
+ auto table_sp = Table::Make(schema_, columns_);
+ const Table& table = *table_sp;
+
+ ASSERT_OK_AND_ASSIGN(auto result, table.RemoveColumn(0));
+
+ auto ex_schema = ::arrow::schema({schema_->field(1), schema_->field(2)});
+ std::vector<std::shared_ptr<ChunkedArray>> ex_columns = {table.column(1),
+ table.column(2)};
+
+ auto expected = Table::Make(ex_schema, ex_columns);
+ ASSERT_TRUE(result->Equals(*expected));
+
+ ASSERT_OK_AND_ASSIGN(result, table.RemoveColumn(1));
+ ex_schema = ::arrow::schema({schema_->field(0), schema_->field(2)});
+ ex_columns = {table.column(0), table.column(2)};
+
+ expected = Table::Make(ex_schema, ex_columns);
+ ASSERT_TRUE(result->Equals(*expected));
+
+ ASSERT_OK_AND_ASSIGN(result, table.RemoveColumn(2));
+ ex_schema = ::arrow::schema({schema_->field(0), schema_->field(1)});
+ ex_columns = {table.column(0), table.column(1)};
+ expected = Table::Make(ex_schema, ex_columns);
+ ASSERT_TRUE(result->Equals(*expected));
+}
+
+TEST_F(TestTable, SetColumn) {
+ const int64_t length = 10;
+ MakeExample1(length);
+
+ auto table_sp = Table::Make(schema_, columns_);
+ const Table& table = *table_sp;
+
+ ASSERT_OK_AND_ASSIGN(auto result,
+ table.SetColumn(0, schema_->field(1), table.column(1)));
+
+ auto ex_schema =
+ ::arrow::schema({schema_->field(1), schema_->field(1), schema_->field(2)});
+
+ auto expected =
+ Table::Make(ex_schema, {table.column(1), table.column(1), table.column(2)});
+ ASSERT_TRUE(result->Equals(*expected));
+}
+
+TEST_F(TestTable, RenameColumns) {
+ MakeExample1(10);
+ auto table = Table::Make(schema_, columns_);
+ EXPECT_THAT(table->ColumnNames(), testing::ElementsAre("f0", "f1", "f2"));
+
+ ASSERT_OK_AND_ASSIGN(auto renamed, table->RenameColumns({"zero", "one", "two"}));
+ EXPECT_THAT(renamed->ColumnNames(), testing::ElementsAre("zero", "one", "two"));
+ ASSERT_OK(renamed->ValidateFull());
+
+ ASSERT_RAISES(Invalid, table->RenameColumns({"hello", "world"}));
+}
+
+TEST_F(TestTable, SelectColumns) {
+ MakeExample1(10);
+ auto table = Table::Make(schema_, columns_);
+
+ ASSERT_OK_AND_ASSIGN(auto subset, table->SelectColumns({0, 2}));
+ ASSERT_OK(subset->ValidateFull());
+
+ auto expexted_schema = ::arrow::schema({schema_->field(0), schema_->field(2)});
+ auto expected = Table::Make(expexted_schema, {table->column(0), table->column(2)});
+ ASSERT_TRUE(subset->Equals(*expected));
+
+ // Out of bounds indices
+ ASSERT_RAISES(Invalid, table->SelectColumns({0, 3}));
+ ASSERT_RAISES(Invalid, table->SelectColumns({-1}));
+}
+
+TEST_F(TestTable, RemoveColumnEmpty) {
+ // ARROW-1865
+ const int64_t length = 10;
+
+ auto f0 = field("f0", int32());
+ auto schema = ::arrow::schema({f0});
+ auto a0 = MakeRandomArray<Int32Array>(length);
+
+ auto table = Table::Make(schema, {std::make_shared<ChunkedArray>(a0)});
+
+ ASSERT_OK_AND_ASSIGN(auto empty, table->RemoveColumn(0));
+
+ ASSERT_EQ(table->num_rows(), empty->num_rows());
+
+ ASSERT_OK_AND_ASSIGN(auto added, empty->AddColumn(0, f0, table->column(0)));
+ ASSERT_EQ(table->num_rows(), added->num_rows());
+}
+
+TEST_F(TestTable, AddColumn) {
+ const int64_t length = 10;
+ MakeExample1(length);
+
+ auto table_sp = Table::Make(schema_, columns_);
+ const Table& table = *table_sp;
+
+ auto f0 = schema_->field(0);
+
+ // Some negative tests with invalid index
+ ASSERT_RAISES(Invalid, table.AddColumn(10, f0, columns_[0]));
+ ASSERT_RAISES(Invalid, table.AddColumn(4, f0, columns_[0]));
+ ASSERT_RAISES(Invalid, table.AddColumn(-1, f0, columns_[0]));
+
+ // Add column with wrong length
+ auto longer_col =
+ std::make_shared<ChunkedArray>(MakeRandomArray<Int32Array>(length + 1));
+ ASSERT_RAISES(Invalid, table.AddColumn(0, f0, longer_col));
+
+ // Add column 0 in different places
+ ASSERT_OK_AND_ASSIGN(auto result, table.AddColumn(0, f0, columns_[0]));
+ auto ex_schema = ::arrow::schema(
+ {schema_->field(0), schema_->field(0), schema_->field(1), schema_->field(2)});
+
+ auto expected = Table::Make(
+ ex_schema, {table.column(0), table.column(0), table.column(1), table.column(2)});
+ ASSERT_TRUE(result->Equals(*expected));
+
+ ASSERT_OK_AND_ASSIGN(result, table.AddColumn(1, f0, columns_[0]));
+ ex_schema = ::arrow::schema(
+ {schema_->field(0), schema_->field(0), schema_->field(1), schema_->field(2)});
+
+ expected = Table::Make(
+ ex_schema, {table.column(0), table.column(0), table.column(1), table.column(2)});
+ ASSERT_TRUE(result->Equals(*expected));
+
+ ASSERT_OK_AND_ASSIGN(result, table.AddColumn(2, f0, columns_[0]));
+ ex_schema = ::arrow::schema(
+ {schema_->field(0), schema_->field(1), schema_->field(0), schema_->field(2)});
+ expected = Table::Make(
+ ex_schema, {table.column(0), table.column(1), table.column(0), table.column(2)});
+ ASSERT_TRUE(result->Equals(*expected));
+
+ ASSERT_OK_AND_ASSIGN(result, table.AddColumn(3, f0, columns_[0]));
+ ex_schema = ::arrow::schema(
+ {schema_->field(0), schema_->field(1), schema_->field(2), schema_->field(0)});
+ expected = Table::Make(
+ ex_schema, {table.column(0), table.column(1), table.column(2), table.column(0)});
+ ASSERT_TRUE(result->Equals(*expected));
+}
+
+class TestTableBatchReader : public TestBase {};
+
+TEST_F(TestTableBatchReader, ReadNext) {
+ ArrayVector c1, c2;
+
+ auto a1 = MakeRandomArray<Int32Array>(10);
+ auto a2 = MakeRandomArray<Int32Array>(20);
+ auto a3 = MakeRandomArray<Int32Array>(30);
+ auto a4 = MakeRandomArray<Int32Array>(10);
+
+ auto sch1 = arrow::schema({field("f1", int32()), field("f2", int32())});
+
+ std::vector<std::shared_ptr<ChunkedArray>> columns;
+
+ std::shared_ptr<RecordBatch> batch;
+
+ std::vector<std::shared_ptr<Array>> arrays_1 = {a1, a4, a2};
+ std::vector<std::shared_ptr<Array>> arrays_2 = {a2, a2};
+ columns = {std::make_shared<ChunkedArray>(arrays_1),
+ std::make_shared<ChunkedArray>(arrays_2)};
+ auto t1 = Table::Make(sch1, columns);
+
+ TableBatchReader i1(*t1);
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_EQ(10, batch->num_rows());
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_EQ(10, batch->num_rows());
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_EQ(20, batch->num_rows());
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_EQ(nullptr, batch);
+
+ arrays_1 = {a1};
+ arrays_2 = {a4};
+ columns = {std::make_shared<ChunkedArray>(arrays_1),
+ std::make_shared<ChunkedArray>(arrays_2)};
+ auto t2 = Table::Make(sch1, columns);
+
+ TableBatchReader i2(*t2);
+
+ ASSERT_OK(i2.ReadNext(&batch));
+ ASSERT_EQ(10, batch->num_rows());
+
+ // Ensure non-sliced
+ ASSERT_EQ(a1->data().get(), batch->column_data(0).get());
+ ASSERT_EQ(a4->data().get(), batch->column_data(1).get());
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_EQ(nullptr, batch);
+}
+
+TEST_F(TestTableBatchReader, Chunksize) {
+ auto a1 = MakeRandomArray<Int32Array>(10);
+ auto a2 = MakeRandomArray<Int32Array>(20);
+ auto a3 = MakeRandomArray<Int32Array>(10);
+
+ auto sch1 = arrow::schema({field("f1", int32())});
+
+ std::vector<std::shared_ptr<Array>> arrays = {a1, a2, a3};
+ auto t1 = Table::Make(sch1, {std::make_shared<ChunkedArray>(arrays)});
+
+ TableBatchReader i1(*t1);
+
+ i1.set_chunksize(15);
+
+ std::shared_ptr<RecordBatch> batch;
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_OK(batch->ValidateFull());
+ ASSERT_EQ(10, batch->num_rows());
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_OK(batch->ValidateFull());
+ ASSERT_EQ(15, batch->num_rows());
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_OK(batch->ValidateFull());
+ ASSERT_EQ(5, batch->num_rows());
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_OK(batch->ValidateFull());
+ ASSERT_EQ(10, batch->num_rows());
+
+ ASSERT_OK(i1.ReadNext(&batch));
+ ASSERT_EQ(nullptr, batch);
+}
+
+} // namespace arrow
projects / ceph.git / blobdiff