From 5b6d4d737a7469885ca0aa6b5a348ef0d3b616c9 Mon Sep 17 00:00:00 2001 From: KaviiSuri Date: Fri, 4 Oct 2024 13:07:13 +0530 Subject: [PATCH] feat: Implement GEOADD and GEODIST commands --- go.mod | 2 +- go.sum | 41 +---- internal/eval/commands.go | 58 ++++--- internal/eval/eval.go | 230 +++++++++++++++++--------- internal/eval/eval_test.go | 117 ++++++++++++- internal/eval/geo/geo.go | 86 ++++++++++ internal/eval/sorted_set.go | 19 --- internal/eval/sortedset/sorted_set.go | 148 +++++++++++++++++ internal/server/utils/round.go | 9 + 9 files changed, 547 insertions(+), 163 deletions(-) create mode 100644 internal/eval/geo/geo.go delete mode 100644 internal/eval/sorted_set.go create mode 100644 internal/eval/sortedset/sorted_set.go create mode 100644 internal/server/utils/round.go diff --git a/go.mod b/go.mod index b79cd31dc..38e68ce3e 100644 --- a/go.mod +++ b/go.mod @@ -27,7 +27,6 @@ require ( github.com/spf13/pflag v1.0.5 // indirect github.com/subosito/gotenv v1.6.0 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect - go.uber.org/atomic v1.11.0 // indirect go.uber.org/multierr v1.11.0 // indirect golang.org/x/arch v0.10.0 // indirect golang.org/x/exp v0.0.0-20240909161429-701f63a606c0 // indirect @@ -47,6 +46,7 @@ require ( github.com/google/btree v1.1.3 github.com/google/go-cmp v0.6.0 github.com/gorilla/websocket v1.5.3 + github.com/mmcloughlin/geohash v0.10.0 github.com/ohler55/ojg v1.24.0 github.com/pelletier/go-toml/v2 v2.2.3 github.com/rs/xid v1.6.0 diff --git a/go.sum b/go.sum index 1a1d21646..aaaed3a20 100644 --- a/go.sum +++ b/go.sum @@ -6,15 +6,11 @@ github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= -github.com/bytedance/sonic v1.12.1 h1:jWl5Qz1fy7X1ioY74WqO0KjAMtAGQs4sYnjiEBiyX24= -github.com/bytedance/sonic v1.12.1/go.mod h1:B8Gt/XvtZ3Fqj+iSKMypzymZxw/FVwgIGKzMzT9r/rk= github.com/bytedance/sonic v1.12.3 h1:W2MGa7RCU1QTeYRTPE3+88mVC0yXmsRQRChiyVocVjU= github.com/bytedance/sonic v1.12.3/go.mod h1:B8Gt/XvtZ3Fqj+iSKMypzymZxw/FVwgIGKzMzT9r/rk= github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= github.com/bytedance/sonic/loader v0.2.0 h1:zNprn+lsIP06C/IqCHs3gPQIvnvpKbbxyXQP1iU4kWM= github.com/bytedance/sonic/loader v0.2.0/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= -github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= -github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cloudwego/base64x v0.1.4 h1:jwCgWpFanWmN8xoIUHa2rtzmkd5J2plF/dnLS6Xd/0Y= @@ -28,12 +24,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 h1:fAjc9m62+UWV/WAFKLNi6ZS0675eEUC9y3AlwSbQu1Y= -github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da h1:aIftn67I1fkbMa512G+w+Pxci9hJPB8oMnkcP3iZF38= github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= -github.com/dgryski/go-metro v0.0.0-20180109044635-280f6062b5bc h1:8WFBn63wegobsYAX0YjD+8suexZDga5CctH4CCTx2+8= -github.com/dgryski/go-metro v0.0.0-20180109044635-280f6062b5bc/go.mod h1:c9O8+fpSOX1DM8cPNSkX/qsBWdkD4yd2dpciOWQjpBw= github.com/dgryski/go-metro v0.0.0-20211217172704-adc40b04c140 h1:y7y0Oa6UawqTFPCDw9JG6pdKt4F9pAhHv0B7FMGaGD0= github.com/dgryski/go-metro v0.0.0-20211217172704-adc40b04c140/go.mod h1:c9O8+fpSOX1DM8cPNSkX/qsBWdkD4yd2dpciOWQjpBw= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= @@ -53,7 +45,6 @@ github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aN github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= -github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4= github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= github.com/klauspost/cpuid/v2 v2.2.8 h1:+StwCXwm9PdpiEkPyzBXIy+M9KUb4ODm0Zarf1kS5BM= github.com/klauspost/cpuid/v2 v2.2.8/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= @@ -64,21 +55,18 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY= github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= -github.com/mattn/go-colorable v0.1.12/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4= github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= -github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= -github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA= github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= +github.com/mmcloughlin/geohash v0.10.0 h1:9w1HchfDfdeLc+jFEf/04D27KP7E2QmpDu52wPbJWRE= +github.com/mmcloughlin/geohash v0.10.0/go.mod h1:oNZxQo5yWJh0eMQEP/8hwQuVx9Z9tjwFUqcTB1SmG0c= github.com/ohler55/ojg v1.24.0 h1:y2AVez6fPTszK/jPhaAYMCAzAoSleConMqSDD5wJKJg= github.com/ohler55/ojg v1.24.0/go.mod h1:gQhDVpQLqrmnd2eqGAvJtn+NfKoYJbe/A4Sj3/Vro4o= -github.com/ohler55/ojg v1.24.1 h1:PaVLelrNgT5/0ppPaUtey54tOVp245z33fkhL2jljjY= -github.com/ohler55/ojg v1.24.1/go.mod h1:gQhDVpQLqrmnd2eqGAvJtn+NfKoYJbe/A4Sj3/Vro4o= github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M= github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -90,12 +78,8 @@ github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/f github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg= github.com/rs/xid v1.6.0 h1:fV591PaemRlL6JfRxGDEPl69wICngIQ3shQtzfy2gxU= github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0= -github.com/rs/zerolog v1.30.0 h1:SymVODrcRsaRaSInD9yQtKbtWqwsfoPcRff/oRXLj4c= -github.com/rs/zerolog v1.30.0/go.mod h1:/tk+P47gFdPXq4QYjvCmT5/Gsug2nagsFWBWhAiSi1w= github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8= github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss= -github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6keLGt6kNQ= -github.com/sagikazarmark/locafero v0.4.0/go.mod h1:Pe1W6UlPYUk/+wc/6KFhbORCfqzgYEpgQ3O5fPuL3H4= github.com/sagikazarmark/locafero v0.6.0 h1:ON7AQg37yzcRPU69mt7gwhFEBwxI6P9T4Qu3N51bwOk= github.com/sagikazarmark/locafero v0.6.0/go.mod h1:77OmuIc6VTraTXKXIs/uvUxKGUXjE1GbemJYHqdNjX0= github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE= @@ -104,8 +88,6 @@ github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9yS github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0= github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8= github.com/spf13/afero v1.11.0/go.mod h1:GH9Y3pIexgf1MTIWtNGyogA5MwRIDXGUr+hbWNoBjkY= -github.com/spf13/cast v1.6.0 h1:GEiTHELF+vaR5dhz3VqZfFSzZjYbgeKDpBxQVS4GYJ0= -github.com/spf13/cast v1.6.0/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= github.com/spf13/cast v1.7.0 h1:ntdiHjuueXFgm5nzDRdOS4yfT43P5Fnud6DH50rz/7w= github.com/spf13/cast v1.7.0/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= @@ -115,7 +97,6 @@ github.com/spf13/viper v1.19.0/go.mod h1:GQUN9bilAbhU/jgc1bKs99f/suXKeUMct8Adx5+ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= -github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= @@ -130,38 +111,20 @@ github.com/twmb/murmur3 v1.1.8 h1:8Yt9taO/WN3l08xErzjeschgZU2QSrwm1kclYq+0aRg= github.com/twmb/murmur3 v1.1.8/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ= github.com/xwb1989/sqlparser v0.0.0-20180606152119-120387863bf2 h1:zzrxE1FKn5ryBNl9eKOeqQ58Y/Qpo3Q9QNxKHX5uzzQ= github.com/xwb1989/sqlparser v0.0.0-20180606152119-120387863bf2/go.mod h1:hzfGeIUDq/j97IG+FhNqkowIyEcD88LrW6fyU3K3WqY= -go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE= -go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= -go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE= -go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= -go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI= -go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ= go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= -golang.org/x/arch v0.0.0-20210923205945-b76863e36670 h1:18EFjUmQOcUvxNYSkA6jO9VAiXCnxFY6NyDX0bHDmkU= -golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= golang.org/x/arch v0.10.0 h1:S3huipmSclq3PJMNe76NGwkBR504WFkQ5dhzWzP8ZW8= golang.org/x/arch v0.10.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= -golang.org/x/crypto v0.25.0 h1:ypSNr+bnYL2YhwoMt2zPxHFmbAN1KZs/njMG3hxUp30= -golang.org/x/crypto v0.25.0/go.mod h1:T+wALwcMOSE0kXgUAnPAHqTLW+XHgcELELW8VaDgm/M= golang.org/x/crypto v0.27.0 h1:GXm2NjJrPaiv/h1tb2UH8QfgC/hOf/+z0p6PT8o1w7A= golang.org/x/crypto v0.27.0/go.mod h1:1Xngt8kV6Dvbssa53Ziq6Eqn0HqbZi5Z6R0ZpwQzt70= -golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI= -golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo= golang.org/x/exp v0.0.0-20240909161429-701f63a606c0 h1:e66Fs6Z+fZTbFBAxKfP3PALWBtpfqks2bwGcexMxgtk= golang.org/x/exp v0.0.0-20240909161429-701f63a606c0/go.mod h1:2TbTHSBQa924w8M6Xs1QcRcFwyucIwBGpK1p2f1YFFY= -golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI= -golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34= golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4= -golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI= golang.org/x/text v0.18.0 h1:XvMDiNzPAl0jr17s6W9lcaIhGUfUORdGCNsuLmPG224= golang.org/x/text v0.18.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/eval/commands.go b/internal/eval/commands.go index 410fa8d34..0d99f7dd2 100644 --- a/internal/eval/commands.go +++ b/internal/eval/commands.go @@ -621,10 +621,10 @@ var ( KeySpecs: KeySpecs{BeginIndex: 1}, } hkeysCmdMeta = DiceCmdMeta{ - Name: "HKEYS", - Info: `HKEYS command is used to retrieve all the keys(or field names) within a hash. Complexity is O(n) where n is the size of the hash.`, - Eval: evalHKEYS, - Arity: 1, + Name: "HKEYS", + Info: `HKEYS command is used to retrieve all the keys(or field names) within a hash. Complexity is O(n) where n is the size of the hash.`, + Eval: evalHKEYS, + Arity: 1, KeySpecs: KeySpecs{BeginIndex: 1}, } hsetnxCmdMeta = DiceCmdMeta{ @@ -911,27 +911,27 @@ var ( Arity: 3, KeySpecs: KeySpecs{BeginIndex: 1}, } - dumpkeyCMmdMeta=DiceCmdMeta{ - Name: "DUMP", - Info: `Serialize the value stored at key in a Redis-specific format and return it to the user. + dumpkeyCMmdMeta = DiceCmdMeta{ + Name: "DUMP", + Info: `Serialize the value stored at key in a Redis-specific format and return it to the user. The returned value can be synthesized back into a Redis key using the RESTORE command.`, - Eval: evalDUMP, - Arity: 1, - KeySpecs: KeySpecs{BeginIndex: 1}, + Eval: evalDUMP, + Arity: 1, + KeySpecs: KeySpecs{BeginIndex: 1}, } - restorekeyCmdMeta=DiceCmdMeta{ - Name: "RESTORE", - Info: `Serialize the value stored at key in a Redis-specific format and return it to the user. + restorekeyCmdMeta = DiceCmdMeta{ + Name: "RESTORE", + Info: `Serialize the value stored at key in a Redis-specific format and return it to the user. The returned value can be synthesized back into a Redis key using the RESTORE command.`, - Eval: evalRestore, - Arity: 2, + Eval: evalRestore, + Arity: 2, KeySpecs: KeySpecs{BeginIndex: 1}, } typeCmdMeta = DiceCmdMeta{ - Name: "TYPE", - Info: `Returns the string representation of the type of the value stored at key. The different types that can be returned are: string, list, set, zset, hash and stream.`, - Eval: evalTYPE, - Arity: 1, + Name: "TYPE", + Info: `Returns the string representation of the type of the value stored at key. The different types that can be returned are: string, list, set, zset, hash and stream.`, + Eval: evalTYPE, + Arity: 1, KeySpecs: KeySpecs{BeginIndex: 1}, } @@ -1040,14 +1040,28 @@ var ( Arity: -4, KeySpecs: KeySpecs{BeginIndex: 1}, } + geoAddCmdMeta = DiceCmdMeta{ + Name: "GEOADD", + Info: `Adds one or more members to a geospatial index. The key is created if it doesn't exist.`, + Arity: -5, + Eval: evalGEOADD, + KeySpecs: KeySpecs{BeginIndex: 1}, + } + geoDistCmdMeta = DiceCmdMeta{ + Name: "GEODIST", + Info: `Returns the distance between two members in the geospatial index.`, + Arity: -4, + Eval: evalGEODIST, + KeySpecs: KeySpecs{BeginIndex: 1}, + } ) func init() { DiceCmds["PING"] = pingCmdMeta DiceCmds["ECHO"] = echoCmdMeta DiceCmds["AUTH"] = authCmdMeta - DiceCmds["DUMP"]=dumpkeyCMmdMeta - DiceCmds["RESTORE"]=restorekeyCmdMeta + DiceCmds["DUMP"] = dumpkeyCMmdMeta + DiceCmds["RESTORE"] = restorekeyCmdMeta DiceCmds["SET"] = setCmdMeta DiceCmds["GET"] = getCmdMeta DiceCmds["MSET"] = msetCmdMeta @@ -1155,6 +1169,8 @@ func init() { DiceCmds["BITFIELD"] = bitfieldCmdMeta DiceCmds["HINCRBYFLOAT"] = hincrbyFloatCmdMeta DiceCmds["HEXISTS"] = hexistsCmdMeta + DiceCmds["GEOADD"] = geoAddCmdMeta + DiceCmds["GEODIST"] = geoDistCmdMeta } // Function to convert DiceCmdMeta to []interface{} diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 30ddfeb08..fe86786cc 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -21,8 +21,8 @@ import ( "unicode" "unsafe" - "github.com/google/btree" - + "github.com/dicedb/dice/internal/eval/geo" + "github.com/dicedb/dice/internal/eval/sortedset" "github.com/dicedb/dice/internal/object" "github.com/rs/xid" @@ -4547,22 +4547,16 @@ func evalZADD(args []string, store *dstore.Store) []byte { key := args[0] obj := store.Get(key) - var tree *btree.BTree - var memberMap map[string]float64 + var ss *sortedset.Set if obj != nil { - if err := object.AssertTypeAndEncoding(obj.TypeEncoding, object.ObjTypeSortedSet, object.ObjEncodingBTree); err != nil { + var err []byte + ss, err = sortedset.FromObject(obj) + if err != nil { return err } - valueSlice, ok := obj.Value.([]interface{}) - if !ok || len(valueSlice) != 2 { - return diceerrors.NewErrWithMessage("Invalid sorted set object") - } - tree = valueSlice[0].(*btree.BTree) - memberMap = valueSlice[1].(map[string]float64) } else { - tree = btree.New(2) - memberMap = make(map[string]float64) + ss = sortedset.New() } added := 0 @@ -4575,24 +4569,14 @@ func evalZADD(args []string, store *dstore.Store) []byte { return diceerrors.NewErrWithMessage(diceerrors.InvalidFloatErr) } - existingScore, exists := memberMap[member] - if exists { - // Remove the existing item from the B-tree - oldItem := &SortedSetItem{Score: existingScore, Member: member} - tree.Delete(oldItem) - } else { - added++ - } - - // Insert the new item into the B-tree - newItem := &SortedSetItem{Score: score, Member: member} - tree.ReplaceOrInsert(newItem) + wasInserted := ss.Upsert(score, member) - // Update the member map - memberMap[member] = score + if wasInserted { + added += 1 + } } - obj = store.NewObj([]interface{}{tree, memberMap}, -1, object.ObjTypeSortedSet, object.ObjEncodingBTree) + obj = store.NewObj(ss, -1, object.ObjTypeSortedSet, object.ObjEncodingBTree) store.Put(key, obj) return clientio.Encode(added, false) @@ -4637,63 +4621,13 @@ func evalZRANGE(args []string, store *dstore.Store) []byte { return clientio.Encode([]string{}, false) } - if err := object.AssertTypeAndEncoding(obj.TypeEncoding, object.ObjTypeSortedSet, object.ObjEncodingBTree); err != nil { - return diceerrors.NewErrWithMessage(diceerrors.WrongTypeErr) - } - - valueSlice, ok := obj.Value.([]interface{}) - if !ok || len(valueSlice) != 2 { - return diceerrors.NewErrWithMessage("Invalid sorted set object") - } - tree := valueSlice[0].(*btree.BTree) - length := tree.Len() - - // Handle negative indices - if start < 0 { - start += length - } - if stop < 0 { - stop += length - } + ss, errMsg := sortedset.FromObject(obj) - if start < 0 { - start = 0 - } - if stop >= length { - stop = length - 1 - } - - if start > stop || start >= length { - return clientio.Encode([]string{}, false) - } - - var result []interface{} - index := 0 - - // iterFunc is the function that will be called for each item in the B-tree. It will append the item to the result if it is within the specified range. - // It will return false if the specified range has been reached. - iterFunc := func(item btree.Item) bool { - if index > stop { - return false - } - if index >= start { - ssi := item.(*SortedSetItem) - result = append(result, ssi.Member) - if withScores { - // Use 'g' format to match Redis's float formatting - scoreStr := strings.ToLower(strconv.FormatFloat(ssi.Score, 'g', -1, 64)) - result = append(result, scoreStr) - } - } - index++ - return true + if errMsg != nil { + return diceerrors.NewErrWithMessage(diceerrors.WrongTypeErr) } - if !reverse { - tree.Ascend(iterFunc) - } else { - tree.Descend(iterFunc) - } + result := ss.GetRange(start, stop, withScores, reverse) return clientio.Encode(result, false) } @@ -4939,3 +4873,135 @@ func evalHINCRBYFLOAT(args []string, store *dstore.Store) []byte { return clientio.Encode(numkey, false) } + +func evalGEOADD(args []string, store *dstore.Store) []byte { + if len(args) < 4 { + return diceerrors.NewErrArity("GEOADD") + } + + key := args[0] + var nx, xx bool + startIdx := 1 + + // Parse options + for startIdx < len(args) { + option := strings.ToUpper(args[startIdx]) + if option == "NX" { + nx = true + startIdx++ + } else if option == "XX" { + xx = true + startIdx++ + } else { + break + } + } + + // Check if we have the correct number of arguments after parsing options + if (len(args)-startIdx)%3 != 0 { + return diceerrors.NewErrArity("GEOADD") + } + + if xx && nx { + return diceerrors.NewErrWithMessage("ERR XX and NX options at the same time are not compatible") + } + + // Get or create sorted set + obj := store.Get(key) + var ss *sortedset.Set + if obj != nil { + var err []byte + ss, err = sortedset.FromObject(obj) + if err != nil { + return err + } + } else { + ss = sortedset.New() + } + + added := 0 + for i := startIdx; i < len(args); i += 3 { + longitude, err := strconv.ParseFloat(args[i], 64) + if err != nil || math.IsNaN(longitude) || longitude < -180 || longitude > 180 { + return diceerrors.NewErrWithMessage("ERR invalid longitude") + } + + latitude, err := strconv.ParseFloat(args[i+1], 64) + if err != nil || math.IsNaN(latitude) || latitude < -85.05112878 || latitude > 85.05112878 { + return diceerrors.NewErrWithMessage("ERR invalid latitude") + } + + member := args[i+2] + _, exists := ss.Get(member) + + // Handle XX option: Only update existing elements + if xx && !exists { + continue + } + + // Handle NX option: Only add new elements + if nx && exists { + continue + } + + hash := geo.EncodeHash(latitude, longitude) + + wasInserted := ss.Upsert(hash, member) + if wasInserted { + added++ + } + } + + obj = store.NewObj(ss, -1, object.ObjTypeSortedSet, object.ObjEncodingBTree) + store.Put(key, obj) + + return clientio.Encode(added, false) +} + +func evalGEODIST(args []string, store *dstore.Store) []byte { + if len(args) < 3 || len(args) > 4 { + return diceerrors.NewErrArity("GEODIST") + } + + key := args[0] + member1 := args[1] + member2 := args[2] + unit := "m" + if len(args) == 4 { + unit = strings.ToLower(args[3]) + } + + // Get the sorted set + obj := store.Get(key) + if obj == nil { + return clientio.RespNIL + } + + ss, err := sortedset.FromObject(obj) + if err != nil { + return err + } + + // Get the scores (geohashes) for both members + score1, ok := ss.Get(member1) + if !ok { + return clientio.RespNIL + } + score2, ok := ss.Get(member2) + if !ok { + return clientio.RespNIL + } + + lat1, lon1 := geo.DecodeHash(score1) + lat2, lon2 := geo.DecodeHash(score2) + + distance := geo.GetDistance(lon1, lat1, lon2, lat2) + + result, err := geo.ConvertDistance(distance, unit) + + if err != nil { + return err + } + + return clientio.Encode(utils.RoundToDecimals(result, 4), false) +} diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index a0b041123..e6ccbb9cd 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -108,6 +108,8 @@ func TestEval(t *testing.T) { testEvalHVALS(t, store) testEvalBitField(t, store) testEvalHINCRBYFLOAT(t, store) + testEvalGEOADD(t, store) + testEvalGEODIST(t, store) } func testEvalPING(t *testing.T, store *dstore.Store) { @@ -4979,7 +4981,6 @@ func testEvalZRANGE(t *testing.T, store *dstore.Store) { runEvalTests(t, tests, evalZRANGE, store) } - func testEvalBitField(t *testing.T, store *dstore.Store) { testCases := map[string]evalTestCase{ "BITFIELD signed SET": { @@ -5264,3 +5265,117 @@ func testEvalDUMP(t *testing.T, store *dstore.Store) { runEvalTests(t, tests, evalDUMP, store) } + +func testEvalGEOADD(t *testing.T, store *dstore.Store) { + tests := map[string]evalTestCase{ + "GEOADD with wrong number of arguments": { + input: []string{"mygeo", "1", "2"}, + output: diceerrors.NewErrArity("GEOADD"), + }, + "GEOADD with non-numeric longitude": { + input: []string{"mygeo", "long", "40.7128", "NewYork"}, + output: diceerrors.NewErrWithMessage("ERR invalid longitude"), + }, + "GEOADD with non-numeric latitude": { + input: []string{"mygeo", "-74.0060", "lat", "NewYork"}, + output: diceerrors.NewErrWithMessage("ERR invalid latitude"), + }, + "GEOADD new member to non-existing key": { + setup: func() {}, + input: []string{"mygeo", "-74.0060", "40.7128", "NewYork"}, + output: clientio.Encode(int64(1), false), + }, + "GEOADD existing member with updated coordinates": { + setup: func() { + evalGEOADD([]string{"mygeo", "-74.0060", "40.7128", "NewYork"}, store) + }, + input: []string{"mygeo", "-73.9352", "40.7304", "NewYork"}, + output: clientio.Encode(int64(0), false), + }, + "GEOADD multiple members": { + setup: func() { + evalGEOADD([]string{"mygeo", "-74.0060", "40.7128", "NewYork"}, store) + }, + input: []string{"mygeo", "-118.2437", "34.0522", "LosAngeles", "-87.6298", "41.8781", "Chicago"}, + output: clientio.Encode(int64(2), false), + }, + "GEOADD with NX option (new member)": { + input: []string{"mygeo", "NX", "-122.4194", "37.7749", "SanFrancisco"}, + output: clientio.Encode(int64(1), false), + }, + "GEOADD with NX option (existing member)": { + setup: func() { + evalGEOADD([]string{"mygeo", "-74.0060", "40.7128", "NewYork"}, store) + }, + input: []string{"mygeo", "NX", "-73.9352", "40.7304", "NewYork"}, + output: clientio.Encode(int64(0), false), + }, + "GEOADD with XX option (new member)": { + input: []string{"mygeo", "XX", "-71.0589", "42.3601", "Boston"}, + output: clientio.Encode(int64(0), false), + }, + "GEOADD with XX option (existing member)": { + setup: func() { + evalGEOADD([]string{"mygeo", "-74.0060", "40.7128", "NewYork"}, store) + }, + input: []string{"mygeo", "XX", "-73.9352", "40.7304", "NewYork"}, + output: clientio.Encode(int64(0), false), + }, + "GEOADD with both NX and XX options": { + input: []string{"mygeo", "NX", "XX", "-74.0060", "40.7128", "NewYork"}, + output: diceerrors.NewErrWithMessage("ERR XX and NX options at the same time are not compatible"), + }, + "GEOADD with invalid option": { + input: []string{"mygeo", "INVALID", "-74.0060", "40.7128", "NewYork"}, + output: diceerrors.NewErrArity("GEOADD"), + }, + "GEOADD to a key of wrong type": { + setup: func() { + store.Put("mygeo", store.NewObj("string_value", -1, object.ObjTypeString, object.ObjEncodingRaw)) + }, + input: []string{"mygeo", "-74.0060", "40.7128", "NewYork"}, + output: []byte("-ERR Existing key has wrong Dice type\r\n"), + }, + "GEOADD with longitude out of range": { + input: []string{"mygeo", "181.0", "40.7128", "Invalid"}, + output: diceerrors.NewErrWithMessage("ERR invalid longitude"), + }, + "GEOADD with latitude out of range": { + input: []string{"mygeo", "-74.0060", "91.0", "Invalid"}, + output: diceerrors.NewErrWithMessage("ERR invalid latitude"), + }, + } + + runEvalTests(t, tests, evalGEOADD, store) +} + +func testEvalGEODIST(t *testing.T, store *dstore.Store) { + tests := map[string]evalTestCase{ + "GEODIST between existing points": { + setup: func() { + evalGEOADD([]string{"points", "13.361389", "38.115556", "Palermo"}, store) + evalGEOADD([]string{"points", "15.087269", "37.502669", "Catania"}, store) + }, + input: []string{"points", "Palermo", "Catania"}, + output: clientio.Encode(float64(166274.1440), false), // Example value + }, + "GEODIST with units (km)": { + setup: func() { + evalGEOADD([]string{"points", "13.361389", "38.115556", "Palermo"}, store) + evalGEOADD([]string{"points", "15.087269", "37.502669", "Catania"}, store) + }, + input: []string{"points", "Palermo", "Catania", "km"}, + output: clientio.Encode(float64(166.2741), false), // Example value + }, + "GEODIST to same point": { + setup: func() { + evalGEOADD([]string{"points", "13.361389", "38.115556", "Palermo"}, store) + }, + input: []string{"points", "Palermo", "Palermo"}, + output: clientio.Encode(float64(0.0000), false), // Expecting distance 0 formatted to 4 decimals + }, + // Add other test cases here... + } + + runEvalTests(t, tests, evalGEODIST, store) +} diff --git a/internal/eval/geo/geo.go b/internal/eval/geo/geo.go new file mode 100644 index 000000000..6db40eaf8 --- /dev/null +++ b/internal/eval/geo/geo.go @@ -0,0 +1,86 @@ +package geo + +import ( + "math" + + "github.com/dicedb/dice/internal/errors" + "github.com/mmcloughlin/geohash" +) + +// Earth's radius in meters +const earthRadius float64 = 6372797.560856 + +// Bit precision for geohash - picked up to match redis +const bitPrecision = 52 + +func DegToRad(deg float64) float64 { + return math.Pi * deg / 180.0 +} + +func RadToDeg(rad float64) float64 { + return 180.0 * rad / math.Pi +} + +func GetDistance( + lon1, + lat1, + lon2, + lat2 float64, +) float64 { + lon1r := DegToRad(lon1) + lon2r := DegToRad(lon2) + v := math.Sin((lon2r - lon1r) / 2) + // if v == 0 we can avoid doing expensive math when lons are practically the same + if v == 0.0 { + return GetLatDistance(lat1, lat2) + } + + lat1r := DegToRad(lat1) + lat2r := DegToRad(lat2) + u := math.Sin((lat2r - lat1r) / 2) + + a := u*u + math.Cos(lat1r)*math.Cos(lat2r)*v*v + + return 2.0 * earthRadius * math.Asin(math.Sqrt(a)) +} + +func GetLatDistance(lat1, lat2 float64) float64 { + return earthRadius * math.Abs(DegToRad(lat2)-DegToRad(lat1)) +} + +// EncodeHash returns a geo hash for a given coordinate, and returns it in float64 so it can be used as score in a zset +func EncodeHash( + latitude, + longitude float64, +) float64 { + h := geohash.EncodeIntWithPrecision(latitude, longitude, bitPrecision) + + return float64(h) +} + +// DecodeHash returns the latitude and longitude from a geo hash +// The hash should be a float64, as it is used as score in a zset +func DecodeHash(hash float64) (lat, lon float64) { + lat, lon = geohash.DecodeIntWithPrecision(uint64(hash), bitPrecision) + + return lat, lon +} + +// ConvertDistance converts a distance from meters to the desired unit +func ConvertDistance( + distance float64, + unit string, +) (converted float64, err []byte) { + switch unit { + case "m": + return distance, nil + case "km": + return distance / 1000, nil + case "mi": + return distance / 1609.34, nil + case "ft": + return distance / 0.3048, nil + default: + return 0, errors.NewErrWithMessage("ERR unsupported unit provided. please use m, km, ft, mi") + } +} diff --git a/internal/eval/sorted_set.go b/internal/eval/sorted_set.go deleted file mode 100644 index 6dfdc3740..000000000 --- a/internal/eval/sorted_set.go +++ /dev/null @@ -1,19 +0,0 @@ -package eval - -import "github.com/google/btree" - -// SortedSetItem represents a member of a sorted set. It includes a score and a member. -type SortedSetItem struct { - btree.Item - Score float64 - Member string -} - -// Less compares two SortedSetItems. Required by the btree.Item interface. -func (a *SortedSetItem) Less(b btree.Item) bool { - other := b.(*SortedSetItem) - if a.Score != other.Score { - return a.Score < other.Score - } - return a.Member < other.Member -} diff --git a/internal/eval/sortedset/sorted_set.go b/internal/eval/sortedset/sorted_set.go new file mode 100644 index 000000000..dac62be9d --- /dev/null +++ b/internal/eval/sortedset/sorted_set.go @@ -0,0 +1,148 @@ +package sortedset + +import ( + "strconv" + "strings" + + diceerrors "github.com/dicedb/dice/internal/errors" + "github.com/dicedb/dice/internal/object" + "github.com/google/btree" +) + +// Item represents a member of a sorted set. It includes a score and a member. +type Item struct { + btree.Item + Score float64 + Member string +} + +// Less compares two Items. Required by the btree.Item interface. +func (a *Item) Less(b btree.Item) bool { + other := b.(*Item) + if a.Score != other.Score { + return a.Score < other.Score + } + return a.Member < other.Member +} + +// is a sorted set data structure that stores members with associated scores. +type Set struct { + // tree is a btree that stores Items. + tree *btree.BTree + // memberMap is a map that stores members and their scores. + memberMap map[string]float64 +} + +// New creates a new . +func New() *Set { + return &Set{ + tree: btree.New(2), + memberMap: make(map[string]float64), + } +} + +func FromObject(obj *object.Obj) (value *Set, err []byte) { + if err := object.AssertTypeAndEncoding(obj.TypeEncoding, object.ObjTypeSortedSet, object.ObjEncodingBTree); err != nil { + return nil, err + } + value, ok := obj.Value.(*Set) + if !ok { + return nil, diceerrors.NewErrWithMessage("Invalid sorted set object") + } + return value, nil +} + +// Add adds a member with a score to the and returns true if the member was added, false if it already existed. +func (ss *Set) Upsert(score float64, member string) bool { + existingScore, exists := ss.memberMap[member] + + if exists { + oldItem := &Item{Score: existingScore, Member: member} + ss.tree.Delete(oldItem) + } + + item := &Item{Score: score, Member: member} + ss.tree.ReplaceOrInsert(item) + ss.memberMap[member] = score + + return !exists +} + +// Remove removes a member from the and returns true if the member was removed, false if it did not exist. +func (ss *Set) Remove(member string) bool { + score, exists := ss.memberMap[member] + if !exists { + return false + } + + item := &Item{Score: score, Member: member} + ss.tree.Delete(item) + delete(ss.memberMap, member) + + return true +} + +// GetRange returns a slice of members with scores between min and max, inclusive. +// it returns the members in ascending order if reverse is false, and descending order if reverse is true. +// If withScores is true, the members will be returned with their scores. +func (ss *Set) GetRange( + start, stop int, + withScores bool, + reverse bool, +) []string { + length := ss.tree.Len() + if start < 0 { + start += length + } + if stop < 0 { + stop += length + } + + if start < 0 { + start = 0 + } + if stop >= length { + stop = length - 1 + } + + if start > stop || start >= length { + return []string{} + } + + var result []string + + index := 0 + + // iterFunc is the function that will be called for each item in the B-tree. It will append the item to the result if it is within the specified range. + // It will return false if the specified range has been reached. + iterFunc := func(item btree.Item) bool { + if index > stop { + return false + } + + if index >= start { + ssi := item.(*Item) + result = append(result, ssi.Member) + if withScores { + // Use 'g' format to match Redis's float formatting + scoreStr := strings.ToLower(strconv.FormatFloat(ssi.Score, 'g', -1, 64)) + result = append(result, scoreStr) + } + } + index++ + return true + } + + if reverse { + ss.tree.Descend(iterFunc) + } else { + ss.tree.Ascend(iterFunc) + } + + return result +} + +func (ss *Set) Get(member string) (float64, bool) { + score, exists := ss.memberMap[member] + return score, exists +} diff --git a/internal/server/utils/round.go b/internal/server/utils/round.go new file mode 100644 index 000000000..467c01eab --- /dev/null +++ b/internal/server/utils/round.go @@ -0,0 +1,9 @@ +package utils + +import "math" + +// RoundToDecimals rounds a float64 or float32 to a specified number of decimal places. +func RoundToDecimals[T float32 | float64](num T, decimals int) T { + pow := math.Pow(10, float64(decimals)) + return T(math.Round(float64(num)*pow) / pow) +}