Skip to content

Commit

Permalink
Add support for --tile-compression on convert operations.
Browse files Browse the repository at this point in the history
  • Loading branch information
lseelenbinder committed Oct 14, 2024
1 parent 429f107 commit 47f3f8c
Show file tree
Hide file tree
Showing 4 changed files with 45 additions and 22 deletions.
17 changes: 16 additions & 1 deletion main.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ var cli struct {
Output string `arg:"" help:"Output PMTiles archive." type:"path"`
Force bool `help:"Force removal."`
NoDeduplication bool `help:"Don't attempt to deduplicate tiles."`
TileCompression string `default:"none" enum:"none,gzip,brotli,zstd" help:"Compression used for tile data (only gzip will be compressed if the source is not compressed; other compressions are assumed to be already compressed)."`
Tmpdir string `help:"An optional path to a folder for tmp data." type:"existingdir"`
} `cmd:"" help:"Convert an MBTiles or older spec version to PMTiles."`

Expand Down Expand Up @@ -197,8 +198,22 @@ func main() {
}
}

var tileCompression pmtiles.Compression
switch cli.Convert.TileCompression {
case "gzip":
tileCompression = pmtiles.Gzip
case "brotli":
tileCompression = pmtiles.Brotli
case "zstd":
tileCompression = pmtiles.Zstd
case "none":
tileCompression = pmtiles.NoCompression
default:
logger.Fatalf("Unknown tile compression: %s", cli.Convert.TileCompression)
}

defer os.Remove(tmpfile.Name())
err := pmtiles.Convert(logger, path, output, !cli.Convert.NoDeduplication, tmpfile)
err := pmtiles.Convert(logger, path, output, !cli.Convert.NoDeduplication, tileCompression, tmpfile)

if err != nil {
logger.Fatalf("Failed to convert %s, %v", path, err)
Expand Down
46 changes: 27 additions & 19 deletions pmtiles/convert.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,15 +28,16 @@ type offsetLen struct {
}

type resolver struct {
deduplicate bool
compress bool
Entries []EntryV3
Offset uint64
OffsetMap map[string]offsetLen
AddressedTiles uint64 // none of them can be empty
compressor *gzip.Writer
compressTmp *bytes.Buffer
hashfunc hash.Hash
deduplicate bool
compress bool
tileCompression Compression
Entries []EntryV3
Offset uint64
OffsetMap map[string]offsetLen
AddressedTiles uint64 // none of them can be empty
compressor *gzip.Writer
compressTmp *bytes.Buffer
hashfunc hash.Hash
}

func (r *resolver) NumContents() uint64 {
Expand Down Expand Up @@ -75,8 +76,8 @@ func (r *resolver) AddTileIsNew(tileID uint64, data []byte) (bool, []byte) {
return false, nil
}
var newData []byte
if !r.compress || (len(data) >= 2 && data[0] == 31 && data[1] == 139) {
// the tile is already compressed
if !r.compress || (len(data) >= 2 && data[0] == 31 && data[1] == 139) || r.tileCompression != Gzip {
// the tile is already compressed or we leave compression unchanged
newData = data
} else {
r.compressTmp.Reset()
Expand All @@ -94,19 +95,22 @@ func (r *resolver) AddTileIsNew(tileID uint64, data []byte) (bool, []byte) {
return true, newData
}

func newResolver(deduplicate bool, compress bool) *resolver {
func newResolver(deduplicate bool, compress bool, tileCompression Compression) *resolver {
b := new(bytes.Buffer)
compressor, _ := gzip.NewWriterLevel(b, gzip.BestCompression)
r := resolver{deduplicate, compress, make([]EntryV3, 0), 0, make(map[string]offsetLen), 0, compressor, b, fnv.New128a()}

r := resolver{deduplicate, compress,
tileCompression,
make([]EntryV3, 0), 0, make(map[string]offsetLen), 0, compressor, b, fnv.New128a()}
return &r
}

// Convert an existing archive on disk to a new PMTiles specification version 3 archive.
func Convert(logger *log.Logger, input string, output string, deduplicate bool, tmpfile *os.File) error {
func Convert(logger *log.Logger, input string, output string, deduplicate bool, tileCompression Compression, tmpfile *os.File) error {
if strings.HasSuffix(input, ".pmtiles") {
return convertPmtilesV2(logger, input, output, deduplicate, tmpfile)
}
return convertMbtiles(logger, input, output, deduplicate, tmpfile)
return convertMbtiles(logger, input, output, deduplicate, tileCompression, tmpfile)
}

func addDirectoryV2Entries(dir directoryV2, entries *[]EntryV3, f *os.File) {
Expand Down Expand Up @@ -186,7 +190,7 @@ func convertPmtilesV2(logger *log.Logger, input string, output string, deduplica
})

// re-use resolve, because even if archives are de-duplicated we may need to recompress.
resolve := newResolver(deduplicate, header.TileType == Mvt)
resolve := newResolver(deduplicate, header.TileType == Mvt, header.TileCompression)

bar := progressbar.Default(int64(len(entries)))
for _, entry := range entries {
Expand Down Expand Up @@ -223,7 +227,7 @@ func convertPmtilesV2(logger *log.Logger, input string, output string, deduplica
return nil
}

func convertMbtiles(logger *log.Logger, input string, output string, deduplicate bool, tmpfile *os.File) error {
func convertMbtiles(logger *log.Logger, input string, output string, deduplicate bool, tileCompression Compression, tmpfile *os.File) error {
start := time.Now()
conn, err := sqlite.OpenConn(input, sqlite.OpenReadOnly)
if err != nil {
Expand Down Expand Up @@ -293,8 +297,12 @@ func convertMbtiles(logger *log.Logger, input string, output string, deduplicate
return fmt.Errorf("no tiles in MBTiles archive")
}

if header.TileType != Mvt {
tileCompression = NoCompression
}

logger.Println("Pass 2: writing tiles")
resolve := newResolver(deduplicate, header.TileType == Mvt)
resolve := newResolver(deduplicate, header.TileType == Mvt, tileCompression)
{
bar := progressbar.Default(int64(tileset.GetCardinality()))
i := tileset.Iterator()
Expand Down Expand Up @@ -393,7 +401,7 @@ func finalize(logger *log.Logger, resolve *resolver, header HeaderV3, tmpfile *o
header.Clustered = true
header.InternalCompression = Gzip
if header.TileType == Mvt {
header.TileCompression = Gzip
header.TileCompression = resolve.tileCompression
}

header.RootOffset = HeaderV3LenBytes
Expand Down
2 changes: 1 addition & 1 deletion pmtiles/convert_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (
)

func TestResolver(t *testing.T) {
resolver := newResolver(true, true)
resolver := newResolver(true, true, Gzip)
resolver.AddTileIsNew(1, []byte{0x1, 0x2})
assert.Equal(t, 1, len(resolver.Entries))
resolver.AddTileIsNew(2, []byte{0x1, 0x3})
Expand Down
2 changes: 1 addition & 1 deletion pmtiles/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ func fakeArchive(t *testing.T, header HeaderV3, metadata map[string]interface{},
keys = append(keys, id)
}
sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
resolver := newResolver(false, false)
resolver := newResolver(false, false, Gzip)
tileDataBytes := make([]byte, 0)
for _, id := range keys {
tileBytes := byTileID[id]
Expand Down

0 comments on commit 47f3f8c

Please sign in to comment.