Skip to content

Commit 5e95651

Browse files
ivanmorozov333ivanmorozov333
authored
Skip default columns on indexation (#6899)
1 parent 1c15550 commit 5e95651

File tree

36 files changed

+424
-173
lines changed

36 files changed

+424
-173
lines changed

‎ydb/core/formats/arrow/arrow_helpers.cpp‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -548,7 +548,7 @@ std::shared_ptr<arrow::Scalar> DefaultScalar(const std::shared_ptr<arrow::DataTy
548548
}
549549
return true;
550550
});
551-
Y_ABORT_UNLESS(out);
551+
AFL_VERIFY(out)("type", type->ToString());
552552
return out;
553553
}
554554

@@ -634,6 +634,19 @@ int ScalarCompare(const std::shared_ptr<arrow::Scalar>& x, const std::shared_ptr
634634
return ScalarCompare(*x, *y);
635635
}
636636

637+
int ScalarCompareNullable(const std::shared_ptr<arrow::Scalar>& x, const std::shared_ptr<arrow::Scalar>& y) {
638+
if (!x && !!y) {
639+
return -1;
640+
}
641+
if (!!x && !y) {
642+
return 1;
643+
}
644+
if (!x && !y) {
645+
return 0;
646+
}
647+
return ScalarCompare(*x, *y);
648+
}
649+
637650
std::shared_ptr<arrow::RecordBatch> SortBatch(const std::shared_ptr<arrow::RecordBatch>& batch,
638651
const std::shared_ptr<arrow::Schema>& sortingKey, const bool andUnique) {
639652
auto sortPermutation = MakeSortPermutation(batch, sortingKey, andUnique);

‎ydb/core/formats/arrow/arrow_helpers.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ std::shared_ptr<arrow::Scalar> GetScalar(const std::shared_ptr<arrow::Array>& ar
9898
bool IsGoodScalar(const std::shared_ptr<arrow::Scalar>& x);
9999
int ScalarCompare(const arrow::Scalar& x, const arrow::Scalar& y);
100100
int ScalarCompare(const std::shared_ptr<arrow::Scalar>& x, const std::shared_ptr<arrow::Scalar>& y);
101+
int ScalarCompareNullable(const std::shared_ptr<arrow::Scalar>& x, const std::shared_ptr<arrow::Scalar>& y);
101102
std::partial_ordering ColumnsCompare(const std::vector<std::shared_ptr<arrow::Array>>& x, const ui32 xRow, const std::vector<std::shared_ptr<arrow::Array>>& y, const ui32 yRow);
102103
bool ScalarLess(const std::shared_ptr<arrow::Scalar>& x, const std::shared_ptr<arrow::Scalar>& y);
103104
bool ScalarLess(const arrow::Scalar& x, const arrow::Scalar& y);

‎ydb/core/formats/arrow/modifier/subset.cpp‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
#include "subset.h"
2+
#include <ydb/library/actors/core/log.h>
3+
4+
namespace NKikimr::NArrow {
5+
6+
TSchemaSubset::TSchemaSubset(const std::set<ui32>& fieldsIdx, const ui32 fieldsCount) {
7+
AFL_VERIFY(fieldsIdx.size() <= fieldsCount);
8+
AFL_VERIFY(fieldsIdx.size());
9+
if (fieldsCount == fieldsIdx.size()) {
10+
return;
11+
}
12+
Exclude = (fieldsCount - fieldsIdx.size()) < fieldsIdx.size();
13+
if (!Exclude) {
14+
FieldIdx = std::vector<ui32>(fieldsIdx.begin(), fieldsIdx.end());
15+
} else {
16+
auto it = fieldsIdx.begin();
17+
for (ui32 i = 0; i < fieldsCount; ++i) {
18+
if (it == fieldsIdx.end() || i < *it) {
19+
FieldIdx.emplace_back(i);
20+
} else if (*it == i) {
21+
++it;
22+
} else {
23+
AFL_VERIFY(false);
24+
}
25+
}
26+
}
27+
}
28+
29+
NKikimrArrowSchema::TSchemaSubset TSchemaSubset::SerializeToProto() const {
30+
NKikimrArrowSchema::TSchemaSubset result;
31+
result.MutableList()->SetExclude(Exclude);
32+
for (auto&& i : FieldIdx) {
33+
result.MutableList()->AddFieldsIdx(i);
34+
}
35+
return result;
36+
}
37+
38+
TConclusionStatus TSchemaSubset::DeserializeFromProto(const NKikimrArrowSchema::TSchemaSubset& proto) {
39+
if (!proto.HasList()) {
40+
return TConclusionStatus::Fail("no schema subset data");
41+
}
42+
Exclude = proto.GetList().GetExclude();
43+
std::vector<ui32> fieldIdx;
44+
for (auto&& i : proto.GetList().GetFieldsIdx()) {
45+
fieldIdx.emplace_back(i);
46+
}
47+
std::swap(fieldIdx, FieldIdx);
48+
return TConclusionStatus::Success();
49+
}
50+
51+
}

‎ydb/core/formats/arrow/modifier/subset.h‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
#pragma once
2+
#include <ydb/core/formats/arrow/protos/fields.pb.h>
3+
#include <ydb/library/conclusion/result.h>
4+
#include <ydb/library/actors/core/log.h>
5+
6+
namespace NKikimr::NArrow {
7+
8+
class TSchemaSubset {
9+
private:
10+
std::vector<ui32> FieldIdx;
11+
TString FieldBits;
12+
bool Exclude = false;
13+
public:
14+
TSchemaSubset() = default;
15+
TSchemaSubset(const std::set<ui32>& fieldsIdx, const ui32 fieldsCount);
16+
17+
template <class T>
18+
std::vector<T> Apply(const std::vector<T>& fullSchema) const {
19+
if (FieldIdx.empty()) {
20+
return fullSchema;
21+
}
22+
std::vector<T> fields;
23+
if (!Exclude) {
24+
for (auto&& i : FieldIdx) {
25+
AFL_VERIFY(i < fullSchema.size());
26+
fields.emplace_back(fullSchema[i]);
27+
}
28+
} else {
29+
auto it = FieldIdx.begin();
30+
for (ui32 i = 0; i < fullSchema.size(); ++i) {
31+
if (it == FieldIdx.end() || i < *it) {
32+
AFL_VERIFY(i < fullSchema.size());
33+
fields.emplace_back(fullSchema[i]);
34+
} else if (i == *it) {
35+
++it;
36+
} else {
37+
AFL_VERIFY(false);
38+
}
39+
}
40+
}
41+
return fields;
42+
}
43+
44+
NKikimrArrowSchema::TSchemaSubset SerializeToProto() const;
45+
[[nodiscard]] TConclusionStatus DeserializeFromProto(const NKikimrArrowSchema::TSchemaSubset& proto);
46+
};
47+
48+
}

‎ydb/core/formats/arrow/modifier/ya.make‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,13 @@ PEERDIR(
44
contrib/libs/apache/arrow
55
ydb/library/conclusion
66
ydb/core/formats/arrow/switch
7+
ydb/core/formats/arrow/protos
78
ydb/library/actors/core
89
)
910

1011
SRCS(
1112
schema.cpp
13+
subset.cpp
1214
)
1315

1416
END()

‎ydb/core/formats/arrow/process_columns.cpp‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "process_columns.h"
22
#include "common/adapter.h"
3+
#include "modifier/subset.h"
34

45
#include <util/string/join.h>
56

@@ -28,16 +29,23 @@ std::shared_ptr<TDataContainer> ExtractColumnsValidateImpl(const std::shared_ptr
2829

2930
template <class TDataContainer>
3031
TConclusion<std::shared_ptr<TDataContainer>> AdaptColumnsImpl(const std::shared_ptr<TDataContainer>& srcBatch,
31-
const std::shared_ptr<arrow::Schema>& dstSchema) {
32+
const std::shared_ptr<arrow::Schema>& dstSchema, TSchemaSubset* subset) {
3233
AFL_VERIFY(srcBatch);
3334
AFL_VERIFY(dstSchema);
3435
std::vector<std::shared_ptr<typename NAdapter::TDataBuilderPolicy<TDataContainer>::TColumn>> columns;
3536
columns.reserve(dstSchema->num_fields());
36-
37+
std::vector<std::shared_ptr<arrow::Field>> fields;
38+
fields.reserve(dstSchema->num_fields());
39+
std::set<ui32> fieldIdx;
40+
ui32 idx = 0;
3741
for (auto& field : dstSchema->fields()) {
3842
const int index = srcBatch->schema()->GetFieldIndex(field->name());
3943
if (index > -1) {
44+
if (subset) {
45+
fieldIdx.emplace(idx);
46+
}
4047
columns.push_back(srcBatch->column(index));
48+
fields.emplace_back(field);
4149
auto srcField = srcBatch->schema()->field(index);
4250
if (field->Equals(srcField)) {
4351
AFL_VERIFY(columns.back()->type()->Equals(field->type()))("event", "cannot_use_incoming_batch")("reason", "invalid_column_type")("column", field->name())
@@ -47,14 +55,17 @@ TConclusion<std::shared_ptr<TDataContainer>> AdaptColumnsImpl(const std::shared_
4755
("column_type", field->ToString(true))("incoming_type", srcField->ToString(true));
4856
return TConclusionStatus::Fail("incompatible column types");
4957
}
50-
} else {
58+
} else if (!subset) {
5159
AFL_ERROR(NKikimrServices::ARROW_HELPER)("event", "not_found_column")("column", field->name())
5260
("column_type", field->type()->ToString())("columns", JoinSeq(",", srcBatch->schema()->field_names()));
5361
return TConclusionStatus::Fail("not found column '" + field->name() + "'");
5462
}
63+
++idx;
5564
}
56-
57-
return NAdapter::TDataBuilderPolicy<TDataContainer>::Build(dstSchema, std::move(columns), srcBatch->num_rows());
65+
if (subset) {
66+
*subset = TSchemaSubset(fieldIdx, dstSchema->num_fields());
67+
}
68+
return NAdapter::TDataBuilderPolicy<TDataContainer>::Build(std::make_shared<arrow::Schema>(fields), std::move(columns), srcBatch->num_rows());
5869
}
5970

6071
template <class TDataContainer, class TStringType>
@@ -114,12 +125,12 @@ std::shared_ptr<arrow::Table> TColumnOperator::Extract(const std::shared_ptr<arr
114125
return ExtractImpl(AbsentColumnPolicy, incoming, columnNames);
115126
}
116127

117-
NKikimr::TConclusion<std::shared_ptr<arrow::RecordBatch>> TColumnOperator::Adapt(const std::shared_ptr<arrow::RecordBatch>& incoming, const std::shared_ptr<arrow::Schema>& dstSchema) {
118-
return AdaptColumnsImpl(incoming, dstSchema);
128+
NKikimr::TConclusion<std::shared_ptr<arrow::RecordBatch>> TColumnOperator::Adapt(const std::shared_ptr<arrow::RecordBatch>& incoming, const std::shared_ptr<arrow::Schema>& dstSchema, TSchemaSubset* subset) {
129+
return AdaptColumnsImpl(incoming, dstSchema, subset);
119130
}
120131

121-
NKikimr::TConclusion<std::shared_ptr<arrow::Table>> TColumnOperator::Adapt(const std::shared_ptr<arrow::Table>& incoming, const std::shared_ptr<arrow::Schema>& dstSchema) {
122-
return AdaptColumnsImpl(incoming, dstSchema);
132+
NKikimr::TConclusion<std::shared_ptr<arrow::Table>> TColumnOperator::Adapt(const std::shared_ptr<arrow::Table>& incoming, const std::shared_ptr<arrow::Schema>& dstSchema, TSchemaSubset* subset) {
133+
return AdaptColumnsImpl(incoming, dstSchema, subset);
123134
}
124135

125136
NKikimr::TConclusion<std::shared_ptr<arrow::RecordBatch>> TColumnOperator::Reorder(const std::shared_ptr<arrow::RecordBatch>& incoming, const std::vector<std::string>& columnNames) {

‎ydb/core/formats/arrow/process_columns.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
namespace NKikimr::NArrow {
77

8+
class TSchemaSubset;
9+
810
class TColumnOperator {
911
public:
1012
enum class EExtractProblemsPolicy {
@@ -36,8 +38,8 @@ class TColumnOperator {
3638
std::shared_ptr<arrow::RecordBatch> Extract(const std::shared_ptr<arrow::RecordBatch>& incoming, const std::vector<TString>& columnNames);
3739
std::shared_ptr<arrow::Table> Extract(const std::shared_ptr<arrow::Table>& incoming, const std::vector<TString>& columnNames);
3840

39-
TConclusion<std::shared_ptr<arrow::RecordBatch>> Adapt(const std::shared_ptr<arrow::RecordBatch>& incoming, const std::shared_ptr<arrow::Schema>& dstSchema);
40-
TConclusion<std::shared_ptr<arrow::Table>> Adapt(const std::shared_ptr<arrow::Table>& incoming, const std::shared_ptr<arrow::Schema>& dstSchema);
41+
TConclusion<std::shared_ptr<arrow::RecordBatch>> Adapt(const std::shared_ptr<arrow::RecordBatch>& incoming, const std::shared_ptr<arrow::Schema>& dstSchema, TSchemaSubset* subset = nullptr);
42+
TConclusion<std::shared_ptr<arrow::Table>> Adapt(const std::shared_ptr<arrow::Table>& incoming, const std::shared_ptr<arrow::Schema>& dstSchema, TSchemaSubset* subset = nullptr);
4143

4244
TConclusion<std::shared_ptr<arrow::RecordBatch>> Reorder(const std::shared_ptr<arrow::RecordBatch>& incoming, const std::vector<std::string>& columnNames);
4345
TConclusion<std::shared_ptr<arrow::Table>> Reorder(const std::shared_ptr<arrow::Table>& incoming, const std::vector<std::string>& columnNames);

‎ydb/core/formats/arrow/protos/fields.proto‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package NKikimrArrowSchema;
2+
3+
message TSchemaSubset {
4+
5+
message TFieldsList {
6+
optional bool Exclude = 1;
7+
repeated uint32 FieldsIdx = 2;
8+
}
9+
10+
oneof Implementation {
11+
TFieldsList List = 1;
12+
string Bits = 2;
13+
14+
}
15+
}

‎ydb/core/formats/arrow/protos/ya.make‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ PROTO_LIBRARY()
22

33
SRCS(
44
ssa.proto
5+
fields.proto
56
)
67

78
PEERDIR(

‎ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3087,6 +3087,13 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
30873087
UNIT_ASSERT_C(it.IsSuccess(), it.GetIssues().ToString());
30883088
}
30893089

3090+
{
3091+
auto it = client.ExecuteQuery(R"(
3092+
UPSERT INTO `/Root/DataShard` (Col3) VALUES ('null');
3093+
)", NYdb::NQuery::TTxControl::BeginTx().CommitTx()).ExtractValueSync();
3094+
UNIT_ASSERT(!it.IsSuccess());
3095+
}
3096+
30903097
{
30913098
auto it = client.StreamExecuteQuery(R"(
30923099
SELECT * FROM `/Root/DataShard` ORDER BY Col1;

0 commit comments

Comments
 (0)