github.com/apache/arrow/go/v16@v16.1.0/arrow/array/json_reader_test.go (about)

     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 array_test
    18  
    19  import (
    20  	"strings"
    21  	"testing"
    22  
    23  	"github.com/apache/arrow/go/v16/arrow"
    24  	"github.com/apache/arrow/go/v16/arrow/array"
    25  	"github.com/apache/arrow/go/v16/arrow/memory"
    26  	"github.com/stretchr/testify/assert"
    27  )
    28  
    29  const jsondata = `
    30  	{"region": "NY", "model": "3", "sales": 742.0}
    31  	{"region": "NY", "model": "S", "sales": 304.125}
    32  	{"region": "NY", "model": "X", "sales": 136.25}
    33  	{"region": "NY", "model": "Y", "sales": 27.5}
    34  	{"region": "CA", "model": "3", "sales": 512}
    35  	{"region": "CA", "model": "S", "sales": 978}
    36  	{"region": "CA", "model": "X", "sales": 1.0}
    37  	{"region": "CA", "model": "Y", "sales": 69}
    38  	{"region": "QC", "model": "3", "sales": 273.5}
    39  	{"region": "QC", "model": "S", "sales": 13}
    40  	{"region": "QC", "model": "X", "sales": 54}
    41  	{"region": "QC", "model": "Y", "sales": 21}
    42  	{"region": "QC", "model": "3", "sales": 152.25}
    43  	{"region": "QC", "model": "S", "sales": 10}
    44  	{"region": "QC", "model": "X", "sales": 42}
    45  	{"region": "QC", "model": "Y", "sales": 37}`
    46  
    47  func TestJSONReader(t *testing.T) {
    48  	schema := arrow.NewSchema([]arrow.Field{
    49  		{Name: "region", Type: arrow.BinaryTypes.String, Nullable: true},
    50  		{Name: "model", Type: arrow.BinaryTypes.String},
    51  		{Name: "sales", Type: arrow.PrimitiveTypes.Float64, Nullable: true},
    52  	}, nil)
    53  
    54  	rdr := array.NewJSONReader(strings.NewReader(jsondata), schema)
    55  	defer rdr.Release()
    56  
    57  	n := 0
    58  	for rdr.Next() {
    59  		n++
    60  		rec := rdr.Record()
    61  		assert.NotNil(t, rec)
    62  		assert.EqualValues(t, 1, rec.NumRows())
    63  		assert.EqualValues(t, 3, rec.NumCols())
    64  	}
    65  
    66  	assert.NoError(t, rdr.Err())
    67  	assert.Equal(t, 16, n)
    68  }
    69  
    70  func TestJSONReaderAll(t *testing.T) {
    71  	schema := arrow.NewSchema([]arrow.Field{
    72  		{Name: "region", Type: arrow.BinaryTypes.String, Nullable: true},
    73  		{Name: "model", Type: arrow.BinaryTypes.String},
    74  		{Name: "sales", Type: arrow.PrimitiveTypes.Float64, Nullable: true},
    75  	}, nil)
    76  
    77  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
    78  	defer mem.AssertSize(t, 0)
    79  
    80  	rdr := array.NewJSONReader(strings.NewReader(jsondata), schema, array.WithAllocator(mem), array.WithChunk(-1))
    81  	defer rdr.Release()
    82  
    83  	assert.True(t, rdr.Next())
    84  	rec := rdr.Record()
    85  	assert.NotNil(t, rec)
    86  	assert.NoError(t, rdr.Err())
    87  
    88  	assert.EqualValues(t, 16, rec.NumRows())
    89  	assert.EqualValues(t, 3, rec.NumCols())
    90  	assert.False(t, rdr.Next())
    91  }
    92  
    93  func TestJSONReaderChunked(t *testing.T) {
    94  	schema := arrow.NewSchema([]arrow.Field{
    95  		{Name: "region", Type: arrow.BinaryTypes.String, Nullable: true},
    96  		{Name: "model", Type: arrow.BinaryTypes.String},
    97  		{Name: "sales", Type: arrow.PrimitiveTypes.Float64, Nullable: true},
    98  	}, nil)
    99  
   100  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
   101  	defer mem.AssertSize(t, 0)
   102  
   103  	rdr := array.NewJSONReader(strings.NewReader(jsondata), schema, array.WithAllocator(mem), array.WithChunk(4))
   104  	defer rdr.Release()
   105  
   106  	n := 0
   107  	for rdr.Next() {
   108  		n++
   109  		rec := rdr.Record()
   110  		assert.NotNil(t, rec)
   111  		assert.NoError(t, rdr.Err())
   112  		assert.EqualValues(t, 4, rec.NumRows())
   113  	}
   114  
   115  	assert.Equal(t, 4, n)
   116  	assert.NoError(t, rdr.Err())
   117  }
   118  
   119  func TestUnmarshalJSON(t *testing.T) {
   120  	schema := arrow.NewSchema([]arrow.Field{
   121  		{Name: "region", Type: arrow.BinaryTypes.String, Nullable: true},
   122  		{Name: "model", Type: arrow.BinaryTypes.String},
   123  		{Name: "sales", Type: arrow.PrimitiveTypes.Float64, Nullable: true},
   124  	}, nil)
   125  
   126  	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
   127  	defer mem.AssertSize(t, 0)
   128  
   129  	recordBuilder := array.NewRecordBuilder(mem, schema)
   130  	defer recordBuilder.Release()
   131  
   132  	jsondata := `{"region": "NY", "model": "3", "sales": 742.0, "extra": 1234}`
   133  
   134  	err := recordBuilder.UnmarshalJSON([]byte(jsondata))
   135  	assert.NoError(t, err)
   136  
   137  	record := recordBuilder.NewRecord()
   138  	defer record.Release()
   139  
   140  	assert.NotNil(t, record)
   141  }