GH: [Go][Parquet]Handle errors to prevent panic

go/arrow/cdata/cdata.go

@@ -189,7 +189,8 @@ func importSchema(schema *CArrowSchema) (ret arrow.Field, err error) {
 			ret.Type = &arrow.DictionaryType{
 				IndexType: ret.Type,
 				ValueType: valueField.Type,
-				Ordered:   schema.dictionary.flags&C.ARROW_FLAG_DICTIONARY_ORDERED != 0}
+				Ordered:   schema.dictionary.flags&C.ARROW_FLAG_DICTIONARY_ORDERED != 0,
+			}
 		}
 
 		return
@@ -395,7 +396,9 @@ func (imp *cimporter) doImportChildren() error {
 		st := imp.dt.(*arrow.StructType)
 		for i, c := range children {
 			imp.children[i].dt = st.Field(i).Type
-			imp.children[i].importChild(imp, c)
+			if err := imp.children[i].importChild(imp, c); err != nil {
+				return err
+			}
 		}
 	case arrow.RUN_END_ENCODED: // import run-ends and values
 		st := imp.dt.(*arrow.RunEndEncodedType)
@@ -416,13 +419,17 @@ func (imp *cimporter) doImportChildren() error {
 		dt := imp.dt.(*arrow.DenseUnionType)
 		for i, c := range children {
 			imp.children[i].dt = dt.Fields()[i].Type
-			imp.children[i].importChild(imp, c)
+			if err := imp.children[i].importChild(imp, c); err != nil {
+				return err
+			}
 		}
 	case arrow.SPARSE_UNION:
 		dt := imp.dt.(*arrow.SparseUnionType)
 		for i, c := range children {
 			imp.children[i].dt = dt.Fields()[i].Type
-			imp.children[i].importChild(imp, c)
+			if err := imp.children[i].importChild(imp, c); err != nil {
+				return err
+			}
 		}
 	}
 
@@ -449,33 +456,28 @@ func (imp *cimporter) doImportArr(src *CArrowArray) error {
 	// and only null columns, then we can release the CArrowArray
 	// struct immediately after import, since we have no imported
 	// memory that we have to track the lifetime of.
+	// On error, we always release regardless of buffer count to avoid leaks.
+	var importErr error
 	defer func() {
-		if imp.alloc.bufCount == 0 {
-			C.ArrowArrayRelease(imp.arr)
-			C.free(unsafe.Pointer(imp.arr))
+		if importErr != nil || imp.alloc.bufCount == 0 {
+			imp.alloc.forceRelease()
 		}
 	}()
 
-	return imp.doImport()
+	importErr = imp.doImport()
+	return importErr
 }
 
 // import is called recursively as needed for importing an array and its children
 // in order to generate array.Data objects
 func (imp *cimporter) doImport() error {
-	// move the array from the src object passed in to the one referenced by
-	// this importer. That way we can set up a finalizer on the created
-	// arrow.ArrayData object so we clean up our Array's memory when garbage collected.
-	defer func(arr *CArrowArray) {
-		// this should only occur in the case of an error happening
-		// during import, at which point we need to clean up the
-		// ArrowArray struct we allocated.
-		if imp.data == nil {
-			C.free(unsafe.Pointer(arr))
-		}
-	}(imp.arr)
-
 	// import any children
 	if err := imp.doImportChildren(); err != nil {
+		for _, c := range imp.children {
+			if c.data != nil {
+				c.data.Release()
+			}
+		}
 		return err
 	}
 
@@ -652,9 +654,7 @@ func (imp *cimporter) importStringLike(offsetByteWidth int64) (err error) {
 		return
 	}
 
-	var (
-		nulls, offsets, values *memory.Buffer
-	)
+	var nulls, offsets, values *memory.Buffer
 	if nulls, err = imp.importNullBitmap(0); err != nil {
 		return
 	}

go/arrow/cdata/cdata_test.go

@@ -1025,3 +1025,101 @@ func TestConfuseGoGc(t *testing.T) {
 
 	wg.Wait()
 }
+
+func TestImportStructWithInvalidSchema(t *testing.T) {
+	mem := mallocator.NewMallocator()
+	defer mem.AssertSize(t, 0)
+
+	arr := createTestStructArr()
+	defer arr.Release()
+
+	carr := createCArr(arr, mem)
+	defer freeTestMallocatorArr(carr, mem)
+
+	sc := testStruct([]string{"+s", "c", "l"}, []string{"", "a", "b"}, []int64{0, flagIsNullable, flagIsNullable})
+	defer freeMallocedSchemas(sc)
+
+	top := (*[1]*CArrowSchema)(unsafe.Pointer(sc))[0]
+	_, err := ImportCRecordBatch(carr, top)
+	assert.Error(t, err)
+}
+
+func TestImportDenseUnionWithInvalidSchema(t *testing.T) {
+	mem := mallocator.NewMallocator()
+	defer mem.AssertSize(t, 0)
+
+	unionArr := createTestDenseUnion()
+	defer unionArr.Release()
+
+	structBld := array.NewStructBuilder(memory.DefaultAllocator, arrow.StructOf(
+		arrow.Field{Name: "union_field", Type: unionArr.DataType(), Nullable: false},
+	))
+	defer structBld.Release()
+
+	unionBld := structBld.FieldBuilder(0).(*array.DenseUnionBuilder)
+	structBld.Append(true)
+	du := unionArr.(*array.DenseUnion)
+	for i := 0; i < du.Len(); i++ {
+		unionBld.Append(du.TypeCode(i))
+		if du.TypeCode(i) == 5 {
+			unionBld.Child(0).(*array.Int32Builder).Append(du.Field(0).(*array.Int32).Value(int(du.ValueOffset(i))))
+		} else {
+			unionBld.Child(1).(*array.Uint8Builder).Append(du.Field(1).(*array.Uint8).Value(int(du.ValueOffset(i))))
+		}
+	}
+
+	structArr := structBld.NewArray()
+	defer structArr.Release()
+
+	carr := createCArr(structArr, mem)
+	defer freeTestMallocatorArr(carr, mem)
+
+	// Create an invalid schema: wrong type for union field (using "i" instead of proper union schema)
+	sc := testStruct([]string{"+s", "i"}, []string{"", "union_field"}, []int64{0, flagIsNullable})
+	defer freeMallocedSchemas(sc)
+
+	top := (*[1]*CArrowSchema)(unsafe.Pointer(sc))[0]
+	_, err := ImportCRecordBatch(carr, top)
+	assert.Error(t, err)
+}
+
+func TestImportSparseUnionWithInvalidSchema(t *testing.T) {
+	mem := mallocator.NewMallocator()
+	defer mem.AssertSize(t, 0)
+
+	unionArr := createTestSparseUnion()
+	defer unionArr.Release()
+
+	structBld := array.NewStructBuilder(memory.DefaultAllocator, arrow.StructOf(
+		arrow.Field{Name: "union_field", Type: unionArr.DataType(), Nullable: false},
+	))
+	defer structBld.Release()
+
+	unionBld := structBld.FieldBuilder(0).(*array.SparseUnionBuilder)
+	structBld.Append(true)
+	su := unionArr.(*array.SparseUnion)
+	for i := 0; i < su.Len(); i++ {
+		unionBld.Append(su.TypeCode(i))
+		if su.TypeCode(i) == 5 {
+			unionBld.Child(0).(*array.Int32Builder).Append(su.Field(0).(*array.Int32).Value(i))
+			unionBld.Child(1).(*array.Uint8Builder).AppendNull()
+		} else {
+			unionBld.Child(0).(*array.Int32Builder).AppendNull()
+			unionBld.Child(1).(*array.Uint8Builder).Append(su.Field(1).(*array.Uint8).Value(i))
+		}
+	}
+
+	structArr := structBld.NewArray()
+	defer structArr.Release()
+
+	carr := createCArr(structArr, mem)
+	defer freeTestMallocatorArr(carr, mem)
+
+	// Create an invalid schema: wrong type for union field (using "u" instead of proper union schema)
+	sc := testStruct([]string{"+s", "u"}, []string{"", "union_field"}, []int64{0, flagIsNullable})
+	defer freeMallocedSchemas(sc)
+
+	top := (*[1]*CArrowSchema)(unsafe.Pointer(sc))[0]
+	_, err := ImportCRecordBatch(carr, top)
+	assert.Error(t, err)
+}

go/arrow/cdata/import_allocator.go

@@ -29,6 +29,7 @@ import "C"
 
 type importAllocator struct {
 	bufCount int64
+	released int32
 
 	arr *CArrowArray
 }
@@ -49,6 +50,12 @@ func (i *importAllocator) Free([]byte) {
 	debug.Assert(atomic.LoadInt64(&i.bufCount) > 0, "too many releases")
 
 	if atomic.AddInt64(&i.bufCount, -1) == 0 {
+		i.forceRelease()
+	}
+}
+
+func (i *importAllocator) forceRelease() {
+	if atomic.CompareAndSwapInt32(&i.released, 0, 1) {
 		defer C.free(unsafe.Pointer(i.arr))
 		C.ArrowArrayRelease(i.arr)
 		if C.ArrowArrayIsReleased(i.arr) != 1 {