zeroshade commented on code in PR #1099:
URL: https://github.com/apache/arrow-go/pull/1099#discussion_r3752105148
##########
parquet/pqarrow/column_readers.go:
##########
@@ -451,6 +451,10 @@ func (lr *listReader) LoadBatch(nrecords int64) error {
}
func (lr *listReader) BuildArray(lenBound int64) (*arrow.Chunked, error) {
+ return lr.buildArray(lenBound, false)
+}
+
+func (lr *listReader) buildArray(lenBound int64, fixedSize bool)
(*arrow.Chunked, error) {
Review Comment:
Minor: the `fixedSize` parameter looks redundant. `fixedSizeListReader`
embeds `*listReader`, and the code this replaces decided the same thing inline
with `lr.field.Type.ID() == arrow.FIXED_SIZE_LIST`. Keeping that type check
here would drop both the bool parameter and the
`fixedSizeListReader.BuildArray` override, for the same behaviour.
Not a blocker, just more indirection than the one-line condition it replaced.
##########
parquet/pqarrow/column_readers.go:
##########
@@ -529,18 +541,85 @@ func (lr *listReader) BuildArray(lenBound int64)
(*arrow.Chunked, error) {
data := array.NewData(lr.field.Type, int(validityIO.Read), buffers,
[]arrow.ArrayData{item}, int(validityIO.NullCount), 0)
defer data.Release()
- if lr.field.Type.ID() == arrow.FIXED_SIZE_LIST {
- defer data.Buffers()[1].Release()
- offsetData :=
arrow.Int32Traits.CastFromBytes(offsetsBuffer.Bytes())
- listSize := lr.field.Type.(*arrow.FixedSizeListType).Len()
- for x := 1; x < data.Len(); x++ {
- size := offsetData[x] - offsetData[x-1]
- if size != listSize {
- return nil, fmt.Errorf("expected all lists to
be of size=%d, but index %d had size=%d", listSize, x, size)
+ out := array.MakeFromData(data)
+ defer out.Release()
+ return arrow.NewChunked(lr.field.Type, []arrow.Array{out}), nil
+}
+
+func (lr *listReader) buildFixedSizeListArray(length int, offsets []int32,
validityBuffer *memory.Buffer,
+ nullCount int64, item arrow.Array) (*arrow.Chunked, error) {
+ listType := lr.field.Type.(*arrow.FixedSizeListType)
+ listSize := int(listType.Len())
+
+ if offsets[0] != 0 {
+ return nil, fmt.Errorf("fixed-size list first offset must be
zero, got %d", offsets[0])
+ }
+ if int64(offsets[length]) != int64(item.Len()) {
+ return nil, fmt.Errorf("fixed-size list final offset %d does
not match decoded child length %d",
+ offsets[length], item.Len())
+ }
+
+ if nullCount == 0 {
+ for idx := 0; idx < length; idx++ {
+ if size := offsets[idx+1] - offsets[idx]; size !=
int32(listSize) {
+ return nil, fmt.Errorf("expected all lists to
be of size=%d, but index %d had size=%d", listSize, idx, size)
+ }
+ }
+ data := array.NewData(lr.field.Type, length,
[]*memory.Buffer{nil},
+ []arrow.ArrayData{item.Data()}, 0, 0)
+ defer data.Release()
+ out := array.MakeFromData(data)
+ defer out.Release()
+ return arrow.NewChunked(lr.field.Type, []arrow.Array{out}), nil
+ }
+
+ pieces := make([]arrow.Array, 0, length)
+ defer func() { releaseArrays(pieces) }()
+
+ for i := 0; i < length; {
+ valid := !lr.field.Nullable ||
bitutil.BitIsSet(validityBuffer.Bytes(), i)
+ end := i + 1
+ for end < length {
+ nextValid := !lr.field.Nullable ||
bitutil.BitIsSet(validityBuffer.Bytes(), end)
+ if nextValid != valid {
+ break
}
+ end++
}
- data.Buffers()[1] = nil
+
+ if valid {
+ for idx := i; idx < end; idx++ {
+ if size := offsets[idx+1] - offsets[idx]; size
!= int32(listSize) {
+ return nil, fmt.Errorf("expected all
lists to be of size=%d, but index %d had size=%d", listSize, idx, size)
+ }
+ }
+ pieces = append(pieces, array.NewSlice(item,
int64(offsets[i]), int64(offsets[end])))
+ } else {
+ for idx := i; idx < end; idx++ {
+ if size := offsets[idx+1] - offsets[idx]; size
!= 0 {
+ return nil, fmt.Errorf("null fixed-size
list at index %d consumed %d child values", idx, size)
+ }
+ }
+ pieces = append(pieces,
array.MakeArrayOfNull(lr.rctx.mem, listType.Elem(), (end-i)*listSize))
+ }
+ i = end
+ }
+
+ if len(pieces) == 0 {
+ pieces = append(pieces, array.MakeArrayOfNull(lr.rctx.mem,
listType.Elem(), 0))
+ }
+ child, err := array.Concatenate(pieces, lr.rctx.mem)
Review Comment:
Question: is this safe for dictionary-encoded or extension child types?
The `nullCount == 0` fast path hands the decoded child array through
untouched, so it never reaches `Concatenate`. But as soon as there's a null
parent, the child goes through this instead — so a fixed-size list of
dictionary values would take a different code path depending only on whether
any parent happened to be null. If `Concatenate` unifies or rejects
dictionaries, that's an asymmetry worth knowing about. I haven't confirmed
either way.
Separately, worth a brief comment noting the cost profile: parents that
alternate null/valid produce one array per run, so this is O(length)
allocations plus the concat in the worst case. Still strictly better than the
previous behaviour of erroring out, but it's the kind of thing that's easier to
reason about later with a note here.
--
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]
