Add GlueingPartitioningOperator + Corresponding changes in window function layer to consume it for MSQ by Akshat-Jain · Pull Request #17038 · apache/druid

Akshat-Jain · 2024-09-12T07:59:41Z

Description

Currently, WindowOperatorQueryFrameProcessor was using the following operator pipeline for window processing: (MSQ shuffling) -> NaiveSortOperator -> NaivePartitioningOperator -> WindowOperator.

Because of this, we had the following issues:

If we wanted to process a subset of rows (in the cases where we could), we had to have the comparison logic (on the PARTITION BY columns) in the WindowOperatorQueryFrameProcessor layer. This is not good because:
1. The partitioning logic would be done again in the NaivePartitioningOperator, hence it's redundant.
2. It's a lot of code, making it unnecessarily difficult to follow the logic.
NaiveSortOperator was a synchronization barrier. It needs to process all rows before sending data to the receiver.

This PR introduces 2 new operators:

GlueingPartitioningOperator: It continuously receives data, and outputs batches of partitioned RACs. It maintains a last-partitioning-boundary of the last-pushed-RAC, and attempts to glue it with the next RAC it receives, ensuring that partitions are handled correctly, even across multiple RACs. You can check GlueingPartitioningOperatorTest for some good examples of the "glueing" work.
PartitionSortOperator: It sorts rows inside partitioned RACs, on the sort columns. The input RACs it receives are expected to be "complete / separate" partitions of data.

With this PR's changes, we are converting the operator pipeline from

(MSQ shuffling) -> NaiveSortOperator -> NaivePartitioningOperator -> WindowOperator

into

(MSQ shuffling) -> GlueingPartitioningOperator -> PartitionSortOperator -> WindowOperator

This allows WindowOperatorQueryFrameProcessor to send RACs of any number of rows into the operator pipeline, without having to do the partitioning on the PARTITION BY columns.

Other notable changes done in the PR

WindowOperatorQueryKit
1. Add translation from NaiveSortOperator -> NaivePartitioningOperator -> WindowOperator into GlueingPartitioningOperator -> PartitionSortOperator -> WindowOperator.
2. Changed logic of creation of window stages. Previously we were creating a single window stage if we had an empty over() clause. But we can't do that now since we rely on MSQ shuffling to cluster the data on the partitioning keys for us, as we don't have NaiveSortOperator in the chain now.
WindowOperatorQueryFrameProcessor
1. Modified the operator so that we don't create the operator chain again and again for new RACs.
2. Modified the operator execution to not run it to completion, and control it manually, since we don't want the completed() to be called.
3. Logic changes to blindly send rows to the operator pipeline whenever number of rows crosses a threshold.
Refactoring work: Created base abstract classes for the following pairs of classes:
1. NaiveSortOperator / PartitionSortOperator
2. NaiveSortOperatorFactory / PartitionSortOperatorFactory
3. NaivePartitioningOperator / GlueingPartitioningOperator
4. NaivePartitioningOperatorFactory / GlueingPartitioningOperatorFactory

Key changed/added classes in this PR

GlueingPartitioningOperator and corresponding factory class
PartitionSortOperator and corresponding factory class
WindowOperatorQueryKit
WindowOperatorQueryFrameProcessor

This PR has:

been self-reviewed.
added documentation for new or modified features or behaviors.
a release note entry in the PR description.
added Javadocs for most classes and all non-trivial methods. Linked related entities via Javadoc links.
added or updated version, license, or notice information in licenses.yaml
added comments explaining the "why" and the intent of the code wherever would not be obvious for an unfamiliar reader.
added unit tests or modified existing tests to cover new code paths, ensuring the threshold for code coverage is met.
added integration tests.
been tested in a test Druid cluster.

…ction layer to consume it for MSQ

kgyrtkirk

thank you for the PR @Akshat-Jain!

…perator pipeline without any extra logic

kgyrtkirk · 2024-09-23T08:36:36Z

