jffwang
/
hetu-core
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Memory connector: Handle null partitions correctly

This commit is contained in:
farhan3 2021-11-17 16:18:45 -05:00
parent 98e7e914c6
commit 46ba508d45
4 changed files with 80 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryDataFragment.java
View File

@ -87,7 +87,7 @@ public class MemoryDataFragment
return MEMORY_DATA_FRAGMENT_CODEC.fromJson(fragment.getBytes());
}
public static Map<String, List<Integer>> getMergedCount(Map<String, List<Integer>> a, Map<String, List<Integer>> b)
public static Map<String, List<Integer>> getMergedPartitionMap(Map<String, List<Integer>> a, Map<String, List<Integer>> b)
{
Map<String, List<Integer>> result = new HashMap<>();
for (Map.Entry<String, List<Integer>> entry : a.entrySet()) {
@ -113,7 +113,7 @@ public class MemoryDataFragment
return new MemoryDataFragment(a.getHostAddress(), a.getRows() + b.getRows(), Math.max(a.getLogicalPartCount(), b.getLogicalPartCount()), Collections.emptyMap());
}
else {
return new MemoryDataFragment(a.getHostAddress(), a.getRows() + b.getRows(), 0, getMergedCount(a.getLogicalPartPartitionMap(), b.getLogicalPartPartitionMap()));
return new MemoryDataFragment(a.getHostAddress(), a.getRows() + b.getRows(), Math.max(a.getLogicalPartCount(), b.getLogicalPartCount()), getMergedPartitionMap(a.getLogicalPartPartitionMap(), b.getLogicalPartPartitionMap()));
}
}
}

9
presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryPageSinkProvider.java
View File

