storj.io/minio@v0.0.0-20230509071714-0cbc90f649b1/pkg/s3select/internal/parquet-go/writer_test.go (about)

     1  /*
     2   * Minio Cloud Storage, (C) 2019 Minio, Inc.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package parquet
    18  
    19  import (
    20  	"os"
    21  	"testing"
    22  
    23  	"storj.io/minio/pkg/s3select/internal/parquet-go/data"
    24  	"storj.io/minio/pkg/s3select/internal/parquet-go/gen-go/parquet"
    25  	"storj.io/minio/pkg/s3select/internal/parquet-go/schema"
    26  )
    27  
    28  func TestWriterWrite(t *testing.T) {
    29  	schemaTree := schema.NewTree()
    30  	{
    31  		one, err := schema.NewElement("one", parquet.FieldRepetitionType_REQUIRED,
    32  			parquet.TypePtr(parquet.Type_INT32), parquet.ConvertedTypePtr(parquet.ConvertedType_INT_16),
    33  			nil, nil, nil)
    34  		if err != nil {
    35  			t.Fatal(err)
    36  		}
    37  
    38  		two, err := schema.NewElement("two", parquet.FieldRepetitionType_REQUIRED,
    39  			parquet.TypePtr(parquet.Type_BYTE_ARRAY), parquet.ConvertedTypePtr(parquet.ConvertedType_UTF8),
    40  			nil, nil, nil)
    41  		if err != nil {
    42  			t.Fatal(err)
    43  		}
    44  
    45  		three, err := schema.NewElement("three", parquet.FieldRepetitionType_REQUIRED,
    46  			parquet.TypePtr(parquet.Type_BOOLEAN), nil, nil, nil, nil)
    47  		if err != nil {
    48  			t.Fatal(err)
    49  		}
    50  
    51  		if err := schemaTree.Set("one", one); err != nil {
    52  			t.Fatal(err)
    53  		}
    54  		if err := schemaTree.Set("two", two); err != nil {
    55  			t.Fatal(err)
    56  		}
    57  		if err := schemaTree.Set("three", three); err != nil {
    58  			t.Fatal(err)
    59  		}
    60  	}
    61  
    62  	file, err := os.Create("test.parquet")
    63  	if err != nil {
    64  		t.Fatal(err)
    65  	}
    66  
    67  	writer, err := NewWriter(file, schemaTree, 100)
    68  	if err != nil {
    69  		t.Fatal(err)
    70  	}
    71  
    72  	oneColumn := data.NewColumn(parquet.Type_INT32)
    73  	oneColumn.AddInt32(100, 0, 0)
    74  
    75  	twoColumn := data.NewColumn(parquet.Type_BYTE_ARRAY)
    76  	twoColumn.AddByteArray([]byte("foo"), 0, 0)
    77  
    78  	threeColumn := data.NewColumn(parquet.Type_BOOLEAN)
    79  	threeColumn.AddBoolean(true, 0, 0)
    80  
    81  	record := map[string]*data.Column{
    82  		"one":   oneColumn,
    83  		"two":   twoColumn,
    84  		"three": threeColumn,
    85  	}
    86  
    87  	err = writer.Write(record)
    88  	if err != nil {
    89  		t.Fatal(err)
    90  	}
    91  
    92  	err = writer.Close()
    93  	if err != nil {
    94  		t.Fatal(err)
    95  	}
    96  }
    97  
    98  func TestWriterWriteJSON(t *testing.T) {
    99  	schemaTree := schema.NewTree()
   100  	{
   101  		one, err := schema.NewElement("one", parquet.FieldRepetitionType_REQUIRED,
   102  			parquet.TypePtr(parquet.Type_INT32), parquet.ConvertedTypePtr(parquet.ConvertedType_INT_16),
   103  			nil, nil, nil)
   104  		if err != nil {
   105  			t.Fatal(err)
   106  		}
   107  
   108  		two, err := schema.NewElement("two", parquet.FieldRepetitionType_REQUIRED,
   109  			parquet.TypePtr(parquet.Type_BYTE_ARRAY), parquet.ConvertedTypePtr(parquet.ConvertedType_UTF8),
   110  			nil, nil, nil)
   111  		if err != nil {
   112  			t.Fatal(err)
   113  		}
   114  
   115  		three, err := schema.NewElement("three", parquet.FieldRepetitionType_REQUIRED,
   116  			parquet.TypePtr(parquet.Type_BOOLEAN), nil, nil, nil, nil)
   117  		if err != nil {
   118  			t.Fatal(err)
   119  		}
   120  
   121  		if err := schemaTree.Set("one", one); err != nil {
   122  			t.Fatal(err)
   123  		}
   124  		if err := schemaTree.Set("two", two); err != nil {
   125  			t.Fatal(err)
   126  		}
   127  		if err := schemaTree.Set("three", three); err != nil {
   128  			t.Fatal(err)
   129  		}
   130  	}
   131  
   132  	file, err := os.Create("test.parquet")
   133  	if err != nil {
   134  		t.Fatal(err)
   135  	}
   136  
   137  	writer, err := NewWriter(file, schemaTree, 100)
   138  	if err != nil {
   139  		t.Fatal(err)
   140  	}
   141  
   142  	record := `{"one": 100, "two": "foo", "three": true}`
   143  	err = writer.WriteJSON([]byte(record))
   144  	if err != nil {
   145  		t.Fatal(err)
   146  	}
   147  
   148  	err = writer.Close()
   149  	if err != nil {
   150  		t.Fatal(err)
   151  	}
   152  }