Skip to content

NodePartitioningManager code cleanup #10293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
findepi merged 7 commits into from
Dec 16, 2021
Merged

NodePartitioningManager code cleanup #10293

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c93b811
Reuse existing variable
findepi Dec 14, 2021
2b253fc
Remove redundant check
findepi Dec 14, 2021
d0b9760
Unwrap from Optional that must bear a value
findepi Dec 14, 2021
0db883f
Rename argument for consistency
findepi Dec 14, 2021
dbb229c
Avoid Optional.get without a check
findepi Dec 14, 2021
bef6667
Remove redundant null-friendliness
findepi Dec 14, 2021
b00716c
empty
findepi Dec 15, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 30 additions & 24 deletions core/trino-main/src/main/java/io/trino/sql/planner/NodePartitioningManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,45 +82,47 @@ public PartitionFunction getPartitionFunction(
PartitioningScheme partitioningScheme,
List<Type> partitionChannelTypes)
{
Optional<int[]> bucketToPartition = partitioningScheme.getBucketToPartition();
checkArgument(bucketToPartition.isPresent(), "Bucket to partition must be set before a partition function can be created");
int[] bucketToPartition = partitioningScheme.getBucketToPartition()
.orElseThrow(() -> new IllegalArgumentException("Bucket to partition must be set before a partition function can be created"));

PartitioningHandle partitioningHandle = partitioningScheme.getPartitioning().getHandle();
if (partitioningHandle.getConnectorHandle() instanceof SystemPartitioningHandle) {
checkArgument(partitioningScheme.getBucketToPartition().isPresent(), "Bucket to partition must be set before a partition function can be created");

return ((SystemPartitioningHandle) partitioningHandle.getConnectorHandle()).getPartitionFunction(
partitionChannelTypes,
partitioningScheme.getHashColumn().isPresent(),
partitioningScheme.getBucketToPartition().get(),
bucketToPartition,
blockTypeOperators);
}

BucketFunction bucketFunction = getBucketFunction(session, partitioningHandle, partitionChannelTypes, bucketToPartition.get().length);
return new BucketPartitionFunction(bucketFunction, partitioningScheme.getBucketToPartition().get());
BucketFunction bucketFunction = getBucketFunction(session, partitioningHandle, partitionChannelTypes, bucketToPartition.length);
return new BucketPartitionFunction(bucketFunction, bucketToPartition);
}

public BucketFunction getBucketFunction(Session session, PartitioningHandle partitioning, List<Type> partitionChannelTypes, int bucketCount)
public BucketFunction getBucketFunction(Session session, PartitioningHandle partitioningHandle, List<Type> partitionChannelTypes, int bucketCount)
{
ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(partitioning.getConnectorId().get());
CatalogName catalogName = partitioningHandle.getConnectorId()
.orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(catalogName);
BucketFunction bucketFunction = partitioningProvider.getBucketFunction(
partitioning.getTransactionHandle().orElse(null),
partitioningHandle.getTransactionHandle().orElseThrow(() -> new IllegalArgumentException("No transactionHandle for partitioning handle: " + partitioningHandle)),
Copy link
Copy Markdown
Member

@losipiuk losipiuk Dec 14, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is transaction handle optional then? What is the case when it is not set? Worth a comment or more descriptive exception message?

Copy link
Copy Markdown
Member Author

@findepi findepi Dec 15, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the empty transaction handle is for SystemPartitioningHandle (internal use)

session.toConnectorSession(),
partitioning.getConnectorHandle(),
partitioningHandle.getConnectorHandle(),
partitionChannelTypes,
bucketCount);
checkArgument(bucketFunction != null, "No bucket function for partitioning: %s", partitioning);
checkArgument(bucketFunction != null, "No bucket function for partitioning: %s", partitioningHandle);
return bucketFunction;
}

public List<ConnectorPartitionHandle> listPartitionHandles(
Session session,
PartitioningHandle partitioningHandle)
{
ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(partitioningHandle.getConnectorId().get());
CatalogName catalogName = partitioningHandle.getConnectorId()
.orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(catalogName);
return partitioningProvider.listPartitionHandles(
partitioningHandle.getTransactionHandle().orElse(null),
session.toConnectorSession(partitioningHandle.getConnectorId().get()),
partitioningHandle.getTransactionHandle().orElseThrow(() -> new IllegalArgumentException("No transactionHandle for partitioning handle: " + partitioningHandle)),
session.toConnectorSession(catalogName),
partitioningHandle.getConnectorHandle());
}

Expand Down Expand Up @@ -197,24 +199,28 @@ private static List<InternalNode> getFixedMapping(ConnectorBucketNodeMap connect
.collect(toImmutableList());
}

public ConnectorBucketNodeMap getConnectorBucketNodeMap(Session session, PartitioningHandle partitioning)
public ConnectorBucketNodeMap getConnectorBucketNodeMap(Session session, PartitioningHandle partitioningHandle)
{
ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(partitioning.getConnectorId().get());
CatalogName catalogName = partitioningHandle.getConnectorId()
.orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(catalogName);
ConnectorBucketNodeMap connectorBucketNodeMap = partitioningProvider.getBucketNodeMap(
partitioning.getTransactionHandle().orElse(null),
session.toConnectorSession(partitioning.getConnectorId().get()),
partitioning.getConnectorHandle());
checkArgument(connectorBucketNodeMap != null, "No partition map %s", partitioning);
partitioningHandle.getTransactionHandle().orElseThrow(() -> new IllegalArgumentException("No transactionHandle for partitioning handle: " + partitioningHandle)),
session.toConnectorSession(catalogName),
partitioningHandle.getConnectorHandle());
checkArgument(connectorBucketNodeMap != null, "No partition map %s", partitioningHandle);
return connectorBucketNodeMap;
}

private ToIntFunction<Split> getSplitToBucket(Session session, PartitioningHandle partitioningHandle)
{
ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(partitioningHandle.getConnectorId().get());
CatalogName catalogName = partitioningHandle.getConnectorId()
.orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
ConnectorNodePartitioningProvider partitioningProvider = getPartitioningProvider(catalogName);

ToIntFunction<ConnectorSplit> splitBucketFunction = partitioningProvider.getSplitBucketFunction(
partitioningHandle.getTransactionHandle().orElse(null),
session.toConnectorSession(partitioningHandle.getConnectorId().get()),
partitioningHandle.getTransactionHandle().orElseThrow(() -> new IllegalArgumentException("No transactionHandle for partitioning handle: " + partitioningHandle)),
session.toConnectorSession(catalogName),
partitioningHandle.getConnectorHandle());
checkArgument(splitBucketFunction != null, "No partitioning %s", partitioningHandle);

Expand Down