This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new b0e8bf246ed7 [SPARK-58258][PYTHON][SQL] Gate the columnar Arrow Python
UDF passthrough on the vector shape matching the declared schema
b0e8bf246ed7 is described below
commit b0e8bf246ed7725cbbdd01e6cb15bb6a70b4165e
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Thu Jul 23 00:35:30 2026 -0700
[SPARK-58258][PYTHON][SQL] Gate the columnar Arrow Python UDF passthrough
on the vector shape matching the declared schema
### What changes were proposed in this pull request?
The columnar Arrow Python UDF input fast path
(`ColumnarArrowPythonInput.writeArrowDirect`) forwards upstream vectors'
buffers verbatim into the worker's IPC stream, whose Schema message was written
once from the declared UDF input schema using the standard interchange
encoding. Its eligibility check (`isArrowBacked`) only verified
`isInstanceOf[ArrowColumnVector]` -- not that each selected vector's physical
shape matches what that schema declares. This PR adds `ArrowUtils.isInterchan
[...]
- the lossless internal struct encodings of nanosecond timestamps and
`CalendarInterval` (SPARK-57975 / SPARK-58005, produced by the Arrow-cached
scan),
- var-width vectors whose offset width disagrees with
`spark.sql.execution.arrow.useLargeVarTypes`,
- the Arrow view types (readable through `ArrowColumnVector` but never
declared by Spark schemas),
and requires it in the passthrough eligibility check. Mismatched batches
take the existing `writeRowByRow` fallback, which re-encodes through
`ArrowWriter` under the declared schema -- including its `DATETIME_OVERFLOW`
guards for values the interchange encoding cannot represent at all.
### Why are the changes needed?
An IPC consumer decodes a RecordBatch body strictly by the shape the
stream's Schema message declared, so forwarding a mismatched vector produces a
stream whose header and body disagree: depending on the column mix that
surfaces as a worker-side decode error or, worse, silently mis-bound buffers
handed to the UDF. This is reachable today: the Arrow cache (SPARK-57268)
always stores standard 32-bit var-width offsets, so caching a string column and
running an Arrow-optimized Python UDF [...]
The safety argument for the passthrough is compositional: every
interchange-shaped vector was produced under the guarded write boundary (its
values already fit the encoding), so the forwarding point only needs to verify
shape -- which this PR makes explicit instead of implicit.
### Does this PR introduce _any_ user-facing change?
No. Shape-matching batches keep the zero-copy fast path; mismatched ones
fall back to the row-based re-encoding, which is the correct-by-construction
slow path.
### How was this patch tested?
- `ArrowUtilsSuite` "isInterchangeShapedField rejects shapes the
interchange schema cannot declare": interchange-encoded fields of every kind
pass (including standard nanos/interval encodings and nested complex types);
the lossless struct encodings are rejected top-level and nested; an untagged
struct with the same child names passes; var-width offset width is checked
against `largeVarTypes` in both directions and nested; view types are rejected.
- `pyspark.sql.tests.arrow.test_arrow_python_udf_cached` gains
`test_udf_on_cached_strings_with_large_var_types`: an Arrow-optimized Python
UDF over an Arrow-cached string column with `useLargeVarTypes=true`. Verified
it fails with a worker-side error without the gate and passes with it; the
existing default-config tests confirm shape-matching batches still take the
fast path and produce correct results.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
This pull request and its description were written by Claude Code.
Closes #57428 from viirya/arrow-udf-passthrough-shape-gate.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../tests/arrow/test_arrow_python_udf_cached.py | 82 ++++++++-
.../org/apache/spark/sql/util/ArrowUtils.scala | 63 +++++++
.../apache/spark/sql/util/ArrowUtilsSuite.scala | 195 ++++++++++++++++++++-
.../python/ColumnarArrowPythonInput.scala | 23 ++-
4 files changed, 358 insertions(+), 5 deletions(-)
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
index 962302730e69..4367f6b959d9 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
@@ -15,6 +15,7 @@
# limitations under the License.
#
+import datetime
import unittest
from pyspark.sql.functions import col, udf
@@ -41,7 +42,9 @@ class ArrowPythonUDFCachedInputTests(ReusedSQLTestCase):
Python runner (ColumnarArrowPythonWithNamedArgumentRunner) end to end --
the path an
Arrow-backed DSv2 source would also take. This pins the init-message
protocol between that
runner and the worker: a mismatched section there desyncs the worker's
init parse and hangs
- the task (SPARK-58241).
+ the task (SPARK-58241). It also pins the columnar fast path's shape gate:
the cached vectors
+ are forwarded verbatim only when their physical shape matches what the
stream's schema
+ declares, and fall back to row-based re-encoding otherwise.
This suite runs in its own module because spark.sql.cache.serializer is a
static conf latched
into a JVM-wide singleton on first cache materialization.
@@ -69,16 +72,52 @@ class ArrowPythonUDFCachedInputTests(ReusedSQLTestCase):
finally:
super().tearDownClass()
- def test_udf_on_cached_strings(self):
+ def _assert_columnar_arrow_eval(self, df, expected_nodes):
+ # Assert the executed plan contains exactly `expected_nodes`
ArrowEvalPythonExec
+ # operators and every one of them takes the columnar input path --
otherwise the
+ # test silently stops exercising the code this suite exists to pin.
+ nodes = []
+ stack = [df._jdf.queryExecution().executedPlan()]
+ while stack:
+ node = stack.pop()
+ name = node.getClass().getSimpleName()
+ if name == "AdaptiveSparkPlanExec":
+ stack.append(node.executedPlan())
+ continue
+ if name == "ArrowEvalPythonExec":
+ nodes.append(node)
+ children = node.children()
+ for i in range(children.size()):
+ stack.append(children.apply(i))
+ self.assertEqual(len(nodes), expected_nodes, "unexpected number of
ArrowEvalPythonExec")
+ for node in nodes:
+ self.assertTrue(node.supportsColumnar(), "expected the columnar
input path")
+
+ def _assert_udf_on_cached_strings(self):
df = self.spark.createDataFrame(
[("hello",), (None,), ("world",)], schema="v string"
).cache()
try:
result = df.select(udf(lambda x: x.upper() if x is not None else
None)(col("v")))
+ self._assert_columnar_arrow_eval(result, expected_nodes=1)
assertDataFrameEqual(result, [Row("HELLO"), Row(None),
Row("WORLD")])
finally:
df.unpersist()
+ def test_udf_on_cached_strings(self):
+ # Cache-scan vectors use standard 32-bit var-width offsets, matching
the default
+ # interchange schema: the columnar fast path forwards them verbatim.
+ self._assert_udf_on_cached_strings()
+
+ def test_udf_on_cached_strings_with_large_var_types(self):
+ # With useLargeVarTypes=true the UDF stream schema declares LargeUtf8
(64-bit offsets)
+ # but the cache always stores standard Utf8, so forwarding the
cache-scan vector
+ # verbatim would produce a stream whose header and body disagree. The
shape gate must
+ # route this through the row-based re-encoding fallback and still
produce correct
+ # results.
+ with self.sql_conf({"spark.sql.execution.arrow.useLargeVarTypes":
"true"}):
+ self._assert_udf_on_cached_strings()
+
def test_udf_on_cached_numeric_column_alongside_unselected_string(self):
# Only the UDF's input columns are serialized to the worker; the
column that is not a
# UDF input rides through the pass-through recombination, so select it
alongside the
@@ -86,10 +125,49 @@ class ArrowPythonUDFCachedInputTests(ReusedSQLTestCase):
df = self.spark.createDataFrame([(1, "a"), (2, "b")], schema="i long,
s string").cache()
try:
result = df.select(col("s"), udf(lambda x: x + 1,
"long")(col("i")))
+ self._assert_columnar_arrow_eval(result, expected_nodes=1)
assertDataFrameEqual(result, [Row("a", 2), Row("b", 3)])
finally:
df.unpersist()
+ def test_udf_on_cached_map_column(self):
+ # Canonical map vectors from the cache pass the gate (their
entry-struct children carry
+ # the canonical "key"/"value" names in order) and ride the columnar
path verbatim. The
+ # rejection of a source whose entry children are swapped is pinned at
the unit level
+ # (ArrowUtilsSuite), since no in-tree producer can build such a vector.
+ df = self.spark.createDataFrame(
+ [({"a": 1},), ({"b": 2, "c": 3},)], schema="m map<string,long>"
+ ).cache()
+ try:
+ result = df.select(udf(lambda m: len(m), "long")(col("m")))
+ self._assert_columnar_arrow_eval(result, expected_nodes=1)
+ assertDataFrameEqual(result, [Row(1), Row(2)])
+ finally:
+ df.unpersist()
+
+ def test_udf_on_timestamps_cached_under_different_time_zone(self):
+ # Pins the timestamp pass-through end to end: the plan takes the
columnar path and
+ # the stored epoch round-trips unchanged across a session-timezone
change between
+ # cache materialization and the UDF query. Note the cache rebuilds its
Arrow schema
+ # at read time with the query's session timezone, so the actual and
declared timezone
+ # labels the gate compares are equal by construction here; the
layout-equal
+ # label-mismatch case (e.g. an external DSv2 Arrow source labeling UTC
while the
+ # stream schema declares the session timezone) has no in-tree columnar
producer and
+ # is pinned by the ArrowUtilsSuite congruence tests instead.
+ ts = datetime.datetime(2025, 1, 6, 12, 30, 45)
+ with self.sql_conf({"spark.sql.session.timeZone": "UTC"}):
+ df = self.spark.createDataFrame([(ts,)], schema="ts
timestamp").cache()
+ df.count() # materialize: the cached vectors are labeled with the
UTC session tz
+ try:
+ with self.sql_conf({"spark.sql.session.timeZone":
"America/Los_Angeles"}):
+ result = df.select(udf(lambda t: t, "timestamp")(col("ts")))
+ self._assert_columnar_arrow_eval(result, expected_nodes=1)
+ # The stored epoch round-trips; naive-datetime wall time is
preserved because
+ # createDataFrame and collect both convert against the
machine-local timezone.
+ assertDataFrameEqual(result, [Row(ts)])
+ finally:
+ df.unpersist()
+
if __name__ == "__main__":
from pyspark.testing import main
diff --git a/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala
b/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala
index 2b026bd4d35d..3877bf2e2c1c 100644
--- a/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala
+++ b/sql/api/src/main/scala/org/apache/spark/sql/util/ArrowUtils.scala
@@ -556,6 +556,69 @@ private[sql] object ArrowUtils {
}
}
+ /**
+ * Whether an Arrow field tree (from an upstream vector) is physically
congruent with the
+ * declared field of the stream it would be forwarded into: same Arrow type,
same child count,
+ * congruent children, and no dictionary encoding. An IPC consumer decodes a
RecordBatch body
+ * strictly by the declared schema, walking a flattened sequence of field
nodes and buffers, so
+ * ANY structural divergence silently shifts the node/buffer assignment of
every following
+ * column or mis-binds buffers within one: the lossless internal struct
encodings where the
+ * declared schema has the interchange types, a var-width vector whose
offset width disagrees
+ * with the declared one, a tagged Variant/Geometry/Geography struct
carrying an extra child
+ * that the recognizers tolerate but the declared canonical field does not
include, or
+ * dictionary-encoded data whose values live in dictionary batches the
stream never writes.
+ *
+ * This compares against the declared field rather than allowlisting
"declarable" shapes: the
+ * declared schema is the exact header the consumer will decode with, so
congruence with it is
+ * the complete correctness condition, not an approximation of it. Field
names, nullability, and
+ * metadata are ignored -- they do not affect the buffer layout -- except
for Map entry
+ * children, whose names carry the key/value semantics (see
mapEntryNamesMatchDeclared). The two
+ * failure directions stay asymmetric: a false reject merely takes the
row-based re-encoding
+ * fallback (slower, still correct, and with the value-domain guards such as
DATETIME_OVERFLOW
+ * reinstated), while a false accept corrupts the stream -- so anything not
provably congruent
+ * must be rejected.
+ */
+ def isCompatibleWithDeclaredField(actual: Field, declared: Field): Boolean =
{
+ actual.getDictionary == null &&
+ sameLayoutArrowType(actual.getType, declared.getType) &&
+ actual.getChildren.size == declared.getChildren.size &&
+ mapEntryNamesMatchDeclared(actual, declared) &&
+ actual.getChildren.asScala.zip(declared.getChildren.asScala).forall { case
(a, d) =>
+ isCompatibleWithDeclaredField(a, d)
+ }
+ }
+
+ // An Arrow Map binds its key/value semantics to the entry-struct children
by name (the spec
+ // puts the key first and MapVector addresses the children as
"key"/"value"), yet Arrow
+ // tolerates vectors whose entry children sit in the opposite order. Such a
vector is
+ // positionally indistinguishable from the canonical one when the key and
value types agree,
+ // and forwarding its buffers verbatim under the declared header silently
swaps keys and
+ // values -- so at Map nodes the entry children must also match the declared
names. Field
+ // names remain ignored everywhere else: Spark's semantics are positional
there and names
+ // carry no meaning.
+ private def mapEntryNamesMatchDeclared(actual: Field, declared: Field):
Boolean =
+ !declared.getType.isInstanceOf[ArrowType.Map] || {
+ val actualNames =
actual.getChildren.asScala.flatMap(_.getChildren.asScala).map(_.getName)
+ val declaredNames =
+
declared.getChildren.asScala.flatMap(_.getChildren.asScala).map(_.getName)
+ actualNames == declaredNames
+ }
+
+ // Physical-layout type comparison, mirroring
ArrowFileReadWrite.checkLayoutMatch: the
+ // timestamp timezone label does not affect the layout (values are int64
epoch numbers either
+ // way) and differs across Spark's own paths (a worker's output comes back
labeled UTC while
+ // the next stream declares the session time zone), so comparing it would
only cause false
+ // rejects. Presence-vs-absence must still match: it distinguishes
TimestampType from
+ // TimestampNTZType, whose identical-looking values are interpreted
differently. Map equality
+ // includes the keysSorted flag, which has no layout impact.
+ private def sameLayoutArrowType(actual: ArrowType, declared: ArrowType):
Boolean =
+ (actual, declared) match {
+ case (a: ArrowType.Timestamp, d: ArrowType.Timestamp) =>
+ a.getUnit == d.getUnit && (a.getTimezone == null) == (d.getTimezone ==
null)
+ case (_: ArrowType.Map, _: ArrowType.Map) => true
+ case (a, d) => a == d
+ }
+
def fromArrowField(field: Field): DataType = {
field.getType match {
case _: ArrowType.Map =>
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/util/ArrowUtilsSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/util/ArrowUtilsSuite.scala
index d60e32896d5c..3d7f5d87596a 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/util/ArrowUtilsSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/util/ArrowUtilsSuite.scala
@@ -22,7 +22,7 @@ import java.time.ZoneId
import scala.jdk.CollectionConverters._
import org.apache.arrow.vector.types.{IntervalUnit, TimeUnit}
-import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType}
+import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding,
Field, FieldType}
import org.apache.spark.{SparkException, SparkFunSuite,
SparkUnsupportedOperationException}
import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.LA
@@ -369,6 +369,199 @@ class ArrowUtilsSuite extends SparkFunSuite {
defaultType.asInstanceOf[ArrowType.Interval].getUnit ===
IntervalUnit.MONTH_DAY_NANO)
}
+ test("isCompatibleWithDeclaredField requires congruence with the declared
field tree") {
+ def declared(schema: StructType, largeVarTypes: Boolean = false): Field = {
+ ArrowUtils.toArrowSchema(schema, "UTC", true,
largeVarTypes).findField("value")
+ }
+ def actual(
+ schema: StructType,
+ lossless: Boolean = false,
+ largeVarTypes: Boolean = false): Field = {
+ ArrowUtils
+ .toArrowSchema(schema, "UTC", true, largeVarTypes,
losslessInternalTypes = lossless)
+ .findField("value")
+ }
+ def compat(a: Field, d: Field): Boolean =
ArrowUtils.isCompatibleWithDeclaredField(a, d)
+
+ // Identical construction is congruent, for every kind of type the schema
can declare,
+ // including the standard nanos/interval encodings, tagged struct types,
and nested trees.
+ Seq[DataType](
+ IntegerType,
+ StringType,
+ TimestampNTZNanosType(9),
+ CalendarIntervalType,
+ GeometryType(4326),
+ VariantType,
+ ArrayType(TimestampNTZNanosType(9)),
+ new StructType().add("i", CalendarIntervalType),
+ MapType(IntegerType, StringType)).foreach { dt =>
+ val s = new StructType().add("value", dt)
+ assert(compat(actual(s), declared(s)), dt.toString)
+ }
+
+ // The lossless internal encodings are structs where the declared schema
has the
+ // interchange types: incongruent in both directions.
+ Seq[DataType](
+ TimestampNTZNanosType(9),
+ CalendarIntervalType,
+ ArrayType(TimestampLTZNanosType(7))).foreach { dt =>
+ val s = new StructType().add("value", dt)
+ assert(!compat(actual(s, lossless = true), declared(s)), dt.toString)
+ }
+
+ // Var-width offset width must agree with the declared field, in both
directions and nested.
+ Seq[DataType](StringType, BinaryType, ArrayType(StringType)).foreach { dt
=>
+ val s = new StructType().add("value", dt)
+ assert(compat(actual(s), declared(s)), dt.toString)
+ assert(!compat(actual(s), declared(s, largeVarTypes = true)),
dt.toString)
+ assert(!compat(actual(s, largeVarTypes = true), declared(s)),
dt.toString)
+ assert(
+ compat(actual(s, largeVarTypes = true), declared(s, largeVarTypes =
true)),
+ dt.toString)
+ }
+
+ // A tagged Geometry struct with an extra child is read back as
GeometryType by the
+ // recognizers (which tolerate extra children), but the declared canonical
field has exactly
+ // two children -- forwarding the three-child body verbatim would shift
every following
+ // column's buffers (the extra child's buffers would be consumed as the
next argument's
+ // data), so it must be incongruent.
+ val geometrySchema = new StructType().add("value", GeometryType(4326))
+ val canonicalGeometry = declared(geometrySchema)
+ val extraChild = new Field(
+ "extra",
+ new FieldType(true, new ArrowType.Int(32, true), null, null),
+ java.util.Collections.emptyList[Field]())
+ val paddedChildren = new
java.util.ArrayList[Field](canonicalGeometry.getChildren)
+ paddedChildren.add(extraChild)
+ val paddedGeometry = new Field(
+ canonicalGeometry.getName,
+ new FieldType(
+ canonicalGeometry.isNullable,
+ canonicalGeometry.getType,
+ null,
+ canonicalGeometry.getMetadata),
+ paddedChildren)
+ assert(ArrowUtils.fromArrowField(paddedGeometry) === GeometryType(4326))
+ assert(!compat(paddedGeometry, canonicalGeometry))
+
+ // List-like types Spark never declares are incongruent with the declared
List.
+ val arraySchema = new StructType().add("value", ArrayType(IntegerType))
+ val declaredList = declared(arraySchema)
+ val intChild = new Field(
+ "element",
+ new FieldType(true, new ArrowType.Int(32, true), null, null),
+ java.util.Collections.emptyList[Field]())
+ Seq[ArrowType](
+ ArrowType.LargeList.INSTANCE,
+ ArrowType.ListView.INSTANCE,
+ ArrowType.LargeListView.INSTANCE,
+ new ArrowType.FixedSizeList(4)).foreach { arrowType =>
+ val listLike = new Field(
+ "value",
+ new FieldType(true, arrowType, null, null),
+ java.util.Collections.singletonList(intChild))
+ assert(!compat(listLike, declaredList), arrowType.toString)
+ }
+
+ // The timestamp timezone LABEL does not affect the physical layout and
differs across
+ // Spark's own paths (a worker's output comes back labeled UTC while the
next stream
+ // declares the session time zone), so it must not cause a reject.
Presence-vs-absence must
+ // still match: it distinguishes TimestampType from TimestampNTZType,
whose values are
+ // interpreted differently. The unit must also match: both units are
int64, but forwarding
+ // one under the other reinterprets every value by a factor of 1000.
+ def tsField(tz: String): Field = new Field(
+ "value",
+ new FieldType(true, new ArrowType.Timestamp(TimeUnit.MICROSECOND, tz),
null, null),
+ java.util.Collections.emptyList[Field]())
+ assert(compat(tsField("UTC"), tsField("America/Los_Angeles")))
+ assert(compat(tsField("America/Los_Angeles"), tsField("UTC")))
+ assert(!compat(tsField(null), tsField("UTC")))
+ assert(!compat(tsField("UTC"), tsField(null)))
+ val tsNanoField = new Field(
+ "value",
+ new FieldType(true, new ArrowType.Timestamp(TimeUnit.NANOSECOND, "UTC"),
null, null),
+ java.util.Collections.emptyList[Field]())
+ assert(!compat(tsNanoField, tsField("UTC")))
+
+ // Map equality includes keysSorted, which has no layout impact and must
not cause a reject.
+ val mapSchema = new StructType().add("value", MapType(IntegerType,
StringType))
+ val declaredMap = declared(mapSchema)
+ val sortedMap = new Field(
+ declaredMap.getName,
+ new FieldType(declaredMap.isNullable, new ArrowType.Map(true), null,
null),
+ declaredMap.getChildren)
+ assert(compat(sortedMap, declaredMap))
+
+ // An Arrow Map binds key/value semantics to the entry-struct children by
name (the key
+ // comes first and MapVector addresses the children as "key"/"value"), yet
Arrow tolerates
+ // vectors whose entry children sit in the opposite order. With int keys
and int values the
+ // swapped tree is positionally identical to the canonical one, so only
the names reveal the
+ // swap -- forwarding the buffers verbatim under the canonical header
would silently decode
+ // {k: v} as {v: k}. Map entry names must therefore match the declared
ones, at every
+ // nesting level.
+ def swapEntries(mapField: Field): Field = {
+ val entries = mapField.getChildren.get(0)
+ val swapped = new Field(
+ entries.getName,
+ new FieldType(entries.isNullable, entries.getType, null,
entries.getMetadata),
+ entries.getChildren.asScala.reverse.asJava)
+ new Field(
+ mapField.getName,
+ new FieldType(mapField.isNullable, mapField.getType, null,
mapField.getMetadata),
+ java.util.Collections.singletonList(swapped))
+ }
+ val intMapSchema = new StructType().add("value", MapType(IntegerType,
IntegerType, false))
+ val declaredIntMap = declared(intMapSchema)
+ assert(compat(declaredIntMap, declaredIntMap))
+ assert(!compat(swapEntries(declaredIntMap), declaredIntMap))
+ val nestedMapSchema = new StructType()
+ .add("value", MapType(IntegerType, MapType(IntegerType, IntegerType,
false), false))
+ val declaredNestedMap = declared(nestedMapSchema)
+ val nestedEntries = declaredNestedMap.getChildren.get(0)
+ val innerSwappedEntries = new Field(
+ nestedEntries.getName,
+ new FieldType(nestedEntries.isNullable, nestedEntries.getType, null,
+ nestedEntries.getMetadata),
+ java.util.Arrays.asList(
+ nestedEntries.getChildren.get(0),
+ swapEntries(nestedEntries.getChildren.get(1))))
+ val innerSwappedMap = new Field(
+ declaredNestedMap.getName,
+ new FieldType(declaredNestedMap.isNullable, declaredNestedMap.getType,
null,
+ declaredNestedMap.getMetadata),
+ java.util.Collections.singletonList(innerSwappedEntries))
+ assert(compat(declaredNestedMap, declaredNestedMap))
+ assert(!compat(innerSwappedMap, declaredNestedMap))
+
+ // A view-typed field is incongruent with the declared Utf8.
+ val stringSchema = new StructType().add("value", StringType)
+ val viewField = new Field(
+ "value",
+ new FieldType(true, ArrowType.Utf8View.INSTANCE, null, null),
+ java.util.Collections.emptyList[Field]())
+ assert(!compat(viewField, declared(stringSchema)))
+
+ // A dictionary-encoded field is incongruent even though its value type
matches the declared
+ // field: its record batches carry indices whose values live in separate
dictionary batches
+ // the UDF stream never writes.
+ val dictField = new Field(
+ "value",
+ new FieldType(
+ true,
+ ArrowType.Utf8.INSTANCE,
+ new DictionaryEncoding(1L, false, new ArrowType.Int(32, true)),
+ null),
+ java.util.Collections.emptyList[Field]())
+ assert(!compat(dictField, declared(stringSchema)))
+ // The same field without the dictionary is congruent, isolating the
dictionary as the
+ // reason. Names and metadata are ignored: only the physical layout
matters.
+ val plainField = new Field(
+ "renamed",
+ new FieldType(true, ArrowType.Utf8.INSTANCE, null, null),
+ java.util.Collections.emptyList[Field]())
+ assert(compat(plainField, declared(stringSchema)))
+ }
+
test("time") {
// Arrow's Time type has no precision field, so TIME(p) precision is
preserved via field
// metadata; the Arrow type itself stays Time(NANOSECOND, 64).
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonInput.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonInput.scala
index 36c401118629..fbc21c0e2bf6 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonInput.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonInput.scala
@@ -28,6 +28,7 @@ import org.apache.arrow.vector.ipc.message.MessageSerializer
import org.apache.spark.api.python.BasePythonRunner
import org.apache.spark.sql.execution.arrow.ArrowWriter
+import org.apache.spark.sql.util.ArrowUtils
import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
/**
@@ -81,10 +82,28 @@ private[python] trait ColumnarArrowPythonInput
}
/**
- * Check whether all selected columns in the batch are Arrow-backed.
+ * Check whether all selected columns in the batch are Arrow-backed AND
physically congruent
+ * with the stream's Schema message. The Schema message was written once
from `root` (built
+ * from the declared UDF input schema), while writeArrowDirect serializes
the upstream
+ * vectors' buffers verbatim -- so each selected vector's field tree is
compared against the
+ * corresponding declared field of that exact header. Any divergence (the
lossless internal
+ * struct encodings from an Arrow-cache scan, a var-width offset width
disagreeing with the
+ * declared one, a tagged struct carrying extra children the declared
canonical field lacks,
+ * dictionary-encoded data, ...) would make the header and body disagree,
which the worker
+ * decodes as silently shifted buffers or rejects. Such batches take
writeRowByRow instead,
+ * which re-encodes through ArrowWriter under the declared schema --
including its
+ * DATETIME_OVERFLOW guards for values the interchange encoding cannot
represent at all.
*/
private def isArrowBacked(batch: ColumnarBatch): Boolean = {
- inputColumnIndices.forall(i =>
batch.column(i).isInstanceOf[ArrowColumnVector])
+ val declaredFields = root.getSchema.getFields
+ inputColumnIndices.indices.forall { i =>
+ batch.column(inputColumnIndices(i)) match {
+ case acv: ArrowColumnVector =>
+ ArrowUtils.isCompatibleWithDeclaredField(
+ acv.getValueVector.getField, declaredFields.get(i))
+ case _ => false
+ }
+ }
}
/**
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
