[SPARK-48356][SQL] Support for FOR statement

[dusantism-db]

What changes were proposed in this pull request?

In this PR, support for FOR statement in SQL scripting is introduced. Examples:

FOR row AS SELECT * FROM t DO
   SELECT row.intCol;
 END FOR;

FOR SELECT * FROM t DO
   SELECT intCol;
 END FOR;

Implementation notes:
As local variables for SQL scripting are currently a work in progress, session variables are used to simulate them.
When FOR begins executing, session variables are declared for each column in the result set, and optionally for the for variable if it is present ("row" in the example above).
On each iteration, these variables are overwritten with the values from the row currently being iterated.
The variables are dropped upon loop completion.

This means that if a session variable which matches the name of a column in the result set already exists, the for statement will drop that variable after completion. If that variable would be referenced after the for statement, the script would fail as the variable would not exist. This limitation is already present in the current iteration of SQL scripting, and will be fixed once local variables are introduced. Also, with local variables the implementation of for statement will be much simpler.

Grammar/parser changes:
forStatement grammar rule
visitForStatement rule visitor
ForStatement logical operator

Why are the changes needed?

FOR statement is an part of SQL scripting control flow logic.

Does this PR introduce any user-facing change?

No

How was this patch tested?

New tests are introduced to all of the three scripting test suites: SqlScriptingParserSuite, SqlScriptingExecutionNodeSuite and SqlScriptingInterpreterSuite.

Was this patch authored or co-authored using generative AI tooling?

No

[davidm-db]

for my knowledge, can you explain what does the case with the Map means exactly, i.e. when will this happen?
also, how did we check that this is the complete list of the relevant cases?

[dusantism-db]

When Map or Struct are in the result set of the query, we can't use Literal(value) to convert them to expressions because Literals don't support them. So for example for Map we recursively convert both keys and values to expressions first, and then create a map expression using CreateMap. The process is similar for structs.

The way i checked is i went through all the spark data types, and for each checked in code of Literal whether it's supported. I only found these two which are not, however I agree we can't be completely sure, and new types will be added to Spark in the future which Literals may or may not support. Probably I should add an error message for currently unsupported type, in case it comes up. Does that make sense to you?

[davidm-db]

Yeah, I would say internal error is fine in this case (i.e. no need to introduce new error for this) since it would mean that we have a bug.
Other than that, this sounds fine to me, but let's wait for Max and/or Wenchen to comment on this if they have any concerns.

[miland-db]

Why do we need to create this every time? Can we fill variablesMap once and then reuse it?

[dusantism-db]

We need to create it every time because the map is different for every row in the result set. You can see we call it on the currRow.

[davidm-db]

Can we rebase first to include already merged changes regarding the label checks, logical plans, etc? And I'll review afterwards again?

[dusantism-db]

@davidm-db @miland-db Rebased, you can review again

[davidm-db]

we don't have drop var statements here due to the fact that they are dropped in handleLeaveStatement?
this is the thing we talked about that will be properly resolved once the proper execution and scopes are introduced?

[dusantism-db]

Yes, that's right. In this case the variables are dropped immediately when the leave statement is encountered, instead of the usual behavior which is to return the dropVariable exec nodes from the iterator.

[davidm-db]

food for thought: does DataFrame have a mechanism to partially collect the data so we don't collect all the results in memory? since we are already using the caching concept, this would be easy to add to the logic of cachedQueryResult.

quickly researching, we can do something like:
sliced_df = df.offset(starting_index).limit(ending_index - starting_index)
but there might be something better...

I wouldn't block the PR on this, but I think we definitely need to consider something like this for a follow-up.

cc: @cloud-fan @MaxGekk

[dusantism-db]

That makes sense, currently the entire result is collected to the driver so it would be problematic if the result size is too large. We should definitely follow up on this

[davidm-db]

Let's see what Wenchen and Max have to say and maybe create a follow-up work item so we don't forget it.

[cloud-fan]

There is a df.toLocalIterator(). Under the hood, it launches jobs for RDD partitions one by one, so at most only the data of one Partition will be collected to the Spark driver at the same time.

[dusantism-db]

@cloud-fan @davidm-db
I refactored to trydf.toLocalIterator()(minimal changes) and it seems to work properly. Could you take a look again?

[davidm-db]

I think we can simplify this a bit further, commented it here.
Otherwise, the idea is really cool. You already have my approval.

[davidm-db]

I've left comments, but in general the approach looks good to me!
This wasn't easy, good job!
Let's resolve the comments and fix tests (seems to be just syntax/style errors).

[cloud-fan]

@miland-db can you confirm whether we should return query or not here? Should we treat it the same as IF/WHILE condition and return it in the iterator?

[cloud-fan]

let's use SparkPlan#executeToIterator(), which returns InternalRow, so that we can save the cost of data conversion in createExpressionFromValue

[cloud-fan]

We can address comments later as they are kind of improvement. Let me merge it first. Thanks!

[dusantism-db]

@cloud-fan Great, thanks. I will create follow up tasks for the improvements you suggested.