path_builder_test.go

Documentation: github.com/apache/arrow/go/v15/parquet/pqarrow

     1  // Licensed to the Apache Software Foundation (ASF) under one
     2  // or more contributor license agreements.  See the NOTICE file
     3  // distributed with this work for additional information
     4  // regarding copyright ownership.  The ASF licenses this file
     5  // to you under the Apache License, Version 2.0 (the
     6  // "License"); you may not use this file except in compliance
     7  // with the License.  You may obtain a copy of the License at
     8  //
     9  // http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  
    17  package pqarrow
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  
    23  	"github.com/apache/arrow/go/v15/arrow"
    24  	"github.com/apache/arrow/go/v15/arrow/array"
    25  	"github.com/apache/arrow/go/v15/arrow/memory"
    26  	"github.com/apache/arrow/go/v15/internal/types"
    27  	"github.com/google/uuid"
    28  	"github.com/stretchr/testify/assert"
    29  	"github.com/stretchr/testify/require"
    30  )
    31  
    32  func TestNonNullableSingleList(t *testing.T) {
    33  	// translates to the following parquet schema:
    34  	// required group bag {
    35  	//   repeated group [unseen] (List) {
    36  	//		 required int64 Entires;
    37  	//	 }
    38  	// }
    39  	// So:
    40  	// def level 0: a null entry
    41  	// def level 1: a non-null entry
    42  	bldr := array.NewListBuilder(memory.DefaultAllocator, arrow.PrimitiveTypes.Int64)
    43  	defer bldr.Release()
    44  
    45  	vb := bldr.ValueBuilder().(*array.Int64Builder)
    46  
    47  	bldr.Append(true)
    48  	vb.Append(1)
    49  
    50  	bldr.Append(true)
    51  	vb.Append(2)
    52  	vb.Append(3)
    53  
    54  	bldr.Append(true)
    55  	vb.Append(4)
    56  	vb.Append(5)
    57  	vb.Append(6)
    58  
    59  	arr := bldr.NewListArray()
    60  	defer arr.Release()
    61  
    62  	mp, err := newMultipathLevelBuilder(arr, false)
    63  	require.NoError(t, err)
    64  	defer mp.Release()
    65  
    66  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
    67  	result, err := mp.write(0, ctx)
    68  	require.NoError(t, err)
    69  
    70  	assert.Equal(t, []int16{2, 2, 2, 2, 2, 2}, result.defLevels)
    71  	assert.Equal(t, []int16{0, 0, 1, 0, 1, 1}, result.repLevels)
    72  	assert.Len(t, result.postListVisitedElems, 1)
    73  	assert.EqualValues(t, 0, result.postListVisitedElems[0].start)
    74  	assert.EqualValues(t, 6, result.postListVisitedElems[0].end)
    75  }
    76  
    77  // next group of tests translates to the following parquet schema:
    78  // optional group bag {
    79  //   repeated group [unseen] (List) {
    80  //		 optional int64 Entires;
    81  //	 }
    82  // }
    83  // So:
    84  // def level 0: a null list
    85  // def level 1: an empty list
    86  // def level 2: a null entry
    87  // def level 3: a non-null entry
    88  
    89  func TestNullableSingleListAllNulls(t *testing.T) {
    90  	bldr := array.NewListBuilder(memory.DefaultAllocator, arrow.PrimitiveTypes.Int64)
    91  	defer bldr.Release()
    92  
    93  	bldr.AppendNull()
    94  	bldr.AppendNull()
    95  	bldr.AppendNull()
    96  	bldr.AppendNull()
    97  
    98  	arr := bldr.NewListArray()
    99  	defer arr.Release()
   100  
   101  	mp, err := newMultipathLevelBuilder(arr, true)
   102  	require.NoError(t, err)
   103  	defer mp.Release()
   104  
   105  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   106  	result, err := mp.write(0, ctx)
   107  	require.NoError(t, err)
   108  
   109  	assert.Equal(t, []int16{0, 0, 0, 0}, result.defLevels)
   110  	assert.Equal(t, []int16{0, 0, 0, 0}, result.repLevels)
   111  }
   112  
   113  func TestNullableSingleListAllEmpty(t *testing.T) {
   114  	bldr := array.NewListBuilder(memory.DefaultAllocator, arrow.PrimitiveTypes.Int64)
   115  	defer bldr.Release()
   116  
   117  	bldr.Append(true)
   118  	bldr.Append(true)
   119  	bldr.Append(true)
   120  	bldr.Append(true)
   121  
   122  	arr := bldr.NewListArray()
   123  	defer arr.Release()
   124  
   125  	mp, err := newMultipathLevelBuilder(arr, true)
   126  	require.NoError(t, err)
   127  	defer mp.Release()
   128  
   129  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   130  	result, err := mp.write(0, ctx)
   131  	require.NoError(t, err)
   132  
   133  	assert.Equal(t, []int16{1, 1, 1, 1}, result.defLevels)
   134  	assert.Equal(t, []int16{0, 0, 0, 0}, result.repLevels)
   135  }
   136  
   137  func TestNullableSingleListAllNullEntries(t *testing.T) {
   138  	bldr := array.NewListBuilder(memory.DefaultAllocator, arrow.PrimitiveTypes.Int64)
   139  	defer bldr.Release()
   140  
   141  	vb := bldr.ValueBuilder().(*array.Int64Builder)
   142  
   143  	bldr.Append(true)
   144  	vb.AppendNull()
   145  	bldr.Append(true)
   146  	vb.AppendNull()
   147  	bldr.Append(true)
   148  	vb.AppendNull()
   149  	bldr.Append(true)
   150  	vb.AppendNull()
   151  
   152  	arr := bldr.NewListArray()
   153  	defer arr.Release()
   154  
   155  	mp, err := newMultipathLevelBuilder(arr, true)
   156  	require.NoError(t, err)
   157  	defer mp.Release()
   158  
   159  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   160  	result, err := mp.write(0, ctx)
   161  	require.NoError(t, err)
   162  
   163  	assert.Equal(t, []int16{2, 2, 2, 2}, result.defLevels)
   164  	assert.Equal(t, []int16{0, 0, 0, 0}, result.repLevels)
   165  	assert.Len(t, result.postListVisitedElems, 1)
   166  	assert.EqualValues(t, 0, result.postListVisitedElems[0].start)
   167  	assert.EqualValues(t, 4, result.postListVisitedElems[0].end)
   168  }
   169  
   170  func TestNullableSingleListAllPresentEntries(t *testing.T) {
   171  	bldr := array.NewListBuilder(memory.DefaultAllocator, arrow.PrimitiveTypes.Int64)
   172  	defer bldr.Release()
   173  
   174  	vb := bldr.ValueBuilder().(*array.Int64Builder)
   175  
   176  	bldr.Append(true)
   177  	bldr.Append(true)
   178  	bldr.Append(true)
   179  	vb.Append(1)
   180  	bldr.Append(true)
   181  	bldr.Append(true)
   182  	vb.Append(2)
   183  	vb.Append(3)
   184  
   185  	arr := bldr.NewListArray()
   186  	defer arr.Release()
   187  
   188  	mp, err := newMultipathLevelBuilder(arr, true)
   189  	require.NoError(t, err)
   190  	defer mp.Release()
   191  
   192  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   193  	result, err := mp.write(0, ctx)
   194  	require.NoError(t, err)
   195  
   196  	assert.Equal(t, []int16{1, 1, 3, 1, 3, 3}, result.defLevels)
   197  	assert.Equal(t, []int16{0, 0, 0, 0, 0, 1}, result.repLevels)
   198  	assert.Len(t, result.postListVisitedElems, 1)
   199  	assert.EqualValues(t, 0, result.postListVisitedElems[0].start)
   200  	assert.EqualValues(t, 3, result.postListVisitedElems[0].end)
   201  }
   202  
   203  func TestNullableSingleListSomeNullEntriesSomeNullLists(t *testing.T) {
   204  	bldr := array.NewListBuilder(memory.DefaultAllocator, arrow.PrimitiveTypes.Int64)
   205  	defer bldr.Release()
   206  
   207  	vb := bldr.ValueBuilder().(*array.Int64Builder)
   208  
   209  	bldr.Append(false)
   210  	bldr.Append(true)
   211  	vb.AppendValues([]int64{1, 2, 3}, nil)
   212  	bldr.Append(true)
   213  	bldr.Append(true)
   214  	bldr.AppendNull()
   215  	bldr.AppendNull()
   216  	bldr.Append(true)
   217  	vb.AppendValues([]int64{4, 5}, nil)
   218  	bldr.Append(true)
   219  	vb.AppendNull()
   220  
   221  	arr := bldr.NewListArray()
   222  	defer arr.Release()
   223  
   224  	mp, err := newMultipathLevelBuilder(arr, true)
   225  	require.NoError(t, err)
   226  	defer mp.Release()
   227  
   228  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   229  	result, err := mp.write(0, ctx)
   230  	require.NoError(t, err)
   231  
   232  	assert.Equal(t, []int16{0, 3, 3, 3, 1, 1, 0, 0, 3, 3, 2}, result.defLevels)
   233  	assert.Equal(t, []int16{0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0}, result.repLevels)
   234  }
   235  
   236  // next group of tests translate to the following parquet schema:
   237  //
   238  // optional group bag {
   239  //   repeated group outer_list (List) {
   240  //     optional group nullable {
   241  //       repeated group inner_list (List) {
   242  //         optional int64 Entries;
   243  //       }
   244  //     }
   245  //   }
   246  // }
   247  // So:
   248  // def level 0: null outer list
   249  // def level 1: empty outer list
   250  // def level 2: null inner list
   251  // def level 3: empty inner list
   252  // def level 4: null entry
   253  // def level 5: non-null entry
   254  
   255  func TestNestedListsWithSomeEntries(t *testing.T) {
   256  	listType := arrow.ListOf(arrow.PrimitiveTypes.Int64)
   257  	bldr := array.NewListBuilder(memory.DefaultAllocator, listType)
   258  	defer bldr.Release()
   259  
   260  	nestedBldr := bldr.ValueBuilder().(*array.ListBuilder)
   261  	vb := nestedBldr.ValueBuilder().(*array.Int64Builder)
   262  
   263  	// produce: [null, [[1, 2, 3], [4, 5]], [[], [], []], []]
   264  
   265  	bldr.AppendNull()
   266  	bldr.Append(true)
   267  	nestedBldr.Append(true)
   268  	vb.AppendValues([]int64{1, 2, 3}, nil)
   269  	nestedBldr.Append(true)
   270  	vb.AppendValues([]int64{4, 5}, nil)
   271  
   272  	bldr.Append(true)
   273  	nestedBldr.Append(true)
   274  	nestedBldr.Append(true)
   275  	nestedBldr.Append(true)
   276  	bldr.Append(true)
   277  
   278  	arr := bldr.NewListArray()
   279  	defer arr.Release()
   280  
   281  	mp, err := newMultipathLevelBuilder(arr, true)
   282  	require.NoError(t, err)
   283  	defer mp.Release()
   284  
   285  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   286  	result, err := mp.write(0, ctx)
   287  	require.NoError(t, err)
   288  
   289  	assert.Equal(t, []int16{0, 5, 5, 5, 5, 5, 3, 3, 3, 1}, result.defLevels)
   290  	assert.Equal(t, []int16{0, 0, 2, 2, 1, 2, 0, 1, 1, 0}, result.repLevels)
   291  }
   292  
   293  func TestNestedListsWithSomeNulls(t *testing.T) {
   294  	listType := arrow.ListOf(arrow.PrimitiveTypes.Int64)
   295  	bldr := array.NewListBuilder(memory.DefaultAllocator, listType)
   296  	defer bldr.Release()
   297  
   298  	nestedBldr := bldr.ValueBuilder().(*array.ListBuilder)
   299  	vb := nestedBldr.ValueBuilder().(*array.Int64Builder)
   300  
   301  	// produce: [null, [[1, null, 3], null, null], [[4, 5]]]
   302  
   303  	bldr.AppendNull()
   304  	bldr.Append(true)
   305  	nestedBldr.Append(true)
   306  	vb.AppendValues([]int64{1, 0, 3}, []bool{true, false, true})
   307  	nestedBldr.AppendNull()
   308  	nestedBldr.AppendNull()
   309  	bldr.Append(true)
   310  	nestedBldr.Append(true)
   311  	vb.AppendValues([]int64{4, 5}, nil)
   312  
   313  	arr := bldr.NewListArray()
   314  	defer arr.Release()
   315  
   316  	mp, err := newMultipathLevelBuilder(arr, true)
   317  	require.NoError(t, err)
   318  	defer mp.Release()
   319  
   320  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   321  	result, err := mp.write(0, ctx)
   322  	require.NoError(t, err)
   323  
   324  	assert.Equal(t, []int16{0, 5, 4, 5, 2, 2, 5, 5}, result.defLevels)
   325  	assert.Equal(t, []int16{0, 0, 2, 2, 1, 1, 0, 2}, result.repLevels)
   326  }
   327  
   328  func TestNestedListsSomeNullsSomeEmpty(t *testing.T) {
   329  	listType := arrow.ListOf(arrow.PrimitiveTypes.Int64)
   330  	bldr := array.NewListBuilder(memory.DefaultAllocator, listType)
   331  	defer bldr.Release()
   332  
   333  	nestedBldr := bldr.ValueBuilder().(*array.ListBuilder)
   334  	vb := nestedBldr.ValueBuilder().(*array.Int64Builder)
   335  
   336  	// produce: [null, [[1, null, 3], [], []], [[4, 5]]]
   337  
   338  	bldr.AppendNull()
   339  	bldr.Append(true)
   340  	nestedBldr.Append(true)
   341  	vb.AppendValues([]int64{1, 0, 3}, []bool{true, false, true})
   342  	nestedBldr.Append(true)
   343  	nestedBldr.Append(true)
   344  	bldr.Append(true)
   345  	nestedBldr.Append(true)
   346  	vb.AppendValues([]int64{4, 5}, nil)
   347  
   348  	arr := bldr.NewListArray()
   349  	defer arr.Release()
   350  
   351  	mp, err := newMultipathLevelBuilder(arr, true)
   352  	require.NoError(t, err)
   353  	defer mp.Release()
   354  
   355  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   356  	result, err := mp.write(0, ctx)
   357  	require.NoError(t, err)
   358  
   359  	assert.Equal(t, []int16{0, 5, 4, 5, 3, 3, 5, 5}, result.defLevels)
   360  	assert.Equal(t, []int16{0, 0, 2, 2, 1, 1, 0, 2}, result.repLevels)
   361  }
   362  
   363  func TestNestedExtensionListsWithSomeNulls(t *testing.T) {
   364  	mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
   365  	defer mem.AssertSize(t, 0)
   366  
   367  	listType := arrow.ListOf(types.NewUUIDType())
   368  	bldr := array.NewListBuilder(mem, listType)
   369  	defer bldr.Release()
   370  
   371  	nestedBldr := bldr.ValueBuilder().(*array.ListBuilder)
   372  	vb := nestedBldr.ValueBuilder().(*types.UUIDBuilder)
   373  
   374  	uuid1 := uuid.New()
   375  	uuid3 := uuid.New()
   376  	uuid4 := uuid.New()
   377  	uuid5 := uuid.New()
   378  
   379  	// produce: [null, [[uuid1, null, uuid3], null, null], [[uuid4, uuid5]]]
   380  
   381  	bldr.AppendNull()
   382  	bldr.Append(true)
   383  	nestedBldr.Append(true)
   384  	vb.Append(uuid1)
   385  	vb.AppendNull()
   386  	vb.Append(uuid3)
   387  	nestedBldr.AppendNull()
   388  	nestedBldr.AppendNull()
   389  	bldr.Append(true)
   390  	nestedBldr.Append(true)
   391  	vb.AppendValues([]uuid.UUID{uuid4, uuid5}, nil)
   392  
   393  	arr := bldr.NewListArray()
   394  	defer arr.Release()
   395  
   396  	mp, err := newMultipathLevelBuilder(arr, true)
   397  	require.NoError(t, err)
   398  	defer mp.Release()
   399  
   400  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   401  	result, err := mp.write(0, ctx)
   402  	require.NoError(t, err)
   403  
   404  	assert.Equal(t, []int16{0, 5, 4, 5, 2, 2, 5, 5}, result.defLevels)
   405  	assert.Equal(t, []int16{0, 0, 2, 2, 1, 1, 0, 2}, result.repLevels)
   406  	assert.Equal(t, result.leafArr.NullN(), 1)
   407  }
   408  
   409  // triplenested translates to parquet:
   410  //
   411  // optional group bag {
   412  //   repeated group outer_list (List) {
   413  //     option group nullable {
   414  //       repeated group middle_list (List) {
   415  //         option group nullable {
   416  //           repeated group inner_list (List) {
   417  //              optional int64 Entries;
   418  //           }
   419  //         }
   420  //       }
   421  //     }
   422  //   }
   423  // }
   424  // So:
   425  // def level 0: a outer list
   426  // def level 1: an empty outer list
   427  // def level 2: a null middle list
   428  // def level 3: an empty middle list
   429  // def level 4: an null inner list
   430  // def level 5: an empty inner list
   431  // def level 6: a null entry
   432  // def level 7: a non-null entry
   433  
   434  func TestTripleNestedAllPresent(t *testing.T) {
   435  	listType := arrow.ListOf(arrow.PrimitiveTypes.Int64)
   436  	nestedListType := arrow.ListOf(listType)
   437  	bldr := array.NewListBuilder(memory.DefaultAllocator, nestedListType)
   438  	defer bldr.Release()
   439  
   440  	dblNestedBldr := bldr.ValueBuilder().(*array.ListBuilder)
   441  	nestedBldr := dblNestedBldr.ValueBuilder().(*array.ListBuilder)
   442  	vb := nestedBldr.ValueBuilder().(*array.Int64Builder)
   443  
   444  	// produce: [ [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9]]] ]
   445  	bldr.Append(true)
   446  	dblNestedBldr.Append(true)
   447  	nestedBldr.Append(true)
   448  	vb.AppendValues([]int64{1, 2, 3}, nil)
   449  	nestedBldr.Append(true)
   450  	vb.AppendValues([]int64{4, 5, 6}, nil)
   451  
   452  	dblNestedBldr.Append(true)
   453  	nestedBldr.Append(true)
   454  	vb.AppendValues([]int64{7, 8, 9}, nil)
   455  
   456  	arr := bldr.NewListArray()
   457  	defer arr.Release()
   458  
   459  	mp, err := newMultipathLevelBuilder(arr, true)
   460  	require.NoError(t, err)
   461  	defer mp.Release()
   462  
   463  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   464  	result, err := mp.write(0, ctx)
   465  	require.NoError(t, err)
   466  
   467  	assert.Equal(t, []int16{7, 7, 7, 7, 7, 7, 7, 7, 7}, result.defLevels)
   468  	assert.Equal(t, []int16{0, 3, 3, 2, 3, 3, 1, 3, 3}, result.repLevels)
   469  }
   470  
   471  func TestTripleNestedSomeNullsSomeEmpty(t *testing.T) {
   472  	listType := arrow.ListOf(arrow.PrimitiveTypes.Int64)
   473  	nestedListType := arrow.ListOf(listType)
   474  	bldr := array.NewListBuilder(memory.DefaultAllocator, nestedListType)
   475  	defer bldr.Release()
   476  
   477  	dblNestedBldr := bldr.ValueBuilder().(*array.ListBuilder)
   478  	nestedBldr := dblNestedBldr.ValueBuilder().(*array.ListBuilder)
   479  	vb := nestedBldr.ValueBuilder().(*array.Int64Builder)
   480  
   481  	// produce: [
   482  	//	  [null, [[1, null, 3], []], []],     first row
   483  	//    [[[]], [[], [1, 2]], null, [[3]]],  second row
   484  	//    null,                               third row
   485  	//    []                                  fourth row
   486  	//  ]
   487  
   488  	// first row
   489  	bldr.Append(true)
   490  	dblNestedBldr.AppendNull()
   491  	dblNestedBldr.Append(true)
   492  	nestedBldr.Append(true)
   493  	vb.AppendValues([]int64{1, 0, 3}, []bool{true, false, true})
   494  	nestedBldr.Append(true)
   495  	dblNestedBldr.Append(true)
   496  
   497  	// second row
   498  	bldr.Append(true)
   499  	dblNestedBldr.Append(true)
   500  	nestedBldr.Append(true)
   501  	dblNestedBldr.Append(true)
   502  	nestedBldr.Append(true)
   503  	nestedBldr.Append(true)
   504  	vb.AppendValues([]int64{1, 2}, nil)
   505  	dblNestedBldr.AppendNull()
   506  	dblNestedBldr.Append(true)
   507  	nestedBldr.Append(true)
   508  	vb.Append(3)
   509  
   510  	// third row
   511  	bldr.AppendNull()
   512  
   513  	// fourth row
   514  	bldr.Append(true)
   515  
   516  	arr := bldr.NewListArray()
   517  	defer arr.Release()
   518  
   519  	mp, err := newMultipathLevelBuilder(arr, true)
   520  	require.NoError(t, err)
   521  	defer mp.Release()
   522  
   523  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   524  	result, err := mp.write(0, ctx)
   525  	require.NoError(t, err)
   526  
   527  	assert.Equal(t, []int16{
   528  		2, 7, 6, 7, 5, 3, // first row
   529  		5, 5, 7, 7, 2, 7, // second row
   530  		0, // third row
   531  		1,
   532  	}, result.defLevels)
   533  	assert.Equal(t, []int16{
   534  		0, 1, 3, 3, 2, 1, // first row
   535  		0, 1, 2, 3, 1, 1, // second row
   536  		0, 0,
   537  	}, result.repLevels)
   538  }
   539  
   540  func TestStruct(t *testing.T) {
   541  	structType := arrow.StructOf(arrow.Field{Name: "list", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64), Nullable: true},
   542  		arrow.Field{Name: "Entries", Type: arrow.PrimitiveTypes.Int64, Nullable: true})
   543  
   544  	bldr := array.NewStructBuilder(memory.DefaultAllocator, structType)
   545  	defer bldr.Release()
   546  
   547  	entryBldr := bldr.FieldBuilder(1).(*array.Int64Builder)
   548  	listBldr := bldr.FieldBuilder(0).(*array.ListBuilder)
   549  	vb := listBldr.ValueBuilder().(*array.Int64Builder)
   550  
   551  	// produce: [ {"Entries": 1, "list": [2, 3]}, {"Entries": 4, "list": [5, 6]}, null]
   552  
   553  	bldr.Append(true)
   554  	entryBldr.Append(1)
   555  	listBldr.Append(true)
   556  	vb.AppendValues([]int64{2, 3}, nil)
   557  
   558  	bldr.Append(true)
   559  	entryBldr.Append(4)
   560  	listBldr.Append(true)
   561  	vb.AppendValues([]int64{5, 6}, nil)
   562  
   563  	bldr.AppendNull()
   564  
   565  	arr := bldr.NewArray()
   566  	defer arr.Release()
   567  
   568  	mp, err := newMultipathLevelBuilder(arr, true)
   569  	require.NoError(t, err)
   570  	defer mp.Release()
   571  
   572  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   573  	result, err := mp.writeAll(ctx)
   574  	require.NoError(t, err)
   575  
   576  	assert.Len(t, result, 2)
   577  	assert.Equal(t, []int16{4, 4, 4, 4, 0}, result[0].defLevels)
   578  	assert.Equal(t, []int16{0, 1, 0, 1, 0}, result[0].repLevels)
   579  
   580  	assert.Equal(t, []int16{2, 2, 0}, result[1].defLevels)
   581  	assert.Nil(t, result[1].repLevels)
   582  }
   583  
   584  func TestFixedSizeListNullableElems(t *testing.T) {
   585  	bldr := array.NewFixedSizeListBuilder(memory.DefaultAllocator, 2, arrow.PrimitiveTypes.Int64)
   586  	defer bldr.Release()
   587  
   588  	vb := bldr.ValueBuilder().(*array.Int64Builder)
   589  	bldr.AppendValues([]bool{false, true, true, false})
   590  	vb.AppendValues([]int64{2, 3, 4, 5}, nil)
   591  
   592  	// produce: [null, [2, 3], [4, 5], null]
   593  
   594  	arr := bldr.NewArray()
   595  	defer arr.Release()
   596  
   597  	mp, err := newMultipathLevelBuilder(arr, true)
   598  	require.NoError(t, err)
   599  	defer mp.Release()
   600  
   601  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   602  	result, err := mp.writeAll(ctx)
   603  	require.NoError(t, err)
   604  
   605  	assert.Len(t, result, 1)
   606  	assert.Equal(t, []int16{0, 3, 3, 3, 3, 0}, result[0].defLevels)
   607  	assert.Equal(t, []int16{0, 0, 1, 0, 1, 0}, result[0].repLevels)
   608  
   609  	// null slots take up space in a fixed size list (they can in variable
   610  	// size lists as well) but the actual written values are only the middle
   611  	// elements
   612  	assert.Len(t, result[0].postListVisitedElems, 1)
   613  	assert.EqualValues(t, 2, result[0].postListVisitedElems[0].start)
   614  	assert.EqualValues(t, 6, result[0].postListVisitedElems[0].end)
   615  }
   616  
   617  func TestFixedSizeListMissingMiddleTwoVisitedRanges(t *testing.T) {
   618  	bldr := array.NewFixedSizeListBuilder(memory.DefaultAllocator, 2, arrow.PrimitiveTypes.Int64)
   619  	defer bldr.Release()
   620  
   621  	vb := bldr.ValueBuilder().(*array.Int64Builder)
   622  	bldr.AppendValues([]bool{true, false, true})
   623  	vb.AppendValues([]int64{0, 1, 2, 3}, nil)
   624  
   625  	// produce: [[0, 1], null, [2, 3]]
   626  
   627  	arr := bldr.NewArray()
   628  	defer arr.Release()
   629  
   630  	mp, err := newMultipathLevelBuilder(arr, true)
   631  	require.NoError(t, err)
   632  	defer mp.Release()
   633  
   634  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   635  	result, err := mp.writeAll(ctx)
   636  	require.NoError(t, err)
   637  
   638  	assert.Len(t, result, 1)
   639  	assert.Equal(t, []int16{3, 3, 0, 3, 3}, result[0].defLevels)
   640  	assert.Equal(t, []int16{0, 1, 0, 0, 1}, result[0].repLevels)
   641  
   642  	// null slots take up space in a fixed size list (they can in variable
   643  	// size lists as well) but the actual written values are only the middle
   644  	// elements
   645  	assert.Len(t, result[0].postListVisitedElems, 2)
   646  	assert.EqualValues(t, 0, result[0].postListVisitedElems[0].start)
   647  	assert.EqualValues(t, 2, result[0].postListVisitedElems[0].end)
   648  
   649  	assert.EqualValues(t, 4, result[0].postListVisitedElems[1].start)
   650  	assert.EqualValues(t, 6, result[0].postListVisitedElems[1].end)
   651  }
   652  
   653  func TestPrimitiveNonNullable(t *testing.T) {
   654  	bldr := array.NewInt64Builder(memory.DefaultAllocator)
   655  	defer bldr.Release()
   656  
   657  	bldr.AppendValues([]int64{1, 2, 3, 4}, nil)
   658  
   659  	arr := bldr.NewArray()
   660  	defer arr.Release()
   661  
   662  	mp, err := newMultipathLevelBuilder(arr, false)
   663  	require.NoError(t, err)
   664  	defer mp.Release()
   665  
   666  	ctx := arrowCtxFromContext(NewArrowWriteContext(context.Background(), nil))
   667  	result, err := mp.write(0, ctx)
   668  	require.NoError(t, err)
   669  
   670  	assert.Nil(t, result.defLevels)
   671  	assert.Nil(t, result.repLevels)
   672  
   673  	assert.Len(t, result.postListVisitedElems, 1)
   674  	assert.EqualValues(t, 0, result.postListVisitedElems[0].start)
   675  	assert.EqualValues(t, 4, result.postListVisitedElems[0].end)
   676  }
   677
View as plain text