Skip to content

Commit 84f194f

Browse files
petrmikheevcopybara-github
authored andcommitted
Return error in DataBag::ImmutableEmptyWithFallbacks if fallbacks are mutable.
Some old usages with mutable fallbacks (including those which explicitly test mutable fallbacks) are switched to a temporary function `DataBag::ImmutableEmptyWithDeprecatedMutableFallbacks`. PiperOrigin-RevId: 862679088 Change-Id: Ibf69f695f5e6951e627e86359bc89ced83da7389
1 parent 2fd4d39 commit 84f194f

18 files changed

+193
-109
lines changed

koladata/adoption_utils.cc

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -119,22 +119,25 @@ absl_nonnull DataBagPtr AdoptionQueue::GetBagWithFallbacks() const {
119119
std::vector<DataBagPtr> fallbacks;
120120
fallbacks.reserve(slices_to_merge_.size() + bags_to_merge_.size());
121121
for (const DataBagPtr& bag : bags_to_merge_) {
122-
if (visited_bags.contains(bag.get())) {
122+
if (bag.get() == nullptr || visited_bags.contains(bag.get())) {
123123
continue;
124124
}
125125
visited_bags.insert(bag.get());
126-
fallbacks.push_back(bag);
126+
fallbacks.push_back(bag->Freeze());
127127
}
128128
for (const DataSlice& slice : slices_to_merge_) {
129129
const DataBagPtr& bag = slice.GetBag();
130-
if (visited_bags.contains(bag.get())) {
130+
if (bag.get() == nullptr || visited_bags.contains(bag.get())) {
131131
continue;
132132
}
133133
visited_bags.insert(bag.get());
134-
fallbacks.push_back(bag);
134+
fallbacks.push_back(bag->Freeze());
135135
}
136136

137-
return DataBag::ImmutableEmptyWithFallbacks(fallbacks);
137+
// Note: ImmutableEmptyWithFallbacks returns an error if fallbacks are
138+
// mutable. Here we explicitly freeze all fallbacks, so we assume that
139+
// the error can not happen.
140+
return *DataBag::ImmutableEmptyWithFallbacks(fallbacks);
138141
}
139142

140143
absl::Status AdoptStub(const DataBagPtr& db, const DataSlice& x) {
@@ -214,7 +217,8 @@ absl::StatusOr<absl_nullable DataBagPtr> WithAdoptedValues(
214217
ASSIGN_OR_RETURN(DataSlice extracted_slice,
215218
extract_utils_internal::Extract(slice));
216219
// NOTE: slices's bag should come first to respect its precedence.
217-
return DataBag::ImmutableEmptyWithFallbacks({extracted_slice.GetBag(), db});
220+
return DataBag::ImmutableEmptyWithDeprecatedMutableFallbacks(
221+
{extracted_slice.GetBag(), db});
218222
}
219223
}
220224

koladata/adoption_utils_test.cc

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -150,8 +150,10 @@ TEST(AdoptionQueueTest, WithFallbacks) {
150150
db3->GetMutableImpl()->SetAttr(obj, "b", internal::DataItem(4.0f)));
151151
ASSERT_OK(db3->GetMutableImpl()->SetAttr(obj, "c", internal::DataItem(3.0f)));
152152

153-
DataBagPtr db_with_fallbacks =
154-
DataBag::ImmutableEmptyWithFallbacks({db2, db3, schema_db});
153+
ASSERT_OK_AND_ASSIGN(
154+
DataBagPtr db_with_fallbacks,
155+
DataBag::ImmutableEmptyWithFallbacks(
156+
{db2->Freeze(), db3->Freeze(), schema_db->Freeze()}));
155157

156158
ASSERT_OK_AND_ASSIGN(
157159
DataSlice slice1,
@@ -180,9 +182,9 @@ TEST(AdoptionQueueTest, WithFallbacks) {
180182
}
181183

182184
TEST(AdoptionQueueTest, TestDbVector) {
183-
auto db1 = DataBag::EmptyMutable();
185+
auto db1 = DataBag::Empty();
184186
DataBagPtr db2;
185-
auto db3 = DataBag::EmptyMutable();
187+
auto db3 = DataBag::Empty();
186188
AdoptionQueue q;
187189
q.Add(db1);
188190
q.Add(db2);

koladata/casting.cc

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,7 @@ bool IsProbablyCastableTo(const internal::DataItem& from_schema,
168168

169169
if (to_schema.holds_value<internal::ObjectId>()) {
170170
// Note: If to_schema and from_schema are both entities, they would be
171-
// traversed futher. If from_schema is Object, compatibility can only be
171+
// traversed further. If from_schema is Object, compatibility can only be
172172
// checked based on data stored.
173173
return from_schema.holds_value<internal::ObjectId>() ||
174174
from_schema == schema::kObject || from_schema == schema::kNone;
@@ -582,10 +582,10 @@ absl::StatusOr<SchemaAlignedSlices> AlignSchemas(
582582
[](const DataSlice& ds) { return ds.GetBag(); });
583583
return ret;
584584
};
585-
ASSIGN_OR_RETURN(
586-
auto common_schema, get_common_schema(),
587-
KodaErrorCausedByNoCommonSchemaError(
588-
_, DataBag::ImmutableEmptyWithFallbacks(get_fallback_db())));
585+
ASSIGN_OR_RETURN(auto common_schema, get_common_schema(),
586+
KodaErrorCausedByNoCommonSchemaError(
587+
_, DataBag::ImmutableEmptyWithDeprecatedMutableFallbacks(
588+
get_fallback_db())));
589589
for (auto& slice : slices) {
590590
// Since we cast to a common schema, we don't need to validate implicit
591591
// compatibility or validate schema (during casting to OBJECT) as no

koladata/casting_test.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1740,7 +1740,7 @@ TEST(Casting, ToObjectEmbedding) {
17401740

17411741
TEST(Casting, ToObjectImmutableBagWithPrimitives) {
17421742
// A primitive slice with an immutable bag should succeed.
1743-
auto db = DataBag::ImmutableEmptyWithFallbacks({});
1743+
auto db = DataBag::ImmutableEmptyWithDeprecatedMutableFallbacks({});
17441744
auto slice = test::DataSlice<int>({1, 2, std::nullopt}, schema::kInt32, db);
17451745
EXPECT_THAT(ToObject(slice),
17461746
IsOkAndHolds(IsEquivalentTo(test::DataSlice<int>(
@@ -1753,7 +1753,7 @@ TEST(Casting, ToObjectErrors) {
17531753
3, internal::DataItem(internal::AllocateExplicitSchema()));
17541754
{
17551755
// Immutable bag.
1756-
auto db = DataBag::ImmutableEmptyWithFallbacks({});
1756+
auto db = DataBag::Empty();
17571757
ASSERT_OK_AND_ASSIGN(
17581758
auto slice,
17591759
DataSlice::Create(object_slice, DataSlice::JaggedShape::FlatFromSize(3),

koladata/check_frozen_test.cc

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,8 @@ TEST(CheckFrozen, Databag) {
4343
}
4444

4545
TEST(CheckFrozen, DatabagWithMutableFallbacks) {
46-
auto db = DataBag::ImmutableEmptyWithFallbacks({DataBag::EmptyMutable()});
46+
auto db = DataBag::ImmutableEmptyWithDeprecatedMutableFallbacks(
47+
{DataBag::EmptyMutable()});
4748
EXPECT_FALSE(db->IsMutable());
4849
EXPECT_THAT(CheckFrozen(arolla::TypedRef::FromValue(db)),
4950
StatusIs(absl::StatusCode::kInvalidArgument,

koladata/data_bag.cc

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include <vector>
2222

2323
#include "absl/base/no_destructor.h"
24+
#include "absl/base/nullability.h"
2425
#include "absl/cleanup/cleanup.h"
2526
#include "absl/container/flat_hash_set.h"
2627
#include "absl/functional/function_ref.h"
@@ -41,15 +42,32 @@
4142

4243
namespace koladata {
4344

44-
DataBagPtr DataBag::ImmutableEmptyWithFallbacks(
45-
absl::Span<const DataBagPtr> fallbacks) {
45+
absl::StatusOr<DataBagPtr> DataBag::ImmutableEmptyWithFallbacks(
46+
absl::Span<const DataBagPtr absl_nullable> fallbacks) {
4647
auto res = DataBagPtr::Make(DataBag::immutable_t());
4748
std::vector<DataBagPtr> non_null_fallbacks;
4849
non_null_fallbacks.reserve(fallbacks.size());
49-
for (int i = 0; i < fallbacks.size(); ++i) {
50-
if (fallbacks[i] != nullptr) {
51-
non_null_fallbacks.push_back(fallbacks[i]);
52-
if (fallbacks[i]->IsMutable() || fallbacks[i]->HasMutableFallbacks()) {
50+
for (const DataBagPtr& fb : fallbacks) {
51+
if (fb != nullptr) {
52+
if (fb->IsMutable() || fb->HasMutableFallbacks()) {
53+
return absl::InvalidArgumentError("fallbacks must be immutable");
54+
}
55+
non_null_fallbacks.push_back(fb);
56+
}
57+
}
58+
res->fallbacks_ = std::move(non_null_fallbacks);
59+
return res;
60+
}
61+
62+
DataBagPtr DataBag::ImmutableEmptyWithDeprecatedMutableFallbacks(
63+
absl::Span<const DataBagPtr absl_nullable> fallbacks) {
64+
auto res = DataBagPtr::Make(DataBag::immutable_t());
65+
std::vector<DataBagPtr> non_null_fallbacks;
66+
non_null_fallbacks.reserve(fallbacks.size());
67+
for (const DataBagPtr& fb : fallbacks) {
68+
if (fb != nullptr) {
69+
non_null_fallbacks.push_back(fb);
70+
if (fb->IsMutable() || fb->HasMutableFallbacks()) {
5371
res->has_mutable_fallbacks_ = true;
5472
}
5573
}
@@ -109,7 +127,11 @@ DataBagPtr DataBag::FreezeWithFallbacks() {
109127
}
110128
}
111129
});
112-
return DataBag::ImmutableEmptyWithFallbacks(std::move(leaf_fallbacks));
130+
// Note: ImmutableEmptyWithFallbacks returns an error if fallbacks are
131+
// mutable. Here we explicitly freeze all fallbacks, so we assume that
132+
// the error can not happen.
133+
return DataBag::ImmutableEmptyWithFallbacks(std::move(leaf_fallbacks))
134+
.value();
113135
}
114136

115137
DataBagPtr DataBag::Freeze() {
@@ -146,7 +168,7 @@ DataBagPtr DataBag::CommonDataBag(absl::Span<const DataBagPtr> databags) {
146168
if (non_null_databags.empty()) {
147169
return nullptr;
148170
}
149-
return ImmutableEmptyWithFallbacks(non_null_databags);
171+
return ImmutableEmptyWithDeprecatedMutableFallbacks(non_null_databags);
150172
}
151173

152174
DataBagPtr DataBag::FromImpl(internal::DataBagImplPtr impl) {

koladata/data_bag.h

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -121,8 +121,14 @@ class DataBag : public arolla::RefcountedBase {
121121
bool HasMutableFallbacks() const { return has_mutable_fallbacks_; }
122122

123123
// Returns a newly created immutable DataBag with fallbacks.
124-
static DataBagPtr ImmutableEmptyWithFallbacks(
125-
absl::Span<const DataBagPtr> fallbacks);
124+
// An error is returned if fallbacks are mutable.
125+
static absl::StatusOr<DataBagPtr> ImmutableEmptyWithFallbacks(
126+
absl::Span<const DataBagPtr absl_nullable> fallbacks);
127+
128+
// Deprecated. Use ImmutableEmptyWithFallbacks instead.
129+
[[deprecated("Use ImmutableEmptyWithFallbacks")]]
130+
static DataBagPtr ImmutableEmptyWithDeprecatedMutableFallbacks(
131+
absl::Span<const DataBagPtr absl_nullable> fallbacks);
126132

127133
// Returns a DataBag that contains all the data its input contain.
128134
// * If they are all the same or only 1 DataBag is non-nullptr, that DataBag

koladata/data_bag_comparison_test.cc

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -50,14 +50,22 @@ TEST(DataBagComparisonTest, ExactlyEqual_Fallbacks) {
5050
auto db2 = DataBag::EmptyMutable();
5151
ASSERT_OK_AND_ASSIGN(internal::DataBagImpl& db2_impl, db2->GetMutableImpl());
5252
ASSERT_OK(db2_impl.SetAttr(ds2, "other", ds1));
53+
db1->UnsafeMakeImmutable();
54+
db2->UnsafeMakeImmutable();
5355

54-
auto db_f1 = DataBag::ImmutableEmptyWithFallbacks({db1});
55-
auto db_f12 = DataBag::ImmutableEmptyWithFallbacks({db1, db2});
56-
auto db_f12_copy = DataBag::ImmutableEmptyWithFallbacks({db1, db2});
57-
auto db_f21 = DataBag::ImmutableEmptyWithFallbacks({db2, db1});
58-
auto db_ff12 = DataBag::ImmutableEmptyWithFallbacks({db_f1, db2});
59-
auto db_ff122 = DataBag::ImmutableEmptyWithFallbacks({db_f12, db2});
60-
auto db_ff212 = DataBag::ImmutableEmptyWithFallbacks({db_f21, db2});
56+
ASSERT_OK_AND_ASSIGN(auto db_f1, DataBag::ImmutableEmptyWithFallbacks({db1}));
57+
ASSERT_OK_AND_ASSIGN(auto db_f12,
58+
DataBag::ImmutableEmptyWithFallbacks({db1, db2}));
59+
ASSERT_OK_AND_ASSIGN(auto db_f12_copy,
60+
DataBag::ImmutableEmptyWithFallbacks({db1, db2}));
61+
ASSERT_OK_AND_ASSIGN(auto db_f21,
62+
DataBag::ImmutableEmptyWithFallbacks({db2, db1}));
63+
ASSERT_OK_AND_ASSIGN(auto db_ff12,
64+
DataBag::ImmutableEmptyWithFallbacks({db_f1, db2}));
65+
ASSERT_OK_AND_ASSIGN(auto db_ff122,
66+
DataBag::ImmutableEmptyWithFallbacks({db_f12, db2}));
67+
ASSERT_OK_AND_ASSIGN(auto db_ff212,
68+
DataBag::ImmutableEmptyWithFallbacks({db_f21, db2}));
6169

6270
EXPECT_TRUE(DataBagComparison::ExactlyEqual(db_f12, db_f12_copy));
6371
EXPECT_FALSE(DataBagComparison::ExactlyEqual(db_f12, db_f21));

koladata/data_bag_repr_test.cc

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,9 @@ TEST(DataBagReprTest,
191191
/*key_schema=*/absl::nullopt, /*value_schema=*/absl::nullopt,
192192
/*item_id=*/uuid));
193193

194-
auto db = DataBag::ImmutableEmptyWithFallbacks({fallback_db1, fallback_db2});
194+
auto db = DataBag::ImmutableEmptyWithFallbacks(
195+
{fallback_db1->Freeze(), fallback_db2->Freeze()})
196+
.value();
195197

196198
// 'a' and 'b' from fallback_db1, and 'c' from fallback_db2.
197199
EXPECT_THAT(DataBagToStr(db), IsOkAndHolds(MatchesRegex(
@@ -248,7 +250,9 @@ TEST(DataBagReprTest,
248250
/*item_schema=*/absl::nullopt,
249251
/*item_id=*/uuid));
250252

251-
auto db = DataBag::ImmutableEmptyWithFallbacks({fallback_db1, fallback_db2});
253+
auto db = DataBag::ImmutableEmptyWithFallbacks(
254+
{fallback_db1->Freeze(), fallback_db2->Freeze()})
255+
.value();
252256

253257
// 'a' and 'b' from fallback_db1, and 'c' from fallback_db2.
254258
EXPECT_THAT(DataBagToStr(db), IsOkAndHolds(MatchesRegex(
@@ -272,7 +276,9 @@ TEST(DataBagReprTest, TestDataBagStringRepresentation_FallbackBags) {
272276
auto ds2, EntityCreator::FromAttrs(fallback_db2, {"b"},
273277
{test::DataItem(123, fallback_db2)}));
274278

275-
auto db = DataBag::ImmutableEmptyWithFallbacks({fallback_db1, fallback_db2});
279+
auto db = DataBag::ImmutableEmptyWithFallbacks(
280+
{fallback_db1->Freeze(), fallback_db2->Freeze()})
281+
.value();
276282
auto ds3 = ds1.WithBag(db);
277283

278284
EXPECT_THAT(
@@ -392,7 +398,9 @@ TEST(DataBagReprTest,
392398
auto ds2, EntityCreator::FromAttrs(fallback_db2, {"b"},
393399
{test::DataItem(123, fallback_db2)}));
394400

395-
auto db = DataBag::ImmutableEmptyWithFallbacks({fallback_db1, fallback_db2});
401+
auto db = DataBag::ImmutableEmptyWithFallbacks(
402+
{fallback_db1->Freeze(), fallback_db2->Freeze()})
403+
.value();
396404
auto ds3 = ds1.WithBag(db);
397405

398406
EXPECT_THAT(
@@ -457,7 +465,9 @@ TEST(DataBagReprTest, TestDataBagStringRepresentation_DuplicatedFallbackBags) {
457465
auto fallback_db = DataBag::EmptyMutable();
458466
ASSERT_OK_AND_ASSIGN(auto ds1, EntityCreator::FromAttrs(
459467
fallback_db, {"a"}, {test::DataItem(42)}));
460-
auto db = DataBag::ImmutableEmptyWithFallbacks({fallback_db, fallback_db});
468+
fallback_db->UnsafeMakeImmutable();
469+
auto db =
470+
DataBag::ImmutableEmptyWithFallbacks({fallback_db, fallback_db}).value();
461471
EXPECT_THAT(
462472
DataBagToStr(db),
463473
IsOkAndHolds(MatchesRegex(
@@ -487,7 +497,9 @@ TEST(DataBagReprTest,
487497
ASSERT_OK(ds2.SetAttr("b", test::DataItem(20, fallback_db2)));
488498
ASSERT_OK(ds2.SetAttr("c", test::DataItem(30, fallback_db2)));
489499

490-
auto db = DataBag::ImmutableEmptyWithFallbacks({fallback_db1, fallback_db2});
500+
auto db = DataBag::ImmutableEmptyWithFallbacks(
501+
{fallback_db1->Freeze(), fallback_db2->Freeze()})
502+
.value();
491503

492504
EXPECT_THAT(DataBagToStr(db), IsOkAndHolds(MatchesRegex(
493505
R"regex(DataBag \$[0-9a-f]{4}:

0 commit comments

Comments
 (0)