diff --git a/presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryDataFragment.java b/presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryDataFragment.java index a526a3399..4ed93f9e3 100644 --- a/presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryDataFragment.java +++ b/presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryDataFragment.java @@ -87,7 +87,7 @@ public class MemoryDataFragment return MEMORY_DATA_FRAGMENT_CODEC.fromJson(fragment.getBytes()); } - public static Map> getMergedCount(Map> a, Map> b) + public static Map> getMergedPartitionMap(Map> a, Map> b) { Map> result = new HashMap<>(); for (Map.Entry> 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())); } } } diff --git a/presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryPageSinkProvider.java b/presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryPageSinkProvider.java index 958d85a65..5e2e9bf81 100644 --- a/presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryPageSinkProvider.java +++ b/presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryPageSinkProvider.java @@ -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> finish() { + // there could be multiple writers in parallel, wait until the last one if (sinkCount.decrementAndGet() == 0) { tablesManager.finishUpdatingTable(tableId); - Map> 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> 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())); diff --git a/presto-memory/src/main/java/io/prestosql/plugin/memory/MemorySplitManager.java b/presto-memory/src/main/java/io/prestosql/plugin/memory/MemorySplitManager.java index e473131f4..208ed6985 100644 --- a/presto-memory/src/main/java/io/prestosql/plugin/memory/MemorySplitManager.java +++ b/presto-memory/src/main/java/io/prestosql/plugin/memory/MemorySplitManager.java @@ -53,53 +53,79 @@ public final class MemorySplitManager List dataFragments = metadata.getDataFragments(table.getId()); + // check if there is a predicate on the partition column + List partitionKeyRanges = new ArrayList<>(); + for (Map.Entry 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 splits = ImmutableList.builder(); + if (partitionKeyRanges.isEmpty()) { + return allSplits(dataFragments, table); + } + for (MemoryDataFragment dataFragment : dataFragments) { Map> 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 partitionKeyRanges = new ArrayList<>(); - - // get the type of the column here so that the partitionKeyValue(Object) can be cast - for (Map.Entry 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> 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> 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> 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 dataFragments, MemoryTableHandle table) + { + ImmutableList.Builder 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()); diff --git a/presto-memory/src/main/java/io/prestosql/plugin/memory/data/LogicalPart.java b/presto-memory/src/main/java/io/prestosql/plugin/memory/data/LogicalPart.java index d5815d894..9b6a817fc 100644 --- a/presto-memory/src/main/java/io/prestosql/plugin/memory/data/LogicalPart.java +++ b/presto-memory/src/main/java/io/prestosql/plugin/memory/data/LogicalPart.java @@ -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; }