github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/actions/lua/encoding/parquet/parquet_test.go (about)

     1  package parquet_test
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"os"
     7  	"testing"
     8  
     9  	"github.com/Shopify/go-lua"
    10  	lualibs "github.com/treeverse/lakefs/pkg/actions/lua"
    11  	"github.com/treeverse/lakefs/pkg/actions/lua/encoding/json"
    12  	"github.com/treeverse/lakefs/pkg/actions/lua/encoding/parquet"
    13  )
    14  
    15  const parquetSchemaRead = `
    16  parquet = require("encoding/parquet")
    17  schema = parquet.get_schema(parquet_content)
    18  
    19  for _, col in pairs(schema) do
    20  	print(col.name .. "\t" .. col.type)
    21  end
    22  `
    23  
    24  const expected = `geoname_id	BYTE_ARRAY
    25  name	BYTE_ARRAY
    26  ascii_name	BYTE_ARRAY
    27  alternate_names	BYTE_ARRAY
    28  feature_class	BYTE_ARRAY
    29  feature_code	BYTE_ARRAY
    30  country_code	BYTE_ARRAY
    31  country_name_en	BYTE_ARRAY
    32  population	INT32
    33  timezone	BYTE_ARRAY
    34  modification_date	INT32
    35  label_en	BYTE_ARRAY
    36  coordinates	BYTE_ARRAY
    37  `
    38  
    39  func TestOpen(t *testing.T) {
    40  	out := bytes.Buffer{}
    41  	l := lua.NewState()
    42  	lualibs.OpenSafe(l, context.Background(), lualibs.OpenSafeConfig{}, &out)
    43  	parquet.Open(l)
    44  	json.Open(l)
    45  
    46  	parquetBytes, err := os.ReadFile("testdata/000.snappy.parquet")
    47  	if err != nil {
    48  		t.Fatal(err)
    49  	}
    50  	l.PushString(string(parquetBytes))
    51  	l.SetGlobal("parquet_content")
    52  
    53  	err = lua.DoString(l, parquetSchemaRead)
    54  	if err != nil {
    55  		t.Fatal(err)
    56  	}
    57  
    58  	printed := out.String()
    59  	if printed != expected {
    60  		t.Fatalf("got unexpected schema: %s", printed)
    61  	}
    62  }