+
+  private boolean needToProcessBatch()
+  {
+    return numRowsInFrameRowsAndCols >= maxRowsMaterialized / 2; // Can this be improved further?


why divide by 2 ? that doesn't give any guarantee that it will be inside bounds
people could set it to half if needed - but I think its easier to document clear things...

We need some threshold to start pushing RACs into the operator pipeline.

We discussed that we should push N rowed RACs into the pipeline. But it's not trivial to create RACs of exact size N.

So I felt it would be better to just have a threshold like push RACs into the pipeline when they cross N rows. I chose N = maxRowsMaterialized / 2 but we can always discuss on better values for this.

that doesn't give any guarantee that it will be inside bounds

I think it does 🤔

convertRowFrameToRowsAndColumns() method enforces the maxRowsMaterialized constraint: ensureMaxRowsInAWindowConstraint(frameRowsAndCols.size() + ldrc.numRows()), hence it won't allow us to accumulate more than maxRowsMaterialized rows.

Thoughts? Let me know if I'm missing something. Thanks!

kgyrtkirk · 2024-09-23T08:39:59Z

+        stageRowSignature = finalWindowStageRowSignature;
+        nextShuffleSpec = finalWindowStageShuffleSpec;
+      } else {
+        nextShuffleSpec = findShuffleSpecForNextWindow(operatorList.get(i + 1), maxWorkerCount);


please open a separate PR and fix the stagebuilder rather than hacking it backwards from here

… to prevent OOM issues

kgyrtkirk · 2024-10-09T11:15:50Z

+
+  protected abstract Iterator<RowsAndColumns> getIteratorForRAC(RowsAndColumns rac);
+
+  protected abstract void handleKeepItGoing(AtomicReference<Signal> signalRef, Iterator<RowsAndColumns> iterator, Receiver receiver);


I don't really see how these methods make the implementing class's simpler...they don't really hide away much complexity...

kgyrtkirk · 2024-10-09T11:21:39Z

+
+    Iterator<RowsAndColumns> partitionsIter = getIteratorForRAC(rac);
+
+    AtomicReference<Signal> keepItGoing = new AtomicReference<>(Signal.GO);


why use an AtomicReference ; the handleKeepItGoing method is void why not use the return value?

It forces us to return null in GlueingPartitioningOperator (as we need to return something):

@Override protected Signal handleKeepItGoing(Iterator<RowsAndColumns> iterator, Receiver receiver) { RowsAndColumns rowsAndColumns = iterator.next(); if (iterator.hasNext()) { return receiver.push(rowsAndColumns); } else { previousRac = rowsAndColumns; return null; } }

And then we have to handle the null specifically, as we don't want to update the signal in the case of null, since handlePush() needs to return the last non-null signal.

Signal keepItGoing = Signal.GO; while (keepItGoing == Signal.GO && partitionsIter.hasNext()) { Signal signal = handleKeepItGoing(partitionsIter, receiver); if (signal != null) { keepItGoing = signal; } }

This didn't seem clean to me, so I ended up going with the AtomicRef approach.

then it maybe return an Optional<Signal> ?

This isn't relevant anymore as I had to move this logic into Glueing/NaivePartitioningOperator layer (away from Abstract layer).

I tried to keep handlePush() and handleKeepItGoing() in Abstract layer, but since we are using static classes now, we can't override the implementation of a static method, so I had to remove these methods and move the logic to the individual classes.

kgyrtkirk · 2024-10-09T11:32:41Z

@@ -110,28 +79,7 @@ public Closeable goOrContinue(Closeable continuation, Receiver receiver)
          @Override
          public Signal push(RowsAndColumns rac)


make the Receiver a static inner class; it should have the Iterator as its field.

handlePush naturally wants to be a method of it

the full if (cont.iter != null) part should also go into the abstract;

put the full body of the while into a method in the abstract

the glueing should override that method and before calling the super() it could check if there are more elements and save that and use that when the Receiver is built

A static receiver wouldn't have access to methods like ensureMaxRowsMaterializedConstraint() which are also called from non-static methods.

Also, can you elaborate why should we make receiver static? Even if we had a StaticReceiver, we would still have to do new StaticReceiver() with the required state.

For example, to add a static receiver in NaivePartitioningOperator, I added the following StaticReceiver class and then replaced new Receiver() {} with new StaticReceiver(receiver, iterHolder):

static class StaticReceiver implements Receiver { private final Receiver delegate; private final AtomicReference<Iterator<RowsAndColumns>> iterHolder; public StaticReceiver( Receiver delegate, AtomicReference<Iterator<RowsAndColumns>> iterHolder ) { this.delegate = delegate; this.iterHolder = iterHolder; } @Override public Signal push(RowsAndColumns rac) { return handlePush(rac, delegate, iterHolder); } @Override public void completed() { if (iterHolder.get() == null) { delegate.completed(); } } }

I think we're not on the same page wrt how the static receiver needs to look like. Can you share your thoughts on the above?

doesn't seem like a dealbreaker...as a last resort ensureMaxRowsMaterializedConstraint can be static and pass the 2 integers for it?

@kgyrtkirk The iterator needs to be created in the receiver's push() method.

Having a static receiver doesn't let us create the iterator there, since getIteratorForRAC() is not static.

getIteratorForRAC() can't be made static since it's supposed to be an abstract method, with different implementations in NaivePartitioningOperator and GlueingPartitioningOperator.

If I remove getIteratorForRAC() from the abstract class, and just have it as a regular non-overridden static method, then how to pass the iterator to the static receiver? I can't pass it in like super.push(rac, iteratorForRac) since it violates the method signature of push().

Also, trying to make everything static also ends up requiring us to pass unnecessary variables as method params, and structure everything just to make the receiver static, which seems unnecessary and not clean to me. We anyway would still have to do new StaticReceiver(field1, field2,...) everytime instead of new Receiver().

Thoughts?

Have made local changes to make receiver and the iterator class static. The code doesn't seem clean though, as I had to pass a bunch of params. Also trying to see what can be moved to Abstract class with those changes. (This is all on my local right now, just updating this info here)

@kgyrtkirk Have pushed the change to the PR to make the relevant classes static. Appreciate if you could take a look again, thanks!

kgyrtkirk · 2024-10-09T12:34:55Z

+public class GlueingPartitioningOperator extends AbstractPartitioningOperator
+{
+  private final int maxRowsMaterialized;
+  private RowsAndColumns previousRac;


it would be better to not have this in Operator scope; I believe it should belong to the Receiver

previousRac is also used in the continuation logic, so needs to be at the Operator scope.

although it works - that's bad design

kgyrtkirk · 2024-10-10T06:47:06Z

+
+    Iterator<RowsAndColumns> partitionsIter = getIteratorForRAC(rac);
+
+    AtomicReference<Signal> keepItGoing = new AtomicReference<>(Signal.GO);


then it maybe return an Optional<Signal> ?

kgyrtkirk · 2024-10-10T06:53:21Z

+        throw new NoSuchElementException();
+      }
+
+      if (!firstPartitionHandled) {


note: you don't necessarily need this boolean; couldn't (previousRac != null) act like it?
or currentIndex==0 ?
having (previousRac != null) at a higher level could also cleanup some conditionals

Changed to using currentIndex == 0

please read my comment again

Have made some change in the way the conditionals are structured: b725332

Please let me know if this aligns with the suggestion.

kgyrtkirk · 2024-10-10T07:11:12Z

+  )
+  {
+    super(partitionColumns, child);
+    this.maxRowsMaterialized = maxRowsMaterialized;


could this be null?

It's int, so it can't be null.

note: unboxing happens at this point as this.maxRowsMaterialized is an int; meanwhile maxRowsMaterialized is an Integer

Makes sense, have added Preconditions.checkNotNull(maxRowsMaterialized, "maxRowsMaterialized cannot be null");, hope that works.

…logic

kgyrtkirk · 2024-10-16T11:02:04Z

+  private RowsAndColumns makeSimpleRac(int... values)
+  {
+    return MapOfColumnsRowsAndColumns.fromMap(
+        ImmutableMap.of("column", new IntArrayColumn(values))
+    );
+  }
+
+  private RowsAndColumnsHelper expectedSimpleRac(int... values)
+  {
+    return new RowsAndColumnsHelper()
+        .expectColumn("column", values)
+        .allColumnsRegistered();
+  }


note: this is not critical; but if you copy these 10 lines into multiple files that creates a maintenance burden after some iterations - its better to try to reuse such things somehow

Moved these methods to RowsAndColumnsHelper class. Hope that works.

kgyrtkirk · 2024-10-16T11:04:54Z

+  )
+  {
+    super(partitionColumns, child);
+    this.maxRowsMaterialized = maxRowsMaterialized;


note: unboxing happens at this point as this.maxRowsMaterialized is an int; meanwhile maxRowsMaterialized is an Integer

kgyrtkirk · 2024-10-16T11:58:33Z

+      ClusteredGroupPartitioner groupPartitioner = rac.as(ClusteredGroupPartitioner.class);
+      if (groupPartitioner == null) {
+        groupPartitioner = new DefaultClusteredGroupPartitioner(rac);
+      }


note: use ClusteredGroupPartitioner.fromRAC

Oh nice, wasn't aware of this. Have made the change in both partitioning operators. Thanks for pointing this out!

kgyrtkirk · 2024-10-16T12:03:41Z

-
-  private final ArrayList<ResultRow> rowsToProcess;
-  private int lastPartitionIndex = -1;
+  private Operator op = null;


can this field have a more verbose name?

Renamed to operator. Any other naming suggestion? 😅

kgyrtkirk · 2024-10-16T12:06:29Z

-        if (frameHasRowsPendingFlush()) {
-          return ReturnOrAwait.runAgain();
+        catch (IOException e) {
+          throw new RuntimeException(e);


note: why the need to catch this?
at line 136 it was okay to throw - why not okay here?

Ah right, have made the change.

kgyrtkirk · 2024-10-16T12:20:46Z

+          return null; // Signal that the operator has completed its work
+        }
+
+        // Return a non-null continuation object to indicate that we want to continue processing.


I don't think this comment and the returned Closeable lives up to the contracts associated with the implemented method...but that didn't worked before this change either.....

Since we have to re-use the same operator, I needed a way to keep it running..

Add GlueingPartitioningOperator + Corresponding changes in window fun…

70c05bb

…ction layer to consume it for MSQ

github-actions Bot added Area - Batch Ingestion Area - Querying Area - MSQ For multi stage queries - https://github.com/apache/druid/issues/12262 labels Sep 12, 2024

kgyrtkirk reviewed Sep 12, 2024

View reviewed changes

Create GluedRACsIterator

dad20ee

Akshat-Jain commented Sep 13, 2024

View reviewed changes

Comment thread ...age-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java Outdated

Akshat-Jain added 2 commits September 19, 2024 19:04

Modify MSQ window function execution to allow sending RACs into the o…

52131ea

…perator pipeline without any extra logic

Revert unnecessary changes

ae2ab9e

Akshat-Jain requested a review from kgyrtkirk September 19, 2024 15:52

cryptoe mentioned this pull request Sep 23, 2024

MSQ window functions: Reject MVDs during window processing #17036

Merged

9 tasks

kgyrtkirk reviewed Sep 23, 2024

View reviewed changes

Akshat-Jain added 3 commits September 25, 2024 11:58

Address review comments - part 1

5d0594b

Address review comments - part 2 - rename base classes to abstract

56c5942

Address review comments - part 3 - add RowsAndColumnsBuilder class

b80fed6

Akshat-Jain marked this pull request as ready for review September 25, 2024 08:17

Akshat-Jain requested a review from kgyrtkirk September 25, 2024 08:19

Akshat-Jain added 7 commits October 3, 2024 10:53

Merge with master

b0f3db8

Fix MSQWindowTest

82e2bb1

Merge with master

7aba378

Remove commented code

929f108

Merge with master

6f14fd0

Make allocator capacity in LazilyDecoratedRowsAndColumns configurable…

348b722

… to prevent OOM issues

Fix MSQWindowTest

d5eab25

Akshat-Jain mentioned this pull request Oct 9, 2024

Refactor WindowOperatorQueryKit to use WindowStage class for representing different window stages #17158

Merged

9 tasks

kgyrtkirk reviewed Oct 9, 2024

View reviewed changes

Akshat-Jain added 2 commits October 9, 2024 22:12

Address review comments

988f3f1

Make the capacity field mandatory in LazilyDecoratedRowsAndColumns

56afb9c

github-advanced-security AI found potential problems Oct 9, 2024

View reviewed changes

Comment thread processing/src/main/java/org/apache/druid/query/operator/GlueingPartitioningOperator.java Fixed

Fix CodeQL

e6167e5

kgyrtkirk reviewed Oct 10, 2024

View reviewed changes

Akshat-Jain added 6 commits October 10, 2024 12:58

Make classes static

d377788

Remove separate method for keepItGoing logic

82a60cb

Address review comments

db19325

Remove firstPartitionHandled boolean, use currentIndex instead

a865fdf

Change isGlueingNeeded() to compare first and last rows of ConcatRAC

3075b96

Pass allocatorCapacity as Long, convert to integer in the ldrc layer

071cfed

Akshat-Jain requested a review from kgyrtkirk October 10, 2024 08:26

github-advanced-security AI found potential problems Oct 10, 2024

View reviewed changes

Comment thread processing/src/main/java/org/apache/druid/query/operator/GlueingPartitioningOperator.java Fixed

Akshat-Jain added 5 commits October 10, 2024 14:42

Fix CodeQL

f947fc7

Address review comment: Change order of conditionals to simplify the …

b725332

…logic

Create AbstractReceiver

00ce202

Move push() method logic to AbstractReceiver

bb8f9fe

Improve some access modifiers

aa2e46f

kgyrtkirk reviewed Oct 16, 2024

View reviewed changes

kgyrtkirk approved these changes Oct 16, 2024

View reviewed changes

Address review comments

b73f472

cryptoe merged commit 450fb01 into apache:master Oct 17, 2024

Akshat-Jain mentioned this pull request Oct 17, 2024

WindowOperatorQueryFrameProcessor: Avoid writing multiple frames to output channel in runIncrementally() #17373

Merged

9 tasks

Akshat-Jain mentioned this pull request Oct 29, 2024

Fix backward compatibility issues in WindowOperatorQueryFrameProcessorFactory and WindowOperatorQueryFrameProcessor #17433

Merged

9 tasks

adarshsanjeev added this to the 32.0.0 milestone Jan 16, 2025


		protected abstract Iterator<RowsAndColumns> getIteratorForRAC(RowsAndColumns rac);

		protected abstract void handleKeepItGoing(AtomicReference<Signal> signalRef, Iterator<RowsAndColumns> iterator, Receiver receiver);


		Iterator<RowsAndColumns> partitionsIter = getIteratorForRAC(rac);

		AtomicReference<Signal> keepItGoing = new AtomicReference<>(Signal.GO);

		@@ -110,28 +79,7 @@ public Closeable goOrContinue(Closeable continuation, Receiver receiver)
		@Override
		public Signal push(RowsAndColumns rac)

Conversation

Akshat-Jain commented Sep 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Other notable changes done in the PR

Key changed/added classes in this PR

Uh oh!

kgyrtkirk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Akshat-Jain commented Sep 12, 2024 •

edited

Loading

kgyrtkirk Oct 16, 2024 •

edited

Loading

kgyrtkirk Oct 16, 2024 •

edited

Loading