Skip to content

[SPARK-61181][CONNECT][PYTHON] Read the local relation settings of the server for local relations of 64 KiB or more, and raise LOCAL_RELATION_SIZE_LIMIT_EXCEEDED, in the Rust and Python clients as PySpark does - #379

Open
HyukjinKwon wants to merge 1 commit into
apache:masterfrom
HyukjinKwon:SPARK-61181-local-relation-settings

Conversation

@HyukjinKwon

Copy link
Copy Markdown
Member

What changes were proposed in this pull request?

SPARK-61181: PySpark Connect's createDataFrame reads spark.sql.session.localRelationCacheThreshold and spark.sql.session.localRelationSizeLimit for every local relation. At the threshold it uploads the relation as cache artifacts. Past the size limit it raises LOCAL_RELATION_SIZE_LIMIT_EXCEEDED.

The Rust client read these settings only for a relation of at least 32 MiB of Arrow data, and sent smaller relations inline without them. This change lowers that bound to 64 KiB. The bound is one constant, LOCAL_RELATION_SETTINGS_BYTES in crates/spark-connect/src/session.rs, with a comment on why it is 64 KiB. It applies on every path that encodes a local relation:

  • local rows (SparkSession::create_dataframe), and the pandas row conversion;
  • create_dataframe_from_columns;
  • the float64, string and int64 pandas fast paths, with an explicit StructType or a DDL schema, through the Arrow IPC bridge (create_dataframe_from_inline_arrow_ipc).

These paths already called arrow_input::create_relation, which applies the threshold and the size limit. Only the bound changes.

The size limit is now checked on the data chunks only, as PySpark does. The schema chunk counts toward the size but is not checked against the limit (crates/spark-connect/src/arrow_input.rs). Before, a schema alone above the limit failed, with an actualSize that PySpark does not report.

Why are the changes needed?

A local relation of 32 MiB or less was sent inline, so it did not fail past the server's size limit. With a session of localRelationCacheThreshold 1 and localRelationSizeLimit 100, a 2,000-row relation of a long and a 100-character string gave:

>>> spark.createDataFrame([(i, "x" * 100) for i in range(2000)], "a long, b string").count()
# before this change: no error; the relation is sent inline
# PySpark 4.3.0:
#   [LOCAL_RELATION_SIZE_LIMIT_EXCEEDED] Local relation size (224760 bytes) exceeds the limit (100 bytes).
# this change:
#   [LOCAL_RELATION_SIZE_LIMIT_EXCEEDED] Local relation size (225496 bytes) exceeds the limit (100 bytes).

The same frame with a StructType schema on the float64 or string fast path also skipped the check. The Arrow importer, used for a DDL schema, did raise.

With the default settings too, a relation of 1 MiB (the default localRelationCacheThreshold) to 32 MiB was sent inline, where PySpark uploads it as cache artifacts. It is now uploaded at the threshold.

Relations below 64 KiB are still sent inline without reading the settings. A 200-row relation of 4,944 bytes stays inline and returns 200 rows, where PySpark 4.3.0 reads the settings and raises LOCAL_RELATION_SIZE_LIMIT_EXCEEDED (checked against PySpark 4.3.0 on the same server). This PR does not change that difference.

Does this PR introduce any user-facing change?

Yes. A local relation of 64 KiB or more under spark.sql.session.localRelationSizeLimit now raises LOCAL_RELATION_SIZE_LIMIT_EXCEEDED on the local rows, pandas, columns and fast paths, as PySpark does. A relation of 64 KiB or more at or above the threshold is uploaded as cache artifacts on those paths too. The size limit is checked on data chunks only, so a schema alone above the limit no longer fails.

How was this patch tested?

  • New tests in python/tests/test_arrow_handoff.py, against the mock Arrow input service, for six paths (local rows, columns, pandas row conversion, float64, string and int64), each with a StructType and a DDL schema:
    • test_local_relations_of_64_kib_read_the_settings_at_the_threshold: the settings are read; the relation is uploaded at a threshold of 1 byte and sent inline at the default threshold.
    • test_local_relations_of_64_kib_past_the_size_limit_raise: LOCAL_RELATION_SIZE_LIMIT_EXCEEDED is raised with the limit, and nothing is uploaded.
    • test_local_relations_of_64_kib_the_same_for_struct_and_ddl: both schema forms have the same outcome.
  • test_small_local_relations_stay_inline_without_reading_settings now uses a 500-row relation, asserted below 64 KiB of Arrow data.
  • Existing tests that encoded the 32 MiB bound now expect the settings to be read for relations of 64 KiB or more: the float64 and string fast-path tests, their hook and iterator tests, the date and timestamp row-input helpers, and the float64 and string cached-upload cases, which now upload at the threshold. Their relations are still sent inline below the default threshold, or uploaded where the test sets the threshold to 1 byte.
  • test_dropin_e2e.py::test_local_relations_of_64_kib_follow_the_size_limit_as_pyspark, against a live server: on a new session with a threshold of 1 and a limit of 100, local rows and a pandas frame, with StructType and DDL, raise LOCAL_RELATION_SIZE_LIMIT_EXCEEDED, and the same data works with the defaults.
  • The new mock and e2e tests fail with the bound at 32 MiB, and pass with this change.
  • Against official PySpark 4.3.0 on the same server, the 2,000-row example above raises LOCAL_RELATION_SIZE_LIMIT_EXCEEDED in both, with the same limit. The actualSize differs by 736 bytes (225496 here, 224760 in PySpark), because the Arrow IPC bytes of the data chunk differ slightly.
  • The lib tests of both crates (609 passed), the build of all tests, also with wasm-udf, the offline Python suites (2273 passed, 8 skipped), and the live drop-in and Arrow suites (215 passed) and live cases (10 passed) against a local Spark 4.3.0 server. Clippy reports no diagnostics in the changed lines.

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

Generated-by: Claude Opus 5.5

This pull request and its description were written by Isaac.

…e server for local relations of 64 KiB or more, and raise LOCAL_RELATION_SIZE_LIMIT_EXCEEDED, in the Rust and Python clients as PySpark does

The Rust client read spark.sql.session.localRelationCacheThreshold and
spark.sql.session.localRelationSizeLimit only for a local relation of at least 32 MiB. For a
smaller one it sent the data inline, so the size limit was not applied. PySpark reads the settings
for every relation and raises LOCAL_RELATION_SIZE_LIMIT_EXCEEDED past the limit. Now a relation of
at least 64 KiB of Arrow data reads them on every createDataFrame path, and is sent as the Arrow
importer sends it. A smaller relation is still sent inline without the configuration RPC. As in
PySpark, the size limit is checked on the data chunks only, not on the schema chunk.

Co-authored-by: Isaac <no-reply@databricks.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant