zeroshade commented on code in PR #1377:
URL: https://github.com/apache/arrow-go/pull/1377#discussion_r4318664457


##########
arrow/compute/datum.go:
##########
@@ -256,50 +276,83 @@ func NewDatum(value interface{}) Datum {
        switch v := value.(type) {
        case Datum:
                return NewDatum(v.data())
+       case arrow.ArrayData:
+               v.Retain()
+               return &ArrayDatum{Value: v}
        case arrow.Array:
                v.Data().Retain()
-               return &ArrayDatum{v.Data()}
+               return &ArrayDatum{Value: v.Data()}
+       case *arrow.Chunked:
+               v.Retain()
+               return &ChunkedDatum{Value: v}
+       case arrow.RecordBatch:
+               v.Retain()
+               return &RecordDatum{Value: v}
+       case arrow.Table:
+               v.Retain()
+               return &TableDatum{Value: v}
        case scalar.Releasable:
                v.Retain()
-               return NewDatumWithoutOwning(v)
+               if sc, ok := v.(scalar.Scalar); ok {
+                       return &ScalarDatum{Value: sc}
+               }
+               return &ScalarDatum{Value: scalar.MakeScalar(v)}
        case scalar.Scalar:
-               return &ScalarDatum{v}
+               return &ScalarDatum{Value: v}
        default:
-               return &ScalarDatum{scalar.MakeScalar(value)}
+               return &ScalarDatum{Value: scalar.MakeScalar(value)}
        }
 }
 
 // NewDatumWithoutOwning is like NewDatum but does not call Retain on the
-// value it is given. The Datum takes over whatever reference the caller
-// supplies, and Datum.Release consumes that reference: it is an ordinary
-// ArrayDatum, ChunkedDatum, RecordDatum, TableDatum or ScalarDatum whose
-// Release calls the wrapped value's Release.
+// value it is given. The returned Datum borrows the underlying value: it
+// does not own any reference to it, so Datum.Release() is a no-op and
+// will NOT release the caller's reference to the value.
+//
+// The caller is responsible for keeping the underlying value alive for as
+// long as the returned Datum is in use. Because the Datum is non-owning,
+// calling Datum.Release on it is safe (no-op) or unnecessary — it will
+// never free the caller's data. This is the key difference from NewDatum,
+// which retains a reference and whose Datum.Release consumes that
+// reference.
+//
+// Behaviour change (breaking): prior versions also documented a

Review Comment:
   This isn't a breaking change for callers following the released docs, so 
please replace this paragraph with a neutral note: `Release` on a Datum from 
`NewDatumWithoutOwning` is a no-op, so use `NewDatum` when the Datum should own 
the value.



