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

     1  package parquet
     2  
     3  import (
     4  	"github.com/Shopify/go-lua"
     5  	"github.com/treeverse/lakefs/pkg/actions/lua/util"
     6  	"github.com/xitongsys/parquet-go-source/buffer"
     7  	"github.com/xitongsys/parquet-go/reader"
     8  )
     9  
    10  func Open(l *lua.State) {
    11  	parquetOpen := func(l *lua.State) int {
    12  		lua.NewLibrary(l, parquetLibrary)
    13  		return 1
    14  	}
    15  	lua.Require(l, "encoding/parquet", parquetOpen, false)
    16  	l.Pop(1)
    17  }
    18  
    19  var parquetLibrary = []lua.RegistryFunction{
    20  	{Name: "get_schema", Function: getSchema},
    21  }
    22  
    23  func check(l *lua.State, err error) {
    24  	if err != nil {
    25  		lua.Errorf(l, "%s", err.Error())
    26  		panic("unreachable")
    27  	}
    28  }
    29  
    30  func getSchema(l *lua.State) int {
    31  	payload := lua.CheckString(l, 1)
    32  	bin := []byte(payload)
    33  	t := struct{}{}
    34  	fp := buffer.NewBufferFileFromBytes(bin)
    35  	r, err := reader.NewParquetReader(fp, &t, 1)
    36  	check(l, err)
    37  	output := make([]map[string]string, 0)
    38  	for i, elem := range r.Footer.GetSchema() {
    39  		if i == 0 {
    40  			continue // root element
    41  		}
    42  		if elem.Type == nil {
    43  			output = append(output, map[string]string{"name": elem.Name, "type": "N/A"})
    44  		} else {
    45  			output = append(output, map[string]string{"name": elem.Name, "type": elem.Type.String()})
    46  		}
    47  	}
    48  	return util.DeepPush(l, output)
    49  }