@ -38,6 +38,7 @@ import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import static com.google.common.base.Preconditions.checkState;
import static io.prestosql.plugin.memory.MemoryErrorCode.MISSING_DATA;
@ -176,9 +177,15 @@ public class MemoryPageSinkProvider
@Override
public CompletableFuture<Collection<Slice>> finish()
{
// there could be multiple writers in parallel, wait until the last one
if (sinkCount.decrementAndGet() == 0) {
tablesManager.finishUpdatingTable(tableId);
Map<String, List<Integer>> logicalPartPartitionMap = tablesManager.getTableLogicalPartPartitionMap(tableId);
// the JSON parser can't handle null keys in a map, so they must be skipped
// see LogicalPart#partitionPage and MemorySplitManager#getSplits for details
Map<String, List<Integer>> logicalPartPartitionMap =
tablesManager.getTableLogicalPartPartitionMap(tableId).entrySet().stream()
.filter(e -> e.getKey() != null)
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
return completedFuture(ImmutableList.of(new MemoryDataFragment(currentHostAddress, addedRows, tablesManager.getTableLogicalPartCount(tableId), logicalPartPartitionMap).toSlice()));
}
return completedFuture(ImmutableList.of(new MemoryDataFragment(currentHostAddress, addedRows, 0, Collections.emptyMap()).toSlice()));

86
presto-memory/src/main/java/io/prestosql/plugin/memory/MemorySplitManager.java
View File

@ -53,53 +53,79 @@ public final class MemorySplitManager
List<MemoryDataFragment> dataFragments = metadata.getDataFragments(table.getId());
// check if there is a predicate on the partition column
List<SortedRangeSet> partitionKeyRanges = new ArrayList<>();
for (Map.Entry<ColumnHandle, Domain> e : table.getPredicate().getDomains().orElse(Collections.emptyMap()).entrySet()) {
if (!e.getKey().isPartitionKey()) {
continue;
}
// in LogicalPart#partitionPage null partition is a special case
// although the partition map in LogicalPart class can have null keys
// when the map is sent to coordinator via MemoryDataFragment the null
// keys are skipped because the JSON parser can't handle null keys in a map
// therefore when query predicate contains a null value, we MUST NOT
// do any partition filtering because we would miss data
//
// e.g. if query is: select * from table where column is null
// then schedule all the splits
//
// see comment in LogicalPart#partitionPage also
if (e.getValue().isNullAllowed()) {
return allSplits(dataFragments, table);
}
SortedRangeSet rangeSet = ((SortedRangeSet) e.getValue().getValues());
partitionKeyRanges.add(rangeSet);
}
ImmutableList.Builder<ConnectorSplit> splits = ImmutableList.builder();
if (partitionKeyRanges.isEmpty()) {
return allSplits(dataFragments, table);
}
for (MemoryDataFragment dataFragment : dataFragments) {
Map<String, List<Integer>> logicalPartPartitionMap = dataFragment.getLogicalPartPartitionMap();
int logicalPartCount = dataFragment.getLogicalPartCount();
long rows = dataFragment.getRows();
// logicalPart ids are 1 based
if (logicalPartPartitionMap.size() == 0) {
int logicalPartCount = dataFragment.getLogicalPartCount();
// logicalPart ids are 1 based
for (int i = 1; i <= logicalPartCount; i++) {
splits.add(new MemorySplit(table.getId(), i, dataFragment.getHostAddress(), rows, OptionalLong.empty()));
}
}
else {
// get the predicate from the MemoryTableHandle
// filter the splits based on the partitionKey and only schedule them
List<SortedRangeSet> partitionKeyRanges = new ArrayList<>();
// get the type of the column here so that the partitionKeyValue(Object) can be cast
for (Map.Entry<ColumnHandle, Domain> e : table.getPredicate().getDomains().orElse(Collections.emptyMap()).entrySet()) {
if (!e.getKey().isPartitionKey()) {
continue;
}
SortedRangeSet rangeSet = ((SortedRangeSet) e.getValue().getValues());
partitionKeyRanges.add(rangeSet);
}
if (partitionKeyRanges.size() > 0) {
for (Map.Entry<String, List<Integer>> entry : logicalPartPartitionMap.entrySet()) {
for (SortedRangeSet rangeSet : partitionKeyRanges) {
Type rangeSetType = rangeSet.getType();
Object value = LogicalPart.deserializeTypedValueFromString(rangeSetType, entry.getKey());
if (rangeSet.containsValue(value)) {
for (Integer i : entry.getValue()) {
splits.add(new MemorySplit(table.getId(), i, dataFragment.getHostAddress(), rows, OptionalLong.empty()));
}
for (Map.Entry<String, List<Integer>> entry : logicalPartPartitionMap.entrySet()) {
for (SortedRangeSet rangeSet : partitionKeyRanges) {
Type rangeSetType = rangeSet.getType();
Object value = LogicalPart.deserializeTypedValueFromString(rangeSetType, entry.getKey());
if (rangeSet.containsValue(value)) {
for (Integer i : entry.getValue()) {
splits.add(new MemorySplit(table.getId(), i, dataFragment.getHostAddress(), rows, OptionalLong.empty()));
}
}
}
}
else {
for (Map.Entry<String, List<Integer>> entry : logicalPartPartitionMap.entrySet()) {
for (Integer i : entry.getValue()) {
splits.add(new MemorySplit(table.getId(), i, dataFragment.getHostAddress(), rows, OptionalLong.empty()));
}
}
}
}
}
return new FixedSplitSource(splits.build());
}
/**
* Schedule entire table
*/
private ConnectorSplitSource allSplits(List<MemoryDataFragment> dataFragments, MemoryTableHandle table)
{
ImmutableList.Builder<ConnectorSplit> splits = ImmutableList.builder();
for (MemoryDataFragment dataFragment : dataFragments) {
int logicalPartCount = dataFragment.getLogicalPartCount();
long rows = dataFragment.getRows();
// logicalPart ids are 1 based
for (int i = 1; i <= logicalPartCount; i++) {
splits.add(new MemorySplit(table.getId(), i, dataFragment.getHostAddress(), rows, OptionalLong.empty()));
}
}
return new FixedSplitSource(splits.build());

15
presto-memory/src/main/java/io/prestosql/plugin/memory/data/LogicalPart.java
View File

@ -951,7 +951,20 @@ public class LogicalPart
int[] retainedPositions = valueAndPosition.getValue().stream().mapToInt(i -> i).toArray();
Object valueKey = valueAndPosition.getKey();
Page subPage = page.getPositions(retainedPositions, 0, retainedPositions.length);
partitions.put(valueKey.toString(), subPage);
// NOTE: null partition key is allowed here in the map
// but when this partition map is sent to coordinator via MemoryDataFragment
// the JSON parser fails and can't handle null keys in the map
// the JSON parser will ignore null keys
// therefore during scheduling if the query predicate is for null
// we MUST NOT do any partition filtering because the partition map
// the coordinator has is missing null partitions
// the coordinator must schedule all splits if the query predicate is null
// see: MemorySplitManager#getSplits
//
// note: the other option is to use an empty string as the null key
// then the JSON parser could send the key to the coordinator
// but then this would cause conflicts with actual empty string values
partitions.put(valueKey == null ? null : valueKey.toString(), subPage);
}
return partitions;
}