##########
arrow/compute/datum_test.go:
##########
@@ -0,0 +1,119 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build go1.18
+
+package compute
+
+import (
+       "testing"
+
+       "github.com/apache/arrow-go/v18/arrow"
+       "github.com/apache/arrow-go/v18/arrow/array"
+       "github.com/apache/arrow-go/v18/arrow/memory"
+       "github.com/apache/arrow-go/v18/arrow/scalar"
+       "github.com/stretchr/testify/assert"
+)
+
+// TestNewDatumWithoutOwningDoesNotLeakOwner verifies that calling Release()
+// on a Datum constructed with NewDatumWithoutOwning is a no-op: it must not
+// free the caller's underlying allocation. A CheckedAllocator tracks bytes
+// before and after so any silent leak or over-release is caught.
+//
+// This is the regression test for issue #1298.
+func TestNewDatumWithoutOwningDoesNotLeakOwner(t *testing.T) {
+       const n = 8
+       const initialAlloc int = 0
+
+       mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+
+       // -- Array --------------------------------------------------------
+       builder := array.NewInt32Builder(mem)
+       builder.AppendValues([]int32{1, 2, 3, 4, 5, 6, 7, 8}, nil)
+       arr := builder.NewArray()
+       defer arr.Release()
+       builder.Release()
+
+       before := mem.CurrentAlloc()
+       assert.Greater(t, before, initialAlloc, "builder should have allocated 
backing buffers")
+       d := NewDatumWithoutOwning(arr)
+       d.Release() // must be no-op for the caller's buffer
+       assert.Equal(t, before, mem.CurrentAlloc(),
+               "ArrayDatum.Release() must not free caller's allocation")
+
+       // -- Chunked ------------------------------------------------------
+       builder2 := array.NewInt32Builder(mem)
+       builder2.AppendValues([]int32{9, 10, 11, 12, 13, 14, 15, 16}, nil)
+       arr2 := builder2.NewArray()
+       defer arr2.Release()
+       builder2.Release()
+
+       chunked := arrow.NewChunked(arr.DataType(), []arrow.Array{arr, arr2})
+       defer chunked.Release()
+
+       before = mem.CurrentAlloc()
+       d = NewDatumWithoutOwning(chunked)
+       d.Release()
+       assert.Equal(t, before, mem.CurrentAlloc(),
+               "ChunkedDatum.Release() must not free caller's allocation")
+
+       // -- RecordBatch --------------------------------------------------
+       schema := arrow.NewSchema([]arrow.Field{{Name: "col", Type: 
arr.DataType()}}, nil)
+       rec := array.NewRecordBatch(schema, []arrow.Array{arr}, n)
+       defer rec.Release()
+
+       before = mem.CurrentAlloc()
+       d = NewDatumWithoutOwning(rec)
+       d.Release()
+       assert.Equal(t, before, mem.CurrentAlloc(),
+               "RecordDatum.Release() must not free caller's allocation")
+
+       // -- Table --------------------------------------------------------
+       tbl := array.NewTable(schema, []arrow.Column{
+               *arrow.NewColumn(schema.Field(0), chunked),
+       }, -1)
+       defer tbl.Release()
+
+       before = mem.CurrentAlloc()
+       d = NewDatumWithoutOwning(tbl)
+       d.Release()
+       assert.Equal(t, before, mem.CurrentAlloc(),
+               "TableDatum.Release() must not free caller's allocation")
+
+       // -- Scalar -------------------------------------------------------
+       sc := scalar.NewInt32Scalar(42)
+
+       before = mem.CurrentAlloc()
+       d = NewDatumWithoutOwning(sc)
+       d.Release()
+       assert.Equal(t, before, mem.CurrentAlloc(),
+               "ScalarDatum.Release() must not free caller's allocation")
+}
+
+// TestNewDatumOwns verifies the owning NewDatum path releases its
+// Retain on Datum.Release(), leaving the caller a single Release to do.
+func TestNewDatumOwns(t *testing.T) {

Review Comment:
   Nit: `TestNewDatumOwns` makes no assertions. With a checked allocator and 
`mem.AssertSize(t, 0)` at the end, it would catch a missing or doubled release.



##########
arrow/compute/executor.go:
##########
@@ -1282,9 +1282,19 @@ func (a *scalarAggExecutor) Execute(ctx context.Context, 
batch *ExecBatch, data
                return fmt.Errorf("%w: scalar aggregate kernel finalized 
without a result", arrow.ErrInvalid)
        }
 
-       // the result owns its value; boxing it without owning hands that single
-       // reference over to the Datum, which the caller releases
-       out := NewDatumWithoutOwning(result.Take())
+       // the result owns its value; boxing it as an owned Datum hands that
+       // single reference over to the Datum, which the caller releases
+       var out Datum
+       switch taken := result.Take().(type) {
+       case nil:

Review Comment:
   Keeping `EmptyDatum` here is fine.



##########
arrow/compute/datum_test.go:
##########
@@ -0,0 +1,119 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build go1.18
+
+package compute
+
+import (
+       "testing"
+
+       "github.com/apache/arrow-go/v18/arrow"
+       "github.com/apache/arrow-go/v18/arrow/array"
+       "github.com/apache/arrow-go/v18/arrow/memory"
+       "github.com/apache/arrow-go/v18/arrow/scalar"
+       "github.com/stretchr/testify/assert"
+)
+
+// TestNewDatumWithoutOwningDoesNotLeakOwner verifies that calling Release()
+// on a Datum constructed with NewDatumWithoutOwning is a no-op: it must not
+// free the caller's underlying allocation. A CheckedAllocator tracks bytes
+// before and after so any silent leak or over-release is caught.
+//
+// This is the regression test for issue #1298.
+func TestNewDatumWithoutOwningDoesNotLeakOwner(t *testing.T) {
+       const n = 8
+       const initialAlloc int = 0
+
+       mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+
+       // -- Array --------------------------------------------------------
+       builder := array.NewInt32Builder(mem)
+       builder.AppendValues([]int32{1, 2, 3, 4, 5, 6, 7, 8}, nil)
+       arr := builder.NewArray()
+       defer arr.Release()
+       builder.Release()
+
+       before := mem.CurrentAlloc()
+       assert.Greater(t, before, initialAlloc, "builder should have allocated 
backing buffers")
+       d := NewDatumWithoutOwning(arr)
+       d.Release() // must be no-op for the caller's buffer
+       assert.Equal(t, before, mem.CurrentAlloc(),
+               "ArrayDatum.Release() must not free caller's allocation")
+
+       // -- Chunked ------------------------------------------------------

Review Comment:
   With each `borrowed` guard disabled in turn, only the Array case fails. 
Chunked, Record and Table still pass, because the caller's arrays keep their 
memory alive. Scalar still passes, because an `Int32` scalar doesn't allocate. 
Could each case give the wrapped value a reference count of 1 that only the 
caller holds, call the Datum's `Release`, then check that the caller's own 
`Release` brings `CurrentAlloc()` to 0, with `defer mem.AssertSize(t, 0)`? For 
the scalar, use one that allocates, such as a list or binary scalar built with 
`mem`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to