Skip to content

Commit

Permalink
Refactor column uniqueness logic
Browse files Browse the repository at this point in the history
  • Loading branch information
rubenada committed Mar 8, 2025
1 parent 802fce3 commit e22a114
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 19 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,6 @@ public Boolean areColumnsUnique(TableScan scan, RelMetadataQuery mq,

public @Nullable Boolean areColumnsUnique(Filter rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
return mq.areColumnsUnique(rel.getInput(), columns, ignoreNulls);
}

Expand Down Expand Up @@ -137,7 +136,6 @@ public Boolean areColumnsUnique(TableScan scan, RelMetadataQuery mq,

public Boolean areColumnsUnique(SetOp rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
// If not ALL then the rows are distinct.
// Therefore the set of all columns is a key.
return !rel.all
Expand All @@ -146,7 +144,6 @@ public Boolean areColumnsUnique(SetOp rel, RelMetadataQuery mq,

public Boolean areColumnsUnique(Intersect rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
if (areColumnsUnique((SetOp) rel, mq, columns, ignoreNulls)) {
return true;
}
Expand All @@ -161,7 +158,6 @@ public Boolean areColumnsUnique(Intersect rel, RelMetadataQuery mq,

public @Nullable Boolean areColumnsUnique(Minus rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
if (areColumnsUnique((SetOp) rel, mq, columns, ignoreNulls)) {
return true;
}
Expand All @@ -174,25 +170,21 @@ public Boolean areColumnsUnique(Intersect rel, RelMetadataQuery mq,
if (maxRowCount != null && maxRowCount <= 1.0d) {
return true;
}
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
return mq.areColumnsUnique(rel.getInput(), columns, ignoreNulls);
}

public @Nullable Boolean areColumnsUnique(TableModify rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
return mq.areColumnsUnique(rel.getInput(), columns, ignoreNulls);
}

public @Nullable Boolean areColumnsUnique(Exchange rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
return mq.areColumnsUnique(rel.getInput(), columns, ignoreNulls);
}

public @Nullable Boolean areColumnsUnique(Correlate rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
switch (rel.getJoinType()) {
case ANTI:
case SEMI:
Expand Down Expand Up @@ -229,7 +221,6 @@ public Boolean areColumnsUnique(Intersect rel, RelMetadataQuery mq,

public @Nullable Boolean areColumnsUnique(Project rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
// LogicalProject maps a set of rows to a different set;
// Without knowledge of the mapping function(whether it
// preserves uniqueness), it is only safe to derive uniqueness
Expand All @@ -243,7 +234,6 @@ public Boolean areColumnsUnique(Intersect rel, RelMetadataQuery mq,

public @Nullable Boolean areColumnsUnique(Calc rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
RexProgram program = rel.getProgram();

return areProjectColumnsUnique(rel, mq, columns, ignoreNulls,
Expand Down Expand Up @@ -295,8 +285,6 @@ public Boolean areColumnsUnique(Intersect rel, RelMetadataQuery mq,

public @Nullable Boolean areColumnsUnique(Join rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);

final RelNode left = rel.getLeft();
final RelNode right = rel.getRight();

Expand Down Expand Up @@ -385,7 +373,6 @@ public Boolean areColumnsUnique(Intersect rel, RelMetadataQuery mq,
return true;
}
if (Aggregate.isSimple(rel) || ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
// group by keys form a unique key
final ImmutableBitSet groupKey = ImmutableBitSet.range(rel.getGroupCount());
final boolean contained = columns.contains(groupKey);
Expand Down Expand Up @@ -430,7 +417,6 @@ public Boolean areColumnsUnique(Values rel, RelMetadataQuery mq,
if (rel.tuples.size() < 2) {
return true;
}
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
final Set<List<Comparable>> set = new HashSet<>();
final List<Comparable> values = new ArrayList<>(columns.cardinality());
for (ImmutableList<RexLiteral> tuple : rel.tuples) {
Expand All @@ -449,13 +435,11 @@ public Boolean areColumnsUnique(Values rel, RelMetadataQuery mq,

public @Nullable Boolean areColumnsUnique(Converter rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
return mq.areColumnsUnique(rel.getInput(), columns, ignoreNulls);
}

public @Nullable Boolean areColumnsUnique(RelSubset rel, RelMetadataQuery mq,
ImmutableBitSet columns, boolean ignoreNulls) {
columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
for (RelNode rel2 : rel.getRels()) {
if (rel2 instanceof Aggregate
|| rel2 instanceof Filter
Expand Down Expand Up @@ -500,7 +484,7 @@ public Boolean areColumnsUnique(Values rel, RelMetadataQuery mq,
* Deduce constant columns from predicates of rel and return the union
* bitsets of checkingColumns and the constant columns.
*/
private static ImmutableBitSet decorateWithConstantColumnsFromPredicates(
static ImmutableBitSet decorateWithConstantColumnsFromPredicates(
ImmutableBitSet checkingColumns, RelNode rel, RelMetadataQuery mq) {
final RelOptPredicateList predicates = mq.getPulledUpPredicates(rel);
if (!RelOptPredicateList.isEmpty(predicates)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
import java.util.function.Supplier;

import static org.apache.calcite.linq4j.Nullness.castNonNull;
import static org.apache.calcite.rel.metadata.RelMdColumnUniqueness.decorateWithConstantColumnsFromPredicates;

/**
* RelMetadataQuery provides a strongly-typed facade on top of
Expand Down Expand Up @@ -605,8 +606,19 @@ public static RelMetadataQuery instance() {
boolean ignoreNulls) {
for (;;) {
try {
return columnUniquenessHandler.areColumnsUnique(rel, this, columns,
ignoreNulls);
Boolean b = columnUniquenessHandler.areColumnsUnique(rel, this, columns, ignoreNulls);
if (b != null && b) {
return true;
}
// Second attempt: try adding the constant columns deduced from the pulled up predicates
ImmutableBitSet decorated = decorateWithConstantColumnsFromPredicates(columns, rel, this);
if (decorated != columns) {
Boolean b2 = columnUniquenessHandler.areColumnsUnique(rel, this, decorated, ignoreNulls);
if (b2 != null && b2) {
return true;
}
}
return b;
} catch (MetadataHandlerProvider.NoHandler e) {
columnUniquenessHandler = revise(BuiltInMetadata.ColumnUniqueness.Handler.class);
}
Expand Down

0 comments on commit e22a114

Please sign in to comment.