schema_test.go

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

     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 schema_test
    18  
    19  import (
    20  	"os"
    21  	"testing"
    22  
    23  	"github.com/apache/arrow/go/v15/parquet"
    24  	format "github.com/apache/arrow/go/v15/parquet/internal/gen-go/parquet"
    25  	"github.com/apache/arrow/go/v15/parquet/schema"
    26  	"github.com/apache/thrift/lib/go/thrift"
    27  	"github.com/stretchr/testify/assert"
    28  	"github.com/stretchr/testify/suite"
    29  )
    30  
    31  func TestColumnPath(t *testing.T) {
    32  	p := parquet.ColumnPath([]string{"toplevel", "leaf"})
    33  	assert.Equal(t, "toplevel.leaf", p.String())
    34  
    35  	p2 := parquet.ColumnPathFromString("toplevel.leaf")
    36  	assert.Equal(t, "toplevel.leaf", p2.String())
    37  
    38  	extend := p2.Extend("anotherlevel")
    39  	assert.Equal(t, "toplevel.leaf.anotherlevel", extend.String())
    40  }
    41  
    42  func NewPrimitive(name string, repetition format.FieldRepetitionType, typ format.Type, fieldID int32) *format.SchemaElement {
    43  	ret := &format.SchemaElement{
    44  		Name:           name,
    45  		RepetitionType: format.FieldRepetitionTypePtr(repetition),
    46  		Type:           format.TypePtr(typ),
    47  	}
    48  	if fieldID >= 0 {
    49  		ret.FieldID = &fieldID
    50  	}
    51  	return ret
    52  }
    53  
    54  func NewGroup(name string, repetition format.FieldRepetitionType, numChildren, fieldID int32) *format.SchemaElement {
    55  	ret := &format.SchemaElement{
    56  		Name:           name,
    57  		RepetitionType: format.FieldRepetitionTypePtr(repetition),
    58  		NumChildren:    &numChildren,
    59  	}
    60  	if fieldID >= 0 {
    61  		ret.FieldID = &fieldID
    62  	}
    63  	return ret
    64  }
    65  
    66  func TestSchemaNodes(t *testing.T) {
    67  	suite.Run(t, new(PrimitiveNodeTestSuite))
    68  	suite.Run(t, new(GroupNodeTestSuite))
    69  	suite.Run(t, new(SchemaConverterSuite))
    70  }
    71  
    72  type PrimitiveNodeTestSuite struct {
    73  	suite.Suite
    74  
    75  	name    string
    76  	fieldID int32
    77  	node    schema.Node
    78  }
    79  
    80  func (p *PrimitiveNodeTestSuite) SetupTest() {
    81  	p.name = "name"
    82  	p.fieldID = 5
    83  }
    84  
    85  func (p *PrimitiveNodeTestSuite) convert(elt *format.SchemaElement) {
    86  	p.node = schema.MustPrimitive(schema.PrimitiveNodeFromThrift(elt))
    87  	p.IsType(&schema.PrimitiveNode{}, p.node)
    88  }
    89  
    90  func (p *PrimitiveNodeTestSuite) TestAttrs() {
    91  	node1 := schema.NewInt32Node("foo" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */)
    92  	node2 := schema.MustPrimitive(schema.NewPrimitiveNodeConverted("bar" /* name */, parquet.Repetitions.Optional, parquet.Types.ByteArray,
    93  		schema.ConvertedTypes.UTF8, 0 /* type len */, 0 /* precision */, 0 /* scale */, -1 /* fieldID */))
    94  
    95  	p.Equal("foo", node1.Name())
    96  	p.Equal(schema.Primitive, node1.Type())
    97  	p.Equal(schema.Primitive, node2.Type())
    98  
    99  	p.Equal(parquet.Repetitions.Repeated, node1.RepetitionType())
   100  	p.Equal(parquet.Repetitions.Optional, node2.RepetitionType())
   101  
   102  	p.Equal(parquet.Types.Int32, node1.PhysicalType())
   103  	p.Equal(parquet.Types.ByteArray, node2.PhysicalType())
   104  
   105  	p.Equal(schema.ConvertedTypes.None, node1.ConvertedType())
   106  	p.Equal(schema.ConvertedTypes.UTF8, node2.ConvertedType())
   107  }
   108  
   109  func (p *PrimitiveNodeTestSuite) TestFromParquet() {
   110  	p.Run("Optional Int32", func() {
   111  		elt := NewPrimitive(p.name, format.FieldRepetitionType_OPTIONAL, format.Type_INT32, p.fieldID)
   112  		p.convert(elt)
   113  
   114  		p.Equal(p.name, p.node.Name())
   115  		p.Equal(p.fieldID, p.node.FieldID())
   116  		p.Equal(parquet.Repetitions.Optional, p.node.RepetitionType())
   117  		p.Equal(parquet.Types.Int32, p.node.(*schema.PrimitiveNode).PhysicalType())
   118  		p.Equal(schema.ConvertedTypes.None, p.node.ConvertedType())
   119  	})
   120  
   121  	p.Run("LogicalType", func() {
   122  		elt := NewPrimitive(p.name, format.FieldRepetitionType_REQUIRED, format.Type_BYTE_ARRAY, p.fieldID)
   123  		elt.ConvertedType = format.ConvertedTypePtr(format.ConvertedType_UTF8)
   124  		p.convert(elt)
   125  
   126  		p.Equal(parquet.Repetitions.Required, p.node.RepetitionType())
   127  		p.Equal(parquet.Types.ByteArray, p.node.(*schema.PrimitiveNode).PhysicalType())
   128  		p.Equal(schema.ConvertedTypes.UTF8, p.node.ConvertedType())
   129  	})
   130  
   131  	p.Run("FixedLenByteArray", func() {
   132  		elt := NewPrimitive(p.name, format.FieldRepetitionType_OPTIONAL, format.Type_FIXED_LEN_BYTE_ARRAY, p.fieldID)
   133  		elt.TypeLength = thrift.Int32Ptr(16)
   134  		p.convert(elt)
   135  
   136  		p.Equal(p.name, p.node.Name())
   137  		p.Equal(p.fieldID, p.node.FieldID())
   138  		p.Equal(parquet.Repetitions.Optional, p.node.RepetitionType())
   139  		p.Equal(parquet.Types.FixedLenByteArray, p.node.(*schema.PrimitiveNode).PhysicalType())
   140  		p.Equal(16, p.node.(*schema.PrimitiveNode).TypeLength())
   141  	})
   142  
   143  	p.Run("convertedtype::decimal", func() {
   144  		elt := NewPrimitive(p.name, format.FieldRepetitionType_OPTIONAL, format.Type_FIXED_LEN_BYTE_ARRAY, p.fieldID)
   145  		elt.ConvertedType = format.ConvertedTypePtr(format.ConvertedType_DECIMAL)
   146  		elt.TypeLength = thrift.Int32Ptr(6)
   147  		elt.Scale = thrift.Int32Ptr(2)
   148  		elt.Precision = thrift.Int32Ptr(12)
   149  
   150  		p.convert(elt)
   151  		p.Equal(parquet.Types.FixedLenByteArray, p.node.(*schema.PrimitiveNode).PhysicalType())
   152  		p.Equal(schema.ConvertedTypes.Decimal, p.node.ConvertedType())
   153  		p.Equal(6, p.node.(*schema.PrimitiveNode).TypeLength())
   154  		p.EqualValues(2, p.node.(*schema.PrimitiveNode).DecimalMetadata().Scale)
   155  		p.EqualValues(12, p.node.(*schema.PrimitiveNode).DecimalMetadata().Precision)
   156  	})
   157  }
   158  
   159  func (p *PrimitiveNodeTestSuite) TestEquals() {
   160  	const fieldID = -1
   161  	node1 := schema.NewInt32Node("foo" /* name */, parquet.Repetitions.Required, fieldID)
   162  	node2 := schema.NewInt64Node("foo" /* name */, parquet.Repetitions.Required, fieldID)
   163  	node3 := schema.NewInt32Node("bar" /* name */, parquet.Repetitions.Required, fieldID)
   164  	node4 := schema.NewInt32Node("foo" /* name */, parquet.Repetitions.Optional, fieldID)
   165  	node5 := schema.NewInt32Node("foo" /* name */, parquet.Repetitions.Required, fieldID)
   166  
   167  	p.True(node1.Equals(node1))
   168  	p.False(node1.Equals(node2))
   169  	p.False(node1.Equals(node3))
   170  	p.False(node1.Equals(node4))
   171  	p.True(node1.Equals(node5))
   172  
   173  	flba1 := schema.MustPrimitive(schema.NewPrimitiveNodeConverted("foo" /* name */, parquet.Repetitions.Required, parquet.Types.FixedLenByteArray,
   174  		schema.ConvertedTypes.Decimal, 12 /* type len */, 4 /* precision */, 2 /* scale */, fieldID))
   175  	flba2 := schema.MustPrimitive(schema.NewPrimitiveNodeConverted("foo" /* name */, parquet.Repetitions.Required, parquet.Types.FixedLenByteArray,
   176  		schema.ConvertedTypes.Decimal, 1 /* type len */, 4 /* precision */, 2 /* scale */, fieldID))
   177  	flba2.SetTypeLength(12)
   178  
   179  	flba3 := schema.MustPrimitive(schema.NewPrimitiveNodeConverted("foo" /* name */, parquet.Repetitions.Required, parquet.Types.FixedLenByteArray,
   180  		schema.ConvertedTypes.Decimal, 1 /* type len */, 4 /* precision */, 2 /* scale */, fieldID))
   181  	flba3.SetTypeLength(16)
   182  
   183  	flba4 := schema.MustPrimitive(schema.NewPrimitiveNodeConverted("foo" /* name */, parquet.Repetitions.Required, parquet.Types.FixedLenByteArray,
   184  		schema.ConvertedTypes.Decimal, 12 /* type len */, 4 /* precision */, 0 /* scale */, fieldID))
   185  	flba5 := schema.MustPrimitive(schema.NewPrimitiveNodeConverted("foo" /* name */, parquet.Repetitions.Required, parquet.Types.FixedLenByteArray,
   186  		schema.ConvertedTypes.None, 12 /* type len */, 4 /* precision */, 0 /* scale */, fieldID))
   187  
   188  	p.True(flba1.Equals(flba2))
   189  	p.False(flba1.Equals(flba3))
   190  	p.False(flba1.Equals(flba4))
   191  	p.False(flba1.Equals(flba5))
   192  }
   193  
   194  func (p *PrimitiveNodeTestSuite) TestPhysicalLogicalMapping() {
   195  	tests := []struct {
   196  		typ       parquet.Type
   197  		cnv       schema.ConvertedType
   198  		typLen    int
   199  		precision int
   200  		scale     int
   201  		shouldErr bool
   202  	}{
   203  		{parquet.Types.Int32, schema.ConvertedTypes.Int32, 0 /* type len */, 0 /* precision */, 0 /* scale */, false},
   204  		{parquet.Types.ByteArray, schema.ConvertedTypes.JSON, 0 /* type len */, 0 /* precision */, 0 /* scale */, false},
   205  		{parquet.Types.Int32, schema.ConvertedTypes.JSON, 0 /* type len */, 0 /* precision */, 0 /* scale */, true},
   206  		{parquet.Types.Int64, schema.ConvertedTypes.TimestampMillis, 0 /* type len */, 0 /* precision */, 0 /* scale */, false},
   207  		{parquet.Types.Int32, schema.ConvertedTypes.Int64, 0 /* type len */, 0 /* precision */, 0 /* scale */, true},
   208  		{parquet.Types.ByteArray, schema.ConvertedTypes.Int8, 0 /* type len */, 0 /* precision */, 0 /* scale */, true},
   209  		{parquet.Types.ByteArray, schema.ConvertedTypes.Interval, 0 /* type len */, 0 /* precision */, 0 /* scale */, true},
   210  		{parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Enum, 0 /* type len */, 0 /* precision */, 0 /* scale */, true},
   211  		{parquet.Types.ByteArray, schema.ConvertedTypes.Enum, 0 /* type len */, 0 /* precision */, 0 /* scale */, false},
   212  		{parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Decimal, 0 /* type len */, 2 /* precision */, 4 /* scale */, true},
   213  		{parquet.Types.Float, schema.ConvertedTypes.Decimal, 0 /* type len */, 2 /* precision */, 4 /* scale */, true},
   214  		{parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Decimal, 0 /* type len */, 4 /* precision */, 0 /* scale */, true},
   215  		{parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Decimal, 10 /* type len */, 4 /* precision */, -1 /* scale */, true},
   216  		{parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Decimal, 10 /* type len */, 2 /* precision */, 4 /* scale */, true},
   217  		{parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Decimal, 10 /* type len */, 6 /* precision */, 4 /* scale */, false},
   218  		{parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Interval, 12 /* type len */, 0 /* precision */, 0 /* scale */, false},
   219  		{parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Interval, 10 /* type len */, 0 /* precision */, 0 /* scale */, true},
   220  	}
   221  	for _, tt := range tests {
   222  		p.Run(tt.typ.String(), func() {
   223  			_, err := schema.NewPrimitiveNodeConverted("foo" /* name */, parquet.Repetitions.Required, tt.typ, tt.cnv, tt.typLen, tt.precision, tt.scale, -1 /* fieldID */)
   224  			if tt.shouldErr {
   225  				p.Error(err)
   226  			} else {
   227  				p.NoError(err)
   228  			}
   229  		})
   230  	}
   231  }
   232  
   233  type GroupNodeTestSuite struct {
   234  	suite.Suite
   235  }
   236  
   237  func (g *GroupNodeTestSuite) fields1() []schema.Node {
   238  	return schema.FieldList{
   239  		schema.NewInt32Node("one" /* name */, parquet.Repetitions.Required, -1 /* fieldID */),
   240  		schema.NewInt64Node("two" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */),
   241  		schema.NewFloat64Node("three" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */),
   242  	}
   243  }
   244  
   245  func (g *GroupNodeTestSuite) fields2() []schema.Node {
   246  	return schema.FieldList{
   247  		schema.NewInt32Node("duplicate" /* name */, parquet.Repetitions.Required, -1 /* fieldID */),
   248  		schema.NewInt64Node("unique" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */),
   249  		schema.NewFloat64Node("duplicate" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */),
   250  	}
   251  }
   252  
   253  func (g *GroupNodeTestSuite) TestAttrs() {
   254  	fields := g.fields1()
   255  
   256  	node1 := schema.MustGroup(schema.NewGroupNode("foo" /* name */, parquet.Repetitions.Repeated, fields, -1 /* fieldID */))
   257  	node2 := schema.MustGroup(schema.NewGroupNodeConverted("bar" /* name */, parquet.Repetitions.Optional, fields, schema.ConvertedTypes.List, -1 /* fieldID */))
   258  
   259  	g.Equal("foo", node1.Name())
   260  	g.Equal(schema.Group, node1.Type())
   261  	g.Equal(len(fields), node1.NumFields())
   262  	g.Equal(parquet.Repetitions.Repeated, node1.RepetitionType())
   263  	g.Equal(parquet.Repetitions.Optional, node2.RepetitionType())
   264  
   265  	g.Equal(schema.ConvertedTypes.None, node1.ConvertedType())
   266  	g.Equal(schema.ConvertedTypes.List, node2.ConvertedType())
   267  }
   268  
   269  func (g *GroupNodeTestSuite) TestEquals() {
   270  	f1 := g.fields1()
   271  	f2 := g.fields1()
   272  
   273  	group1 := schema.Must(schema.NewGroupNode("group" /* name */, parquet.Repetitions.Repeated, f1, -1 /* fieldID */))
   274  	group2 := schema.Must(schema.NewGroupNode("group" /* name */, parquet.Repetitions.Repeated, f2, -1 /* fieldID */))
   275  	group3 := schema.Must(schema.NewGroupNode("group2" /* name */, parquet.Repetitions.Repeated, f2, -1 /* fieldID */))
   276  
   277  	f2 = append(f2, schema.NewFloat32Node("four" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */))
   278  	group4 := schema.Must(schema.NewGroupNode("group" /* name */, parquet.Repetitions.Repeated, f2, -1 /* fieldID */))
   279  	group5 := schema.Must(schema.NewGroupNode("group" /* name */, parquet.Repetitions.Repeated, g.fields1(), -1 /* fieldID */))
   280  
   281  	g.True(group1.Equals(group1))
   282  	g.True(group1.Equals(group2))
   283  	g.False(group1.Equals(group3))
   284  	g.False(group1.Equals(group4))
   285  	g.False(group5.Equals(group4))
   286  }
   287  
   288  func (g *GroupNodeTestSuite) TestFieldIndex() {
   289  	fields := g.fields1()
   290  	group := schema.MustGroup(schema.NewGroupNode("group" /* name */, parquet.Repetitions.Required, fields, -1 /* fieldID */))
   291  	for idx, field := range fields {
   292  		f := group.Field(idx)
   293  		g.Same(field, f)
   294  		g.Equal(idx, group.FieldIndexByField(f))
   295  		g.Equal(idx, group.FieldIndexByName(field.Name()))
   296  	}
   297  
   298  	// Non field nodes
   299  	nonFieldAlien := schema.NewInt32Node("alien" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   300  	nonFieldFamiliar := schema.NewInt32Node("one" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */)
   301  	g.Less(group.FieldIndexByField(nonFieldAlien), 0)
   302  	g.Less(group.FieldIndexByField(nonFieldFamiliar), 0)
   303  }
   304  
   305  func (g *GroupNodeTestSuite) TestFieldIndexDuplicateName() {
   306  	fields := g.fields2()
   307  	group := schema.MustGroup(schema.NewGroupNode("group" /* name */, parquet.Repetitions.Required, fields, -1 /* fieldID */))
   308  	for idx, field := range fields {
   309  		f := group.Field(idx)
   310  		g.Same(f, field)
   311  		g.Equal(idx, group.FieldIndexByField(f))
   312  	}
   313  }
   314  
   315  type SchemaConverterSuite struct {
   316  	suite.Suite
   317  
   318  	name string
   319  	node schema.Node
   320  }
   321  
   322  func (s *SchemaConverterSuite) SetupSuite() {
   323  	s.name = "parquet_schema"
   324  }
   325  
   326  func (s *SchemaConverterSuite) convert(elems []*format.SchemaElement) {
   327  	s.node = schema.Must(schema.FromParquet(elems))
   328  	s.Equal(schema.Group, s.node.Type())
   329  }
   330  
   331  func (s *SchemaConverterSuite) checkParentConsistency(groupRoot *schema.GroupNode) bool {
   332  	// each node should have the group as parent
   333  	for i := 0; i < groupRoot.NumFields(); i++ {
   334  		field := groupRoot.Field(i)
   335  		if field.Parent() != groupRoot {
   336  			return false
   337  		}
   338  		if field.Type() == schema.Group {
   339  			if !s.checkParentConsistency(field.(*schema.GroupNode)) {
   340  				return false
   341  			}
   342  		}
   343  	}
   344  	return true
   345  }
   346  
   347  func (s *SchemaConverterSuite) TestNestedExample() {
   348  	elements := make([]*format.SchemaElement, 0)
   349  	elements = append(elements,
   350  		NewGroup(s.name, format.FieldRepetitionType_REPEATED, 2 /* numChildren */, 0 /* fieldID */),
   351  		NewPrimitive("a" /* name */, format.FieldRepetitionType_REQUIRED, format.Type_INT32, 1 /* fieldID */),
   352  		NewGroup("bag" /* name */, format.FieldRepetitionType_OPTIONAL, 1 /* numChildren */, 2 /* fieldID */))
   353  	elt := NewGroup("b" /* name */, format.FieldRepetitionType_REPEATED, 1 /* numChildren */, 3 /* fieldID */)
   354  	elt.ConvertedType = format.ConvertedTypePtr(format.ConvertedType_LIST)
   355  	elements = append(elements, elt, NewPrimitive("item" /* name */, format.FieldRepetitionType_OPTIONAL, format.Type_INT64, 4 /* fieldID */))
   356  
   357  	s.convert(elements)
   358  
   359  	// construct the expected schema
   360  	fields := make([]schema.Node, 0)
   361  	fields = append(fields, schema.NewInt32Node("a" /* name */, parquet.Repetitions.Required, 1 /* fieldID */))
   362  
   363  	// 3-level list encoding
   364  	item := schema.NewInt64Node("item" /* name */, parquet.Repetitions.Optional, 4 /* fieldID */)
   365  	list := schema.MustGroup(schema.NewGroupNodeConverted("b" /* name */, parquet.Repetitions.Repeated, schema.FieldList{item}, schema.ConvertedTypes.List, 3 /* fieldID */))
   366  	bag := schema.MustGroup(schema.NewGroupNode("bag" /* name */, parquet.Repetitions.Optional, schema.FieldList{list}, 2 /* fieldID */))
   367  	fields = append(fields, bag)
   368  
   369  	sc := schema.MustGroup(schema.NewGroupNode(s.name, parquet.Repetitions.Repeated, fields, 0 /* fieldID */))
   370  	s.True(sc.Equals(s.node))
   371  	s.Nil(s.node.Parent())
   372  	s.True(s.checkParentConsistency(s.node.(*schema.GroupNode)))
   373  }
   374  
   375  func (s *SchemaConverterSuite) TestZeroColumns() {
   376  	elements := []*format.SchemaElement{NewGroup("schema" /* name */, format.FieldRepetitionType_REPEATED, 0 /* numChildren */, 0 /* fieldID */)}
   377  	s.NotPanics(func() { s.convert(elements) })
   378  }
   379  
   380  func (s *SchemaConverterSuite) TestInvalidRoot() {
   381  	// According to the Parquet spec, the first element in the list<SchemaElement>
   382  	// is a group whose children (and their descendants) contain all of the rest of
   383  	// the flattened schema elements. If the first element is not a group, it is malformed
   384  	elements := []*format.SchemaElement{NewPrimitive("not-a-group" /* name */, format.FieldRepetitionType_REQUIRED,
   385  		format.Type_INT32, 0 /* fieldID */), format.NewSchemaElement()}
   386  	s.Panics(func() { s.convert(elements) })
   387  
   388  	// While the parquet spec indicates that the root group should have REPEATED
   389  	// repetition type, some implementations may return REQUIRED or OPTIONAL
   390  	// groups as the first element. These tests check that this is okay as a
   391  	// practicality matter
   392  	elements = []*format.SchemaElement{
   393  		NewGroup("not-repeated" /* name */, format.FieldRepetitionType_REQUIRED, 1 /* numChildren */, 0 /* fieldID */),
   394  		NewPrimitive("a" /* name */, format.FieldRepetitionType_REQUIRED, format.Type_INT32, 1 /* fieldID */)}
   395  	s.NotPanics(func() { s.convert(elements) })
   396  
   397  	elements[0] = NewGroup("not-repeated" /* name */, format.FieldRepetitionType_OPTIONAL, 1 /* numChildren */, 0 /* fieldID */)
   398  	s.NotPanics(func() { s.convert(elements) })
   399  }
   400  
   401  func (s *SchemaConverterSuite) TestNotEnoughChildren() {
   402  	s.Panics(func() {
   403  		s.convert([]*format.SchemaElement{NewGroup(s.name, format.FieldRepetitionType_REPEATED, 2 /* numChildren */, 0 /* fieldID */)})
   404  	})
   405  }
   406  
   407  func TestColumnDesc(t *testing.T) {
   408  	n := schema.MustPrimitive(schema.NewPrimitiveNodeConverted("name" /* name */, parquet.Repetitions.Optional, parquet.Types.ByteArray,
   409  		schema.ConvertedTypes.UTF8, 0 /* type len */, 0 /* precision */, 0 /* scale */, -1 /* fieldID */))
   410  	descr := schema.NewColumn(n, 4, 1)
   411  
   412  	assert.Equal(t, "name", descr.Name())
   413  	assert.EqualValues(t, 4, descr.MaxDefinitionLevel())
   414  	assert.EqualValues(t, 1, descr.MaxRepetitionLevel())
   415  	assert.Equal(t, parquet.Types.ByteArray, descr.PhysicalType())
   416  	assert.Equal(t, -1, descr.TypeLength())
   417  
   418  	expectedDesc := `column descriptor = {
   419    name: name,
   420    path: ,
   421    physical_type: BYTE_ARRAY,
   422    converted_type: UTF8,
   423    logical_type: String,
   424    max_definition_level: 4,
   425    max_repetition_level: 1,
   426  }`
   427  	assert.Equal(t, expectedDesc, descr.String())
   428  
   429  	n = schema.MustPrimitive(schema.NewPrimitiveNodeConverted("name" /* name */, parquet.Repetitions.Optional, parquet.Types.FixedLenByteArray, schema.ConvertedTypes.Decimal, 12 /* type len */, 10 /* precision */, 4 /* scale */, -1 /* fieldID */))
   430  	descr2 := schema.NewColumn(n, 4, 1)
   431  
   432  	assert.Equal(t, parquet.Types.FixedLenByteArray, descr2.PhysicalType())
   433  	assert.Equal(t, 12, descr2.TypeLength())
   434  
   435  	expectedDesc = `column descriptor = {
   436    name: name,
   437    path: ,
   438    physical_type: FIXED_LEN_BYTE_ARRAY,
   439    converted_type: DECIMAL,
   440    logical_type: Decimal(precision=10, scale=4),
   441    max_definition_level: 4,
   442    max_repetition_level: 1,
   443    length: 12,
   444    precision: 10,
   445    scale: 4,
   446  }`
   447  	assert.Equal(t, expectedDesc, descr2.String())
   448  }
   449  
   450  func TestSchemaDescriptor(t *testing.T) {
   451  	t.Run("Equals", func(t *testing.T) {
   452  		inta := schema.NewInt32Node("a" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   453  		intb := schema.NewInt64Node("b" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */)
   454  		intb2 := schema.NewInt64Node("b2" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */)
   455  		intc := schema.NewByteArrayNode("c" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */)
   456  
   457  		item1 := schema.NewInt64Node("item1" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   458  		item2 := schema.NewBooleanNode("item2" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */)
   459  		item3 := schema.NewInt32Node("item3" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */)
   460  		list := schema.MustGroup(schema.NewGroupNodeConverted("records" /* name */, parquet.Repetitions.Repeated, schema.FieldList{item1, item2, item3}, schema.ConvertedTypes.List, -1 /* fieldID */))
   461  
   462  		bag := schema.MustGroup(schema.NewGroupNode("bag" /* name */, parquet.Repetitions.Optional, schema.FieldList{list}, -1 /* fieldID */))
   463  		bag2 := schema.MustGroup(schema.NewGroupNode("bag" /* name */, parquet.Repetitions.Required, schema.FieldList{list}, -1 /* fieldID */))
   464  
   465  		descr1 := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, schema.FieldList{inta, intb, intc, bag}, -1 /* fieldID */)))
   466  		assert.True(t, descr1.Equals(descr1))
   467  
   468  		descr2 := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, schema.FieldList{inta, intb, intc, bag2}, -1 /* fieldID */)))
   469  		assert.False(t, descr1.Equals(descr2))
   470  
   471  		descr3 := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, schema.FieldList{inta, intb2, intc, bag}, -1 /* fieldID */)))
   472  		assert.False(t, descr1.Equals(descr3))
   473  
   474  		descr4 := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("SCHEMA" /* name */, parquet.Repetitions.Repeated, schema.FieldList{inta, intb, intc, bag}, -1 /* fieldID */)))
   475  		assert.True(t, descr1.Equals(descr4))
   476  
   477  		descr5 := schema.NewSchema(schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, schema.FieldList{inta, intb, intc, bag, intb2}, -1 /* fieldID */)))
   478  		assert.False(t, descr1.Equals(descr5))
   479  
   480  		col1 := schema.NewColumn(inta, 5 /* maxDefLvl */, 1 /* maxRepLvl */)
   481  		col2 := schema.NewColumn(inta, 6 /* maxDefLvl */, 1 /* maxRepLvl */)
   482  		col3 := schema.NewColumn(inta, 5 /* maxDefLvl */, 2 /* maxRepLvl */)
   483  
   484  		assert.True(t, col1.Equals(col1))
   485  		assert.False(t, col1.Equals(col2))
   486  		assert.False(t, col2.Equals(col3))
   487  	})
   488  
   489  	t.Run("BuildTree", func(t *testing.T) {
   490  		inta := schema.NewInt32Node("a" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   491  		fields := schema.FieldList{inta}
   492  		fields = append(fields,
   493  			schema.NewInt64Node("b" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */),
   494  			schema.NewByteArrayNode("c" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */))
   495  
   496  		item1 := schema.NewInt64Node("item1" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   497  		item2 := schema.NewBooleanNode("item2" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */)
   498  		item3 := schema.NewInt32Node("item3" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */)
   499  		list := schema.MustGroup(schema.NewGroupNodeConverted("records" /* name */, parquet.Repetitions.Repeated, schema.FieldList{item1, item2, item3}, schema.ConvertedTypes.List, -1 /* fieldID */))
   500  		bag := schema.MustGroup(schema.NewGroupNode("bag" /* name */, parquet.Repetitions.Optional, schema.FieldList{list}, -1 /* fieldID */))
   501  		fields = append(fields, bag)
   502  
   503  		sc := schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, fields, -1 /* fieldID */))
   504  		descr := schema.NewSchema(sc)
   505  
   506  		const nleaves = 6
   507  		assert.Equal(t, nleaves, descr.NumColumns())
   508  
   509  		//                             mdef mrep
   510  		// required int32 a            0    0
   511  		// optional int64 b            1    0
   512  		// repeated byte_array c       1    1
   513  		// optional group bag          1    0
   514  		//   repeated group records    2    1
   515  		//     required int64 item1    2    1
   516  		//     optional boolean item2  3    1
   517  		//     repeated int32 item3    3    2
   518  		var (
   519  			exMaxDefLevels = [...]int16{0, 1, 1, 2, 3, 3}
   520  			exMaxRepLevels = [...]int16{0, 0, 1, 1, 1, 2}
   521  		)
   522  
   523  		for i := 0; i < nleaves; i++ {
   524  			col := descr.Column(i)
   525  			assert.Equal(t, exMaxDefLevels[i], col.MaxDefinitionLevel())
   526  			assert.Equal(t, exMaxRepLevels[i], col.MaxRepetitionLevel())
   527  		}
   528  
   529  		assert.Equal(t, "a", descr.Column(0).Path())
   530  		assert.Equal(t, "b", descr.Column(1).Path())
   531  		assert.Equal(t, "c", descr.Column(2).Path())
   532  		assert.Equal(t, "bag.records.item1", descr.Column(3).Path())
   533  		assert.Equal(t, "bag.records.item2", descr.Column(4).Path())
   534  		assert.Equal(t, "bag.records.item3", descr.Column(5).Path())
   535  
   536  		for i := 0; i < nleaves; i++ {
   537  			col := descr.Column(i)
   538  			assert.Equal(t, i, descr.ColumnIndexByNode(col.SchemaNode()))
   539  		}
   540  
   541  		nonColumnAlien := schema.NewInt32Node("alien" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   542  		nonColumnFamiliar := schema.NewInt32Node("a" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */)
   543  		assert.Less(t, descr.ColumnIndexByNode(nonColumnAlien), 0)
   544  		assert.Less(t, descr.ColumnIndexByNode(nonColumnFamiliar), 0)
   545  
   546  		assert.Same(t, inta, descr.ColumnRoot(0))
   547  		assert.Same(t, bag, descr.ColumnRoot(3))
   548  		assert.Same(t, bag, descr.ColumnRoot(4))
   549  		assert.Same(t, bag, descr.ColumnRoot(5))
   550  
   551  		assert.Same(t, sc, descr.Root())
   552  	})
   553  
   554  	t.Run("HasRepeatedFields", func(t *testing.T) {
   555  		inta := schema.NewInt32Node("a" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   556  		fields := schema.FieldList{inta}
   557  		fields = append(fields,
   558  			schema.NewInt64Node("b" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */),
   559  			schema.NewByteArrayNode("c" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */))
   560  
   561  		sc := schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, fields, -1 /* fieldID */))
   562  		descr := schema.NewSchema(sc)
   563  		assert.True(t, descr.HasRepeatedFields())
   564  
   565  		item1 := schema.NewInt64Node("item1" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   566  		item2 := schema.NewBooleanNode("item2" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */)
   567  		item3 := schema.NewInt32Node("item3" /* name */, parquet.Repetitions.Repeated, -1 /* fieldID */)
   568  		list := schema.MustGroup(schema.NewGroupNodeConverted("records" /* name */, parquet.Repetitions.Repeated, schema.FieldList{item1, item2, item3}, schema.ConvertedTypes.List, -1 /* fieldID */))
   569  		bag := schema.MustGroup(schema.NewGroupNode("bag" /* name */, parquet.Repetitions.Optional, schema.FieldList{list}, -1 /* fieldID */))
   570  		fields = append(fields, bag)
   571  
   572  		sc = schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, fields, -1 /* fieldID */))
   573  		descr = schema.NewSchema(sc)
   574  		assert.True(t, descr.HasRepeatedFields())
   575  
   576  		itemKey := schema.NewInt64Node("key" /* name */, parquet.Repetitions.Required, -1 /* fieldID */)
   577  		itemValue := schema.NewBooleanNode("value" /* name */, parquet.Repetitions.Optional, -1 /* fieldID */)
   578  		sc = schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, append(fields, schema.FieldList{
   579  			schema.MustGroup(schema.NewGroupNode("my_map" /* name */, parquet.Repetitions.Optional, schema.FieldList{
   580  				schema.MustGroup(schema.NewGroupNodeConverted("map" /* name */, parquet.Repetitions.Repeated, schema.FieldList{itemKey, itemValue}, schema.ConvertedTypes.Map, -1 /* fieldID */)),
   581  			}, -1 /* fieldID */)),
   582  		}...), -1 /* fieldID */))
   583  		descr = schema.NewSchema(sc)
   584  		assert.True(t, descr.HasRepeatedFields())
   585  	})
   586  }
   587  
   588  func ExamplePrintSchema() {
   589  	fields := schema.FieldList{schema.NewInt32Node("a" /* name */, parquet.Repetitions.Required, 1 /* fieldID */)}
   590  	item1 := schema.NewInt64Node("item1" /* name */, parquet.Repetitions.Optional, 4 /* fieldID */)
   591  	item2 := schema.NewBooleanNode("item2" /* name */, parquet.Repetitions.Required, 5 /* fieldID */)
   592  	list := schema.MustGroup(schema.NewGroupNodeConverted("b" /* name */, parquet.Repetitions.Repeated, schema.FieldList{item1, item2}, schema.ConvertedTypes.List, 3 /* fieldID */))
   593  	bag := schema.MustGroup(schema.NewGroupNode("bag" /* name */, parquet.Repetitions.Optional, schema.FieldList{list}, 2 /* fieldID */))
   594  	fields = append(fields, bag)
   595  
   596  	fields = append(fields,
   597  		schema.MustPrimitive(schema.NewPrimitiveNodeConverted("c" /* name */, parquet.Repetitions.Required, parquet.Types.Int32, schema.ConvertedTypes.Decimal, 0 /* type len */, 3 /* precision */, 2 /* scale */, 6 /* fieldID */)),
   598  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("d" /* name */, parquet.Repetitions.Required, schema.NewDecimalLogicalType(10 /* precision */, 5 /* scale */), parquet.Types.Int64, -1 /* type len */, 7 /* fieldID */)))
   599  
   600  	sc := schema.MustGroup(schema.NewGroupNode("schema" /* name */, parquet.Repetitions.Repeated, fields, 0 /* fieldID */))
   601  	schema.PrintSchema(sc, os.Stdout, 2)
   602  
   603  	// Output:
   604  	// repeated group field_id=0 schema {
   605  	//   required int32 field_id=1 a;
   606  	//   optional group field_id=2 bag {
   607  	//     repeated group field_id=3 b (List) {
   608  	//       optional int64 field_id=4 item1;
   609  	//       required boolean field_id=5 item2;
   610  	//     }
   611  	//   }
   612  	//   required int32 field_id=6 c (Decimal(precision=3, scale=2));
   613  	//   required int64 field_id=7 d (Decimal(precision=10, scale=5));
   614  	// }
   615  }
   616  
   617  func TestPanicSchemaNodeCreation(t *testing.T) {
   618  	assert.Panics(t, func() {
   619  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("map" /* name */, parquet.Repetitions.Required, schema.MapLogicalType{}, parquet.Types.Int64, -1 /* type len */, -1 /* fieldID */))
   620  	}, "nested logical type on non-group node")
   621  
   622  	assert.Panics(t, func() {
   623  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("string" /* name */, parquet.Repetitions.Required, schema.StringLogicalType{}, parquet.Types.Boolean, -1 /* type len */, -1 /* fieldID */))
   624  	}, "incompatible primitive type")
   625  
   626  	assert.Panics(t, func() {
   627  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("interval" /* name */, parquet.Repetitions.Required, schema.IntervalLogicalType{}, parquet.Types.FixedLenByteArray, 11 /* type len */, -1 /* fieldID */))
   628  	}, "incompatible primitive length")
   629  
   630  	assert.Panics(t, func() {
   631  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("decimal" /* name */, parquet.Repetitions.Required, schema.NewDecimalLogicalType(16, 6), parquet.Types.Int32, -1 /* type len */, -1 /* fieldID */))
   632  	}, "primitive too small for given precision")
   633  
   634  	assert.Panics(t, func() {
   635  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("uuid" /* name */, parquet.Repetitions.Required, schema.UUIDLogicalType{}, parquet.Types.FixedLenByteArray, 64 /* type len */, -1 /* fieldID */))
   636  	}, "incompatible primitive length")
   637  
   638  	assert.Panics(t, func() {
   639  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("float16" /* name */, parquet.Repetitions.Required, schema.Float16LogicalType{}, parquet.Types.FixedLenByteArray, 4 /* type len */, -1 /* fieldID */))
   640  	}, "incompatible primitive length")
   641  
   642  	assert.Panics(t, func() {
   643  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("negative_len" /* name */, parquet.Repetitions.Required, schema.NoLogicalType{}, parquet.Types.FixedLenByteArray, -16 /* type len */, -1 /* fieldID */))
   644  	}, "non-positive length for fixed length binary")
   645  
   646  	assert.Panics(t, func() {
   647  		schema.MustPrimitive(schema.NewPrimitiveNodeLogical("zero_len" /* name */, parquet.Repetitions.Required, schema.NoLogicalType{}, parquet.Types.FixedLenByteArray, 0 /* type len */, -1 /* fieldID */))
   648  	}, "non-positive length for fixed length binary")
   649  
   650  	assert.Panics(t, func() {
   651  		schema.MustGroup(schema.NewGroupNodeLogical("list" /* name */, parquet.Repetitions.Repeated, schema.FieldList{}, schema.JSONLogicalType{}, -1 /* fieldID */))
   652  	}, "non-nested logical type on group node")
   653  }
   654  
   655  func TestNullLogicalConvertsToNone(t *testing.T) {
   656  	var (
   657  		empty schema.LogicalType
   658  		n     schema.Node
   659  	)
   660  	assert.NotPanics(t, func() {
   661  		n = schema.MustPrimitive(schema.NewPrimitiveNodeLogical("value" /* name */, parquet.Repetitions.Required, empty, parquet.Types.Double, -1 /* type len */, -1 /* fieldID */))
   662  	})
   663  	assert.True(t, n.LogicalType().IsNone())
   664  	assert.Equal(t, schema.ConvertedTypes.None, n.ConvertedType())
   665  	assert.NotPanics(t, func() {
   666  		n = schema.MustGroup(schema.NewGroupNodeLogical("items" /* name */, parquet.Repetitions.Repeated, schema.FieldList{}, empty, -1 /* fieldID */))
   667  	})
   668  	assert.True(t, n.LogicalType().IsNone())
   669  	assert.Equal(t, schema.ConvertedTypes.None, n.ConvertedType())
   670  }
   671
View as plain text