Skip to content
This repository was archived by the owner on Feb 21, 2024. It is now read-only.

Commit 6b75988

Browse files
travisturnertravisturner
committed
add key translation to exports
1 parent 060e5e3 commit 6b75988

File tree

2 files changed

+204
-6
lines changed

2 files changed

+204
-6
lines changed

‎api.go

+39-6
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,18 @@ func (api *API) ExportCSV(_ context.Context, indexName string, fieldName string,
339339
return ErrClusterDoesNotOwnShard
340340
}
341341

342+
// Find index.
343+
index := api.holder.Index(indexName)
344+
if index == nil {
345+
return newNotFoundError(ErrIndexNotFound)
346+
}
347+
348+
// Find field from the index.
349+
field := index.Field(fieldName)
350+
if field == nil {
351+
return newNotFoundError(ErrFieldNotFound)
352+
}
353+
342354
// Find the fragment.
343355
f := api.holder.fragment(indexName, fieldName, viewStandard, shard)
344356
if f == nil {
@@ -348,13 +360,34 @@ func (api *API) ExportCSV(_ context.Context, indexName string, fieldName string,
348360
// Wrap writer with a CSV writer.
349361
cw := csv.NewWriter(w)
350362

363+
// Define the function to write each bit as a string,
364+
// translating to keys where necessary.
365+
fn := func(rowID, columnID uint64) error {
366+
var rowStr string
367+
var colStr string
368+
var err error
369+
370+
if field.keys() {
371+
if rowStr, err = api.holder.translateFile.TranslateRowToString(index.Name(), field.Name(), rowID); err != nil {
372+
return errors.Wrap(err, "translating row")
373+
}
374+
} else {
375+
rowStr = strconv.FormatUint(rowID, 10)
376+
}
377+
378+
if index.Keys() {
379+
if colStr, err = api.holder.translateFile.TranslateColumnToString(index.Name(), columnID); err != nil {
380+
return errors.Wrap(err, "translating column")
381+
}
382+
} else {
383+
colStr = strconv.FormatUint(columnID, 10)
384+
}
385+
386+
return cw.Write([]string{rowStr, colStr})
387+
}
388+
351389
// Iterate over each column.
352-
if err := f.forEachBit(func(rowID, columnID uint64) error {
353-
return cw.Write([]string{
354-
strconv.FormatUint(rowID, 10),
355-
strconv.FormatUint(columnID, 10),
356-
})
357-
}); err != nil {
390+
if err := f.forEachBit(fn); err != nil {
358391
return errors.Wrap(err, "writing CSV")
359392
}
360393

‎http/client_test.go

+165
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
package http_test
1616

1717
import (
18+
"bufio"
19+
"bytes"
1820
"context"
1921
"fmt"
2022
gohttp "net/http"
@@ -168,6 +170,169 @@ func TestClient_MultiNode(t *testing.T) {
168170
}
169171
}
170172

173+
// Ensure client can export data.
174+
func TestClient_Export(t *testing.T) {
175+
cmd := test.MustRunCluster(t, 1)[0]
176+
host := cmd.URL()
177+
178+
cmd.MustCreateIndex(t, "keyed", pilosa.IndexOptions{Keys: true})
179+
cmd.MustCreateIndex(t, "unkeyed", pilosa.IndexOptions{Keys: false})
180+
181+
cmd.MustCreateField(t, "keyed", "keyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000), pilosa.OptFieldKeys())
182+
cmd.MustCreateField(t, "keyed", "unkeyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000))
183+
cmd.MustCreateField(t, "unkeyed", "keyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000), pilosa.OptFieldKeys())
184+
cmd.MustCreateField(t, "unkeyed", "unkeyedf", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 1000))
185+
186+
c := MustNewClient(host, http.GetHTTPClient(nil))
187+
188+
data := []pilosa.Bit{
189+
{RowID: 1, ColumnID: 100, RowKey: "row1", ColumnKey: "col100"},
190+
{RowID: 1, ColumnID: 101, RowKey: "row1", ColumnKey: "col101"},
191+
{RowID: 1, ColumnID: 102, RowKey: "row1", ColumnKey: "col102"},
192+
{RowID: 1, ColumnID: 103, RowKey: "row1", ColumnKey: "col103"},
193+
{RowID: 2, ColumnID: 200, RowKey: "row2", ColumnKey: "col200"},
194+
{RowID: 2, ColumnID: 201, RowKey: "row2", ColumnKey: "col201"},
195+
{RowID: 2, ColumnID: 202, RowKey: "row2", ColumnKey: "col202"},
196+
{RowID: 2, ColumnID: 203, RowKey: "row2", ColumnKey: "col203"},
197+
}
198+
199+
t.Run("Import unkeyed,unkeyedf", func(t *testing.T) {
200+
// Populate data.
201+
for _, bit := range data {
202+
_, err := c.Query(context.Background(), "unkeyed", &pilosa.QueryRequest{
203+
Query: fmt.Sprintf(`Set(%d, unkeyedf=%d)`, bit.ColumnID, bit.RowID),
204+
Remote: false,
205+
})
206+
if err != nil {
207+
t.Fatal(err)
208+
}
209+
}
210+
211+
buf := bytes.NewBuffer(nil)
212+
bw := bufio.NewWriter(buf)
213+
214+
// Send export request.
215+
if err := c.ExportCSV(context.Background(), "unkeyed", "unkeyedf", 0, bw); err != nil {
216+
t.Fatal(err)
217+
}
218+
219+
got := buf.String()
220+
221+
// Expected output.
222+
exp := ""
223+
for _, bit := range data {
224+
exp += fmt.Sprintf("%d,%d\n", bit.RowID, bit.ColumnID)
225+
}
226+
227+
// Verify data.
228+
if got != exp {
229+
t.Fatalf("unexpected export data: %s", got)
230+
}
231+
})
232+
233+
t.Run("Import unkeyed,keyedf", func(t *testing.T) {
234+
// Populate data.
235+
for _, bit := range data {
236+
_, err := c.Query(context.Background(), "unkeyed", &pilosa.QueryRequest{
237+
Query: fmt.Sprintf(`Set(%d, keyedf=%s)`, bit.ColumnID, bit.RowKey),
238+
Remote: false,
239+
})
240+
if err != nil {
241+
t.Fatal(err)
242+
}
243+
}
244+
245+
buf := bytes.NewBuffer(nil)
246+
bw := bufio.NewWriter(buf)
247+
248+
// Send export request.
249+
if err := c.ExportCSV(context.Background(), "unkeyed", "keyedf", 0, bw); err != nil {
250+
t.Fatal(err)
251+
}
252+
253+
got := buf.String()
254+
255+
// Expected output.
256+
exp := ""
257+
for _, bit := range data {
258+
exp += fmt.Sprintf("%s,%d\n", bit.RowKey, bit.ColumnID)
259+
}
260+
261+
// Verify data.
262+
if got != exp {
263+
t.Fatalf("unexpected export data: %s", got)
264+
}
265+
})
266+
267+
t.Run("Import keyed,unkeyedf", func(t *testing.T) {
268+
// Populate data.
269+
for _, bit := range data {
270+
_, err := c.Query(context.Background(), "keyed", &pilosa.QueryRequest{
271+
Query: fmt.Sprintf(`Set("%s", unkeyedf=%d)`, bit.ColumnKey, bit.RowID),
272+
Remote: false,
273+
})
274+
if err != nil {
275+
t.Fatal(err)
276+
}
277+
}
278+
279+
buf := bytes.NewBuffer(nil)
280+
bw := bufio.NewWriter(buf)
281+
282+
// Send export request.
283+
if err := c.ExportCSV(context.Background(), "keyed", "unkeyedf", 0, bw); err != nil {
284+
t.Fatal(err)
285+
}
286+
287+
got := buf.String()
288+
289+
// Expected output.
290+
exp := ""
291+
for _, bit := range data {
292+
exp += fmt.Sprintf("%d,%s\n", bit.RowID, bit.ColumnKey)
293+
}
294+
295+
// Verify data.
296+
if got != exp {
297+
t.Fatalf("unexpected export data: %s", got)
298+
}
299+
})
300+
301+
t.Run("Import keyed,keyedf", func(t *testing.T) {
302+
// Populate data.
303+
for _, bit := range data {
304+
_, err := c.Query(context.Background(), "keyed", &pilosa.QueryRequest{
305+
Query: fmt.Sprintf(`Set("%s", keyedf=%s)`, bit.ColumnKey, bit.RowKey),
306+
Remote: false,
307+
})
308+
if err != nil {
309+
t.Fatal(err)
310+
}
311+
}
312+
313+
buf := bytes.NewBuffer(nil)
314+
bw := bufio.NewWriter(buf)
315+
316+
// Send export request.
317+
if err := c.ExportCSV(context.Background(), "keyed", "keyedf", 0, bw); err != nil {
318+
t.Fatal(err)
319+
}
320+
321+
got := buf.String()
322+
323+
// Expected output.
324+
exp := ""
325+
for _, bit := range data {
326+
exp += fmt.Sprintf("%s,%s\n", bit.RowKey, bit.ColumnKey)
327+
}
328+
329+
// Verify data.
330+
if got != exp {
331+
t.Fatalf("unexpected export data: %s", got)
332+
}
333+
})
334+
}
335+
171336
// Ensure client can bulk import data.
172337
func TestClient_Import(t *testing.T) {
173338
cmd := test.MustRunCluster(t, 1)[0]

0 commit comments

Comments
 (0)