Skip to content

feat: support Spark-compatible stack function - #26266

Open
comphead wants to merge 11 commits into
apache:mainfrom
comphead:spark-stack
Open

comphead wants to merge 11 commits into
apache:mainfrom
comphead:spark-stack

Conversation

@comphead

@comphead comphead commented Oct 13, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

Spark's stack(n, expr1, ..., exprk) separates the values into n rows. It is the standard way to unpivot in Spark SQL, and Comet falls back to Spark for it.

What changes are included in this PR?

  • stack in datafusion-spark. A ScalarUDF cannot emit rows, so like json_tuple it returns the n rows as a List<Struct<col0, col1, ...>>, which is the data type of Spark's Generator. As in Spark, the list is not nullable while its elements and fields are. The caller (Comet, or unnest in SQL) expands it.
  • The arguments are validated as in Spark's Stack.checkInputDataTypes, with the same messages. n is coerced to INT. The values of a column must have the same Spark type, and NULLs, Arrow types of one Spark type (such as Utf8 and Utf8View, or lists and maps whose field names differ) and integer literals that fit an INT column are cast to it.
  • Primitive columns are copied with a strided loop, and string columns of only scalars (such as unpivot keys) are tiled with a few large concat calls. Other types use take over the concatenated values. On the new stack bench this is about 3x to 5x faster than interleave.

What is the testing strategy for this PR?

spark/generator/stack.slt ports the cases of Spark's GeneratorFunctionSuite, GeneratorExpressionSuite and TypeCoercionSuite, and of the generators, table-valued-functions and join-lateral golden files. It also covers each kernel path. Unit tests cover an empty batch, the row limit and the type widening, which SQL cannot reach.

Are there any user-facing changes?

A new Spark-compatible stack function. Unlike Spark, n must be an integer literal, so a foldable expression such as 3 - 1 is rejected.

The rows are materialized per batch, so a call that generates more than 16777216 rows (n times the batch size) fails with a resources error. Spark streams the rows and has no limit. With 8192-row batches this rejects n above 2048 for a column argument.

@comphead
comphead marked this pull request as draft October 13, 2026 15:31
comphead

This comment was marked as resolved.

@github-actions github-actions Bot added sqllogictest SQL Logic Tests (.slt) spark labels Oct 13, 2026
@codecov-commenter

codecov-commenter commented Oct 14, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 90.60773% with 34 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.95%. Comparing base (390d210) to head (fb1b29b).
⚠️ Report is 93 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/spark/src/function/generator/stack.rs 90.58% 17 Missing and 17 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #26266      +/-   ##
==========================================
+ Coverage   82.80%   82.95%   +0.14%     
==========================================
  Files        1149     1156       +7     
  Lines      454480   458762    +4282     
  Branches   454480   458762    +4282     
==========================================
+ Hits       376347   380544    +4197     
+ Misses      55267    55151     -116     
- Partials    22866    23067     +201     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@comphead
comphead marked this pull request as ready for review October 15, 2026 19:58
@andygrove

Copy link
Copy Markdown
Member

Thanks for the PR. I reviewed this by reading the diff and surrounding code. I did not build or run it, so the failure scenarios below are from code reading and not reproduced.

Main concerns

  1. Unbounded output (stack.rs ~L514). parse_n accepts any n up to i32::MAX, and the only guard is that the i32 offsets fit. SELECT stack(2000000000, 1) looks like it would allocate roughly 16 GB for an Int64 column, plus a large null bitmap and a ~2e9-iteration padding loop. stack_take would also build an 8 GB Vec<u32> of indices. A modest n (e.g. 100000) times an 8192-row batch is large too. Constant folding of an all-scalar call can reach this at plan time. Spark streams these rows lazily, and there is no memory accounting here. Could we cap the output size, or at least fail cleanly?
  2. Type unification depends on argument order (~L356). An integer literal only unifies with INT when the INT value comes first. stack(2, 0, a) with a INT fails with an Int64 vs Int32 mismatch, while stack(2, a, 0) works. Spark's integer literals are INT, so common unpivots like stack(2, 'x', 1, 'y', int_col) break depending on order. The slt test for this case is commented out as a known gap. Preferring the non-literal type for a column, or treating integer literals that fit in INT as Int32, would fix it.
  3. n must be a bare literal (~L266). stack(1 + 1, ...), stack(2::INT, ...) and stack(CAST(2 AS INT), ...) fail at planning. The error says the expression is not foldable when it is.

Smaller items

  1. same_field compares nullability (~L376). Spark's sameType ignores nullability, so stack(2, arr_a, arr_b) with List<Int64 not null> and List<Int64 nullable> is rejected. The same applies to struct fields and map values. FixedSizeList vs List also never matches.
  2. Strict execution-time cast (~L443). invoke_with_args uses cast_to with safe: false. A caller that supplies its own return_field (e.g. Comet) with an Int32 column and Int64 data above i32::MAX fails the whole batch with a generic Arrow cast error.
  3. Broadcast is inferred from length (~L519). A source counts as per-row when source.len() > 1, and the check is repeated in stack_take, the strided loops and the all-len-1 check that gates tile. If an array's length differs from number_rows, zip silently truncates and stack_take indexing can go out of bounds. Tracking scalar-ness once, e.g. with a broadcast flag computed in stack_columns, would remove the implicit invariant.
  4. Three kernel paths (~L592). tile (doubling + concat for all-scalar Utf8/Binary columns) sits beside stack_take and stack_primitive for one benchmark shape. It has one slt case, and the 256-copy block boundary is only covered indirectly. A single path for scalar-only columns would be easier to maintain.

