@@ -149,22 +149,111 @@ func GetDEnumAsStringFromConstantExpr(expr opt.Expr) (enumAsString string, ok bo
149149 return "" , false
150150}
151151
152- // BuildLookupJoinLookupTableDistribution builds the Distribution that results
153- // from performing lookups of a LookupJoin, if that distribution can be
154- // statically determined. If crdbRegionColID is non-zero, it is the column ID
155- // of the input REGIONAL BY ROW table holding the crdb_region column, and
156- // inputDistribution is the distribution of the operation on that table
157- // (Scan or LocalityOptimizedSearch).
158- // The distribution of the lookup join is returned, plus the first lookup index
159- // column, as matched in the lookup with crdbRegionColID (if it can be
160- // determined, otherwise zero).
152+ // getCRBDRegionColSetFromInput finds the set of column ids in the input to a
153+ // lookup join which are known to all be `crdb_region` columns, either in a Scan
154+ // or by being a `crdb_region` lookup table key column in a lookup join or chain
155+ // of lookup joins, all connected by `crdb_region` key column lookups. This
156+ // column set plus the distribution of the lookup join's input relation is
157+ // returned.
158+ func getCRBDRegionColSetFromInput (
159+ ctx context.Context ,
160+ evalCtx * eval.Context ,
161+ join * memo.LookupJoinExpr ,
162+ required * physical.Required ,
163+ maybeGetBestCostRelation func (grp memo.RelExpr , required * physical.Required ) (best memo.RelExpr , ok bool ),
164+ ) (crdbRegionColSet opt.ColSet , inputDistribution physical.Distribution ) {
165+ var needRemap bool
166+ var setOpCols opt.ColSet
167+
168+ if bestCostInputRel , ok := maybeGetBestCostRelation (join .Input , required ); ok {
169+ maybeScan := bestCostInputRel
170+ var projectExpr * memo.ProjectExpr
171+ if projectExpr , ok = maybeScan .(* memo.ProjectExpr ); ok {
172+ maybeScan , ok = maybeGetBestCostRelation (projectExpr .Input , required )
173+ if ! ok {
174+ return crdbRegionColSet , physical.Distribution {}
175+ }
176+ }
177+ if selectExpr , ok := maybeScan .(* memo.SelectExpr ); ok {
178+ maybeScan , ok = maybeGetBestCostRelation (selectExpr .Input , required )
179+ if ! ok {
180+ return crdbRegionColSet , physical.Distribution {}
181+ }
182+ }
183+ if indexJoinExpr , ok := maybeScan .(* memo.IndexJoinExpr ); ok {
184+ maybeScan , ok = maybeGetBestCostRelation (indexJoinExpr .Input , required )
185+ if ! ok {
186+ return crdbRegionColSet , physical.Distribution {}
187+ }
188+ }
189+ if lookupJoinExpr , ok := maybeScan .(* memo.LookupJoinExpr ); ok {
190+ crdbRegionColSet , inputDistribution =
191+ BuildLookupJoinLookupTableDistribution (
192+ ctx , evalCtx , lookupJoinExpr , required , maybeGetBestCostRelation )
193+ return crdbRegionColSet , inputDistribution
194+ }
195+ if localityOptimizedScan , ok := maybeScan .(* memo.LocalityOptimizedSearchExpr ); ok {
196+ maybeScan = localityOptimizedScan .Local
197+ needRemap = true
198+ setOpCols = localityOptimizedScan .Relational ().OutputCols
199+ }
200+ scanExpr , ok := maybeScan .(* memo.ScanExpr )
201+ if ! ok {
202+ return crdbRegionColSet , physical.Distribution {}
203+ }
204+ tab := maybeScan .Memo ().Metadata ().Table (scanExpr .Table )
205+ if ! tab .IsRegionalByRow () {
206+ return crdbRegionColSet , physical.Distribution {}
207+ }
208+ inputDistribution =
209+ BuildProvided (ctx , evalCtx , scanExpr , & required .Distribution )
210+ index := tab .Index (scanExpr .Index )
211+ crdbRegionColID := scanExpr .Table .IndexColumnID (index , 0 )
212+ if needRemap {
213+ scanCols := scanExpr .Relational ().OutputCols
214+ if scanCols .Len () == setOpCols .Len () {
215+ destCol , _ := setOpCols .Next (0 )
216+ for srcCol , ok := scanCols .Next (0 ); ok ; srcCol , ok = scanCols .Next (srcCol + 1 ) {
217+ if srcCol == crdbRegionColID {
218+ crdbRegionColID = destCol
219+ break
220+ }
221+ destCol , _ = setOpCols .Next (destCol + 1 )
222+ }
223+ }
224+ }
225+ if projectExpr != nil {
226+ if ! projectExpr .Passthrough .Contains (crdbRegionColID ) {
227+ return crdbRegionColSet , physical.Distribution {}
228+ }
229+ }
230+ crdbRegionColSet .Add (crdbRegionColID )
231+ }
232+ return crdbRegionColSet , inputDistribution
233+ }
234+
235+ // BuildLookupJoinLookupTableDistribution builds and returns the Distribution of
236+ // a lookup table in a lookup join if that distribution can be statically
237+ // determined. It also finds the set of column ids in the input to the lookup
238+ // join which are known to all be `crdb_region` columns, either in a Scan or by
239+ // being a `crdb_region` lookup table key column in a lookup join or chain of
240+ // lookup joins, all connected by `crdb_region` key column lookups. This column
241+ // set plus the distribution of the lookup join's lookup table is returned.
242+ // Parameter `maybeGetBestCostRelation` is expected to be a function that looks
243+ // up the best cost relation in the `grp` relation's memo group. Typically this
244+ // is `xform.Coster.MaybeGetBestCostRelation`. It's passed as a function since
245+ // `BuildLookupJoinLookupTableDistribution` may be called from packages not
246+ // allowed to import the `xform` package due to import cycles.
161247func BuildLookupJoinLookupTableDistribution (
162248 ctx context.Context ,
163249 evalCtx * eval.Context ,
164250 lookupJoin * memo.LookupJoinExpr ,
165- crdbRegionColID opt.ColumnID ,
166- inputDistribution physical.Distribution ,
167- ) (firstLookupIndexCol opt.ColumnID , provided physical.Distribution ) {
251+ required * physical.Required ,
252+ maybeGetBestCostRelation func (grp memo.RelExpr , required * physical.Required ) (best memo.RelExpr , ok bool ),
253+ ) (crdbRegionColSet opt.ColSet , provided physical.Distribution ) {
254+ localCrdbRegionColSet , inputDistribution :=
255+ getCRBDRegionColSetFromInput (ctx , evalCtx , lookupJoin , required , maybeGetBestCostRelation )
256+
168257 lookupTableMeta := lookupJoin .Memo ().Metadata ().TableMeta (lookupJoin .Table )
169258 lookupTable := lookupTableMeta .Table
170259
@@ -175,13 +264,13 @@ func BuildLookupJoinLookupTableDistribution(
175264
176265 if lookupJoin .LocalityOptimized || lookupJoin .ChildOfLocalityOptimizedSearch {
177266 provided .FromLocality (evalCtx .Locality )
178- return 0 /* firstLookupIndexCol */ , provided
267+ return crdbRegionColSet , provided
179268 } else if lookupTable .IsGlobalTable () {
180269 provided .FromLocality (evalCtx .Locality )
181- return 0 /* firstLookupIndexCol */ , provided
270+ return crdbRegionColSet , provided
182271 } else if homeRegion , ok := lookupTable .HomeRegion (); ok {
183272 provided .Regions = []string {homeRegion }
184- return 0 /* firstLookupIndexCol */ , provided
273+ return crdbRegionColSet , provided
185274 } else if lookupTable .IsRegionalByRow () {
186275 if len (lookupJoin .KeyCols ) > 0 {
187276 inputExpr := lookupJoin .Input
@@ -190,46 +279,56 @@ func BuildLookupJoinLookupTableDistribution(
190279 if filterExpr , ok := invertedJoinExpr .GetConstExprFromFilter (firstKeyColID ); ok {
191280 if homeRegion , ok = GetDEnumAsStringFromConstantExpr (filterExpr ); ok {
192281 provided .Regions = []string {homeRegion }
193- return colIDOfFirstLookupIndexColumn , provided
282+ crdbRegionColSet .UnionWith (localCrdbRegionColSet )
283+ crdbRegionColSet .Add (colIDOfFirstLookupIndexColumn )
284+ return crdbRegionColSet , provided
194285 }
195286 }
196287 } else if projectExpr , ok := inputExpr .(* memo.ProjectExpr ); ok {
197288 regionName := projectExpr .GetProjectedEnumConstant (firstKeyColID )
198289 if regionName != "" {
199290 provided .Regions = []string {regionName }
200- return colIDOfFirstLookupIndexColumn , provided
291+ crdbRegionColSet .UnionWith (localCrdbRegionColSet )
292+ crdbRegionColSet .Add (colIDOfFirstLookupIndexColumn )
293+ return crdbRegionColSet , provided
201294 }
202295 }
203- if crdbRegionColID == firstKeyColID {
296+ if localCrdbRegionColSet . Contains ( firstKeyColID ) {
204297 provided .FromIndexScan (ctx , evalCtx , lookupTableMeta , lookupJoin .Index , nil )
205298 if ! inputDistribution .Any () &&
206299 (provided .Any () || len (provided .Regions ) > len (inputDistribution .Regions )) {
207- return colIDOfFirstLookupIndexColumn , inputDistribution
300+ crdbRegionColSet .UnionWith (localCrdbRegionColSet )
301+ crdbRegionColSet .Add (colIDOfFirstLookupIndexColumn )
302+ return crdbRegionColSet , inputDistribution
208303 }
209- return 0 /* firstLookupIndexCol */ , provided
304+ return crdbRegionColSet , provided
210305 }
211306 } else if len (lookupJoin .LookupJoinPrivate .LookupExpr ) > 0 {
212307 if filterIdx , ok := lookupJoin .GetConstPrefixFilter (lookupJoin .Memo ().Metadata ()); ok {
213308 firstIndexColEqExpr := lookupJoin .LookupJoinPrivate .LookupExpr [filterIdx ].Condition
214309 if firstIndexColEqExpr .Op () == opt .EqOp {
215310 if regionName , ok := GetDEnumAsStringFromConstantExpr (firstIndexColEqExpr .Child (1 )); ok {
216311 provided .Regions = []string {regionName }
217- return colIDOfFirstLookupIndexColumn , provided
312+ crdbRegionColSet .UnionWith (localCrdbRegionColSet )
313+ crdbRegionColSet .Add (colIDOfFirstLookupIndexColumn )
314+ return crdbRegionColSet , provided
218315 }
219316 }
220- } else if lookupJoin .ColIsEquivalentWithLookupIndexPrefix (lookupJoin .Memo ().Metadata (), crdbRegionColID ) {
317+ } else if lookupJoin .LookupIndexPrefixIsEquatedWithColInColSet (lookupJoin .Memo ().Metadata (), localCrdbRegionColSet ) {
221318 // We have a `crdb_region = crdb_region` term in `LookupJoinPrivate.LookupExpr`.
222319 provided .FromIndexScan (ctx , evalCtx , lookupTableMeta , lookupJoin .Index , nil )
223320 if ! inputDistribution .Any () &&
224321 (provided .Any () || len (provided .Regions ) > len (inputDistribution .Regions )) {
225- return colIDOfFirstLookupIndexColumn , inputDistribution
322+ crdbRegionColSet .UnionWith (localCrdbRegionColSet )
323+ crdbRegionColSet .Add (colIDOfFirstLookupIndexColumn )
324+ return crdbRegionColSet , inputDistribution
226325 }
227- return 0 /* firstLookupIndexCol */ , provided
326+ return crdbRegionColSet , provided
228327 }
229328 }
230329 }
231330 provided .FromIndexScan (ctx , evalCtx , lookupTableMeta , lookupJoin .Index , nil )
232- return 0 /* firstLookupIndexCol */ , provided
331+ return crdbRegionColSet , provided
233332}
234333
235334// BuildInvertedJoinLookupTableDistribution builds the Distribution that results
0 commit comments