diff --git a/python/pyspark/sql/tests/arrow/test_arrow.py b/python/pyspark/sql/tests/arrow/test_arrow.py index a29408980faba..76a29afd4063e 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow.py +++ b/python/pyspark/sql/tests/arrow/test_arrow.py @@ -1853,6 +1853,41 @@ def test_toArrow_with_compression_codec_large_dataset(self): self.assertEqual(t.num_rows, 10000) self.assertEqual(t.column_names, ["id", "str_col", "mod_col"]) + def test_toPandas_double_nested_array_empty_outer(self): + schema = StructType([StructField("data", ArrayType(ArrayType(StringType())))]) + df = self.spark.createDataFrame([Row(data=[])], schema=schema) + pdf = df.toPandas() + self.assertEqual(len(pdf), 1) + self.assertEqual(len(pdf["data"][0]), 0) + + def test_toPandas_array_of_map_empty_outer(self): + schema = StructType([StructField("data", ArrayType(MapType(StringType(), StringType())))]) + df = self.spark.createDataFrame([Row(data=[])], schema=schema) + pdf = df.toPandas() + self.assertEqual(len(pdf), 1) + self.assertEqual(len(pdf["data"][0]), 0) + + def test_toPandas_triple_nested_array_empty_outer(self): + # SPARK-55056: This triggers SIGSEGV without the fix. + # When the outer array is empty, the second-level ArrayWriter is never + # invoked, so its count stays 0. Arrow format requires ListArray offset + # buffer to have N+1 entries even when N=0, but getBufferSizeFor(0) + # returns 0 and the buffer is omitted in IPC serialization. + schema = StructType([StructField("data", ArrayType(ArrayType(ArrayType(StringType()))))]) + df = self.spark.createDataFrame([Row(data=[])], schema=schema) + pdf = df.toPandas() + self.assertEqual(len(pdf), 1) + self.assertEqual(len(pdf["data"][0]), 0) + + def test_toPandas_nested_array_with_map_empty_outer(self): + schema = StructType( + [StructField("data", ArrayType(ArrayType(MapType(StringType(), StringType()))))] + ) + df = self.spark.createDataFrame([Row(data=[])], schema=schema) + pdf = df.toPandas() + self.assertEqual(len(pdf), 1) + self.assertEqual(len(pdf["data"][0]), 0) + @unittest.skipIf( not have_pandas or not have_pyarrow, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala index 8d68e74dbf874..cb7e407cf819d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala @@ -387,6 +387,11 @@ private[arrow] class ArrayWriter( val valueVector: ListVector, val elementWriter: ArrowFieldWriter) extends ArrowFieldWriter { + // SPARK-55056: Arrow format requires ListArray offset buffer to have N+1 entries. + // Even when N=0, the buffer must contain [0]. Initialize offset buffer at construction + // to ensure it exists even if no elements are written. + valueVector.getOffsetBuffer.setInt(0, 0) + override def setNull(): Unit = { } @@ -408,6 +413,8 @@ private[arrow] class ArrayWriter( override def reset(): Unit = { super.reset() + // Re-initialize offset buffer after reset (see constructor comment) + valueVector.getOffsetBuffer.setInt(0, 0) elementWriter.reset() } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala index 2c0c0494bbacf..bb8b8bc28e06b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/arrow/ArrowWriterSuite.scala @@ -875,4 +875,46 @@ class ArrowWriterSuite extends SparkFunSuite { assert(map2.keyArray().array().mkString(",") == Array(1).mkString(",")) assert(stringRepr(map2) == Array("bob", "40").mkString(",")) } + + test("SPARK-55056: triple nested array with empty outer array") { + // Schema: array>> + // This triggers SIGSEGV without the fix. When the outer array is empty, + // the second-level ArrayWriter is never invoked, so its count stays 0. + // Arrow format requires ListArray offset buffer to have N+1 entries even + // when N=0, but getBufferSizeFor(0) returns 0 and the buffer is omitted. + val schema = new StructType() + .add("data", ArrayType(ArrayType(ArrayType(StringType)))) + val writer = ArrowWriter.create(schema, null) + assert(writer.schema === schema) + + // Write a row with an empty outer array + writer.write(InternalRow(ArrayData.toArrayData(Array.empty))) + writer.finish() + + val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0)) + val array0 = reader.getArray(0) + assert(array0.numElements() === 0) + + writer.root.close() + } + + test("SPARK-55056: nested array with map inside empty outer array") { + // Schema: array>> + // Regression test - two-level array with map does not trigger the issue, + // but we keep this test to ensure the fix doesn't break normal cases. + val schema = new StructType() + .add("data", ArrayType(ArrayType(MapType(StringType, StringType)))) + val writer = ArrowWriter.create(schema, null) + assert(writer.schema === schema) + + // Write a row with an empty outer array + writer.write(InternalRow(ArrayData.toArrayData(Array.empty))) + writer.finish() + + val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0)) + val array0 = reader.getArray(0) + assert(array0.numElements() === 0) + + writer.root.close() + } }