SNOW-1882075,SNOW-1990301,SNOW-1990302,SNOW-1821290: AST finalizations for stateless batches, and DataframeWriter,DataframeReader new APIs

[sfc-gh-heshah]

1. Which Jira issue is this PR addressing? Make sure that there is an accompanying issue to your PR.

   This PR includes all changes required to complete the following tickets.

   • SNOW-1882075: AST for DataFrameWriter.insertInto/insert_into
   • SNOW-1990301: AST for DataframerWriter.{format,save}
   • SNOW-1990302: AST for DataframeReader.{format,load}
   • SNOW-1821290: Finalize the AST schema v1

   It also includes changes in support of SNOW-1997136: Ensure AST batch is stateless.

2. Fill out the following pre-review checklist:

   • I am adding a new automated test(s) to verify correctness of my new code
     • If this test skips Local Testing mode, I'm requesting review from @snowflakedb/local-testing
   • I am adding new logging messages
   • I am adding a new telemetry message
   • I am adding new credentials
   • I am adding a new dependency
   • If this is a new feature/behavior, I'm adding the Local Testing parity changes.
   • I acknowledge that I have ensured my changes to be thread-safe. Follow the link for more information: Thread-safe Developer Guidelines
   • If adding any arguments to public Snowpark APIs or creating new public Snowpark APIs, I acknowledge that I have ensured my changes include AST support.

3. Please describe how your code solves the related issue.

   AST generation for DataframeReader APIs which return the same DataframeReader instance have been updated to follow the same pattern as DataframeWriter. Since all such methods are setters, it is only necessary to have the last set attributes in the AST sent to the serveer. This is in support of the new format API in both Snowpark classes.

   AST generation was added for the new APIs DataframeReader.{format,load}, and DataframeWriter.{format,read,insert_into} along with expectation test updates.

   Changes were also made in support of AST batch statelessness in the IR, by changing fields typed VarId to be Expr typed. This allows more consistent usage of _set_ast_ref as a helper method across all Snowpark APIs.

   This PR is accompanied by a server-side PR.

[sfc-gh-snowflakedb-snyk-sa]

🎉 Snyk checks have passed. No issues have been found so far.

✅ security/snyk check is complete. No issues have been found. (View Details)

✅ license/snyk check is complete. No issues have been found. (View Details)

[sfc-gh-heshah]

• SNOW-1997136: Ensure transitive closure in AST Request messages #3201
• SNOW-1821290: Finalize AST schema v1 by refactoring Assign to Bind, VarId to BindId #3183
• SNOW-1882075,SNOW-1990301,SNOW-1990302,SNOW-1821290: AST finalizations for stateless batches, and DataframeWriter,DataframeReader new APIs #3180 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[sfc-gh-vbudati]

Thanks Hemit!

[sfc-gh-lspiegelberg]

I recall there used to be some issue adding this to the constructor, do we have tests covering on/off for using Column objects (i.e. calling this API, and also col)?

[sfc-gh-heshah]

Could you clarify, do you mean on/off in terms of our global AST enabled flag?

[sfc-gh-lspiegelberg]

oh like with AST flag enabled/disabled.

[sfc-gh-heshah]

Oh, yes we do have tests in both cases then (public calls made by tests, and internal calls to the constructor) if I'm understanding correctly.

[sfc-gh-lspiegelberg]

No questions re the minor changes, but confused about _flush_ast. What's the role of it? Why do we need it? Could we get away with not adding yet another parameter?

[sfc-gh-heshah]

save_as_table and copy_into_location (to which the _flush_ast parameter was added) are both public APIs with eager evaluation and internal uses. However, most other DataframeWriter APIs rely on them to perform the eager evaluation (in terms of Snowpark execution), and to flush the AST.

Only using _emit_ast meant that internal callers can't ask save_as_table and copy_into_location to both suppress their AST statement generation and flush the AST batch, because the latter was still guarded by our if _emit_ast check.

Using the combination means that internal callers can set _emit_ast to False to suppress the AST generation from these APIs, while also setting _flush_ast to true to indicate that the eager request sent through the underlying connection should still include the AST batch.

It's a unique case: public APIs, eager (in terms of actually performing the execution within the function itself rather than in a helper), and used internally by other public APIs. I don't think combination exists anywhere else in the Snowpark codebase unless I'm mistaken.

Elsewhere we seem to pass an _ast argument containing a pre-built proto.Expr itself to be inherited by the callee, so we do use another parameter. In this case I think _flush_ast makes more sense to indicate what's actually happening, however, we can follow the common pattern and just use an _ast parameter with the built proto.Expr() for save_as_table and copy_into_table to add to the Request before flushing.

[sfc-gh-lspiegelberg]

save_as_table and copy_into_location (to which the _flush_ast parameter was added) are both public APIs with eager evaluation and internal uses. However, most other DataframeWriter APIs rely on them to perform the eager evaluation (in terms of Snowpark execution), and to flush the AST.

Only using _emit_ast meant that internal callers can't ask save_as_table and copy_into_location to both suppress their AST statement generation and flush the AST batch, because the latter was still guarded by our if _emit_ast check.

Using the combination means that internal callers can set _emit_ast to False to suppress the AST generation from these APIs, while also setting _flush_ast to true to indicate that the eager request sent through the underlying connection should still include the AST batch.

It's a unique case: public APIs, eager (in terms of actually performing the execution within the function itself rather than in a helper), and used internally by other public APIs. I don't think combination exists anywhere else in the Snowpark codebase unless I'm mistaken.

Elsewhere we seem to pass an _ast argument containing a pre-built proto.Expr itself to be inherited by the callee, so we do use another parameter. In this case I think _flush_ast makes more sense to indicate what's actually happening, however, we can follow the common pattern and just use an _ast parameter with the built proto.Expr() for save_as_table and copy_into_table to add to the Request before flushing.

Hm, internal use of save_as_table should not generate AST/not flush it. My mind model here is that all of this moves to the server-side, i.e. what is right now internal will simply execute server-side.