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
17package arrow
18
19import (
20	"fmt"
21	"sort"
22	"strings"
23)
24
25type Metadata struct {
26	keys   []string
27	values []string
28}
29
30func NewMetadata(keys, values []string) Metadata {
31	if len(keys) != len(values) {
32		panic("arrow: len mismatch")
33	}
34
35	n := len(keys)
36	if n == 0 {
37		return Metadata{}
38	}
39
40	md := Metadata{
41		keys:   make([]string, n),
42		values: make([]string, n),
43	}
44	copy(md.keys, keys)
45	copy(md.values, values)
46	return md
47}
48
49func MetadataFrom(kv map[string]string) Metadata {
50	md := Metadata{
51		keys:   make([]string, 0, len(kv)),
52		values: make([]string, 0, len(kv)),
53	}
54	for k := range kv {
55		md.keys = append(md.keys, k)
56	}
57	sort.Strings(md.keys)
58	for _, k := range md.keys {
59		md.values = append(md.values, kv[k])
60	}
61	return md
62}
63
64func (md Metadata) Len() int         { return len(md.keys) }
65func (md Metadata) Keys() []string   { return md.keys }
66func (md Metadata) Values() []string { return md.values }
67
68func (md Metadata) String() string {
69	o := new(strings.Builder)
70	fmt.Fprintf(o, "[")
71	for i := range md.keys {
72		if i > 0 {
73			fmt.Fprintf(o, ", ")
74		}
75		fmt.Fprintf(o, "%q: %q", md.keys[i], md.values[i])
76	}
77	fmt.Fprintf(o, "]")
78	return o.String()
79}
80
81// FindKey returns the index of the key-value pair with the provided key name,
82// or -1 if such a key does not exist.
83func (md Metadata) FindKey(k string) int {
84	for i, v := range md.keys {
85		if v == k {
86			return i
87		}
88	}
89	return -1
90}
91
92func (md Metadata) clone() Metadata {
93	if len(md.keys) == 0 {
94		return Metadata{}
95	}
96
97	o := Metadata{
98		keys:   make([]string, len(md.keys)),
99		values: make([]string, len(md.values)),
100	}
101	copy(o.keys, md.keys)
102	copy(o.values, md.values)
103
104	return o
105}
106
107// Schema is a sequence of Field values, describing the columns of a table or
108// a record batch.
109type Schema struct {
110	fields []Field
111	index  map[string][]int
112	meta   Metadata
113}
114
115// NewSchema returns a new Schema value from the slice of fields and metadata.
116//
117// NewSchema panics if there is a field with an invalid DataType.
118func NewSchema(fields []Field, metadata *Metadata) *Schema {
119	sc := &Schema{
120		fields: make([]Field, 0, len(fields)),
121		index:  make(map[string][]int, len(fields)),
122	}
123	if metadata != nil {
124		sc.meta = metadata.clone()
125	}
126	for i, field := range fields {
127		if field.Type == nil {
128			panic("arrow: field with nil DataType")
129		}
130		sc.fields = append(sc.fields, field)
131		sc.index[field.Name] = append(sc.index[field.Name], i)
132	}
133	return sc
134}
135
136func (sc *Schema) Metadata() Metadata { return sc.meta }
137func (sc *Schema) Fields() []Field    { return sc.fields }
138func (sc *Schema) Field(i int) Field  { return sc.fields[i] }
139
140func (sc *Schema) FieldsByName(n string) ([]Field, bool) {
141	indices, ok := sc.index[n]
142	if !ok {
143		return nil, ok
144	}
145	fields := make([]Field, 0, len(indices))
146	for _, v := range indices {
147		fields = append(fields, sc.fields[v])
148	}
149	return fields, ok
150}
151
152// FieldIndices returns the indices of the named field or nil.
153func (sc *Schema) FieldIndices(n string) []int {
154	return sc.index[n]
155}
156
157func (sc *Schema) HasField(n string) bool { return len(sc.FieldIndices(n)) > 0 }
158func (sc *Schema) HasMetadata() bool      { return len(sc.meta.keys) > 0 }
159
160// Equal returns whether two schema are equal.
161// Equal does not compare the metadata.
162func (sc *Schema) Equal(o *Schema) bool {
163	switch {
164	case sc == o:
165		return true
166	case sc == nil || o == nil:
167		return false
168	case len(sc.fields) != len(o.fields):
169		return false
170	}
171
172	for i := range sc.fields {
173		if !sc.fields[i].Equal(o.fields[i]) {
174			return false
175		}
176	}
177	return true
178}
179
180func (s *Schema) String() string {
181	o := new(strings.Builder)
182	fmt.Fprintf(o, "schema:\n  fields: %d\n", len(s.Fields()))
183	for i, f := range s.Fields() {
184		if i > 0 {
185			o.WriteString("\n")
186		}
187		fmt.Fprintf(o, "    - %v", f)
188	}
189	if meta := s.Metadata(); meta.Len() > 0 {
190		fmt.Fprintf(o, "\n  metadata: %v", meta)
191	}
192	return o.String()
193}
194