I'd treat 1-3 as the ones to address before merging, especially 1 and 3.

Cap the generated rows per batch, unify integer literals with an INT value
in any position, ignore nullability when matching types, check array lengths
and cover the `tile` block boundaries.
@comphead

Copy link
Copy Markdown
Contributor Author

Thanks, pushed ac51d5844.

  1. Output is now capped at 2^24 generated rows per batch and fails with a resources error, including during constant folding.
  2. An integer literal that fits INT now unifies with an Int32 value in any position.
  3. n is still limited to a bare literal. The return type is planned before folding, and scalar_arguments deliberately excludes expression casts. The error message and docs now say so.

Nullability is ignored when matching types. Array lengths are checked against number_rows. The tile block boundaries have a unit test. I kept the strict cast, because overflowing the planned type should be an error. I kept tile, which measured about 3.5x faster than take for unpivot.

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the quick turnaround, @comphead. The integer literal unification and the array length check look good, and the slt cases cover them. Your reasons for keeping the strict cast, tile and the literal-only n make sense to me. For Comet, Spark folds a foldable n to a literal before we see the plan anyway.

There are two things inline, and one of them is my mistake from the last round. Spark's Stack.checkInputDataTypes compares column types with !=, so nullability does matter there. My item 4 was wrong, sorry about that.

Smaller items:

  • The cap is per batch, so whether stack(4096, col) succeeds depends on the batch size. It fails with 8192-row batches and works with 4096. Spark has no limit, and Comet uses 8192-row batches, so any n above 2048 with a column argument fails there. Could the user-facing changes section mention the limit?
  • TINYINT and SMALLINT literals are also widened in an INT column, since the check is is_integer(). Spark rejects those. Limiting it to Int64 literals would match the comment.
  • The cap error reads "1 rows of 2000000000 generated rows exceed the limit". Something like "generating 2000000000 rows for each of 1 input rows exceeds the limit of 16777216 rows per batch" would read more clearly.

I checked both inline cases with unit tests against ac51d5844, not end to end.

return internal_err!("stack: `n` must be a scalar");
};
let n = parse_n(n)?;
if values.iter().all(|v| matches!(v, ColumnarValue::Scalar(_))) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The cap is only checked against the real batch size on the array path. When every value is a literal, this calls stack(.., 1) and returns one scalar list, and the caller then expands it to number_rows copies through ScalarValue::to_array_of_size with no check. With n = 2^20 and number_rows = 8192, invoke_with_args returns Ok, and the expanded result would be about 8.6 billion rows against a cap of 16.8 million. Something like SELECT stack(1048576, 1) FROM generate_series(1, 8192) should reach it, and so does constant folding, which turns the call into a literal at plan time.

Could the number_rows * n check run before this shortcut? That covers execution, including Comet. Folded literals would still get through, but at least that gap would be known.

/// Whether two Arrow types are the same Spark type. Types of one logical type, such
/// as `Utf8` and `Utf8View`, match, and list elements and map entries match
/// regardless of their field names, which Spark types do not have. As in Spark's
/// `sameType`, nullability is ignored.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comes from my last review, and I was wrong about it. Spark's Stack.checkInputDataTypes compares each value's type with the column type using != (in both 3.5.8 and 4.1.1), so containsNull, valueContainsNull and StructField.nullable all count. StackTypeCoercion only retypes a NULL literal.

Ignoring nullability now lets a query plan and then fail. The column takes the first value's type, so when that's List(non-null Int64) and a later value is a List(Int64) with a NULL element, the cast at execution fails with Non-nullable field of ListArray "item" cannot contain nulls. Something like SELECT stack(2, arrow_cast([1], 'List(non-null Int64)'), [2, NULL]) should show it. The existing Nullability does not make types differ case passes only because [2] has no NULL.

I'd go back to matching nullability exactly, as Spark does. If you'd rather stay lenient, the column type's nested fields would need to be nullable whenever the values differ.

Check the generated rows limit before the all-scalar shortcut, match
nullability exactly as Spark does, widen only BIGINT literals and clarify the
limit error.
@comphead

Copy link
Copy Markdown
Contributor Author

Thanks, pushed fb1b29b9c.

  • The generated rows check now runs before the all-scalar shortcut, with the real batch size. A unit test calls invoke_with_args with n = 2^20 and 8192 rows. The folded literal case still gets through, as you noted.
  • Nullability is matched exactly again, as in Spark. The slt case now expects the plan-time error for List(non-null Int64) followed by [2, NULL].
  • Only Int64 literals are widened in an INT column, so TINYINT and SMALLINT literals are rejected as in Spark. A unit test covers it.
  • The limit error now reads "generating 2000000000 rows for each of 1 input rows exceeds the limit of 16777216 rows per batch".
  • The limit and the n > 2048 consequence for 8192-row batches are in the user-facing changes section of the PR description.

@comphead
comphead requested a review from andygrove October 21, 2026 16:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

spark sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants